#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ GPU显存清理脚本（强化版） 在任务完成后调用，彻底释放GPU显存和内存 支持强制杀死占用GPU的Python进程 """ import sys import subprocess import os import gc import time def get_gpu_processes(): """获取占用GPU的进程列表""" try: result = subprocess.run( "nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv,noheader", shell=True, capture_output=True, text=True, timeout=5 ) processes = [] if result.returncode == 0 and result.stdout.strip(): for line in result.stdout.strip().split('\n'): if line.strip(): parts = [p.strip() for p in line.split(',')] if len(parts) >= 3: pid, name, memory = parts processes.append({ 'pid': int(pid), 'name': name, 'memory': memory }) return processes except Exception as e: print("[GPU清理] 获取GPU进程失表: {}".format(e)) return [] def get_all_python_processes(): """获取所有Python进程列表（包括孤儿进程）""" try: result = subprocess.run( "tasklist /FI \"IMAGENAME eq python.exe\" /FO CSV /NH", shell=True, capture_output=True, text=True, timeout=5 ) processes = [] if result.returncode == 0 and result.stdout.strip(): for line in result.stdout.strip().split('\n'): if line.strip(): parts = [p.strip('"') for p in line.split(',')] if len(parts) >= 2: name = parts[0] pid = parts[1] try: processes.append({ 'pid': int(pid), 'name': name }) except: pass return processes except Exception as e: print("[GPU清理] 获取Python进程失败: {}".format(e)) return [] def should_kill_process(pid, model_paths=None): """判断进程是否应该被清理（只清理AI模型相关进程） Args: pid: 进程ID model_paths: 模型路径列表，如果提供，只清理这些路径下的进程 """ try: # 获取进程命令行 cmdline_result = subprocess.run( "wmic process where ProcessId={} get CommandLine /format:list".format(pid), shell=True, capture_output=True, text=True, timeout=2 ) if cmdline_result.returncode != 0 or not cmdline_result.stdout: return False cmdline = cmdline_result.stdout.lower() # 白名单（包含这些关键词的进程不清理） whitelist_keywords = [ '直播伴侣', 'finderliveobs', 'xwechat', 'obs', 'obs-studio', 'streamlabs', 'douyin', 'tiktok', 'bilibili', '抖音', 'kuaishou', '快手', 'pycharm', 'vscode', 'visual studio code', 'jupyter', 'spyder', 'anaconda', 'conda', 'sublime', 'notepad++', 'wechat', 'weixin', '微信', 'qq', 'asr', 'funasr', 'aigcpanel', 'whisper', 'cosyvoice', ] # 检查白名单 for keyword in whitelist_keywords: if keyword in cmdline: # 提取进程名称（从命令行中获取） process_name = "未知" if "python" in cmdline: process_name = "Python进程" elif "finderliveobs" in cmdline or "xwechat" in cmdline: process_name = "视频号直播伴侣" elif "obs" in cmdline: process_name = "OBS直播软件" elif "wechat" in cmdline or "weixin" in cmdline: process_name = "微信" print("[GPU清理] [保护] 跳过白名单进程 PID={} ({}, 匹配关键词: {})".format(pid, process_name, keyword)) return False # 🔧 必须提供模型路径才进行清理（安全模式） if model_paths and len(model_paths) > 0: for model_path in model_paths: # 标准化路径格式（统一使用小写和反斜杠） normalized_path = model_path.lower().replace('/', '\\') if normalized_path in cmdline: print("[GPU清理] 匹配模型路径 PID={} (路径: {})".format(pid, model_path)) return True # 如果提供了路径但都不匹配，不清理 return False else: # 🔧 安全修复：没有提供路径时，不清理任何进程（防止误杀） # 之前的关键词匹配模式已移除，确保只清理模型路径下的进程 print("[GPU清理] 警告: 未提供模型路径，跳过进程 PID={}".format(pid)) return False except Exception as e: print("[GPU清理] 检查进程 PID={} 失败: {}".format(pid, e)) return False def cleanup_gpu(force_kill=True, model_paths=None): """彻底清理GPU显存和内存（默认强力模式） Args: force_kill: 是否强制清理 model_paths: 模型路径列表，只清理这些路径下的进程 """ try: print("[GPU清理] ========== 开始GPU显存和内存清理 ==========") if model_paths and len(model_paths) > 0: print("[GPU清理] [模式] 路径过滤：只清理模型路径下的Python进程") print("[GPU清理] [路径] 共 {} 个模型路径:".format(len(model_paths))) for path in model_paths: print("[GPU清理] - {}".format(path)) else: print("[GPU清理] [模式] 智能清理：只清理AI模型进程，保护其他Python应用") # ==================== 步骤1：PyTorch GPU深度清理 ==================== print("[GPU清理] [步骤1] PyTorch GPU深度清理...") try: import torch if torch.cuda.is_available(): # 显示清理前的显存状态 allocated_before = torch.cuda.memory_allocated() / 1024**3 reserved_before = torch.cuda.memory_reserved() / 1024**3 print("[GPU清理] 清理前: 已分配={:.2f}GB, 已保留={:.2f}GB".format(allocated_before, reserved_before)) # 1. 清空CUDA缓存（多次循环确保彻底清理） print("[GPU清理] 正在清空CUDA缓存...") for i in range(5): torch.cuda.empty_cache() gc.collect() time.sleep(0.1) # 2. 同步所有CUDA设备 print("[GPU清理] 同步CUDA设备...") torch.cuda.synchronize() # 3. 重置内存统计 torch.cuda.reset_peak_memory_stats() torch.cuda.reset_accumulated_memory_stats() # 4. 尝试释放所有未使用的缓存内存 try: # 清理内存分配器的缓存 torch.cuda.memory.empty_cache() except: pass # 5. 设置内存分配器配置（优化内存使用） try: os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512' except: pass # 显示清理后的显存状态 allocated_after = torch.cuda.memory_allocated() / 1024**3 reserved_after = torch.cuda.memory_reserved() / 1024**3 freed_memory = (reserved_before - reserved_after) print("[GPU清理] 清理后: 已分配={:.2f}GB, 已保留={:.2f}GB".format(allocated_after, reserved_after)) print("[GPU清理] 释放显存: {:.2f}GB".format(freed_memory)) print("[GPU清理] [OK] PyTorch GPU清理完成") else: print("[GPU清理] CUDA不可用，跳过GPU清理") except ImportError: print("[GPU清理] PyTorch未安装，跳过GPU清理") except Exception as e: print("[GPU清理] PyTorch清理出错: {}".format(e)) # ==================== 步骤2：强制系统内存清理 ==================== print("[GPU清理] [步骤2] 执行强制系统内存清理...") # 1. 多轮垃圾回收（清理不同代的对象） total_collected = 0 for i in range(5): collected = gc.collect(generation=2) # 清理所有代 total_collected += collected time.sleep(0.05) print("[GPU清理] 垃圾回收: 清理了 {} 个对象".format(total_collected)) # 2. 清理Python内部缓存 try: import ctypes # 尝试释放Python未使用的内存回操作系统 if hasattr(ctypes, 'windll'): # Windows平台 ctypes.windll.kernel32.SetProcessWorkingSetSize(-1, -1, -1) except: pass print("[GPU清理] [OK] 系统内存清理完成") # ==================== 步骤3：显示当前状态 ==================== print("[GPU清理] [步骤3] 检查GPU状态...") try: result = subprocess.run( "nvidia-smi --query-gpu=index,memory.used,memory.free,memory.total --format=csv,noheader", shell=True, capture_output=True, text=True, timeout=5 ) if result.returncode == 0: print("[GPU清理] ========== GPU 显存状态 ==========") for line in result.stdout.strip().split('\n'): if line.strip(): print("[GPU清理] {}".format(line.strip())) except: pass # ==================== 步骤4：清理所有Python进程（包括孤儿进程） ==================== print("[GPU清理] [步骤4] 检查并清理所有Python进程...") # 获取所有Python进程（包括GPU和非GPU进程） all_python_processes = get_all_python_processes() if all_python_processes: print("[GPU清理] 检测到 {} 个Python进程（包括AI任务和孤儿进程）".format(len(all_python_processes))) # 过滤需要清理的进程 print("[GPU清理] [智能过滤] 正在识别需要清理的进程...") processes_to_kill = [] skipped_count = 0 for proc in all_python_processes: if should_kill_process(proc['pid'], model_paths): processes_to_kill.append(proc) print("[GPU清理] PID={}, 将被清理 (AI模型进程)".format(proc['pid'])) else: skipped_count += 1 print("[GPU清理] 识别完成: {} 个进程将被清理, {} 个进程已跳过".format( len(processes_to_kill), skipped_count )) # 清理识别出的AI模型进程 if processes_to_kill: print("[GPU清理] [智能清理] 开始清理AI模型进程...") killed_count = 0 failed_pids = [] for proc in processes_to_kill: try: # 使用 PowerShell 强制杀死进程（更可靠） result = subprocess.run( 'powershell -Command "Stop-Process -Id {} -Force"'.format(proc['pid']), shell=True, capture_output=True, timeout=3 ) if result.returncode == 0: print("[GPU清理] [OK] 已清理进程 PID={}".format(proc['pid'])) killed_count += 1 else: failed_pids.append(proc['pid']) print("[GPU清理] [FAIL] 清理进程 PID={} 失败".format(proc['pid'])) time.sleep(0.2) except Exception as e: failed_pids.append(proc['pid']) print("[GPU清理] [ERROR] 清理进程PID={} 异常: {}".format(proc['pid'], str(e))) print("[GPU清理] [完成] 成功清理 {} 个进程, 失败 {} 个".format(killed_count, len(failed_pids))) else: print("[GPU清理] [OK] 没有需要清理的AI模型进程") # 清理后再次检查 time.sleep(1) remaining = get_all_python_processes() remaining_ai_processes = [p for p in remaining if should_kill_process(p['pid'], model_paths)] if not remaining_ai_processes: print("[GPU清理] [OK] AI模型进程已清理完成") else: print("[GPU清理] [提示] 仍有 {} 个AI模型进程存在".format(len(remaining_ai_processes))) for proc in remaining_ai_processes: print("[GPU清理] 残留进程 PID={}".format(proc['pid'])) else: print("[GPU清理] [OK] 没有检测到Python进程") # ==================== 步骤5：最终GPU显存清理 ==================== print("[GPU清理] [步骤5] 最终GPU显存清理...") try: import torch if torch.cuda.is_available(): # 进程清理后，再次清理GPU缓存 for i in range(3): torch.cuda.empty_cache() gc.collect() time.sleep(0.1) # 显示最终状态 final_allocated = torch.cuda.memory_allocated() / 1024**3 final_reserved = torch.cuda.memory_reserved() / 1024**3 print("[GPU清理] 最终状态: 已分配={:.2f}GB, 已保留={:.2f}GB".format(final_allocated, final_reserved)) print("[GPU清理] [OK] 最终清理完成") else: print("[GPU清理] CUDA不可用") except: pass print("[GPU清理] ========== GPU和内存清理完成 ==========\n") return True except Exception as e: print("[GPU清理] [ERROR] 清理失败: {}".format(str(e))) import traceback traceback.print_exc() return False if __name__ == "__main__": # 解析命令行参数获取模型路径 model_paths = [] if len(sys.argv) > 1: # 从命令行参数获取模型路径（以逗号分隔） paths_arg = sys.argv[1] if paths_arg and paths_arg.strip(): model_paths = [p.strip() for p in paths_arg.split(',') if p.strip()] if model_paths: print("[GPU清理] [信息] 路径过滤模式：只清理指定模型路径下的Python进程") print("[GPU清理] [信息] 受保护的应用：直播伴侣、OBS、IDE等不会被清理") print("[GPU清理] [信息] 模型路径数: {}\n".format(len(model_paths))) else: print("[GPU清理] [警告] 未提供模型路径参数！") print("[GPU清理] [说明] 脚本仅在提供模型路径时才进行进程清理") print("[GPU清理] [说明] 用法: python cleanup_gpu.py \"path1,path2,path3\"") print("[GPU清理] [信息] 将执行GPU显存和系统内存清理（不清理任何进程）\n") cleanup_gpu(force_kill=True, model_paths=model_paths) sys.exit(0)