Add cuda-entrypoint.sh and update Dockerfile-cuda

Author: alvarobartt

Dockerfile-cuda

@@ -126,25 +126,26 @@ 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 \
     libssl-dev \
     curl \
-    cuda-compat-13-1 \
+    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"]

cuda-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 "$@"