From 00a1983515804ed08b2f328a250aab6fd42f5fb5 Mon Sep 17 00:00:00 2001
From: Yorick van Pelt <yorick@yorickvanpelt.nl>
Date: Mon, 15 Apr 2024 21:46:50 +0200
Subject: [PATCH] fix rpaths in trtllm-backend

---
 nix/tensorrt-llm.nix   |  2 +-
 nix/trtllm-backend.nix | 12 +++++++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/nix/tensorrt-llm.nix b/nix/tensorrt-llm.nix
index 458c833..c717893 100644
--- 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
index d160044..4dab9b2 100644
--- 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'
   '';
 }