NVIDIA深入理解之pynvml库

一、前言

写在前面

该文章是对我之前文章《Fedora上安装NVIDIA闭源显卡驱动》的一个拓展，正好寒假闲的没事干不如加深一下对NVIDIA的了解。Python是当前非常流行的一门编程语言，它以kiss为设计思想，能封装就能封装，给用户提供比较良好且便于理解的编程体验，那么我们尽量要了解的这个库叫做pynvml,它是Python的一个第三方库，提供了对NVIDIA的管理库（NVML）的接口，使得开发者在正确安装NVIDIA闭源驱动（无论是run包还是包管理器封装的驱动）后，都可以使用该库与NVIDIA GPU进行交互，以获取GPU当前的一些状态信息和参数。为了简便称呼，我以下提及的“GPU”专门指的是NVIDIA显卡，本文章默认你已经安装好了闭源驱动，如果没有正确安装，可以去我之前的文章看一看。

在Windows操作系统中，我们去查看GPU信息是一件非常简单且合理的事情，实际上只要打开Windows任务管理器，里面就已经有GPU的各种信息，有一个图形界面就是好。但是在Linux中，查看GPU信息真是一件貌似不太容易的事情，实际上也是非常容易的，你看完我的文章就会清楚了。

实际上我们并不一定要通过pynvml这个库去了解GPU状态信息，实际上等你安装好闭源驱动之后，它本身就提供了一个脚本去查看一些信息，比如nvidia-smi。这个我们在之前的文章里就已经演示过了，这里不再赘述（你在terminal直接敲这个命令就好了）。这里我再推荐若干好用的工具帮你查看GPU信息。

当你学习并了解了NVIDIA显卡驱动的API使用，你也许会好奇AMD显卡和Intel显卡驱动是否也提供若干API以供用户参考和使用，很抱歉我没有搭载AMD显卡的设备，不能实操说明，不过你可以搜索rocm-smi这个工具，并且py3nvml库虽然是主要为NVIDIA显卡设计的，但是它也貌似支持AMD显卡的基本信息显示，你可以试一试，然后告诉我。

工具介绍

1、Linux风格的工具 -- nvtop

大家知道在Linux中查看进程/内存占用等信息的工具叫top,那么nvtop就是专门用来查看GPU显卡信息的工具，是非常*iux的工具。

sudo dnf install nvtop

效果如图所示，大家在使用Linux时一定要使用这种terminal风格的信息显示，特别是服务器管理的时候很少是安装了图形界面的，因此命令行才是最高效的查看和管理服务器的方法。

2、Windows风格的工具 -- Mission Center

Mission Center | Flathub

这是flathub上一个非常有名的模仿Windows风格任务管理器的资源查看器。

只有在安装NVIDIA闭源驱动之后GPU0那一项才会显示，否则是没有那一项的，请大家注意。

3、KDE桌面自带了资源显示功能

上图中左下角黑框框里面的就是GPU信息，可以放在桌面上实时查看，是不是非常炫酷呢？

二、正文

简单介绍pynvml的含义

好了，以上都是前言部分，下面我们开始正式介绍pynvml库（是不是都快忘记主题了）。

首先我们必须清楚pynvml为什么叫这个名字，我一开始了解的时候也感觉很困惑，这个名字太奇怪了，老是敲错这个nvml.NVML是NVIDIA Manager Lib的简写，它是由NVIDIA官方使用c语言编写的、用于管理和控制NVIDIA硬件的一组程序。NVML提供一组API,开发者可以通过API来查询和控制NVIDIA显卡的状态和配置，比如电源管理、温度控制、性能监控等。NVML和CUDA都是NVIDIA提供的工具，它们通常被显卡开发者同时使用以发挥GPU的最大性能。

NVIDIA Management Library (NVML) | NVIDIA Developer

NVML的官方解释在这里

虽然NVIDIA驱动是闭源的，但是这个驱动对外提供的接口是开放的，NVIDIA官网也有给出详细的函数说明，不过我们使用c/c++去查看显卡信息确实是小题大做了，没有这个必要，这是GPU驱动开发人员应该做的事情。那么对于非NVIDIA驱动开发人员来说，不需要深入挖掘NVIDIA的功能。pynvml这个命名也是一个经典组合，它意味着该库是NVML的Python绑定，除此之外还有PYQT、pygtk等库也是这个命名。

pynvml编程实践

1、安装pynvml

pip install pynvml

Linux环境中如果没有pip,需要先安装pip。

我这里需要说明一下pynvml是比较早期的一个库了，它是兼容Python2和Python3的，我们目前都是使用Python3居多，因此有一个名叫py3nvml的新库只支持Python3,大家也可以使用这个库。

pip install py3nvml

pynvml库开源项目

py3nvml库开源项目

有感兴趣开发可以关注这两个开源项目，可以提交pr给开源做贡献。

