与官方脚本的主要区别和升级

📋 核心区别对比
特性 官方脚本 修改后脚本
缓存机制 ❌ 无缓存,每次都下载 ✅ 支持本地缓存,优先使用本地文件
默认缓存目录 ❌ 无 ✅ ~/Downloads/ollama-cache/
下载信息显示 ❌ 只显示"Downloading" ✅ 显示完整下载地址、保存位置
镜像地址支持 ❌ 无 ✅ 支持 OLLAMA_MIRROR 环境变量
离线安装 ❌ 不支持 ✅ 支持,只需预先下载文件到缓存目录
用户提示 ❌ 无提示 ✅ 多处友好提示(默认目录、镜像使用等)
🚀 主要升级功能

1. 智能缓存系统 ⭐⭐⭐

官方脚本:每次都下载

curl … | tar -xf -

修改后脚本:先检查缓存

if [ -f “TEMPDIR/TEMP_DIR/TEMPDIR/{filename}.tar.zst” ]; then
使用缓存文件
else
下载文件
fi

优势:

  • 避免重复下载相同文件
  • 支持离线安装
  • 节省带宽和时间

2. 默认缓存目录 ⭐⭐

默认使用 ~/Downloads/ollama-cache/

用户可以自定义:

./install.sh --cache-dir=~/my-cache

或修改脚本中的 DEFAULT_CACHE_DIR 变量

优势:

  • 开箱即用,无需配置
  • 目录固定,方便管理
  • 支持自定义目录

3. 镜像地址支持 ⭐⭐⭐

支持官方镜像环境变量

export OLLAMA_MIRROR=“https://gh.llkk.cc/https://github.com/ollama/ollama/releases/latest/download/”
./install.sh
优势:

  • 解决国内网络问题
  • 兼容官方脚本的使用方式
  • 支持各种镜像源

4. 透明下载信息 ⭐⭐

官方脚本:只显示

Downloading ollama-linux-amd64.tar.zst

修改后脚本:完整信息

开始下载 ollama-linux-amd64.tar.zst
下载地址: https://ollama.com/download/ollama-linux-amd64.tar.zst
保存位置: /home/sunny/Downloads/ollama-cache/ollama-linux-amd64.tar.zst
优势:

  • 让用户心中有数
  • 便于调试网络问题
  • 清楚文件保存位置
  1. 智能格式选择 ⭐

优先使用 .tar.zst(新格式), 检查失败时回退到 .tgz(旧格式),并提示用户可能的原因

优势:

  • 兼容不同版本的 Ollama
  • 提供清晰的错误提示

6. 多种运行模式 ⭐⭐

模式1:标准模式(使用缓存)

./install.sh

模式2:自定义缓存目录

./install.sh --cache-dir=~/my-cache

模式3:传统模式(不使用缓存)

./install.sh --no-cache

模式4:使用镜像地址

export OLLAMA_MIRROR=“…”

📊 使用场景对比

场景1:网络环境差
官方脚本:下载失败,无法安装
修改后脚本:使用镜像地址或预先下载文件
场景2:多次安装/重装
官方脚本:每次都重新下载
修改后脚本:第二次开始直接使用缓存,秒级安装
场景3:批量部署
官方脚本:每台机器都要下载
修改后脚本:下载一次,复制到多台机器安装
场景4:离线环境
官方脚本:无法安装
修改后脚本:预先下载文件后可离线安装

🎯 技术改进

代码质量提升

  1. 逻辑更清晰:分离缓存检查和下载逻辑
  2. 错误处理更好:添加用户友好的提示
  3. 兼容性更强:支持官方环境变量和用法
  4. 可维护性更高:清晰的注释和结构
    用户体验优化
  5. 信息透明:显示完整的下载信息
  6. 操作简单:开箱即用,无需复杂配置
  7. 提示友好:遇到问题时给出解决方案
  8. 灵活性高:支持多种使用方式

📝 代码行数对比

  • 官方脚本:455 行
  • 修改后脚本:约 620 行
  • 增加代码:约 165 行(主要是缓存功能和用户提示)

💡 兼容性说明

修改后的脚本完全兼容官方脚本:

  • 支持官方的所有环境变量(OLLAMA_VERSION, OLLAMA_MIRROR 等)
  • 支持官方的所有命令行选项
  • 安装过程和结果与官方脚本一致
  • 只是增加了缓存功能,不影响原有功能

