fix rpaths in trtllm-backend

replicate · Apr 15, 2024 · 00a1983 · 00a1983
1 parent 57b92cf
commit 00a1983
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/nix/tensorrt-llm.nix b/nix/tensorrt-llm.nix
@@ -140,7 +140,7 @@ stdenv.mkDerivation (o: {
       mkdir -p $out/cpp/build/tensorrt_llm/plugins
       pushd tensorrt_llm
       cp ./libtensorrt_llm.so $out/cpp/build/tensorrt_llm/
-      patchelf --add-needed 'libcudnn.so.8' $out/cpp/build/tensorrt_llm/libtensorrt_llm.so
+      patchelf --add-needed 'libcudnn.so.8' --add-rpath ${cudaPackages.cudnn.lib}/lib $out/cpp/build/tensorrt_llm/libtensorrt_llm.so
       cp ./plugins/libnvinfer_plugin_tensorrt_llm.so* $out/cpp/build/tensorrt_llm/plugins/
       for f in $out/cpp/build/tensorrt_llm/plugins/*.so*; do
         if [ ! -L "$f" ]; then

diff --git a/nix/trtllm-backend.nix b/nix/trtllm-backend.nix
@@ -82,9 +82,19 @@ oldGccStdenv.mkDerivation rec {
     "-DTRT_INCLUDE_DIR=${tensorrt-src}/include"
     "-DTRTLLM_DIR=${tensorrt-llm}"
   ];
+  postInstall = ''
+    mkdir -p $out/backends/tensorrtllm
+    cp libtriton_*.so triton_tensorrtllm_worker $out/backends/tensorrtllm
+    rm -r /build/source/inflight_batcher_llm/build/_deps/repo-core-build
+    rm -r /build/source/inflight_batcher_llm/build/libtriton_tensorrtllm_common.so
+  '';
   # buildInputs = [ tensorrt-llm ];
   postFixup = ''
     patchelf $out/backends/tensorrtllm/libtriton_tensorrtllm.so \
-      --add-rpath ${trt_lib_dir}:${tensorrt-llm}/cpp/build/tensorrt_llm:${tensorrt-llm}/cpp/build/tensorrt_llm/plugins
+      --add-rpath '$ORIGIN:${trt_lib_dir}:${tensorrt-llm}/cpp/build/tensorrt_llm:${tensorrt-llm}/cpp/build/tensorrt_llm/plugins:${cudaPackages.cudnn.lib}/lib'
+    patchelf $out/backends/tensorrtllm/libtriton_tensorrtllm_common.so \
+      --add-rpath '$ORIGIN:${trt_lib_dir}:${tensorrt-llm}/cpp/build/tensorrt_llm:${tensorrt-llm}/cpp/build/tensorrt_llm/plugins:${cudaPackages.cudnn.lib}/lib'
+    patchelf $out/backends/tensorrtllm/triton_tensorrtllm_worker \
+      --add-rpath '$ORIGIN:${trt_lib_dir}:${tensorrt-llm}/cpp/build/tensorrt_llm:${tensorrt-llm}/cpp/build/tensorrt_llm/plugins:${cudaPackages.cudnn.lib}/lib'
   '';
 }