ollama更新加速脚本

#!/bin/sh
# This script installs Ollama on Linux.
# It detects the current operating system architecture and installs the appropriate version of Ollama.

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} $*"; }

TEMP_DIR=$(mktemp -d)
cleanup() { rm -rf $TEMP_DIR; }
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
}

[ "$(uname -s)" = "Linux" ] || error 'This script is intended to run on Linux only.'

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

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  2'" ;;
    *) ;;
esac

# GitHub 代理加速
GITHUB_PROXY="${GITHUB_PROXY:-https://gh.xmly.dev/}"

# 获取最新版本号
if [ -z "${OLLAMA_VERSION:-}" ]; then
    status "Getting latest Ollama version..."
    OLLAMA_VERSION=$(curl -s https://api.github.com/repos/ollama/ollama/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/')
    if [ -z "$OLLAMA_VERSION" ]; then
        error "Failed to fetch latest version. Please set OLLAMA_VERSION manually."
    fi
    status "Latest version: $OLLAMA_VERSION"
fi

SUDO=
if [ "$(id -u)" -ne 0 ]; then
    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"

    status "Downloading \({filename} via proxy: \){GITHUB_PROXY}"

    # Check if .tar.zst is available
    if curl --fail --silent --head --location "\({GITHUB_PROXY}\){url_base}/${filename}.tar.zst" >/dev/null 2>&1; then
        if ! available zstd; then
            error "This version requires zstd for extraction. Please install zstd and try again:
  - Debian/Ubuntu: sudo apt-get install zstd
  - RHEL/CentOS/Fedora: sudo dnf install zstd
  - Arch: sudo pacman -S zstd"
        fi

        status "Downloading ${filename}.tar.zst"
        curl --fail --show-error --location --progress-bar \
            "\({GITHUB_PROXY}\){url_base}/${filename}.tar.zst" | \
            zstd -d | \(SUDO tar -xf - -C "\){dest_dir}"
        return 0
    fi

    status "Downloading ${filename}.tgz"
    curl --fail --show-error --location --progress-bar \
        "\({GITHUB_PROXY}\){url_base}/${filename}.tgz" | \
        \(SUDO tar -xzf - -C "\){dest_dir}"
}

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 \(OLLAMA_VERSION 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 "https://github.com/ollama/ollama/releases/download/\({OLLAMA_VERSION}" "\)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 "https://github.com/ollama/ollama/releases/download/\({OLLAMA_VERSION}" "\)OLLAMA_INSTALL_DIR" "ollama-linux-${ARCH}-jetpack6"
    elif grep R35 /etc/nv_tegra_release > /dev/null ; then
        download_and_extract "https://github.com/ollama/ollama/releases/download/\({OLLAMA_VERSION}" "\)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

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"
Environment="OLLAMA_HOST=0.0.0.0:11434"
[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

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

if [ -f /etc/nv_tegra_release ] ; then
    status "NVIDIA JetPack ready."
    install_success
    exit 0
fi

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() {
    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 "https://github.com/ollama/ollama/releases/download/\({OLLAMA_VERSION}" "\)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/"

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...'
            \(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
}

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

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