Python 获取某个进程和系统gpu使用情况的脚本

以下是一个使用Python获取某个进程和系统GPU使用情况的示例脚本：

import subprocess

def get_process_gpu_usage(process_name):
    try:
        # 使用nvidia-smi命令获取GPU信息
        output = subprocess.check_output(['nvidia-smi', '--query-gpu=utilization.gpu', '--format=csv'])
        gpu_info = output.decode('utf-8').strip().split('\n')[1:]  # 去除标题行
        gpu_usage = [int(line.split(',')[0].strip().strip('%')) for line in gpu_info]
        
        # 使用psutil库获取进程信息
        import psutil
        process_list = [p for p in psutil.process_iter(['pid', 'name']) if process_name in p.info['name']]
        process_gpu_usage = []
        for process in process_list:
            pid = process.info['pid']
            try:
                # 使用nvidia-smi命令获取进程GPU使用情况
                output = subprocess.check_output(['nvidia-smi', '--query-compute-apps=pid,utilization.gpu', '--format=csv,nounits'])
                process_info = output.decode('utf-8').strip().split('\n')[1:]  # 去除标题行
                process_gpu_usage.extend([int(line.split(',')[1].strip().strip('%')) for line in process_info if int(line.split(',')[0]) == pid])
            except subprocess.CalledProcessError:
                pass
        
        return gpu_usage, process_gpu_usage
    except Exception as e:
        print(f"Error: {e}")
        return None, None

# 例子：获取名为"example_process"的进程和系统GPU使用情况
gpu_usage, process_gpu_usage = get_process_gpu_usage("example_process")

# 打印结果
print(f"System GPU Usage: {gpu_usage}%")
print(f"Process GPU Usage: {process_gpu_usage}%")

请注意，此脚本依赖于nvidia-smi和psutil库。确保您的系统中已安装这些库并且可从命令行访问nvidia-smi