我们先介绍pynvml,这个库比较经典，相关的文章也比较多。

2、编程思路

使用pynvml库的基本思路如下：

注意这里说的是基本思路，如果您是驱动开发者，可以不遵守基本思路。

①导包

Python库使用前必须import，导包使用固定格式就好

②初始化

第二步和第五步是成双成对的，由于pynvml会调用显卡驱动程序，因此pynvml在使用前要建立和NVIDIA驱动的连接，在连接的时候是要占用一些系统资源的（比如内存缓冲区、CPU时间），然后使用完了之后需要释放这些资源，因此我们要遵守开发规范。

③获取设备句柄（Handle）

有很多计算机设备是不止一个GPU的，比如一台计算机有两个GPU（核显与独显），甚至多个GPU组成阵列也是非常正常的，因此你得告诉pynvml你要观察的是哪一块GPU,不能张冠李戴啊。

④调用函数

当你获取到Handle之后，也就是告诉pynvml你要观察的GPU号，然后就可以获取信息了，那么获取信息要调用库函数。

⑤释放NVML

有申请就有释放。

那么以下是一个符合开发规范的简单的例子：

编程环境：

OS：Fedora Linux

IDE：vscode

解释器版本：Python3.12.1

from pynvml import * #导包
import humanfriendly

nvmlInit()  #初始化

handle = nvml.nvmlDeviceGetHandleByIndex(0) #获取句柄

#获取GPU温度信息
temperature = nvmlDeviceGetTemperature(handle,NVML_TEMPERATURE_GPU)
print("GPU Temperature:",temperature)

#获取GPU显存信息
memory_info = nvmlDeviceGetMemoryInfo(handle)
print("Tota memory",humanfriendly.format_size(memory_info.total))
print("Free memory",humanfriendly.format_size(memory_info.free))
print("Used memroy",humanfriendly.format_size(memory_info.used))

nvmlShutdown()  #释放

除了pynvml库我们还使用了humanfriendly库，这个库也是需要安装的。

输出结果：

GPU Temperature: 41
Tota memory 6.44 GB
Free memory 5.32 GB
Used memroy 1.12 GB

当前GPU的温度是41摄氏度，显存是6.44GB,已经使用了1.12GB，还剩5.32GB。

我的代码中，直接使用了nvmlDeviceGetHandleByIndex(0)来获取句柄，这是因为我只有一块显卡，所以我知道它就是GPU0,计算机从0开始计数，如果你不知道你的GPU编号的话，可以遍历一遍所有GPU，以确定你要观察哪一块GPU.

from pynvml import *

nvmlInit()

GPU_count = nvmlDeviceGetCount()
print(GPU_count)

for i in range(GPU_count):
    handle = nvmlDeviceGetHandleByIndex(i)
    name = nvmlDeviceGetName(handle=handle)
    print(name.encode('utf-8'))

nvmlShutdown()

输出结果：

1
b'NVIDIA GeForce GTX 1660 Ti'

也就是我总共只有一块GPU,并且这块GPU的名字是GTX 1660Ti 。

我首先获取了我的GPU设备的数量，然后遍历这些GPU设备分别获取它们的句柄，再通过函数调用获得显卡名字，最后打印出来，思路应该比较清晰。

相信大家也看出来了，导包、初始化、获取句柄、释放，这几个操作几乎都是固定搭配，不用修改的，就是函数调用这一步需要学习，那么我们就重点讲函数调用。

我们分析一下这几个函数的名字特点：

其实它们虽然看起来名字长很复杂，其实拆分开来看是相当有规律的，是小驼峰命名法。

nvml + Device + Get + 你要获取什么信息

比如你要获取内存信息：nvmlDeviceGetMemoryInfo

然后函数调用的时候必须携带句柄信息。

了解到这些原则之后编程就简单起来了。

我列出一下一些函数，以供大家参考：

`nvmlDeviceGetCount()`	获取系统中的GPU设备数量
nvmlDeviceGetHandleByIndex(index)	根据设备的索引号获取设备的句柄
nvmlDeviceGetName(handle)	根据设备的句柄获取设备的名称
nvmlDeviceGetMemoryInfo(handle)	根据设备的句柄获取设备的内存信息
nvmlDeviceGetTemperature(handle, sensorType)	根据设备的句柄和传感器类型获取设备的温度
nvmlDeviceGetFanSpeed(handle)	根据设备的句柄获取设备的风扇速度
nvmlDeviceGetPowerState(handle)	根据设备的句柄获取设备的电源状态
nvmlDeviceGetUtilizationRates(handle)	根据设备的句柄获取设备的使用率
nvmlDeviceGetPerformanceState(handle)	根据设备的句柄获取设备的性能状态
nvmlDeviceGetPowerManagementMode(handle)	根据设备的句柄获取设备的电源管理模式
nvmlDeviceGetPowerUsage(handle)	根据设备的句柄获取设备的电源使用情况
nvmlDeviceGetPowerLimit(handle)	根据设备的句柄获取设备的电源限制

