From 08c7b172980d707324ee7545e6f8b5be8dbddf4b Mon Sep 17 00:00:00 2001 From: Enrico Ros Date: Mon, 25 Mar 2024 10:36:18 -0700 Subject: [PATCH] Fix NVIDIA VRAM detection on WSL2 environments (#1894) * NVIDIA VRAM detection on WSL2 environments More robust single NVIDIA GPU memory detection, following the improved NVIDIA WSL2 detection patch yesterday #1891. Tested and working on WSL2, Linux. Signed-off-by: Enrico Ros * Update aio/entrypoint.sh Signed-off-by: Ettore Di Giacinto --------- Signed-off-by: Enrico Ros Signed-off-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto --- aio/entrypoint.sh | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/aio/entrypoint.sh b/aio/entrypoint.sh index aeb5e4de..795cb86a 100755 --- a/aio/entrypoint.sh +++ b/aio/entrypoint.sh @@ -57,29 +57,33 @@ function detect_gpu() { } function detect_gpu_size() { - if [ "$GPU_ACCELERATION" = true ]; then - GPU_SIZE=gpu-8g - fi - # Attempting to find GPU memory size for NVIDIA GPUs - if echo "$gpu_model" | grep -iq nvidia; then + if [ "$GPU_ACCELERATION" = true ] && [ "$GPU_VENDOR" = "nvidia" ]; then echo "NVIDIA GPU detected. Attempting to find memory size..." - nvidia_sm=($(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits)) + # Using head -n 1 to get the total memory of the 1st NVIDIA GPU detected. + # If handling multiple GPUs is required in the future, this is the place to do it + nvidia_sm=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -n 1) if [ ! -z "$nvidia_sm" ]; then - echo "Total GPU Memory: ${nvidia_sm[0]} MiB" + echo "Total GPU Memory: $nvidia_sm MiB" + # if bigger than 8GB, use 16GB + #if [ "$nvidia_sm" -gt 8192 ]; then + # GPU_SIZE=gpu-16g + #else + GPU_SIZE=gpu-8g + #fi else - echo "Unable to determine NVIDIA GPU memory size." + echo "Unable to determine NVIDIA GPU memory size. Falling back to CPU." + GPU_SIZE=gpu-8g fi - # if bigger than 8GB, use 16GB - #if [ "$nvidia_sm" -gt 8192 ]; then - # GPU_SIZE=gpu-16g - #fi - else - echo "Non-NVIDIA GPU detected. GPU memory size detection for non-NVIDIA GPUs is not supported in this script." - fi + + # Default to a generic GPU size until we implement GPU size detection for non NVIDIA GPUs + elif [ "$GPU_ACCELERATION" = true ]; then + echo "Non-NVIDIA GPU detected. Specific GPU memory size detection is not implemented." + GPU_SIZE=gpu-8g # default to cpu if GPU_SIZE is not set - if [ -z "$GPU_SIZE" ]; then + else + echo "GPU acceleration is not enabled or supported. Defaulting to CPU." GPU_SIZE=cpu fi }