Adding install instruction to make GPUs usable

setup.sh

@@ -24,6 +24,53 @@ done
 IMAGE_NAME="local-ai"; CONTAINER_NAME="local-ai"; PORT=11434
 NOTES=()   # machine-specific remarks, echoed to the terminal and RUN_AIDER.md
 
+# Distro-specific commands to make an NVIDIA GPU usable from Docker.
+# $1 = need_driver (1/0)
+gpu_install_steps() {
+  local need_driver=$1 id="" like=""
+  [ -r /etc/os-release ] && { id=$(. /etc/os-release; echo "${ID:-}"); like=$(. /etc/os-release; echo "${ID_LIKE:-}"); }
+  case " $id $like " in
+    *" ubuntu "*|*" debian "*)
+      if [ "$need_driver" = 1 ]; then
+        if [[ " $id $like " == *" ubuntu "* ]]; then
+          echo "# 1. NVIDIA driver (then reboot; with Secure Boot you must enroll the MOK key when prompted)"
+          echo "sudo ubuntu-drivers install"
+        else
+          echo "# 1. NVIDIA driver (needs the non-free component enabled; then reboot)"
+          echo "sudo apt install nvidia-driver"
+        fi
+        echo "sudo reboot"
+        echo "nvidia-smi                       # after reboot: should list your GPU"
+        echo
+      fi
+      echo "# 2. NVIDIA container toolkit (lets Docker use the GPU)"
+      echo "curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
+      echo "curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list"
+      echo "sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit"
+      echo "sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker"
+      ;;
+    *" fedora "*|*" rhel "*|*" centos "*)
+      if [ "$need_driver" = 1 ]; then
+        echo "# 1. NVIDIA driver from RPM Fusion (enable RPM Fusion first: https://rpmfusion.org/Configuration), then reboot"
+        echo "sudo dnf install akmod-nvidia xorg-x11-drv-nvidia-cuda && sudo reboot"
+        echo
+      fi
+      echo "# 2. NVIDIA container toolkit"
+      echo "curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo"
+      echo "sudo dnf install -y nvidia-container-toolkit"
+      echo "sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker"
+      ;;
+    *)
+      echo "# Install the NVIDIA driver with your distribution's package manager, then the container toolkit:"
+      echo "# https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html"
+      echo "# and finish with: sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker"
+      ;;
+  esac
+  echo
+  echo "# 3. Verify, then re-run ./setup.sh"
+  echo "docker run --rm --gpus all ubuntu nvidia-smi"
+}
+
 # ---------------------------------------------------------------- detection
 cpu_model=$(grep -m1 'model name' /proc/cpuinfo | cut -d: -f2- | xargs || true)
 cpu_threads=$(nproc)
@@ -31,7 +78,7 @@ ram_gb=$(awk '/MemTotal/ {printf "%.1f", $2/1024/1024}' /proc/meminfo)
 disk_gb=$(df -Pk . | awk 'NR==2 {printf "%.0f", $4/1024/1024}')
 grep -qw avx2 /proc/cpuinfo && avx2=yes || avx2=no
 
-gpu_name=""; vram_gb=0; gpu_usable=0; nodriver_gpu=0; est_vram=0
+gpu_name=""; vram_gb=0; gpu_usable=0; nodriver_gpu=0; est_vram=0; toolkit_only=0; GPU_ADVICE=""; GPU_STEPS=""
 if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
   gpu_name=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
   vram_mb=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
@@ -58,7 +105,7 @@ if command -v docker >/dev/null 2>&1; then
   fi
 fi
 if [ "$gpu_usable" = 1 ] && [ "$docker_nvidia" = no ]; then
-  NOTES+=("GPU $gpu_name found but Docker has no NVIDIA runtime (install nvidia-container-toolkit). Running on CPU.")
+  toolkit_only=1
   gpu_usable=0
 fi
 [ "$FORCE_GPU" = 1 ] && gpu_usable=1
@@ -106,23 +153,39 @@ done
 [ "$rec" -ge 0 ] || rec=0
 
 # ------------------------------------------- would a GPU driver help?
