如何在 tensorflow 服务 (gpu) 中加载自定义动态库 (*.so)？

Question

我写了自己的cudaMelloc，如下所示，我计划将其应用到tensorflow服务（GPU）中，通过LD_PRELOAD机制跟踪cudaMelloc调用（可用于限制每个tf服务的GPU使用容器也经过适当修改）。

typedef cudaError_t (*cu_malloc)(void **, size_t);

/* cudaMalloc wrapper function */
cudaError_t cudaMalloc(void **devPtr, size_t size)
{
    //cudaError_t (*cu_malloc)(void **devPtr, size_t size);
    cu_malloc real_cu_malloc = NULL;
    char *error;

    real_cu_malloc = (cu_malloc)dlsym(RTLD_NEXT, "cudaMalloc");
    if ((error = dlerror()) != NULL) {
        fputs(error, stderr);
        exit(1);
    }
    cudaError_t res = real_cu_malloc(devPtr, size);
    printf("cudaMalloc(%d) = %p\n", (int)size, devPtr);
    return res;
}

我使用以下命令将上述代码编译成动态库文件：

nvcc --compiler-options "-DRUNTIME -shared -fpic" --cudart=shared -o libmycudaMalloc.so mycudaMalloc.cu -ldl

当应用于使用命令nvcc -g --cudart=shared -o vector_add_dynamic vector_add.cu 编译的vector_add 程序时，它运行良好：

root@ubuntu:~# LD_PRELOAD=./libmycudaMalloc.so ./vector_add_dynamic 
cudaMalloc(800000) = 0x7ffe22ce1580
cudaMalloc(800000) = 0x7ffe22ce1588
cudaMalloc(800000) = 0x7ffe22ce1590

但是当我使用以下命令将其应用于 tensorflow 服务时，cudaMelloc 调用并不引用我编写的动态库。

root@ubuntu:~# LD_PRELOAD=/root/libmycudaMalloc.so ./tensorflow_model_server --port=8500 --rest_api_port=8501 --model_name=resnet --model_base_path=/models/resnet

这是我的问题：

1. 是不是因为 tensorflow-serving 是以完全静态的方式构建的，所以 tf-serving 指的是libcudart_static.a 而不是libcudart.so？

2. 如果是这样，我如何构建 tf-serving 以启用动态链接？

Answer 1

是不是因为 tensorflow-serving 是以完全静态的方式构建的，所以 tf-serving 是指 libcudart_static.a 而不是 libcudart.so？

它可能不是完全静态构建的。你可以通过运行来查看是否是：

readelf -d tensorflow_model_server | grep NEEDED

但它可能 与libcudart_static.a 链接。您可以通过以下方式查看是否：

readelf -Ws tensorflow_model_server | grep ' cudaMalloc$'

如果您看到未解析的 (U) 符号（就像您看到 vector_add_dynamic 二进制文件一样），那么 LD_PRELOAD 应该可以工作。但您可能会看到一个已定义的（T 或 t）符号。

如果是这样，我如何构建 tf-serving 以启用动态链接？

当然：它是开源的。您所要做的就是弄清楚如何构建它，然后在没有libcudart_static.a 的情况下如何构建它，然后弄清楚当您这样做时会出现什么问题（如果有的话）。