From 91ed362e74bc1b2135cf7481787774f27734fdad Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Fri, 13 Dec 2024 16:59:19 +0000
Subject: [PATCH] fix: update trtllm dockefile after rebase

---
 Dockerfile_trtllm | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/Dockerfile_trtllm b/Dockerfile_trtllm
index 0bd271d8..4ac11d9a 100644
--- a/Dockerfile_trtllm
+++ b/Dockerfile_trtllm
@@ -1,5 +1,5 @@
 ARG CUDA_ARCH_LIST="75-real;80-real;86-real;89-real;90-real"
-ARG OMPI_VERSION="4.1.6"
+ARG OMPI_VERSION="4.1.7rc1"
 
 # Build dependencies resolver stage
 FROM lukemathwalker/cargo-chef:latest AS chef
@@ -10,7 +10,7 @@ COPY . .
 RUN cargo chef prepare --recipe-path recipe.json
 
 # CUDA dependent dependencies resolver stage
-FROM nvidia/cuda:12.6.1-cudnn-devel-ubuntu22.04 AS cuda-builder
+FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu24.04 AS cuda-builder
 
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
@@ -18,18 +18,13 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     build-essential \
     cmake \
     curl \
-    gcc  \
-    g++ \
+    gcc-14  \
+    g++-14 \
     git \
     git-lfs \
     libssl-dev \
+    libucx-dev \
     ninja-build \
-    pkg-config \
-    python3 \
-    python3-dev \
-    python3-setuptools \
-    tar \
-    wget \
     ffmpeg \
     libavcodec-dev \
     libavfilter-dev \
@@ -38,8 +33,13 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     libavutil-dev \
     libswscale-dev \
     pkg-config \
-    libclang-dev \
-    clang
+    pipx \
+    python3 \
+    python3-dev \
+    python3-setuptools \
+    tar \
+    wget && \
+    pipx ensurepath
 
 ENV TGI_INSTALL_PREFIX=/usr/local/tgi
 ENV TENSORRT_INSTALL_PREFIX=/usr/local/tensorrt
@@ -93,13 +93,15 @@ RUN mkdir $TGI_INSTALL_PREFIX && mkdir "$TGI_INSTALL_PREFIX/include" && mkdir "$
     cd backends/trtllm && \
     CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX cargo build --release
 
-FROM nvidia/cuda:12.6.1-cudnn-runtime-ubuntu22.04 AS runtime
-RUN apt update && apt install -y python3-minimal python3-dev python3-pip && \
+FROM nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04 AS runtime
+RUN apt update && apt install -y libucx0 pipx python3-minimal python3-dev python3-pip python3-venv && \
     rm -rf /var/lib/{apt,dpkg,cache,log}/ && \
-    python3 -m pip install transformers tokenizers
+    pipx ensurepath && \
+    pipx install --include-deps transformers tokenizers
 
 WORKDIR /usr/local/tgi/bin
 
+ENV PATH=/root/.local/share/pipx/venvs/transformers/bin/:$PATH
 ENV LD_LIBRARY_PATH="/usr/local/tgi/lib:/usr/local/mpi/lib:/usr/local/tensorrt/lib:/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH"
 ENV TOKENIZERS_PARALLELISM=false
 ENV OMPI_MCA_plm_rsh_agent=""
@@ -108,6 +110,9 @@ COPY --from=mpi-builder /usr/local/mpi /usr/local/mpi
 COPY --from=trt-builder /usr/local/tensorrt /usr/local/tensorrt
 COPY --from=tgi-builder /usr/local/tgi /usr/local/tgi
 COPY --from=tgi-builder /usr/src/text-generation-inference/target/release/text-generation-backends-trtllm /usr/local/tgi/bin/text-generation-launcher
+# Copy the ffmpeg libraries
+COPY --from=cuda-builder /usr/lib/x86_64-linux-gnu/* /usr/lib/x86_64-linux-gnu-copy/
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu-copy"
 
 FROM runtime
 
@@ -115,4 +120,4 @@ LABEL co.huggingface.vendor="Hugging Face Inc."
 LABEL org.opencontainers.image.authors="hardware@hf.co"
 
 ENTRYPOINT ["./text-generation-launcher"]
-CMD ["--executor-worker", "/usr/local/tgi/bin/executorWorker"]
+CMD ["--executor-worker", "/usr/local/tgi/bin/executorWorker"]
\ No newline at end of file