build: Support rebuilding and replacing trtllm backend (#99)

container/Dockerfile

@@ -69,11 +69,12 @@ RUN pip install "git+https://github.com/triton-inference-server/perf_analyzer.gi
 
 # Backend & Framework Specific Installation
 ARG FRAMEWORK="STANDARD"
-ARG TENSORRTLLM_BACKEND_COMMIT=
+ARG TENSORRTLLM_BACKEND_REPO_TAG=
+ARG TENSORRTLLM_BACKEND_REBUILD=
 ENV FRAMEWORK=${FRAMEWORK}
 RUN --mount=type=bind,source=./container/deps/requirements.tensorrtllm.txt,target=/tmp/requirements.txt \
     --mount=type=bind,source=./container/deps/clone_tensorrtllm.sh,target=/tmp/clone_tensorrtllm.sh \
-    if [[ "$FRAMEWORK" == "TENSORRTLLM" ]] ; then pip install --timeout=2000 -r /tmp/requirements.txt; /tmp/clone_tensorrtllm.sh --tensorrtllm-backend-commit ${TENSORRTLLM_BACKEND_COMMIT} ; fi
+    if [[ "$FRAMEWORK" == "TENSORRTLLM" ]] ; then pip install --timeout=2000 -r /tmp/requirements.txt; /tmp/clone_tensorrtllm.sh --tensorrtllm-backend-repo-tag ${TENSORRTLLM_BACKEND_REPO_TAG} --tensorrtllm-backend-rebuild ${TENSORRTLLM_BACKEND_REBUILD} ; fi
 
 RUN --mount=type=bind,source=./container/deps/requirements.vllm.txt,target=/tmp/requirements.txt \
     if [[ "$FRAMEWORK" == "VLLM" ]] ; then pip install --timeout=2000 -r /tmp/requirements.txt ; fi
@@ -98,7 +99,7 @@ RUN --mount=type=bind,source=./container/deps/requirements.standard.txt,target=/
 ARG TENSORRTLLM_FRAMEWORK
 ENV FRAMEWORK_LD_LIBRARY_PATH=${TENSORRTLLM_FRAMEWORK:+/opt/tritonserver/backends/tensorrtllm/}
 ENV LD_LIBRARY_PATH=${FRAMEWORK_LD_LIBRARY_PATH}:${LD_LIBRARY_PATH}
-ENV TENSORRTLLM_BACKEND_COMMIT=$TENSORRTLLM_BACKEND_COMMIT
+ENV TENSORRTLLM_BACKEND_REPO_TAG=$TENSORRTLLM_BACKEND_REPO_TAG
 ENV TRTLLM_USE_MPI_KVCACHE=${TENSORRTLLM_FRAMEWORK:+"1"}
 
 # TODO set VLLM Version

container/build.sh

@@ -59,8 +59,13 @@ STANDARD_BASE_IMAGE_TAG=${STANDARD_BASE_VERSION}-py3
 TENSORRTLLM_BASE_VERSION=24.12
 TENSORRTLLM_BASE_IMAGE=nvcr.io/nvidia/tritonserver
 TENSORRTLLM_BASE_IMAGE_TAG=${TENSORRTLLM_BASE_VERSION}-trtllm-python-py3
-# IMPORTANT NOTE: Ensure the commit matches the TRTLLM backend version used in the base image above
-TENSORRTLLM_BACKEND_COMMIT=v0.16.0
+# IMPORTANT NOTE: Ensure the repo tag complies with the TRTLLM backend version
+# used in the base image above.
+TENSORRTLLM_BACKEND_REPO_TAG=v0.16.0
+# Set this as 1 to rebuild and replace trtllm backend bits in the container.
+# This will allow building triton distributed container image with custom
+# trt-llm backend repo branch.
+TENSORRTLLM_BACKEND_REBUILD=0
 
 # vllm installation is done later in the Dockerfile so it will overwrite the
 # vllm version installed in the base image.
@@ -91,7 +96,7 @@ get_options() {
 		missing_requirement $1
             fi
             ;;
-        --tensorrtllm-backend-commit)
+        --tensorrtllm-backend-repo-tag)
             if [ "$2" ]; then
                 TRTLLM_BACKEND_COMMIT=$2
                 shift
@@ -99,6 +104,14 @@ get_options() {
 		missing_requirement $1
             fi
             ;;
+        --tensorrtllm-backend-rebuild)
+            if [ "$2" ]; then
+                TRTLLM_BACKEND_REBUILD=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
         --base-image)
             if [ "$2" ]; then
                 BASE_IMAGE=$2
