如何拦截nvidia-smi对NVML的调用并添加自定义逻辑？

环境配置 #

• Ubuntu 20.04
• NVIDIA-SMI 515.65.01
• Driver Version: 515.65.01
• CUDA Version: 11.7

需求与尝试过程 #

想要拦截nvidia-smi对NVML的调用，每次触发调用时打印I am here，采用dlsym+LD_PRELOAD的方法实现：

自定义拦截代码 #

nvmlReturn_t nvmlDeviceGetMemoryInfo_v2(nvmlDevice_t device, nvmlMemory_v2_t *memory){
    void* handle = dlopen("libnvidia-ml.so", RTLD_LOCAL | RTLD_LAZY);
    nvmlReturn_t (*fn) (nvmlDevice_t, nvmlMemory_v2_t *) = (nvmlReturn_t (*) (nvmlDevice_t, nvmlMemory_v2_t *))dlsym(handle, "nvmlDeviceGetMemoryInfo_v2");
    nvmlReturn_t current_return = fn(device, memory);
    printf("i am here");
    return current_return;
}

编译命令 #

gcc -I/usr/local/cuda/include -shared -fPIC -o mylib.so mylib.c -ldl -L/usr/lib/x86_64-linux-gnu/ -lnvidia-ml

加载执行命令 #

LD_PRELOAD=./mylib.so nvidia-smi

异常现象 #

执行后未打印I am here，拦截失败；但在自定义的小型测试项目（包含<nvml.h>）中，该预加载库可成功拦截NVML函数调用。

疑问 #

为什么nvidia-smi不受预加载影响？是否函数包装错误、库构建/加载存在问题？有没有其他实现方法？

核心原因 #

1. 静态链接限制：部分版本的nvidia-smi并非动态链接libnvidia-ml.so，而是将NVML代码静态编译进自身二进制文件。这种情况下，LD_PRELOAD无法拦截内部静态函数调用，因为预加载仅作用于动态链接的符号。
2. 函数符号不匹配：你包装的nvmlDeviceGetMemoryInfo_v2可能不是nvidia-smi实际调用的函数。它可能使用了旧版无后缀的nvmlDeviceGetMemoryInfo，或者nvmlInit、nvmlDeviceGetCount等基础初始化/设备枚举函数。
3. 库加载与绑定问题：编译时链接libnvidia-ml.so会导致预加载库提前绑定原库符号，而非运行时动态查找；同时dlopen使用相对路径可能出现查找失败的情况。

验证与修复步骤 #

1. 检查nvidia-smi的链接方式
   执行命令查看动态依赖：

   ldd /usr/bin/nvidia-smi
   

   如果输出中没有libnvidia-ml.so，说明它是静态链接NVML，LD_PRELOAD方法无效。

2. 确认nvidia-smi调用的NVML函数
   使用strace追踪动态符号调用：

   strace -e trace=call /usr/bin/nvidia-smi 2>&1 | grep nvml
   

   或查看动态符号表（仅动态链接版本有效）：

   nm -D /usr/bin/nvidia-smi | grep nvml
   

   根据输出确定需要拦截的函数名称。

3. 修正预加载库实现

   • 编译时去掉-lnvidia-ml，避免符号提前绑定：

     gcc -I/usr/local/cuda/include -shared -fPIC -o mylib.so mylib.c -ldl
     

   • 使用绝对路径dlopen原库，避免路径查找错误：

     void* handle = dlopen("/usr/lib/x86_64-linux-gnu/libnvidia-ml.so", RTLD_NOW | RTLD_GLOBAL);
     

   • 同时包装多个核心NVML函数，提升拦截命中概率：

     #include <nvml.h>
     #include <dlfcn.h>
     #include <stdio.h>
     
     nvmlReturn_t nvmlInit_v2(void) {
         static nvmlReturn_t (*orig_fn)(void);
         if (!orig_fn) {
             void* handle = dlopen("/usr/lib/x86_64-linux-gnu/libnvidia-ml.so", RTLD_NOW | RTLD_GLOBAL);
             orig_fn = dlsym(handle, "nvmlInit_v2");
         }
         printf("I am here: nvmlInit_v2\n");
         return orig_fn();
     }
     
     // 同理添加nvmlDeviceGetCount、nvmlDeviceGetHandleByIndex等函数的拦截...
     

替代方案 #

• 若nvidia-smi是静态链接，可尝试使用Ptrace工具动态注入代码，或修改二进制文件（注意版权与风险）；
• 编写独立监控脚本，通过定时调用NVML接口或解析nvidia-smi输出来实现自定义逻辑，替代拦截方案。

内容的提问来源于stack exchange，提问作者Pierpaolo Sorbellini