[#6102][fix] support non-system python installation (#7763)

Signed-off-by: Yuan Tong <[email protected]>

Author: tongyuantongyu

cpp/tensorrt_llm/CMakeLists.txt

@@ -294,7 +294,10 @@ if(TARGET ${NIXL_WRAPPER_TARGET})
 endif()
 
 if(NOT WIN32)
-  set_target_properties(${SHARED_TARGET} PROPERTIES BUILD_RPATH "$ORIGIN")
+  # Load libraries at $PREFIX/lib from
+  # $PREFIX/lib/python3.12/site-packages/tensorrt_llm/libs
+  set_target_properties(${SHARED_TARGET}
+                        PROPERTIES BUILD_RPATH "$ORIGIN;$ORIGIN/../../../..")
 endif()
 
 if(BUILD_PYT)

cpp/tensorrt_llm/nanobind/CMakeLists.txt

@@ -41,8 +41,6 @@ endif()
 target_link_libraries(
   ${TRTLLM_NB_MODULE}
   PUBLIC ${SHARED_TARGET}
-         ${UNDEFINED_FLAG}
-         ${NO_AS_NEEDED_FLAG}
          ${Python3_LIBRARIES}
          ${TORCH_LIBRARIES}
          torch_python
@@ -54,10 +52,14 @@ target_compile_definitions(
                              PYBIND11_DETAILED_ERROR_MESSAGES=1)
 
 if(NOT WIN32)
-  set_target_properties(
-    ${TRTLLM_NB_MODULE}
-    PROPERTIES
-      LINK_FLAGS
-      "-Wl,-rpath,'$ORIGIN/libs' -Wl,-rpath,'$ORIGIN/../nvidia/nccl/lib' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}"
+  set(TRTLLM_NB_MODULE_RPATH_LIST
+      "$ORIGIN/libs" # TRTLLM libraries
+      "$ORIGIN/../../.." # Shared libraries under $PREFIX/lib
+      "$ORIGIN/../nvidia/nccl/lib" # NCCL libraries
   )
+  set_target_properties(${TRTLLM_NB_MODULE}
+                        PROPERTIES BUILD_RPATH "${TRTLLM_NB_MODULE_RPATH_LIST}")
+  set_target_properties(
+    ${TRTLLM_NB_MODULE} PROPERTIES LINK_FLAGS
+                                   "${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}")
 endif()

cpp/tensorrt_llm/pybind/CMakeLists.txt

@@ -43,8 +43,6 @@ endif()
 target_link_libraries(
   ${TRTLLM_PYBIND_MODULE}
   PUBLIC ${SHARED_TARGET}
-         ${UNDEFINED_FLAG}
-         ${NO_AS_NEEDED_FLAG}
          ${Python3_LIBRARIES}
          ${TORCH_LIBRARIES}
          torch_python
@@ -56,10 +54,15 @@ target_compile_definitions(
                                  PYBIND11_DETAILED_ERROR_MESSAGES=1)
 
 if(NOT WIN32)
-  set_target_properties(
-    ${TRTLLM_PYBIND_MODULE}
-    PROPERTIES
-      LINK_FLAGS
-      "-Wl,-rpath,'$ORIGIN/libs' -Wl,-rpath,'$ORIGIN/../nvidia/nccl/lib' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}"
+  set(TRTLLM_PYBIND_MODULE_RPATH_LIST
+      "$ORIGIN/libs" # TRTLLM libraries
+      "$ORIGIN/../../.." # Shared libraries under $PREFIX/lib
+      "$ORIGIN/../nvidia/nccl/lib" # NCCL libraries
   )
+  set_target_properties(
+    ${TRTLLM_PYBIND_MODULE} PROPERTIES BUILD_RPATH
+                                       "${TRTLLM_PYBIND_MODULE_RPATH_LIST}")
+  set_target_properties(
+    ${TRTLLM_PYBIND_MODULE} PROPERTIES LINK_FLAGS
+                                       "${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}")
 endif()

tensorrt_llm/__init__.py

@@ -27,6 +27,35 @@ def _add_trt_llm_dll_directory():
 
 _add_trt_llm_dll_directory()
 
+
+def _preload_python_lib():
+    """
+    Preload Python library.
+
+    On Linux, the python executable links to libpython statically,
+    so the dynamic library `libpython3.x.so` is not loaded.
+    When using virtual environment on top of non-system Python installation,
+    our libraries installed under `$VENV_PREFIX/lib/python3.x/site-packages/`
+    have difficulties loading `$PREFIX/lib/libpython3.x.so.1.0` on their own,
+    since venv does not symlink `libpython3.x.so` into `$VENV_PREFIX/lib/`,
+    and the relative path from `$VENV_PREFIX` to `$PREFIX` is arbitrary.
+
+    We preload the libraries here since the Python executable under `$PREFIX/bin`
+    can easily find the library.
+    """
+    import platform
+    on_linux = platform.system() == "Linux"
+    if on_linux:
+        import sys
+        from ctypes import cdll
+        v_major, v_minor, *_ = sys.version_info
+        pythonlib = f'libpython{v_major}.{v_minor}.so'
+        _ = cdll.LoadLibrary(pythonlib + '.1.0')
+        _ = cdll.LoadLibrary(pythonlib)
+
+
+_preload_python_lib()
+
 import sys
 
 # Need to import torch before tensorrt_llm library, otherwise some shared binary files