ExternalVendorCode/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2025-01-30 08:04:13 +00:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Fix NVIDIA VRAM detection on WSL2 environments (#1894)

Browse Source 
* 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 <enrico.ros@gmail.com>

* Update aio/entrypoint.sh

Signed-off-by: Ettore Di Giacinto <mudler@users.noreply.github.com>

---------

Signed-off-by: Enrico Ros <enrico.ros@gmail.com>
Signed-off-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
This commit is contained in:
Enrico Ros 2024-03-25 10:36:18 -07:00 committed by GitHub
parent 5e12382524
commit 08c7b17298
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 20 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
aio/entrypoint.sh
View File

@ -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
}