Script to manage Nvidia driver versions on proxmox hosts and in lxc containers.

nvidia_pve_passthrough.sh

#!/usr/bin/env bash

SHOW_VERSION_COUNT=20
USE_VERSION=""
RECOMMENDED_VERSION="580.105.08"

# Ingest Args (FUTURE NOT WORKING)
process_args() {
    echo "Processing arguments for function: $@"
    # Iterate while arguments exist
    while [ "$#" -gt 0 ]; do
        case "$1" in
            -a|--all)
                echo "Showing all versions."
				SHOW_VERSION_COUNT=9999
                shift # Move to the next argument
                ;;
            -v|--version)
                echo "Using version '$2' if it exists."
				USE_VERSION=$2
                shift 2 # Move past the option and its value
                ;;
            *)
                echo "Unknown option or argument: $1"
                shift
                ;;
        esac
    done
}

# Function to check if running inside an LXC container
is_lxc() {
    if command -v systemd-detect-virt &> /dev/null; then
        if [ $(systemd-detect-virt) = "lxc" ]; then
            return 0 # True, it's LXC
        fi
    fi
    if grep -q "lxc" /proc/1/cgroup; then
        return 0 # True, it's LXC
    fi
    return 1 # False, not LXC
}

# Check for and install missing dependency
confirm_depend () {
	if ! command -v "$1" &> /dev/null; then
		echo "[DEPEND] Missing dependency '$1' installing now."
		apt-get update
		apt install -y "$1"
	fi
}

# Handle arguments
#process_args

# Get run scope
IS_LXC=$(is_lxc && echo "true" || echo "false")

# Detect current version
CURRENT_VERSION="NOT INSTALLED"
if command -v nvidia-smi &> /dev/null; then
	CURRENT_VERSION_RAW=$(nvidia-smi --version | grep 'DRIVER version')
	CURRENT_VERSION="${CURRENT_VERSION_RAW#*:}"
	CURRENT_VERSION=$(echo "$CURRENT_VERSION" | xargs)
fi

PVE_LXC_DIR="/etc/pve/lxc/"
INSTALL_MODE="HOST"
INSTALL_TARGET="SELF"
if [[ "$IS_LXC" == "true" ]]; then
    echo "[CONTAINER] Detected script is running in an LXC container."
	INSTALL_MODE="LXC"
else
    echo "[HOST] Detected script is running on bare metal."
	INSTALL_MODE="HOST"
	
	# Check if proxmox host and if lxc directory exists
	if [ -d "$PVE_LXC_DIR" ]; then
		# Check if host drivers are installed
		if [[ "$CURRENT_VERSION" == "NOT INSTALLED" ]]; then
			echo "Detected Proxmox host, install drivers on host to setup passthrough in containers."
			echo "Defaulting to HOST installation."
		else
			# Get LXC Configs
			PVE_LXC_CONFIGS_RAW=( $PVE_LXC_DIR/*.conf )
			
			PVE_LXC_CONFIGS=()
			for i in "${!PVE_LXC_CONFIGS_RAW[@]}"; do
				PVE_LXC_CONFIGS[i]="$(basename ${PVE_LXC_CONFIGS_RAW[i]%*/} .conf)"
			done
			
			echo "Select [HOST] or Container ID to install NVIDIA Drivers:
['h' for HOST] ['q' to exit]"
			select item in "${PVE_LXC_CONFIGS[@]}"; do
				if [ "$REPLY" == "q" ]; then
					echo "Exiting."
					exit 98
				elif [ "$REPLY" == "h" ]; then
					INSTALL_MODE="HOST"
					break
				elif [ -n "$item" ]; then
					INSTALL_MODE="CONTAINER"
					INSTALL_TARGET=$item
					break
				else
					echo "Invalid choice, please try again."
				fi
			done
		fi
	fi
fi

echo "Install details:
INSTALL MODE:     $INSTALL_MODE
INSTALL TARGET:   $INSTALL_TARGET"

