使用java获取nvidia显卡信息

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

基于nivdia-smi工具

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

使用java获取nvidia显卡信息_第1张图片
我们可以看到上图的输出,表面了当前的驱动版本、CUDA版本、显卡数量、显存、显卡利用率,当前使用进程等信息。但是明显是给人来看的,我们想放到程序中则需要进行解析。好在该工具提供了xml格式输出:
使用java获取nvidia显卡信息_第2张图片
于是有了思路:

  1. 通过进程调用nvidia-smi,并获取xml输出结果
  2. 解析xml格式,获取需要的内容
  3. 将需要的内容转为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.publicid>
  <url>https://maven.maxhenkel.de/repository/publicurl>
repository>

然后使用进行依赖引入:

<dependency>
  <groupId>de.bommel24.nvmljgroupId>
  <artifactId>nvmljartifactId>
  <version>1.0.2version>
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)。具体哪个管用试试就知道了,不过通常是会出现再usrlib/lib32/lib64的子目录或者子子(子…)目录下了。

你可能感兴趣的:(java,nvml,nvidia-smi)