这里要提醒一下，并不是所有设备都支持以上的函数，越新的设备会有更多的传感器，会支持更多的函数，旧设备也许因为一些问题，比如没有对应的传感器或缺少部分驱动导致一些函数调用失败，那么如果失败的话是会报错了，大家在编程的时候要注意报错的可能性。

Traceback (most recent call last):
File "/home/april_zhao/文档/practice/Python/nvidia/test2.py", line 19, in <module>
print(nvmlDeviceGetFanSpeed(handle))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/april_zhao/.local/lib/python3.12/site-packages/pynvml/nvml.py", line 2275, in nvmlDeviceGetFanSpeed
_nvmlCheckReturn(ret)
File "/home/april_zhao/.local/lib/python3.12/site-packages/pynvml/nvml.py", line 833, in _nvmlCheckReturn
raise NVMLError(ret)
pynvml.nvml.NVMLError_NotSupported: Not Supported

比如我的这张1660Ti显卡就不支持获取显卡风扇转速，也许是因为笔记本显卡的原因，硬件上可能没有单独的GPU风扇，散热归主板管了，显卡驱动没有能力查看其他硬件的风扇情况。因此这里是直接报错的，"Not Supported"，表明不支持该函数，可以使用try语句去解决这个问题。

除了打印GPU整体信息外，我们还可以打印占用GPU的进程的信息。

from pynvml import *
import humanfriendly

nvmlInit()

handle = nvmlDeviceGetHandleByIndex(0)

processes = nvmlDeviceGetComputeRunningProcesses(handle=handle)
for process in processes:
    print("Process ID:",process.pid)
    print("Memory Used:",humanfriendly.format_size(process.usedGpuMemory))


nvmlShutdown()

它可以去检查占用当前GPU的进程的信息。

输出：

Process ID: 3811
Memory Used: 232.57 MB

那么PID为3811是哪一个进程呢？

原来是我们的Chrome浏览器正在占用GPU,因为我现在正在写文章，所以Chrome肯定是开着的。

以上是NVML低级的绑定，它还有一个对nvidia-smi命令的绑定。

from pynvml.smi import nvidia_smi


nvsmi = nvidia_smi.getInstance()
dict = nvsmi.DeviceQuery('memory.free, memory.total,memory.used')

print("GPU total memory:{0}GB".format(dict['gpu'][0]['fb_memory_usage']['total'] / 1024.0))
print("GPU free memory:{0}GB".format(dict['gpu'][0]['fb_memory_usage']['free'] / 1024.0))
print("GPU usage memory:{0}GB".format(dict['gpu'][0]['fb_memory_usage']['used'] / 1024.0))

输出结果如下：

GPU total memory:6.0GB
GPU free memory:5.05352783203125GB
GPU usage memory:0.94647216796875GB

接下来是py3nvml的例子，实际上在GitHub还是py3nvnml的star数量比较多，说明对它的关注度更大。

from py3nvml.py3nvml import *
import humanfriendly

nvmlInit()
print("Driver Version: {}".format(nvmlSystemGetDriverVersion()))

deviceCount = nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = nvmlDeviceGetHandleByIndex(i)
    print("Device {}: {}".format(i, nvmlDeviceGetName(handle)))
    print("GPU Temp:{}".format(nvmlDeviceGetTemperature(handle,NVML_TEMPERATURE_GPU)))
    processes = nvmlDeviceGetComputeRunningProcesses(handle) # 获取正在运行的计算进程
    for process in processes:
        print("pid={}, used_memory={}".format(process.pid, humanfriendly.format_size(process.usedGpuMemory)))

nvmlShutdown()

py3nvml也有对nvisia-smi的封装

import py3nvml.nvidia_smi as smi

print(smi.XmlDeviceQuery())

这一步操作相当于nvidia-smi -q -x命令，获取xml格式的输出，你可以使用xpath等库解析这个xml文档。

实际上，我们可以获取NV显卡的实时信息然后用图形界面的方式显示出来，进行二次封装，这就是一个不错的点子。

除了我上述提到的函数调用之外，还有很多其他函数能够获取更多的信息，大家在IDE里输入nvmlDeviceGet之后就会给出提示，可以多尝试看输出什么东西出来。

from py3nvml.py3nvml import *
import humanfriendly

nvmlInit()
print("Driver Version: {}".format(nvmlSystemGetDriverVersion()))

deviceCount = nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = nvmlDeviceGetHandleByIndex(i)
    try:
        test = nvmlDeviceGetCurrentDriverModel(handle)
        print(test)
    except:
        print("Not Support")
    
nvmlShutdown()

可以使用try语句挨个尝试IDE提示有的函数。