🏆 总结

官方脚本:简洁、可靠,适合网络环境好的用户
修改后脚本:功能更强、体验更好,特别适合:

  • 国内网络环境
  • 需要多次安装
  • 离线安装场景
  • 批量部署需求
    推荐使用修改后的脚本,因为它保留了官方脚本的所有优点,同时增加了实用的缓存功能。

以下是新修改的脚本内容:

#!/bin/sh
# This script installs Ollama on Linux and macOS.
# It detects the current operating system architecture and installs the appropriate version of Ollama.
#
# 默认缓存目录: ~/Downloads/ollama-cache/
# 如需修改默认目录,请编辑脚本中的 DEFAULT_CACHE_DIR 变量
#
# 用法: ./install.sh [选项]
#   --cache-dir=路径    指定缓存和工作目录(替代临时目录)
#   --no-cache         禁用缓存功能,使用临时目录(传统方式)
#   -h, --help         显示帮助信息
#
# 示例:
#   ./install.sh                    # 使用默认缓存目录 ~/Downloads/ollama-cache/
#   ./install.sh --cache-dir=~/my-cache  # 使用指定缓存目录
#   ./install.sh --no-cache        # 禁用缓存,使用临时目录

# Wrap script in main function so that a truncated partial download doesn't end
# up executing half a script.
main() {

set -eu

red="$( (/usr/bin/tput bold || :; /usr/bin/tput setaf 1 || :) 2>&-)"
plain="$( (/usr/bin/tput sgr0 || :) 2>&-)"

status() { echo ">>> $*" >&2; }
error() { echo "${red}ERROR:${plain} $*"; exit 1; }
warning() { echo "${red}WARNING:${plain} $*"; }

# Parse command line arguments
CACHE_DIR=""
USE_CACHE=true
while [ $# -gt 0 ]; do
    case "$1" in
        --cache-dir|--local-dir)
            CACHE_DIR="$2"
            shift 2
            ;;
        --cache-dir=*|--local-dir=*)
            CACHE_DIR="${1#*=}"
            shift
            ;;
        --no-cache)
            USE_CACHE=false
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo "Install Ollama on Linux or macOS"
            echo ""
            echo "Options:"
            echo "  --cache-dir=PATH    指定缓存和工作目录(替代临时目录)"
            echo "                      脚本会先检查该目录是否有需要的安装文件"
            echo "                      如果没有则下载并保存到该目录"
            echo "                      下次运行可直接使用已下载的文件"
            echo "                      支持的文件:"
            echo "                      - macOS: Ollama-darwin.zip"
            echo "                      - Linux: ollama-linux-\${ARCH}.tar.zst 或 .tgz"
            echo "                      - 可选: ollama-linux-\${ARCH}-rocm.tar.zst/.tgz (AMD GPU)"
            echo "                      - 可选: ollama-linux-\${ARCH}-jetpack*.tar.zst/.tgz (Jetson)"
            echo "  --no-cache          禁用缓存功能,使用临时目录(传统方式)"
            echo "  -h, --help         显示帮助信息"
            echo ""
            echo "示例:"
            echo "  $0                    # 使用默认缓存目录 ~/Download/ollama-cache/"
            echo "  $0 --cache-dir=~/my-cache  # 使用指定缓存目录"
            echo "  $0 --no-cache        # 禁用缓存,使用临时目录"
            exit 0
            ;;
        *)
            error "未知选项: $1"
            ;;
    esac
done

# 设置默认缓存目录(用户可以修改这里的默认值)
# 注意:将 ~/Downloads/ 修改为您想要的默认目录
DEFAULT_CACHE_DIR="$HOME/Downloads/ollama-cache"

# 设置工作目录
if [ "$USE_CACHE" = true ]; then
    # 如果没有指定缓存目录,使用默认目录
    if [ -z "$CACHE_DIR" ]; then
        CACHE_DIR="$DEFAULT_CACHE_DIR"
        status "使用默认缓存目录: $CACHE_DIR"
        status "提示:您可以通过 --cache-dir 参数指定其他目录,或修改脚本中的 DEFAULT_CACHE_DIR 变量"
    else
        status "使用指定缓存目录: $CACHE_DIR"
    fi
    
    # 如果目录不存在则创建
    if [ ! -d "$CACHE_DIR" ]; then
        status "创建缓存目录: $CACHE_DIR"
        mkdir -p "$CACHE_DIR" || error "无法创建目录: $CACHE_DIR"
    fi
    TEMP_DIR="$CACHE_DIR"
    
    # 不自动清理用户指定的目录
    cleanup() { :; }