-if [ "$nodriver_gpu" = 1 ]; then
+if [ "$nodriver_gpu" = 1 ] || [ "$toolkit_only" = 1 ]; then
+  [ "$nodriver_gpu" = 1 ] && gvram=$est_vram || gvram=$vram_gb
   best=-1
-  if awk -v v="$est_vram" 'BEGIN{exit !(v>0)}'; then
+  if awk -v v="$gvram" 'BEGIN{exit !(v>0)}'; then
     for i in "${!rows[@]}"; do
       mem=$(cut -f3 <<<"${rows[$i]}")
-      awk -v m="$mem" -v v="$est_vram" 'BEGIN{exit !(m<=v*0.9)}' && best=$i
+      awk -v m="$mem" -v v="$gvram" 'BEGIN{exit !(m<=v*0.9)}' && best=$i
     done
   fi
-  drv="install the NVIDIA driver and nvidia-container-toolkit, then re-run ./setup.sh"
+  if [ "$nodriver_gpu" = 1 ]; then
-  if [ "$est_vram" = 0 ]; then
+    what="the NVIDIA driver and nvidia-container-toolkit"
-    NOTES+=("NVIDIA GPU ($gpu_name) found but not usable (no driver). Installing the driver may speed things up a lot if it has enough VRAM: $drv.")
+    have="has no working NVIDIA driver"
-  elif [ "$best" -gt "$rec" ]; then
-    NOTES+=("Worth installing the NVIDIA driver: with ~${est_vram} GB VRAM, ${tags[$best]} would run fully on the GPU (better than the current CPU pick ${tags[$rec]}). $drv.")
-  elif [ "$best" -ge 0 ]; then
-    NOTES+=("Optional: an NVIDIA driver would make models up to ${tags[$best]} run fully on the ~${est_vram} GB GPU (much faster) and partly speed up larger ones such as ${tags[$rec]} via partial offload, but ${est_vram} GB VRAM does not unlock bigger models than ${tags[$rec]}. Worth it for speed: $drv.")
   else
-    NOTES+=("A driver would not help much: ~${est_vram} GB VRAM is too small to hold even the smallest model.")
+    what="nvidia-container-toolkit"
+    have="is not usable from Docker (no NVIDIA runtime)"
+  fi
+  head="Your GPU ($gpu_name) $have, so the model runs on the CPU."
+  if [ "$gvram" = 0 ]; then
+    GPU_ADVICE="$head If it has enough VRAM, installing $what may speed things up a lot."
+  elif [ "$best" -gt "$rec" ]; then
+    GPU_ADVICE="$head Recommended: install $what. With ~${gvram} GB VRAM, ${tags[$best]} would run fully on the GPU (better than the current CPU pick ${tags[$rec]})."
+  elif [ "$best" -ge 0 ]; then
+    GPU_ADVICE="$head Optional, for speed: install $what. With ~${gvram} GB VRAM, models up to ${tags[$best]} would run fully on the GPU and larger ones such as ${tags[$rec]} partly (partial offload), but this VRAM does not unlock bigger models than ${tags[$rec]}."
+  else
+    GPU_ADVICE=""   # VRAM too small to matter
+  fi
+  if [ -n "$GPU_ADVICE" ]; then
+    GPU_STEPS=$(gpu_install_steps "$nodriver_gpu")
+    echo "=== GPU hint ==="
+    echo "  $GPU_ADVICE"
+    echo "  To do so (optional; then re-run ./setup.sh):"
+    sed 's/^/    /' <<<"$GPU_STEPS"
+    echo
   fi
 fi
 
@@ -226,6 +289,11 @@ model: ollama_chat/$MODEL
 - Nothing is sent off the machine once the image is built.
 - To switch model or after changing hardware, re-run \`./setup.sh\` then \`./build-ai-image.sh\`.
 EOF
+if [ -n "$GPU_ADVICE" ]; then
+  echo; echo "## Optional: use your GPU"; echo
+  echo "$GPU_ADVICE"; echo
+  echo '```bash'; echo "$GPU_STEPS"; echo '```'
+fi
 if [ "${#NOTES[@]}" -gt 0 ]; then
   echo; echo "## Notes about this machine"; echo
   for n in "${NOTES[@]}"; do echo "- $n"; done