Fix support for containers w/ CUDA 13.0+ (#831)

Author: alvarobartt

Cargo.lock

@@ -22,7 +22,7 @@ default-members = [
 resolver = "2"
 
 [workspace.package]
-version = "1.9.0"
+version = "1.9.1"
 edition = "2021"
 authors = ["Olivier Dehaene", "Nicolas Patry", "Alvaro Bartolome"]
 homepage = "https://github.com/huggingface/text-embeddings-inference"

Cargo.toml

@@ -126,7 +126,7 @@ ARG DEFAULT_USE_FLASH_ATTENTION=True
 ENV HUGGINGFACE_HUB_CACHE=/data \
     PORT=80 \
     USE_FLASH_ATTENTION=$DEFAULT_USE_FLASH_ATTENTION \
-    LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH}"
+    LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     ca-certificates \
@@ -135,16 +135,17 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
     cuda-compat-12-9 \
     && rm -rf /var/lib/apt/lists/*
 
+COPY --chmod=775 cuda-entrypoint.sh entrypoint.sh
+
 FROM base AS grpc
 
 COPY --from=grpc-builder /usr/src/target/release/text-embeddings-router /usr/local/bin/text-embeddings-router
 
-ENTRYPOINT ["text-embeddings-router"]
+ENTRYPOINT ["./entrypoint.sh"]
 CMD ["--json-output"]
 
 FROM base
 
 COPY --from=http-builder /usr/src/target/release/text-embeddings-router /usr/local/bin/text-embeddings-router
-
-ENTRYPOINT ["text-embeddings-router"]
+ENTRYPOINT ["./entrypoint.sh"]
 CMD ["--json-output"]

Dockerfile-cuda

@@ -115,7 +115,7 @@ ARG DEFAULT_USE_FLASH_ATTENTION=True
 ENV HUGGINGFACE_HUB_CACHE=/data \
     PORT=80 \
     USE_FLASH_ATTENTION=$DEFAULT_USE_FLASH_ATTENTION \
-    LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH}"
+    LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     ca-certificates \

Dockerfile-cuda-all

@@ -5,6 +5,27 @@ if ! command -v nvidia-smi &>/dev/null; then
     exit 1
 fi
 
+# NOTE: Given that we need to support CUDA versions earlier than CUDA 12.9.1, we
+# need to include the `cuda-compat-12-9` in `LD_LIBRARY_PATH` when the host CUDA
+# version is lower than that; whilst we shouldn't include that when CUDA is 13.0+
+# as otherwise it will fail due to it.
+if [ -d /usr/local/cuda/compat ]; then
+    DRIVER_CUDA=$(nvidia-smi 2>/dev/null | awk '/CUDA Version/ {print $3; exit}')
+
+    IFS='.' read -r MAJ MIN PATCH <<EOF
+${DRIVER_CUDA:-0.0.0}
+EOF
+    : "${MIN:=0}"
+    : "${PATCH:=0}"
+
+    DRIVER_INT=$((10#${MAJ} * 10000 + 10#${MIN} * 100 + 10#${PATCH}))
+    TARGET_INT=$((12 * 10000 + 9 * 100 + 1))
+
+    if [ "$DRIVER_INT" -lt "$TARGET_INT" ]; then
+        export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH}"
+    fi
+fi
+
 compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv | sed -n '2p' | sed 's/\.//g')
 
 if [ ${compute_cap} -eq 75 ]; then

cuda-all-entrypoint.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+if ! command -v nvidia-smi &>/dev/null; then
+    echo "Error: 'nvidia-smi' command not found."
+    exit 1
+fi
+
+# NOTE: Given that we need to support CUDA versions earlier than CUDA 12.9.1, we
+# need to include the `cuda-compat-12-9` in `LD_LIBRARY_PATH` when the host CUDA
+# version is lower than that; whilst we shouldn't include that when CUDA is 13.0+
+# as otherwise it will fail due to it.
+if [ -d /usr/local/cuda/compat ]; then
+    DRIVER_CUDA=$(nvidia-smi 2>/dev/null | awk '/CUDA Version/ {print $3; exit}')
+
+    IFS='.' read -r MAJ MIN PATCH <<EOF
+${DRIVER_CUDA:-0.0.0}
+EOF
+    : "${MIN:=0}"
+    : "${PATCH:=0}"
+
+    DRIVER_INT=$((10#${MAJ} * 10000 + 10#${MIN} * 100 + 10#${PATCH}))
+    TARGET_INT=$((12 * 10000 + 9 * 100 + 1))
+
+    if [ "$DRIVER_INT" -lt "$TARGET_INT" ]; then
+        export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH}"
+    fi
+fi
+
+exec text-embeddings-router "$@"

cuda-entrypoint.sh

@@ -10,7 +10,7 @@
       "name": "Apache 2.0",
       "url": "https://www.apache.org/licenses/LICENSE-2.0"
     },
-    "version": "1.9.0"
+    "version": "1.9.1"
   },
   "paths": {
     "/decode": {