else
    # 禁用缓存,使用临时目录
    TEMP_DIR=$(mktemp -d)
    cleanup() { rm -rf $TEMP_DIR; }
    status "禁用缓存功能,使用临时目录"
fi
trap cleanup EXIT

available() { command -v $1 >/dev/null; }
require() {
    local MISSING=''
    for TOOL in $*; do
        if ! available $TOOL; then
            MISSING="$MISSING $TOOL"
        fi
    done

    echo $MISSING
}

OS="$(uname -s)"
ARCH=$(uname -m)
case "$ARCH" in
    x86_64) ARCH="amd64" ;;
    aarch64|arm64) ARCH="arm64" ;;
    *) error "Unsupported architecture: $ARCH" ;;
esac

VER_PARAM="${OLLAMA_VERSION:+?version=$OLLAMA_VERSION}"

# 支持镜像地址(兼容官方脚本)
OLLAMA_DOWNLOAD_BASE="${OLLAMA_MIRROR:-https://ollama.com/download}"

###########################################
# macOS
###########################################

if [ "$OS" = "Darwin" ]; then
    NEEDS=$(require curl unzip)
    if [ -n "$NEEDS" ]; then
        status "ERROR: The following tools are required but missing:"
        for NEED in $NEEDS; do
            echo "  - $NEED"
        done
        exit 1
    fi

    DOWNLOAD_URL="${OLLAMA_DOWNLOAD_BASE}/Ollama-darwin.zip${VER_PARAM}"

    if pgrep -x Ollama >/dev/null 2>&1; then
        status "Stopping running Ollama instance..."
        pkill -x Ollama 2>/dev/null || true
        sleep 2
    fi

    if [ -d "/Applications/Ollama.app" ]; then
        status "Removing existing Ollama installation..."
        rm -rf "/Applications/Ollama.app"
    fi

    # 检查缓存目录中是否有文件
    if [ -f "$TEMP_DIR/Ollama-darwin.zip" ]; then
        status "使用缓存中的 Ollama-darwin.zip"
    else
        # 显示下载信息,让用户心中有数
        status "开始下载 Ollama for macOS..."
        status "下载地址: $DOWNLOAD_URL"
        status "保存位置: $TEMP_DIR/Ollama-darwin.zip"
        
        curl --fail --show-error --location --progress-bar \
            -o "$TEMP_DIR/Ollama-darwin.zip" "$DOWNLOAD_URL"
        
        # 如果是缓存目录模式,文件已经保存在缓存目录中
        if [ -n "$CACHE_DIR" ]; then
            status "下载完成,文件已保存到缓存目录: $CACHE_DIR/Ollama-darwin.zip"
        else
            status "下载完成(临时目录模式)"
        fi
    fi

    status "Installing Ollama to /Applications..."
    unzip -q "$TEMP_DIR/Ollama-darwin.zip" -d "$TEMP_DIR"
    mv "$TEMP_DIR/Ollama.app" "/Applications/"

    if [ ! -L "/usr/local/bin/ollama" ] || [ "$(readlink "/usr/local/bin/ollama")" != "/Applications/Ollama.app/Contents/Resources/ollama" ]; then
        status "Adding 'ollama' command to PATH (may require password)..."
        mkdir -p "/usr/local/bin" 2>/dev/null || sudo mkdir -p "/usr/local/bin"
        ln -sf "/Applications/Ollama.app/Contents/Resources/ollama" "/usr/local/bin/ollama" 2>/dev/null || \
            sudo ln -sf "/Applications/Ollama.app/Contents/Resources/ollama" "/usr/local/bin/ollama"
    fi

    if [ -z "${OLLAMA_NO_START:-}" ]; then
        status "Starting Ollama..."
        open -a Ollama --args hidden
    fi

    status "Install complete. You can now run 'ollama'."
    exit 0
fi

###########################################
# Linux
###########################################

[ "$OS" = "Linux" ] || error 'This script is intended to run on Linux and macOS only.'

IS_WSL2=false

KERN=$(uname -r)
case "$KERN" in
    *icrosoft*WSL2 | *icrosoft*wsl2) IS_WSL2=true;;
    *icrosoft) error "Microsoft WSL1 is not currently supported. Please use WSL2 with 'wsl --set-version <distro> 2'" ;;
    *) ;;
