Adds ability to build llgtrt based on TRTLLM built from source (#18)

* Adds ability to build llgtrt based on TRTLLM built from source

* Recommented a couple stray uncommented lines

---------

Co-authored-by: Lucas Pickup <[email protected]>

Author: dans-msft

README.md

@@ -48,6 +48,17 @@ To build a container, use:
 
 The build script will initialize submodules if they are missing. It takes about 15 minutes on a GitHub runner and should typically be faster on a local machine.
 
+#### Optional: Building TensorRT-LLM from source
+
+The build process above uses prebuilt binaries from a release version of TensorRT-LLM.  It is also possible to build your own version of TensorRT-LLM from source and create a build of llgtrt based on that.  This can be used to build a version of llgtrt that will work with versions of TensorRT-LLM newer than the released versions in nVidia's repositories.
+
+To do so, first build TensorRT-LLM from source following the instructions in https://nvidia.github.io/TensorRT-LLM/installation/build-from-source-linux.html
+
+Now, build llgtrt based on the Docker image you built above
+```bash
+./docker/build.sh --trtllm tensorrt_llm/release
+```
+
 ### Building the TensorRT-LLM Engine
 
 This is based on the [TensorRT-LLM Quick-start](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html).

docker/Dockerfile

@@ -1,4 +1,16 @@
-FROM nvcr.io/nvidia/tensorrt:24.12-py3 AS llgtrt_dev
+# This argument allows you to override the base image with your own custom build of TensorRT or TensorRT-LLM
+ARG BASE_IMAGE=nvcr.io/nvidia/tensorrt:24.12-py3
+
+# Set this to false if your image already includes TensorRT-LLM
+ARG INSTALL_TRTLLM=true
+
+# Set this to 1 to enable the CXX11_ABI during native builds
+ARG USE_CXX11_ABI=0
+
+FROM ${BASE_IMAGE} as llgtrt_dev
+
+ARG INSTALL_TRTLLM
+ARG USE_CXX11_ABI
 
 RUN apt-get update
 RUN DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
@@ -32,10 +44,14 @@ COPY TensorRT-LLM/examples/llama/requirements.txt /tmp/reqs/
 # COPY TensorRT-LLM/docker/common/install_mpi4py.sh /tmp/reqs/
 # COPY TensorRT-LLM/docker/common/install_tensorrt.sh /tmp/reqs/
 
-# RUN bash /tmp/reqs/install_mpi4py.sh
-# RUN bash /tmp/reqs/install_tensorrt.sh
+#RUN bash /tmp/reqs/install_mpi4py.sh
+#RUN bash /tmp/reqs/install_tensorrt.sh
 
-RUN cd /tmp/reqs && pip install --upgrade -r requirements.txt
+# Install tensorrt_llm and its backing c++ libraries from pip unless disabled
+RUN echo "INSTALL_TRTLLM is set to: ${INSTALL_TRTLLM}" && \
+    if [ "${INSTALL_TRTLLM}" = "true" ]; then \
+        cd /tmp/reqs && pip install --upgrade -r requirements.txt; \
+    fi
 
 # more packages for this image
 RUN pip install \
@@ -45,7 +61,7 @@ RUN pip install \
  pandas matplotlib plotly wheel
 RUN pip uninstall -y guidance
 
-# RUN pip install --upgrade transformers
+RUN pip install --upgrade transformers
 
 RUN cd /usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/ && \
     ln -s libnvinfer_plugin_tensorrt_llm.so libnvinfer_plugin_tensorrt_llm.so.10
@@ -65,7 +81,8 @@ COPY . .
 
 # link stub, so that it builds without nvidia-runtime
 RUN cd /usr/local/cuda/lib64 && ln -s stubs/libnvidia-ml.so libnvidia-ml.so.1
-RUN ./scripts/build.sh --clean
+RUN bash -c 'env'
+RUN ./scripts/build.sh --clean --cxx11abi $USE_CXX11_ABI
 # remove stub just in case
 RUN rm /usr/local/cuda/lib64/libnvidia-ml.so.1
 

docker/build.sh

@@ -1,11 +1,49 @@
-#!/bin/sh
+#!/bin/bash
 
-cd $(dirname $0)/..
+# Default values
+DEV_MODE=false
+BASE_IMAGE=""
+INSTALL_TRTLLM=true
+USE_CXX11_ABI=0
 
-if [ ! -f TensorRT-LLM/README.md ] ; then git submodule update --init ; fi
+# Parse command-line arguments
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --dev)
+            DEV_MODE=true
+            shift
+            ;;
+        --trtllm)
+            if [[ -n "$2" && "$2" != --* ]]; then
+                BASE_IMAGE="$2"
+                INSTALL_TRTLLM=false
+                USE_CXX11_ABI=1
+                shift 2
+            else
+                echo "Error: --trtllm requires an argument (image name)."
+                exit 1
+            fi
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
 
