Align cuda-{all-entrypoint,entrypoint}.sh

Signed-off-by: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>

Author: alvarobartt

cuda-all-entrypoint.sh

@@ -1,26 +1,27 @@
 #!/bin/bash
 
+if ! command -v nvidia-smi &>/dev/null; then
+    echo "Error: 'nvidia-smi' command not found."
+    exit 1
+fi
+
 # GKE injects /usr/local/nvidia from the host but does not add it to PATH or
 # run ldconfig for it, unlike the nvidia-container-runtime on AWS which does both.
 if [ -d /usr/local/nvidia/bin ]; then
     export PATH="${PATH}:/usr/local/nvidia/bin"
 fi
 if [ -d /usr/local/nvidia/lib64 ]; then
-    echo /usr/local/nvidia/lib64 > /etc/ld.so.conf.d/nvidia-host.conf
+    echo /usr/local/nvidia/lib64 >/etc/ld.so.conf.d/nvidia-host.conf
     ldconfig
 fi
 
-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}')
+    # Match both "CUDA Version" and "CUDA UMD Version" (on `driver-6xx` onwards)
+    DRIVER_CUDA=$(nvidia-smi 2>/dev/null | grep -oE 'CUDA[[:space:]]+([A-Za-z]+[[:space:]])?Version:[[:space:]]*[0-9]+(\.[0-9]+)+' | grep -oE '[0-9]+(\.[0-9]+)+' | head -n1)
 
     IFS='.' read -r MAJ MIN PATCH <<EOF
 ${DRIVER_CUDA:-0.0.0}

cuda-entrypoint.sh

@@ -1,20 +1,20 @@
 #!/bin/bash
 
+if ! command -v nvidia-smi &>/dev/null; then
+    echo "Error: 'nvidia-smi' command not found."
+    exit 1
+fi
+
 # GKE injects /usr/local/nvidia from the host but does not add it to PATH or
 # run ldconfig for it, unlike the nvidia-container-runtime on AWS which does both.
 if [ -d /usr/local/nvidia/bin ]; then
     export PATH="${PATH}:/usr/local/nvidia/bin"
 fi
 if [ -d /usr/local/nvidia/lib64 ]; then
-    echo /usr/local/nvidia/lib64 > /etc/ld.so.conf.d/nvidia-host.conf
+    echo /usr/local/nvidia/lib64 >/etc/ld.so.conf.d/nvidia-host.conf
     ldconfig
 fi
 
-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+