使用java获取nvidia显卡信息

news2024/12/26 0:15:20

前言
AI开发通常使用到GPU，但通常使用的是python、c++等语言，java用的则非常少。这也导致了java在gpu相关的库比较少。现在的需求是要获取nvidia显卡的使用情况，如剩余显存。这里给出两种较简单的解决方案。

基于nivdia-smi工具

显卡是硬件，当然需要操作系统支持相关的驱动程序。首先要确保的正确安装了nvidia的相关驱动程序。输入命令nvidis-smi后有类似下面的输出，即证明当前机器已安装好了驱动程序。我们这个方案也是基于nvidia-smi这个命令来做的。

在这里插入图片描述
我们可以看到上图的输出，表面了当前的驱动版本、CUDA版本、显卡数量、显存、显卡利用率，当前使用进程等信息。但是明显是给人来看的，我们想放到程序中则需要进行解析。好在该工具提供了xml格式输出：
在这里插入图片描述
于是有了思路：

通过进程调用nvidia-smi，并获取xml输出结果
解析xml格式，获取需要的内容
将需要的内容转为java bean以供使用

具体的xml格式使用"nvidia-smi -q -x"一看便知，我在这里也就不赘述如何解析了。这里给出一个简单的库，目前只包含了显存信息的解析，其余内容没有实现：https://gitee.com/angelhand/vel-koz

基于nvml绑定库

上面提到的nvidia-smi工具实际上是英伟达基于工具nvml实现的，这里有对该工具的一些介绍：https://developer.nvidia.com/nvidia-management-library-nvml。

该工具是一个c编写的管理工具，上面文档有介绍具体的API。工具可以通过GPU Deployment Kit一并安装，GPU Deployment Kit又是CUDA Toolkit的一部分。不过CUDA8.0之后，GPU Deployment Kit不再单独提供，而是合并到了CUDA Toolkit里了。什么意思呢，就是CUDA8.0之后，要使用nvml，只需要安装CUDA Toolkit：https://developer.nvidia.com/cuda-downloads。使用这个方案是一定要安装的哦。不过Windows11环境下，使用系统更新自动安装好的驱动就可以了，不需要额外下载（40系显卡31版本驱动，旧版本不保证）。

环境安装好之后使用https://github.com/henkelmax/nvmlj这个绑定库。这里有同学要问绑定库是什么意思，简单解释一下。前面提到的nvml工具是c写的，提供了一个nvml.h的头文件，其他语言要使用就需要对这些功能做一个映射（绑定）。这个头文件就类似于java中的接口。java对c的映射可以使用JNI技术，另外还有比较简单的JNA。我们这个库就是基于JNA实现的。下面来展示如何使用。

首先修改一下源：

<repository>
  <id>henkelmax.public</id>
  <url>https://maven.maxhenkel.de/repository/public</url>
</repository>

然后使用进行依赖引入：

<dependency>
  <groupId>de.bommel24.nvmlj</groupId>
  <artifactId>nvmlj</artifactId>
  <version>1.0.2</version>
</dependency>

使用起来也很简单：

public class Main {
    public static void main(String[] args) throws NVMLJException {
        // The path to the nvml library (Defaults to C:/Program Files/NVIDIA Corporation/NVSMI/nvml.dll)
        // windows系统需要这样设置
        System.setProperty("nvml.path", "C:\\Windows\\System32\\nvml.dll");
        // 这是我的wsl找到的，实体机可以whereis一下
        System.setProperty("nvml.path", "/usr/lib/wsl/lib/libnvidia-ml.so.1");
        NVMLJ.nvmlInit();
        System.out.println(NVMLJ.nvmlDeviceGetCount());
        // 不要忘记关闭资源
        NVMLJ.nvmlShutdown();
    }
}

这里有个点需要注意一下，就是代码里那行System.setProperty()。set的内容实际上就是nvml的动态链接库。Windows较新版本的目录就是上面提到的目录，可以找一下是否存在。Linux下，我使用的是wsl ubuntu22.04系统，在上面那个目录里。如果你使用的是其他版本的系统，可以whereis libnvidia-mi.so.1来查找一下库的具体位置。你可能会查出来多个位置，或者没有。那么再查一下libnvidia-ml.so（后面没有.1）。具体哪个管用试试就知道了，不过通常是会出现再usr的lib/lib32/lib64的子目录或者子子（子…）目录下了。