esac

SUDO=
if [ "$(id -u)" -ne 0 ]; then
    # Running as root, no need for sudo
    if ! available sudo; then
        error "This script requires superuser permissions. Please re-run as root."
    fi

    SUDO="sudo"
fi

NEEDS=$(require curl awk grep sed tee xargs)
if [ -n "$NEEDS" ]; then
    status "ERROR: The following tools are required but missing:"
    for NEED in $NEEDS; do
        echo "  - $NEED"
    done
    exit 1
fi

# Function to download and extract with fallback from zst to tgz
download_and_extract() {
    local url_base="$1"
    local dest_dir="$2"
    local filename="$3"
    local file_ext=""
    local cache_file=""
    
    # 第一步:先检查缓存中是否有任何格式的文件(优先使用本地文件)
    # 检查 .tar.zst 格式(新版本)
    if [ -f "$TEMP_DIR/${filename}.tar.zst" ]; then
        file_ext=".tar.zst"
        cache_file="$TEMP_DIR/${filename}.tar.zst"
        status "使用缓存文件: ${filename}${file_ext}"
        
        # 检查是否有 zstd 工具
        if ! available zstd; then
            error "缓存文件 ${filename}.tar.zst 需要 zstd 解压。请安装 zstd:
  - Debian/Ubuntu: sudo apt-get install zstd
  - RHEL/CentOS/Fedora: sudo dnf install zstd
  - Arch: sudo pacman -S zstd"
        fi
        
        cat "$cache_file" | zstd -d | $SUDO tar -xf - -C "${dest_dir}"
        return 0
    # 检查 .tgz 格式(旧版本)
    elif [ -f "$TEMP_DIR/${filename}.tgz" ]; then
        file_ext=".tgz"
        cache_file="$TEMP_DIR/${filename}.tgz"
        status "使用缓存文件: ${filename}${file_ext}"
        
        cat "$cache_file" | $SUDO tar -xzf - -C "${dest_dir}"
        return 0
    fi
    
    # 第二步:缓存中没有文件,需要下载
    # 检查 .tar.zst 是否可用(新版本格式)
    if curl --fail --silent --head --location "${url_base}/${filename}.tar.zst${VER_PARAM}" >/dev/null 2>&1; then
        # zst 文件存在 - 检查是否有 zstd 工具
        if ! available zstd; then
            error "此版本需要 zstd 解压。请安装 zstd:
  - Debian/Ubuntu: sudo apt-get install zstd
  - RHEL/CentOS/Fedora: sudo dnf install zstd
  - Arch: sudo pacman -S zstd"
        fi

        file_ext=".tar.zst"
        cache_file="$TEMP_DIR/${filename}${file_ext}"
        local download_url="${url_base}/${filename}${file_ext}${VER_PARAM}"
        
        status "开始下载 ${filename}${file_ext}"
        status "下载地址: $download_url"
        status "保存位置: $cache_file"
        
        # 如果是缓存目录模式,先下载到文件
        if [ -n "$CACHE_DIR" ]; then
            curl --fail --show-error --location --progress-bar \
                -o "$cache_file" "$download_url"
            status "下载完成,文件已保存到缓存目录: $cache_file"
            cat "$cache_file" | zstd -d | $SUDO tar -xf - -C "${dest_dir}"
        else
            # 临时目录模式,流式传输直接解压
            status "临时目录模式,下载后直接解压(不保存文件)"
            curl --fail --show-error --location --progress-bar \
                "$download_url" | \
                zstd -d | $SUDO tar -xf - -C "${dest_dir}"
        fi
        return 0
    fi

    # 回退到 .tgz 格式(旧版本)
    file_ext=".tgz"
    cache_file="$TEMP_DIR/${filename}${file_ext}"
    local download_url="${url_base}/${filename}${file_ext}${VER_PARAM}"
    
    status "注意: .tar.zst 格式检查失败,使用 .tgz 格式(旧版本)"
    status "提示: 如果网络有问题,可以设置 OLLAMA_MIRROR 环境变量使用镜像地址"
    status "例如: export OLLAMA_MIRROR='https://gh.llkk.cc/https://github.com/ollama/ollama/releases/latest/download/'"
    status "开始下载 ${filename}${file_ext}"
    status "下载地址: $download_url"
    status "保存位置: $cache_file"
    
    # 如果是缓存目录模式,先下载到文件
    if [ -n "$CACHE_DIR" ]; then
        curl --fail --show-error --location --progress-bar \
            -o "$cache_file" "$download_url"
        status "下载完成,文件已保存到缓存目录: $cache_file"
        cat "$cache_file" | $SUDO tar -xzf - -C "${dest_dir}"
    else
        # 临时目录模式,流式传输直接解压
        status "临时目录模式,下载后直接解压(不保存文件)"
        curl --fail --show-error --location --progress-bar \
            "$download_url" | \
            $SUDO tar -xzf - -C "${dest_dir}"
    fi
}