# Prepare installation proceedure
if [[ "$INSTALL_MODE" != "CONTAINER" ]]; then
	# Depends
	confirm_depend "curl"
	confirm_depend "pup"

	# Get latest version information
	VERSION_LATEST_RAW=$(curl -s https://download.nvidia.com/XFree86/Linux-x86_64/latest.txt)
	VERSION_LATEST="${VERSION_LATEST_RAW%% *}"

	# Get all versions
	VERSION_LIST=$(curl -s https://download.nvidia.com/XFree86/Linux-x86_64/ | pup 'ul li span.dir a[href] attr{href}')

	# Create array and count items
	mapfile -t VERSION_LIST_ASC <<<"$VERSION_LIST"
	VERSION_COUNT=${#VERSION_LIST_ASC[@]}

	# Remove trailing slash
	for i in "${!VERSION_LIST_ASC[@]}"; do
		VERSION_LIST_ASC[i]="${VERSION_LIST_ASC[i]%*/}"
	done

	# Sort NEW -> OLD
	VERSION_LIST_DESC=()
	for ((i = VERSION_COUNT - 1; i >= 0; i--)); do
	  VERSION_LIST_DESC+=("${VERSION_LIST_ASC[i]}")
	done

	# Show information
	echo "NVIDIA DRIVER INSTALL TOOL
DRIVER COUNT:          $VERSION_COUNT
LATEST VERSION:        $VERSION_LATEST
RECOMMENDED VERSION:   $RECOMMENDED_VERSION
CURRENT VERSION:       $CURRENT_VERSION

['l' for latest] ['r' for recommended] ['q' to quit]
Enter the number for the driver to install:
	"

	# Show selections
	select item in "${VERSION_LIST_DESC[@]:0:$SHOW_VERSION_COUNT}"; do
		if [ "$REPLY" == "q" ]; then
			echo "Exiting."
			exit 99
		elif [ "$REPLY" == "l" ]; then
			USE_VERSION=$VERSION_LATEST
			break
		elif [ "$REPLY" == "r" ]; then
			USE_VERSION=$RECOMMENDED_VERSION
			break
		elif [ -n "$item" ]; then
			USE_VERSION=$item
			break
		else
			echo "Invalid choice, please try again."
		fi
	done

	echo "Selected version: $USE_VERSION"

	# Remove old version if required
	if command -v nvidia-uninstall &> /dev/null; then
		echo "[OLD VERSION] Found old version of nvidia drivers, removing them first..."
		nvidia-uninstall --silent
	fi

	# Download Drivers and make Executable
	wget https://download.nvidia.com/XFree86/Linux-x86_64/$USE_VERSION/NVIDIA-Linux-x86_64-$USE_VERSION.run -P /tmp
	chmod +x /tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run

	# Install Drivers
	if [[ "$INSTALL_MODE" == "HOST" ]]; then
		if [ -d "$PVE_LXC_DIR" ]; then
			confirm_depend "pve-headers"
			confirm_depend "pve-nvidia-vgpu-helper"
			pve-nvidia-vgpu-helper setup
		fi
		/tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run -a -q --silent --kernel-module-type=proprietary
	else
		/tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run -a -q --silent --no-kernel-module --kernel-module-type=proprietary
	fi

	# Cleanup
	rm -f /tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run

	# Reboot if not host
	if [[ "$INSTALL_MODE" == "HOST" ]]; then
		echo "!!! === Please reboot when as soon as is reasonable. === !!!"
	else
		reboot
	fi
else
	# Container version should be the same as the host version
	echo "Using host version for LXC container: $CURRENT_VERSION"
	USE_VERSION=$CURRENT_VERSION

	# Configure the container to allow the setup of the GPU
	if ! grep -qF "/dev/nvidia0" "$PVE_LXC_DIR/${INSTALL_TARGET}.conf"; then
		echo "Updating LXC Config: $INSTALL_TARGET"
		
		# Stop the container
		pct shutdown $INSTALL_TARGET
		
		# Set Passthrough Configs
		pct set $INSTALL_TARGET --dev0 "/dev/nvidia0"
		pct set $INSTALL_TARGET --dev1 "/dev/nvidiactl"
		pct set $INSTALL_TARGET --dev2 "/dev/nvidia-modeset"
		pct set $INSTALL_TARGET --dev3 "/dev/nvidia-uvm"
		pct set $INSTALL_TARGET --dev4 "/dev/nvidia-uvm-tools"
		pct set $INSTALL_TARGET --dev5 "/dev/nvram"
		
		# Update Tags
		TAGS=$(pct config $INSTALL_TARGET | grep tags | cut -d' ' -f2- | sed 's/;/A/g' | sed 's/A/,/g')
		TAGS+=",GPU"
		pct set $INSTALL_TARGET --tags "$TAGS"
		
		# Start the container
		pct start $INSTALL_TARGET
	fi

	TMPFILE=$(mktemp)
	# Install in container
	echo "Generating LXC install script."
	echo "#!/usr/bin/env bash

# TEMPORARY GENERATED SCRIPT TO INSTALL NVIDIA drivers
# VERSION: $USE_VERSION
$ CONTAINER ID: $INSTALL_TARGET
	
if command -v nvidia-uninstall &> /dev/null; then
	nvidia-uninstall --silent
fi

wget https://download.nvidia.com/XFree86/Linux-x86_64/$USE_VERSION/NVIDIA-Linux-x86_64-$USE_VERSION.run -P /tmp
chmod +x /tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run
/tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run -a -q --silent --no-kernel-module --kernel-module-type=proprietary
rm -f /tmp/NVIDIA-Linux-x86_64-$USE_VERSION.run
reboot
	" > $TMPFILE
	
	echo "Pushing script to LXC container: $INSTALL_TARGET"
	pct push $INSTALL_TARGET $TMPFILE $TMPFILE
	
	echo "Executing script inside LXC container: $INSTALL_TARGET"
	pct exec $INSTALL_TARGET -- chmod +x $TMPFILE
	pct exec $INSTALL_TARGET -- $TMPFILE
fi