TensorFlow Serving（GPU）中加载自定义动态库及LD_PRELOAD拦截cudaMalloc失效问题排查与解决方案咨询

你猜的没错——TensorFlow Serving默认构建确实大概率静态链接了libcudart_static.a，这直接导致LD_PRELOAD的动态库没法拦截cudaMalloc调用，因为静态链接的符号会被直接打包进二进制文件，不走动态链接器的符号解析流程。下面是解决这个问题的完整方案，以及一些额外的排查细节：

首先可以用两个命令确认TF Serving的链接方式：

1. 用ldd检查动态依赖：

ldd ./tensorflow_model_server | grep cudart

如果没有输出任何包含libcudart.so的行，基本坐实了静态链接的猜测。
2. 再用nm查看符号类型：

nm ./tensorflow_model_server | grep cudaMalloc

如果输出里的cudaMalloc符号标记为T（表示在文本段，属于静态链接的符号），就100%确认了。

要让TF Serving使用动态CUDA runtime库，需要在编译时显式指定相关选项，推荐两种方式：

1. 基于官方Docker构建（最稳妥） #

TensorFlow Serving的官方Docker构建流程支持自定义编译参数，步骤如下：

• 先克隆源码仓库：

git clone https://github.com/tensorflow/serving.git
cd serving

• 找到GPU版本的Dockerfile（比如tensorflow_serving/tools/docker/Dockerfile.gpu），修改其中的bazel build命令，添加--define=dynamic_link_cudart=true参数，修改后的命令类似：

bazel build --config=cuda --define=dynamic_link_cudart=true //tensorflow_serving/model_servers:tensorflow_model_server

• 按照官方文档完成Docker镜像构建，之后从镜像中提取编译好的tensorflow_model_server二进制文件即可。

2. 本地源码编译 #

如果不需要Docker，直接在本地环境编译：

• 确保已安装和TF Serving依赖版本匹配的CUDA Toolkit，且libcudart.so在系统库路径中（比如/usr/local/cuda/lib64）。
• 执行bazel编译命令时加上动态链接参数和ABI兼容选项：

bazel build --config=cuda --define=dynamic_link_cudart=true --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" //tensorflow_serving/model_servers:tensorflow_model_server

注：-D_GLIBCXX_USE_CXX11_ABI=0是为了匹配TensorFlow默认的ABI设置，避免出现链接错误。

1. 确保环境变量传递到位：TF Serving启动时可能会fork子进程，建议直接用env LD_PRELOAD=/root/libmycudaMalloc.so ./tensorflow_model_server ...的方式启动，或者先export LD_PRELOAD=/root/libmycudaMalloc.so再启动，避免变量传递失败。
2. 检查库版本兼容性：你的libmycudaMalloc.so必须用和TF Serving相同的CUDA版本编译，否则会出现符号不兼容的问题，导致拦截失效。
3. 优化日志输出：TF Serving的日志可能会覆盖stdout，建议把你的wrapper日志写入文件（比如fprintf(fopen("/tmp/cuda_malloc_log.txt", "a"), "cudaMalloc(%d) = %p\n", (int)size, *devPtr);），这样能更准确判断是否被调用。

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