for BINDIR in /usr/local/bin /usr/bin /bin; do
    echo $PATH | grep -q $BINDIR && break || continue
done
OLLAMA_INSTALL_DIR=$(dirname ${BINDIR})

if [ -d "$OLLAMA_INSTALL_DIR/lib/ollama" ] ; then
    status "Cleaning up old version at $OLLAMA_INSTALL_DIR/lib/ollama"
    $SUDO rm -rf "$OLLAMA_INSTALL_DIR/lib/ollama"
fi
status "Installing ollama to $OLLAMA_INSTALL_DIR"
$SUDO install -o0 -g0 -m755 -d $BINDIR
$SUDO install -o0 -g0 -m755 -d "$OLLAMA_INSTALL_DIR/lib/ollama"
download_and_extract "$OLLAMA_DOWNLOAD_BASE" "$OLLAMA_INSTALL_DIR" "ollama-linux-${ARCH}"

if [ "$OLLAMA_INSTALL_DIR/bin/ollama" != "$BINDIR/ollama" ] ; then
    status "Making ollama accessible in the PATH in $BINDIR"
    $SUDO ln -sf "$OLLAMA_INSTALL_DIR/ollama" "$BINDIR/ollama"
fi

# Check for NVIDIA JetPack systems with additional downloads
if [ -f /etc/nv_tegra_release ] ; then
    if grep R36 /etc/nv_tegra_release > /dev/null ; then
        download_and_extract "$OLLAMA_DOWNLOAD_BASE" "$OLLAMA_INSTALL_DIR" "ollama-linux-${ARCH}-jetpack6"
    elif grep R35 /etc/nv_tegra_release >/dev/null ; then
        download_and_extract "$OLLAMA_DOWNLOAD_BASE" "$OLLAMA_INSTALL_DIR" "ollama-linux-${ARCH}-jetpack5"
    else
        warning "Unsupported JetPack version detected.  GPU may not be supported"
    fi
fi

install_success() {
    status 'The Ollama API is now available at 127.0.0.1:11434.'
    status 'Install complete. Run "ollama" from the command line.'
}
trap install_success EXIT

# Everything from this point onwards is optional.

configure_systemd() {
    if ! id ollama >/dev/null 2>&1; then
        status "Creating ollama user..."
        $SUDO useradd -r -s /bin/false -U -m -d /usr/share/ollama ollama
    fi
    if getent group render >/dev/null 2>&1; then
        status "Adding ollama user to render group..."
        $SUDO usermod -a -G render ollama
    fi
    if getent group video >/dev/null 2>&1; then
        status "Adding ollama user to video group..."
        $SUDO usermod -a -G video ollama
    fi

    status "Adding current user to ollama group..."
    $SUDO usermod -a -G ollama $(whoami)

    status "Creating ollama systemd service..."
    cat <<EOF | $SUDO tee /etc/systemd/system/ollama.service >/dev/null
[Unit]
Description=Ollama Service
After=network-online.target

[Service]
ExecStart=$BINDIR/ollama serve
User=ollama
Group=ollama
Restart=always
RestartSec=3
Environment="PATH=$PATH"

[Install]
WantedBy=default.target
EOF
    SYSTEMCTL_RUNNING="$(systemctl is-system-running || true)"
    case $SYSTEMCTL_RUNNING in
        running|degraded)
            status "Enabling and starting ollama service..."
            $SUDO systemctl daemon-reload
            $SUDO systemctl enable ollama

            start_service() { $SUDO systemctl restart ollama; }
            trap start_service EXIT
            ;;
        *)
            warning "systemd is not running"
            if [ "$IS_WSL2" = true ]; then
                warning "see https://learn.microsoft.com/en-us/windows/wsl/systemd#how-to-enable-systemd to enable it"
            fi
            ;;
    esac
}

