diff --git a/README.md b/README.md index 0331020..2469633 100644 --- a/README.md +++ b/README.md @@ -49,14 +49,17 @@ Requirements: # bash ./patch.sh -h SYNOPSIS - patch.sh [OPTION]... + patch.sh [-s] [-r|-h|-c VERSION|-l] DESCRIPTION - The patch for Nvidia drivers to increase encoder sessions + The patch for Nvidia drivers to remove NVENC session limit - -s Silent mode (No output) - -r Rollback to original (Restore lib from backup) - -h Print this help message + -s Silent mode (No output) + -r Rollback to original (Restore lib from backup) + -h Print this help message + -c VERSION Check if version VERSION supported by this patch. + Returns true exit code (0) if version is supported. + -l List supported driver versions ``` diff --git a/patch.sh b/patch.sh index 9201064..d89d07f 100755 --- a/patch.sh +++ b/patch.sh @@ -6,29 +6,36 @@ set -euo pipefail ; # <- this semicolon and comment make options apply backup_path="/opt/nvidia/libnvidia-encode-backup" silent_flag='' -rollback_flag='' print_usage() { printf ' SYNOPSIS - patch.sh [OPTION]... + patch.sh [-s] [-r|-h|-c VERSION|-l] DESCRIPTION - The patch for Nvidia drivers to increase encoder sessions + The patch for Nvidia drivers to remove NVENC session limit - -s Silent mode (No output) - -r Rollback to original (Restore lib from backup) - -h Print this help message + -s Silent mode (No output) + -r Rollback to original (Restore lib from backup) + -h Print this help message + -c VERSION Check if version VERSION supported by this patch. + Returns true exit code (0) if version is supported. + -l List supported driver versions ' } -while getopts 'rsh' flag; do - case "${flag}" in - r) rollback_flag='true' ;; - s) silent_flag='true' ;; - *) print_usage - exit 1 ;; - esac +# shellcheck disable=SC2209 +opmode="patch" + +while getopts 'rshc:l' flag; do + case "${flag}" in + r) opmode="${opmode}rollback" ;; + s) silent_flag='true' ;; + h) opmode="${opmode}help" ;; + c) opmode="${opmode}checkversion" ; checked_version="$OPTARG" ;; + l) opmode="${opmode}listversions" ;; + *) echo "Incorrect option specified in command line" ; exit 2 ;; + esac done if [[ $silent_flag ]]; then @@ -105,56 +112,84 @@ declare -A object_list=( ["435.21"]='libnvcuvid.so' ) -NVIDIA_SMI="$(command -v nvidia-smi)" +check_version_supported () { + local ver="$1" + [[ "${patch_list[$ver]+isset}" && "${object_list[$ver]+isset}" ]] +} -if ! driver_version=$("$NVIDIA_SMI" --query-gpu=driver_version --format=csv,noheader,nounits | head -n 1) ; then - echo 'Something went wrong. Check nvidia driver' - exit 1; -fi - -echo "Detected nvidia driver version: $driver_version" - -if [[ ! "${patch_list[$driver_version]+isset}" || ! "${object_list[$driver_version]+isset}" ]]; then - echo "Patch for this ($driver_version) nvidia driver not found." 1>&2 - echo "Available patches for: " 1>&2 +get_supported_versions () { for drv in "${!patch_list[@]}"; do - echo "$drv" 1>&2 - done - exit 1; -fi + [[ "${object_list[$drv]+isset}" ]] && echo "$drv" + done | sort -t. -n + return 0 +} -patch="${patch_list[$driver_version]}" -object="${object_list[$driver_version]}" - -declare -a driver_locations=( - '/usr/lib/x86_64-linux-gnu' - '/usr/lib/x86_64-linux-gnu/nvidia/current/' - '/usr/lib64' - "/usr/lib/nvidia-${driver_version%%.*}" -) - -dir_found='' -for driver_dir in "${driver_locations[@]}" ; do - if [[ -e "$driver_dir/$object.$driver_version" ]]; then - dir_found='true' - break +patch_common () { + NVIDIA_SMI="$(command -v nvidia-smi || true)" + if [[ ! "$NVIDIA_SMI" ]] ; then + echo 'nvidia-smi utility not found. Probably driver is not installed.' + exit 1 fi -done -[[ "$dir_found" ]] || { echo "ERROR: cannot detect driver directory"; exit 1; } + if ! driver_version=$("$NVIDIA_SMI" --query-gpu=driver_version --format=csv,noheader,nounits | head -n 1) ; then + echo 'Something went wrong. Check nvidia driver' + exit 1 + fi -if [[ $rollback_flag ]]; then + echo "Detected nvidia driver version: $driver_version" + + if ! check_version_supported "$driver_version" ; then + echo "Patch for this ($driver_version) nvidia driver not found." + echo "Patch is available for versions: " + get_supported_versions + exit 1 + fi + + patch="${patch_list[$driver_version]}" + object="${object_list[$driver_version]}" + + declare -a driver_locations=( + '/usr/lib/x86_64-linux-gnu' + '/usr/lib/x86_64-linux-gnu/nvidia/current/' + '/usr/lib64' + "/usr/lib/nvidia-${driver_version%%.*}" + ) + + dir_found='' + for driver_dir in "${driver_locations[@]}" ; do + if [[ -e "$driver_dir/$object.$driver_version" ]]; then + dir_found='true' + break + fi + done + + [[ "$dir_found" ]] || { echo "ERROR: cannot detect driver directory"; exit 1; } + +} + +rollback () { + patch_common if [[ -f "$backup_path/$object.$driver_version" ]]; then cp -p "$backup_path/$object.$driver_version" \ "$driver_dir/$object.$driver_version" echo "Restore from backup $object.$driver_version" else echo "Backup not found. Try to patch first." - exit 1; + exit 1 fi -else - if [[ ! -f "$backup_path/$object.$driver_version" ]]; then - echo "Attention! Backup not found. Copy current $object to backup." +} + +patch () { + patch_common + if [[ -f "$backup_path/$object.$driver_version" ]]; then + bkp_hash="$(sha1sum "$backup_path/$object.$driver_version" | cut -f1 -d\ )" + drv_hash="$(sha1sum "$driver_dir/$object.$driver_version" | cut -f1 -d\ )" + if [[ "$bkp_hash" != "$drv_hash" ]] ; then + echo "Backup exists and driver file differ from backup. Skipping patch." + return 0 + fi + else + echo "Attention! Backup not found. Copying current $object to backup." mkdir -p "$backup_path" cp -p "$driver_dir/$object.$driver_version" \ "$backup_path/$object.$driver_version" @@ -165,4 +200,28 @@ else sha1sum "${PATCH_OUTPUT_DIR-$driver_dir}/$object.$driver_version" ldconfig echo "Patched!" -fi +} + +query_version_support () { + if check_version_supported "$checked_version" ; then + echo "SUPPORTED" + exit 0 + else + echo "NOT SUPPORTED" + exit 1 + fi +} + +list_supported_versions () { + get_supported_versions +} + +case "${opmode}" in + patch) patch ;; + patchrollback) rollback ;; + patchhelp) print_usage ; exit 2 ;; + patchcheckversion) query_version_support ;; + patchlistversions) list_supported_versions ;; + *) echo "Incorrect combination of flags. Use option -h to get help." + exit 2 ;; +esac