@@ -216,7 +229,8 @@ show_image_options() {
     echo "   Base: '${BASE_IMAGE}'"
     echo "   Base_Image_Tag: '${BASE_IMAGE_TAG}'"
     if [[ $FRAMEWORK == "TENSORRTLLM" ]]; then
-	echo "   Tensorrtllm Backend Commit: '${TENSORRTLLM_BACKEND_COMMIT}'"
+	    echo "   Tensorrtllm Backend Repo Tag: '${TENSORRTLLM_BACKEND_REPO_TAG}'"
+        echo "   Tensorrtllm Backend Rebuild: '${TENSORRTLLM_BACKEND_REBUILD}'"
     fi
     echo "   Build Context: '${BUILD_CONTEXT}'"
     echo "   Build Arguments: '${BUILD_ARGS}'"
@@ -230,7 +244,8 @@ show_help() {
     echo "  [--base-imge-tag base image tag]"
     echo "  [--platform platform for docker build"
     echo "  [--framework framework one of ${!FRAMEWORKS[@]}]"
-    echo "  [--tensorrtllm-backend-commit commit or tag]"
+    echo "  [--tensorrtllm-backend-repo-tag commit or tag]"
+    echo "  [--tensorrtllm-backend-rebuild whether or not to rebuild the backend]"
     echo "  [--build-arg additional build args to pass to docker build]"
     echo "  [--tag tag for image]"
     echo "  [--no-cache disable docker build cache]"
@@ -262,8 +277,9 @@ if [ ! -z ${GITLAB_TOKEN} ]; then
     BUILD_ARGS+=" --build-arg GITLAB_TOKEN=${GITLAB_TOKEN} "
 fi
 
-if [[ $FRAMEWORK == "TENSORRTLLM" ]] && [ ! -z ${TENSORRTLLM_BACKEND_COMMIT} ]; then
-    BUILD_ARGS+=" --build-arg TENSORRTLLM_BACKEND_COMMIT=${TENSORRTLLM_BACKEND_COMMIT} "
+if [[ $FRAMEWORK == "TENSORRTLLM" ]] && [ ! -z ${TENSORRTLLM_BACKEND_REPO_TAG} ]; then
+    BUILD_ARGS+=" --build-arg TENSORRTLLM_BACKEND_REPO_TAG=${TENSORRTLLM_BACKEND_REPO_TAG} "
+    BUILD_ARGS+=" --build-arg TENSORRTLLM_BACKEND_REBUILD=${TENSORRTLLM_BACKEND_REBUILD} "
 fi
 
 if [ ! -z ${HF_TOKEN} ]; then

container/deps/clone_tensorrtllm.sh

@@ -14,7 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-TENSORRTLLM_BACKEND_COMMIT=
+TENSORRTLLM_BACKEND_REPO_TAG=
+TENSORRTLLM_BACKEND_REBUILD=
 GIT_TOKEN=
 GIT_REPO=
 
@@ -25,9 +26,17 @@ get_options() {
             show_help
             exit
             ;;
-    --tensorrtllm-backend-commit)
+    --tensorrtllm-backend-repo-tag)
             if [ "$2" ]; then
-                TENSORRTLLM_BACKEND_COMMIT=$2
+                TENSORRTLLM_BACKEND_REPO_TAG=$2
+                shift
+            else
+		missing_requirement $1
+            fi
+            ;;
+    --tensorrtllm-backend-rebuild)
+            if [ "$2" ]; then
+                TENSORRTLLM_BACKEND_REBUILD=$2
                 shift
             else
 		missing_requirement $1
@@ -68,14 +77,16 @@ show_options() {
     echo ""
     echo "Getting TENSORRTLLM Backend Repo"
     echo ""
-    echo "   TENSORRTLLM Backend Commit: '${TENSORRTLLM_BACKEND_COMMIT}'"
+    echo "   Tensorrtllm Backend Repo Tag: '${TENSORRTLLM_BACKEND_REPO_TAG}'"
+    echo "   Tensorrtllm Backend Rebuild: '${TENSORRTLLM_BACKEND_REBUILD}'"
     echo ""
 }
 
 
 show_help() {
     echo "usage: clone_tensorrtllm.sh"
-    echo "  [--tensorrtllm-backend-commit commit]"
+    echo "  [--tensorrtllm-backend-repo-tag commit]"
+    echo "  [--tensorrtllm-backend-rebuild whether to rebuild backend]"
     echo "  [--git-token git-token]"
     echo "  [--git-repo git-repo]"
     exit 0
@@ -106,9 +117,29 @@ show_options
 
 git clone ${GIT_REPO}
 cd tensorrtllm_backend
-git reset --hard ${TENSORRTLLM_BACKEND_COMMIT}
+git checkout -b ${TENSORRTLLM_BACKEND_REPO_TAG}
 git submodule update --init --recursive
 git lfs install
 git lfs pull
+
+if [ -z ${TENSORRTLLM_BACKEND_REBUILD} ]; then
+    # Install cmake
+    apt update -q=2 \
+	    && apt install -y gpg wget \
+        && wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - |  tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null \
+	    && . /etc/os-release \
+	    && echo "deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ $UBUNTU_CODENAME main" | tee /etc/apt/sources.list.d/kitware.list >/dev/null \
+	    && apt-get update -q=2 \
+	    && apt-get install -y --no-install-recommends cmake=3.28.3* cmake-data=3.28.3* \
+        && cmake --version
+
+    # Build the backend
+    (cd inflight_batcher_llm/src \
+        && cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DUSE_CXX11_ABI=1 .. \
+        && make install \
+        && cp libtriton_tensorrtllm.so /opt/tritonserver/backends/tensorrtllm/ \
+        && cp trtllmExecutorWorker /opt/tritonserver/backends/tensorrtllm/ \
+    )
+fi
 cd ..
 mv tensorrtllm_backend /