chore: map compute_cap from GPU name (#276)

Author: haixiw

sagemaker-entrypoint-cuda-all.sh

@@ -1,5 +1,26 @@
 #!/bin/bash
 
+verlte() {
+    [ "$1" = "$2" ] && return 1 || [ "$2" = "$(echo -e "$1\n$2" | sort -V | head -n1)" ]
+}
+
+if [ -f /usr/local/cuda/compat/libcuda.so.1 ]; then
+    CUDA_COMPAT_MAX_DRIVER_VERSION=$(readlink /usr/local/cuda/compat/libcuda.so.1 | cut -d"." -f 3-)
+    echo "CUDA compat package requires Nvidia driver ≤${CUDA_COMPAT_MAX_DRIVER_VERSION}"
+    cat /proc/driver/nvidia/version
+    NVIDIA_DRIVER_VERSION=$(sed -n 's/^NVRM.*Kernel Module *\([0-9.]*\).*$/\1/p' /proc/driver/nvidia/version 2>/dev/null || true)
+    echo "Current installed Nvidia driver version is ${NVIDIA_DRIVER_VERSION}"
+    if [ $(verlte "$CUDA_COMPAT_MAX_DRIVER_VERSION" "$NVIDIA_DRIVER_VERSION") ]; then
+        echo "Setup CUDA compatibility libs path to LD_LIBRARY_PATH"
+        export LD_LIBRARY_PATH=/usr/local/cuda/compat:$LD_LIBRARY_PATH
+        echo $LD_LIBRARY_PATH
+    else
+        echo "Skip CUDA compat libs setup as newer Nvidia driver is installed"
+    fi
+else
+    echo "Skip CUDA compat libs setup as package not found"
+fi
+
 if [[ -z "${HF_MODEL_ID}" ]]; then
   echo "HF_MODEL_ID must be set"
   exit 1
@@ -15,9 +36,37 @@ if ! command -v nvidia-smi &> /dev/null; then
     exit 1
 fi
 
+# Query GPU name using nvidia-smi
+gpu_name=$(nvidia-smi --query-gpu=gpu_name --format=csv | awk 'NR==2')
+if [ $? -ne 0 ]; then
+    echo "Error: $gpu_name"
+    echo "Query gpu_name failed"
+else
+    echo "Query gpu_name succeeded. Printing output: $gpu_name"
+fi
+
+# Function to get compute capability based on GPU name
+get_compute_cap() {
+    gpu_name="$1"
+
+    # Check if the GPU name contains "A10G"
+    if [[ "$gpu_name" == *"A10G"* ]]; then
+        echo "86"
+    # Check if the GPU name contains "A100"
+    elif [[ "$gpu_name" == *"A100"* ]]; then
+        echo "80"
+    # Check if the GPU name contains "H100"
+    elif [[ "$gpu_name" == *"H100"* ]]; then
+        echo "90"
+    else
+        echo "80"  # Default compute capability
+    fi
+}
+
 if [[ -z "${CUDA_COMPUTE_CAP}" ]]
 then
-    compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv | sed -n '2p' | sed 's/\.//g')
+    compute_cap=$(get_compute_cap "$gpu_name")
+    echo "the compute_cap is $compute_cap"
 else
     compute_cap=$CUDA_COMPUTE_CAP
 fi