PyTorch/CUDA 多GPU环境管理：从nvidia-smi到pynvml的5个高级脚本

Mu Tian

326人浏览 · 2026-07-04 10:40:17

Mu Tian · 2026-07-04 10:40:17 发布

PyTorch/CUDA 多GPU环境管理：从nvidia-smi到pynvml的5个高级脚本

在共享GPU集群的开发环境中，高效管理多GPU资源往往比单纯的技术实现更具挑战性。当多个研究团队共用有限的GPU设备时，开发者常面临三大痛点：无法快速识别空闲GPU、难以监控长期任务的显存泄漏、以及僵尸进程导致的资源死锁。传统解决方案依赖人工执行nvidia-smi命令，这种被动响应模式在复杂场景下显得力不从心。

1. 智能GPU选择器：动态分配空闲设备

常规的CUDA_VISIBLE_DEVICES指定方式存在明显局限——它要求开发者预先知道哪些GPU可用。我们开发的自适应选择脚本能实时分析设备状态，自动分配符合计算需求的GPU资源。

import pynvml
import torch

def select_idle_gpu(min_memory=1024, max_retry=3):
    pynvml.nvmlInit()
    for _ in range(max_retry):
        available_gpus = []
        for i in range(pynvml.nvmlDeviceGetCount()):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            util = pynvml.nvmlDeviceGetUtilizationRates(handle)
            
            if mem_info.free/1024**2 > min_memory and util.gpu < 50:
                available_gpus.append(str(i))
        
        if available_gpus:
            torch.cuda.set_device(int(available_gpus[0]))
            return f"CUDA_VISIBLE_DEVICES={','.join(available_gpus)}"
    
    raise RuntimeError("No available GPU meets the criteria")

# 使用示例
os.environ.update(eval(select_idle_gpu(min_memory=2048)))

该脚本实现了三个关键功能：

动态阈值筛选 ：同时考虑显存余量（默认>1GB）和GPU利用率（<50%）
优先级策略 ：返回满足条件的首个设备，避免随机选择的不确定性
异常重试机制 ：在竞争激烈的环境中自动进行多次尝试

提示：将min_memory参数设置为模型预估显存占用的1.2倍，可有效防止内存不足错误

2. 显存监控器：实时追踪进程内存变化

长期运行的训练任务可能出现显存缓慢增长的问题。以下脚本以1秒间隔记录指定进程的显存变化，帮助开发者定位内存泄漏。

import time
from collections import deque

def monitor_gpu_memory(pid, duration=3600, interval=1):
    pynvml.nvmlInit()
    history = deque(maxlen=duration//interval)
    
    try:
        while True:
            current_mem = 0
            for i in range(pynvml.nvmlDeviceGetCount()):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
                for p in procs:
                    if p.pid == pid:
                        current_mem += p.usedGpuMemory/1024**2
            
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
            history.append((timestamp, current_mem))
            print(f"[{timestamp}] PID {pid} using {current_mem:.2f}MB")
            
            if len(history) > 1 and current_mem - history[0][1] > 500:
                print(f"Warning: Memory increased by {current_mem-history[0][1]:.2f}MB")
            
            time.sleep(interval)
    except KeyboardInterrupt:
        return list(history)

该监控器具有以下特点：

功能	实现方式	典型应用场景
多GPU聚合	累加进程在所有设备上的显存使用	分布式训练监控
趋势预警	滑动窗口比较内存变化量	早期发现内存泄漏
时间戳记录	标准化时间格式输出	与日志系统集成

3. GPU压力测试工具：模拟高负载场景

在部署新模型前，需要验证GPU在极端条件下的稳定性。下面的脚本可以模拟不同级别的显存和计算负载：

def gpu_stress_test(device_index, mem_percent=0.8, duration=60):
    device = torch.device(f'cuda:{device_index}')
    total_mem = torch.cuda.get_device_properties(device).total_memory
    block_size = int(total_mem * mem_percent / 10)
    
    # 创建内存压力
    blocks = []
    try:
        for i in range(10):
            blocks.append(torch.randn(block_size, device=device))
        
        # 创建计算压力
        start = time.time()
        while time.time() - start < duration:
            x = torch.randn(10000, 10000, device=device)
            torch.mm(x, x.t())
        
    finally:
        del blocks
        torch.cuda.empty_cache()

参数调节建议：

mem_percent ：0.5-0.9，模拟不同内存占用率
duration ：测试持续时间（秒）
block_size ：将内存分配拆分为多个块，避免单次分配失败

4. 僵尸进程清理工具

被异常终止的PyTorch进程可能继续占用GPU资源。以下脚本自动识别并清理这些"僵尸"：

def clean_zombie_processes():
    pynvml.nvmlInit()
    zombie_found = False
    
    for i in range(pynvml.nvmlDeviceGetCount()):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
        
        for p in procs:
            try:
                os.kill(p.pid, 0)
            except OSError:
                print(f"Killing zombie process {p.pid} on GPU {i}")
                os.kill(p.pid, 9)
                zombie_found = True
    
    if not zombie_found:
        print("No zombie processes detected")
    return zombie_found

清理策略说明：

遍历所有GPU设备上运行的进程
尝试向进程发送信号0（无操作）检测存活状态
对无响应的进程发送SIGKILL(9)
返回是否发现僵尸进程的布尔值

5. GPU使用报告生成器

定期生成资源使用报告有助于优化集群调度策略。这个脚本生成包含关键指标的HTML报告：

def generate_gpu_report(output_file="gpu_report.html"):
    pynvml.nvmlInit()
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    
    gpu_data = []
    for i in range(pynvml.nvmlDeviceGetCount()):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        name = pynvml.nvmlDeviceGetName(handle)
        mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        util = pynvml.nvmlDeviceGetUtilizationRates(handle)
        temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
        
        gpu_data.append({
            "index": i,
            "name": name.decode(),
            "mem_used": mem_info.used/1024**3,
            "mem_total": mem_info.total/1024**3,
            "util_gpu": util.gpu,
            "util_mem": util.memory,
            "temp": temp
        })
    
    # 生成HTML报告
    html = f"""<html><head><title>GPU Report {timestamp}</title>
             <style>table {{border-collapse: collapse;}} td, th {{border: 1px solid #ddd; padding: 8px;}}</style></head>
             <body><h1>GPU Utilization Report</h1><p>Generated at {timestamp}</p>
             <table><tr><th>GPU</th><th>Name</th><th>Mem Used</th><th>Util %</th><th>Temp °C</th></tr>"""
    
    for gpu in gpu_data:
        mem_pct = gpu["mem_used"] / gpu["mem_total"] * 100
        html += f"""<tr><td>{gpu['index']}</td><td>{gpu['name']}</td>
                 <td>{gpu['mem_used']:.1f}G/{gpu['mem_total']:.1f}G ({mem_pct:.1f}%)</td>
                 <td>GPU: {gpu['util_gpu']}% | Mem: {gpu['util_mem']}%</td>
                 <td>{gpu['temp']}</td></tr>"""
    
    html += "</table></body></html>"
    
    with open(output_file, "w") as f:
        f.write(html)
    return output_file

报告包含的关键指标：