diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 2e625f46879..4c97c8c1803 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/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) diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index 2cc544227c3..5eaa3c15fad 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/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() diff --git a/cpp/tensorrt_llm/pybind/CMakeLists.txt b/cpp/tensorrt_llm/pybind/CMakeLists.txt index 14aedca40ea..723dec0511d 100755 --- a/cpp/tensorrt_llm/pybind/CMakeLists.txt +++ b/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() diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 487888f34e3..fa757e60981 100644 --- a/tensorrt_llm/__init__.py +++ b/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