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

Question

我写了自己的 cudaMelloc 如下，我打算将其应用在 tensorflow serving (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;
}

我使用以下命令将上面的代码编译成一个动态lib文件：

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

所以这是我的问题：

是不是因为tensorflow-serving是完全静态构建的，所以tf-serving指的是libcudart_static.a而不是libcudart.so？
如果是这样，我如何构建 tf-serving 以启用动态链接？

Answer 1

Is it because that tensorflow-serving is built in a fully static manner, such that tf-serving refers to the libcudart_static.a instead of 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）符号。

If so, how could I build tf-serving to enable dynamic linking?

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

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

How to load customized dynamic libs (*.so) in tensorflow serving (gpu)?

gpu

dynamic-linking

tensorflow

tensorflow-serving