if available systemctl; then
    configure_systemd
fi

# WSL2 only supports GPUs via nvidia passthrough
# so check for nvidia-smi to determine if GPU is available
if [ "$IS_WSL2" = true ]; then
    if available nvidia-smi && [ -n "$(nvidia-smi | grep -o "CUDA Version: [0-9]*\.[0-9]*")" ]; then
        status "Nvidia GPU detected."
    fi
    install_success
    exit 0
fi

# Don't attempt to install drivers on Jetson systems
if [ -f /etc/nv_tegra_release ] ; then
    status "NVIDIA JetPack ready."
    install_success
    exit 0
fi

# Install GPU dependencies on Linux
if ! available lspci && ! available lshw; then
    warning "Unable to detect NVIDIA/AMD GPU. Install lspci or lshw to automatically detect and install GPU dependencies."
    exit 0
fi

check_gpu() {
    # Look for devices based on vendor ID for NVIDIA and AMD
    case $1 in
        lspci)
            case $2 in
                nvidia) available lspci && lspci -d '10de:' | grep -q 'NVIDIA' || return 1 ;;
                amdgpu) available lspci && lspci -d '1002:' | grep -q 'AMD' || return 1 ;;
            esac ;;
        lshw)
            case $2 in
                nvidia) available lshw && $SUDO lshw -c display -numeric -disable network | grep -q 'vendor: .* \[10DE\]' || return 1 ;;
                amdgpu) available lshw && $SUDO lshw -c display -numeric -disable network | grep -q 'vendor: .* \[1002\]' || return 1 ;;
            esac ;;
        nvidia-smi) available nvidia-smi || return 1 ;;
    esac
}

if check_gpu nvidia-smi; then
    status "NVIDIA GPU installed."
    exit 0
fi

if ! check_gpu lspci nvidia && ! check_gpu lshw nvidia && ! check_gpu lspci amdgpu && ! check_gpu lshw amdgpu; then
    install_success
    warning "No NVIDIA/AMD GPU detected. Ollama will run in CPU-only mode."
    exit 0
fi

if check_gpu lspci amdgpu || check_gpu lshw amdgpu; then
    download_and_extract "$OLLAMA_DOWNLOAD_BASE" "$OLLAMA_INSTALL_DIR" "ollama-linux-${ARCH}-rocm"

    install_success
    status "AMD GPU ready."
    exit 0
fi

CUDA_REPO_ERR_MSG="NVIDIA GPU detected, but your OS and Architecture are not supported by NVIDIA.  Please install the CUDA driver manually https://docs.nvidia.com/cuda/cuda-installation-guide-linux/"
# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#rhel-7-centos-7
# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#rhel-8-rocky-8
# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#rhel-9-rocky-9
# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#fedora
install_cuda_driver_yum() {
    status 'Installing NVIDIA repository...'
    
    case $PACKAGE_MANAGER in
        yum)
            $SUDO $PACKAGE_MANAGER -y install yum-utils
            if curl -I --silent --fail --location "https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo" >/dev/null ; then
                $SUDO $PACKAGE_MANAGER-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo
            else
                error $CUDA_REPO_ERR_MSG
            fi
            ;;
        dnf)
            if curl -I --silent --fail --location "https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo" >/dev/null ; then
                $SUDO $PACKAGE_MANAGER config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo
            else
                error $CUDA_REPO_ERR_MSG
            fi
            ;;
    esac

    case $1 in
        rhel)
            status 'Installing EPEL repository...'
            # EPEL is required for third-party dependencies such as dkms and libvdpau
            $SUDO $PACKAGE_MANAGER -y install https://dl.fedoraproject.org/pub/epel/epel-release-latest-$2.noarch.rpm || true
            ;;
    esac

    status 'Installing CUDA driver...'

    if [ "$1" = 'centos' ] || [ "$1$2" = 'rhel7' ]; then
        $SUDO $PACKAGE_MANAGER -y install nvidia-driver-latest-dkms
    fi

    $SUDO $PACKAGE_MANAGER -y install cuda-drivers
}

# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#ubuntu
# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#debian
install_cuda_driver_apt() {
    status 'Installing NVIDIA repository...'
    if curl -I --silent --fail --location "https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-keyring_1.1-1_all.deb" >/dev/null ; then
        curl -fsSL -o $TEMP_DIR/cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-keyring_1.1-1_all.deb
    else
        error $CUDA_REPO_ERR_MSG
    fi

    case $1 in
        debian)
            status 'Enabling contrib sources...'
            $SUDO sed 's/main/contrib/' < /etc/apt/sources.list | $SUDO tee /etc/apt/sources.list.d/contrib.list > /dev/null
            if [ -f "/etc/apt/sources.list.d/debian.sources" ]; then
                $SUDO sed 's/main/contrib/' < /etc/apt/sources.list.d/debian.sources | $SUDO tee /etc/apt/sources.list.d/contrib.sources > /dev/null
            fi
            ;;
    esac

    status 'Installing CUDA driver...'
    $SUDO dpkg -i $TEMP_DIR/cuda-keyring.deb
    $SUDO apt-get update

    [ -n "$SUDO" ] && SUDO_E="$SUDO -E" || SUDO_E=
    DEBIAN_FRONTEND=noninteractive $SUDO_E apt-get -y install cuda-drivers -q
}

if [ ! -f "/etc/os-release" ]; then
    error "Unknown distribution. Skipping CUDA installation."
fi

. /etc/os-release

OS_NAME=$ID
OS_VERSION=$VERSION_ID

PACKAGE_MANAGER=
for PACKAGE_MANAGER in dnf yum apt-get; do
    if available $PACKAGE_MANAGER; then
        break
    fi
done

if [ -z "$PACKAGE_MANAGER" ]; then
    error "Unknown package manager. Skipping CUDA installation."
fi

if ! check_gpu nvidia-smi || [ -z "$(nvidia-smi | grep -o "CUDA Version: [0-9]*\.[0-9]*")" ]; then
    case $OS_NAME in
        centos|rhel) install_cuda_driver_yum 'rhel' $(echo $OS_VERSION | cut -d '.' -f 1) ;;
        rocky) install_cuda_driver_yum 'rhel' $(echo $OS_VERSION | cut -c1) ;;
        fedora) [ $OS_VERSION -lt '39' ] && install_cuda_driver_yum $OS_NAME $OS_VERSION || install_cuda_driver_yum $OS_NAME '39';;
        amzn) install_cuda_driver_yum 'fedora' '37' ;;
        debian) install_cuda_driver_apt $OS_NAME $OS_VERSION ;;
        ubuntu) install_cuda_driver_apt $OS_NAME $(echo $OS_VERSION | sed 's/\.//') ;;
        *) exit ;;
    esac
fi

if ! lsmod | grep -q nvidia || ! lsmod | grep -q nvidia_uvm; then
    KERNEL_RELEASE="$(uname -r)"
    case $OS_NAME in
        rocky) $SUDO $PACKAGE_MANAGER -y install kernel-devel kernel-headers ;;
        centos|rhel|amzn) $SUDO $PACKAGE_MANAGER -y install kernel-devel-$KERNEL_RELEASE kernel-headers-$KERNEL_RELEASE ;;
        fedora) $SUDO $PACKAGE_MANAGER -y install kernel-devel-$KERNEL_RELEASE ;;
        debian|ubuntu) $SUDO apt-get -y install linux-headers-$KERNEL_RELEASE ;;
        *) exit ;;
    esac

    NVIDIA_CUDA_VERSION=$($SUDO dkms status | awk -F: '/added/ { print $1 }')
    if [ -n "$NVIDIA_CUDA_VERSION" ]; then
        $SUDO dkms install $NVIDIA_CUDA_VERSION
    fi

    if lsmod | grep -q nouveau; then
        status 'Reboot to complete NVIDIA CUDA driver install.'
        exit 0
    fi

    $SUDO modprobe nvidia
    $SUDO modprobe nvidia_uvm
fi

# make sure the NVIDIA modules are loaded on boot with nvidia-persistenced
if available nvidia-persistenced; then
    $SUDO touch /etc/modules-load.d/nvidia.conf
    MODULES="nvidia nvidia-uvm"
    for MODULE in $MODULES; do
        if ! grep -qxF "$MODULE" /etc/modules-load.d/nvidia.conf; then
            echo "$MODULE" | $SUDO tee -a /etc/modules-load.d/nvidia.conf > /dev/null
        fi
    done
fi

status "NVIDIA GPU ready."
install_success
}

main "$@"

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