-docker build --progress=plain -t llgtrt/llgtrt:dev --target llgtrt_dev . -f docker/Dockerfile
+# Set the default base image if --trtllm was not provided
+if [[ -z "$BASE_IMAGE" ]]; then
+    BASE_IMAGE="nvcr.io/nvidia/tensorrt:24.12-py3"
+fi
 
-if [ "$1" != "--dev" ] ; then
-    docker build --progress=plain -t llgtrt/llgtrt:latest --target llgtrt_prod . -f docker/Dockerfile
+if $DEV_MODE; then
+    TARGET="--target llgtrt_dev"
+else
+    TARGET="--target llgtrt_prod"
 fi
+
+# Build the Docker image with the appropriate arguments
+DOCKER_BUILD_ARGS="--progress=plain --build-arg BASE_IMAGE=$BASE_IMAGE --build-arg INSTALL_TRTLLM=$INSTALL_TRTLLM --build-arg USE_CXX11_ABI=$USE_CXX11_ABI $TARGET"
+
+echo "Building Docker image $TARGET with arguments: $DOCKER_BUILD_ARGS"
+docker build $DOCKER_BUILD_ARGS . -f docker/Dockerfile

scripts/build.sh

@@ -1,5 +1,6 @@
 #!/bin/bash
 
+USE_CXX11_ABI=0
 
 while test $# -gt 0; do
     case "$1" in
@@ -8,6 +9,15 @@ while test $# -gt 0; do
             rm -rf target/release/* 2>/dev/null || :
             shift
             ;;
+        --cxx11abi)
+            if [[ -n "$2" && "$2" =~ ^[01]$ ]]; then
+                USE_CXX11_ABI="$2"
+                shift 2
+            else
+                echo "Error: --cxx11abi requires an argument (0 or 1)."
+                exit 1
+            fi
+            ;;
         *)
             echo "Unknown option $1"
             exit 1
@@ -25,7 +35,8 @@ fi
 
 mkdir -p trtllm-c/build
 cd trtllm-c/build
-cmake ..
+cmake -DUSE_CXX11_ABI=$USE_CXX11_ABI ..
 make -j
 cd ../../llgtrt
+export RUSTC_LOG=rustc_codegen_ssa::back::link=info
 cargo build --release

trtllm-c/CMakeLists.txt

@@ -7,7 +7,7 @@ set(TRTLLM_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../TensorRT-LLM")
 set(TRTLLM_INCLUDE_DIR "${TRTLLM_DIR}/cpp/include")
 set(TRTLLM_INCLUDE_COMMON_DIR "${TRTLLM_DIR}/cpp")
 
-set(USE_CXX11_ABI 0)
+option(USE_CXX11_ABI "Enable CXX11 ABI" 0)
 message(STATUS "Use CXX11 ABI: ${USE_CXX11_ABI}")
 add_compile_options("-D_GLIBCXX_USE_CXX11_ABI=${USE_CXX11_ABI}")