ansible-devops/scripts/cuda.sh

#!/bin/bash
set -euo pipefail

# 全局变量
CUDA_VERSION="12.6"
DRIVER_VERSION="560.35.05"
PACKAGE_NAME="cuda_${CUDA_VERSION}_${DRIVER_VERSION}_linux.run"
PACKAGE_PATH="/opt/${PACKAGE_NAME}"
INTERNAL_BASE_URL="http://10.101.0.51:5588/cuda-linux"
OFFICIAL_BASE_URL="https://developer.download.nvidia.com/compute/cuda/${CUDA_VERSION}/local_installers"
TEMP_DIR="/tmp/cuda_temp"
CUDA_MAJOR_MINOR=$(echo "$CUDA_VERSION" | cut -d. -f1-2)
CUDA_INSTALL_DIR="/usr/local/cuda-${CUDA_MAJOR_MINOR}"
ENV_PROFILE="/etc/profile"
LOG_FILE="/var/log/cuda_manager_$(date +%Y%m%d%H%M%S).log"

# 颜色定义
GREEN='\033[1;32m'
RED='\033[1;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # 重置颜色

# 日志函数
log() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] $*"
    echo "$msg" >> "$LOG_FILE"
}

# 步骤提示
step() {
    local msg="==> $1"
    echo -e "${GREEN}$msg${NC}"
    log "$msg"
}

# 警告提示
warning() {
    local msg="警告: $1"
    echo -e "${YELLOW}$msg${NC}"
    log "$msg"
}

# 错误提示
error() {
    local msg="错误: $1"
    echo -e "${RED}$msg${NC}"
    log "$msg"
    exit 1
}

# 执行命令并记录日志
run_cmd() {
    step "执行: $1"
    eval "$1" &>> "$LOG_FILE" || {
        error "命令执行失败: $1"
    }
}

# 检查命令是否存在
check_cmd() {
    command -v "$1" &>/dev/null || error "未找到命令: $1"
}

# 参数解析
ACTION=""
FORCE=0

while [[ $# -gt 0 ]]; do
    case "$1" in
        --install) ACTION="install"; shift ;;
        --uninstall) ACTION="uninstall"; shift ;;
        --version)
            CUSTOM_VERSION="$2"
            if [[ $CUSTOM_VERSION =~ ^([0-9]+\.[0-9]+(\.[0-9]+)?)(_([0-9]+\.[0-9]+\.[0-9]+))$ ]]; then
                CUDA_VERSION="${BASH_REMATCH[1]}"
                DRIVER_VERSION="${BASH_REMATCH[4]}"
                PACKAGE_NAME="cuda_${CUDA_VERSION}_${DRIVER_VERSION}_linux.run"
                PACKAGE_PATH="/opt/${PACKAGE_NAME}"
                CUDA_MAJOR_MINOR=$(echo "$CUDA_VERSION" | cut -d. -f1-2)
                CUDA_INSTALL_DIR="/usr/local/cuda-${CUDA_MAJOR_MINOR}"
                OFFICIAL_BASE_URL="https://developer.download.nvidia.com/compute/cuda/${CUDA_VERSION}/local_installers"
            else
                error "版本格式错误，应为 x.x.x_y.y.y 或 x.x_y.y.y"
            fi
            shift 2 ;;
        --force) FORCE=1; shift ;;
        *) error "未知参数: $1" ;;
    esac
done

[[ -z "$ACTION" ]] && error "必须指定 --install 或 --uninstall"


# 下载安装包
download_package() {
    step "检查安装包: $PACKAGE_PATH"
    if [[ -f "$PACKAGE_PATH" ]]; then
        step "使用本地安装包"
        return 0
    fi

    step "本地包不存在，开始下载"
    mkdir -p "$(dirname "$PACKAGE_PATH")"
    
    local urls=(
        "${INTERNAL_BASE_URL}/${PACKAGE_NAME}"
        "${OFFICIAL_BASE_URL}/${PACKAGE_NAME}"
    )
    
    for url in "${urls[@]}"; do
        step "尝试从 $url 下载"
        if wget -q -O "$PACKAGE_PATH" "$url"; then
            step "下载成功，文件大小: $(du -h "$PACKAGE_PATH" | cut -f1)"
            return 0
        else
            warning "从 $url 下载失败"
        fi
    done
    
    error "无法下载安装包，请检查网络或手动放置到 $PACKAGE_PATH"
}


# 安装 CUDA
install_cuda() {
    step "开始安装 CUDA ${CUDA_VERSION}"
    
    # 下载安装包
    download_package
    
    # 创建临时目录
    mkdir -p "$TEMP_DIR"
    trap 'rm -rf "$TEMP_DIR"' EXIT

    # 检查安装包参数
    step "检查安装包支持的参数"
    sh "$PACKAGE_PATH" --help > "$TEMP_DIR/help.txt" 2>&1
    
    # 检查是否支持 --toolkit 参数
    if grep -q -- '--toolkit' "$TEMP_DIR/help.txt"; then
        step "安装包支持 --toolkit 参数"
        run_cmd "sh $PACKAGE_PATH --silent --toolkit"
    else
        warning "安装包不支持 --toolkit 参数，尝试完整安装"
        run_cmd "sh $PACKAGE_PATH --silent"
    fi

    # 配置环境变量
    step "配置 CUDA 环境变量"
    if ! grep -q "cuda-${CUDA_MAJOR_MINOR}" "$ENV_PROFILE"; then
        cat >> "$ENV_PROFILE" << EOF

# CUDA ${CUDA_VERSION}
export PATH=${CUDA_INSTALL_DIR}/bin:\$PATH
export LD_LIBRARY_PATH=${CUDA_INSTALL_DIR}/lib64:\$LD_LIBRARY_PATH
EOF
        step "已添加环境变量到 $ENV_PROFILE"
    else
        step "环境变量已存在，跳过添加"
    fi
    
    # 生效环境变量
    export LC_BYOBU=0 && source "/etc/profile" #临时解决
    
    # 验证安装
    step "验证 CUDA 安装"
    if command -v nvcc &>/dev/null; then
        nvcc_version=$(nvcc -V | grep release | awk '{print $5}' | tr -d ',')
        if [[ "$nvcc_version" == *"${CUDA_MAJOR_MINOR}"* ]]; then
            step "CUDA ${CUDA_VERSION} 安装成功"
        else
            error "CUDA 版本不匹配，期望 ${CUDA_VERSION}，实际 $nvcc_version"
        fi
    else
        error "nvcc 命令未找到，安装失败"
    fi
}


# 卸载 CUDA
uninstall_cuda() {
    step "开始卸载 CUDA ${CUDA_VERSION}"
    
    # 创建临时目录
    mkdir -p "$TEMP_DIR"
    trap 'rm -rf "$TEMP_DIR"' EXIT

    # 检查官方卸载脚本
    OFFICIAL_UNINSTALLER="${CUDA_INSTALL_DIR}/bin/cuda-uninstaller"
    if [[ -x "$OFFICIAL_UNINSTALLER" ]]; then
        step "找到官方卸载脚本: $OFFICIAL_UNINSTALLER"
        
        # 执行官方卸载脚本
        step "执行官方卸载程序"
        run_cmd "$OFFICIAL_UNINSTALLER --silent"
    else
        warning "未找到官方卸载脚本，尝试其他方法"
        
        # 检查安装包是否存在
        if [[ -f "$PACKAGE_PATH" ]]; then
            step "找到安装包: $PACKAGE_PATH"
        else
            if [[ $FORCE -eq 1 ]]; then
                warning "未找到安装包，继续强制卸载"
            else
                step "未找到安装包，开始下载"
                download_package
            fi
        fi

        # 检查安装包是否支持 --uninstall 参数
        step "检查安装包是否支持 --uninstall 参数"
        sh "$PACKAGE_PATH" --help > "$TEMP_DIR/help.txt" 2>&1
        
        if grep -q -- '--uninstall' "$TEMP_DIR/help.txt"; then
            step "安装包支持 --uninstall 参数"
            run_cmd "sh $PACKAGE_PATH --silent --uninstall"
        else
            step "安装包不支持 --uninstall 参数，尝试解压查找卸载脚本"
            
            # 解压安装包
            step "解压安装包到 $TEMP_DIR"
            run_cmd "sh $PACKAGE_PATH --extract=$TEMP_DIR"
            
            # 查找卸载脚本
            UNINSTALL_SCRIPT=$(find "$TEMP_DIR" -name "uninstall_cuda*" -type f | head -n1)
            
            if [[ -n "$UNINSTALL_SCRIPT" ]]; then
                step "找到卸载脚本: $UNINSTALL_SCRIPT"
                run_cmd "sh $UNINSTALL_SCRIPT"
            else
                warning "未找到卸载脚本，继续手动清理"
            fi
        fi
    fi

    # 清理环境变量
    step "清理环境变量"
    if grep -q "cuda-${CUDA_MAJOR_MINOR}" "$ENV_PROFILE"; then
        run_cmd "sed -i '/cuda-${CUDA_MAJOR_MINOR}/d' $ENV_PROFILE"
        step "已从 $ENV_PROFILE 移除 CUDA 环境变量"
    else
        step "环境变量已清理"
    fi
    
    # 清理安装目录
    if [[ -d "$CUDA_INSTALL_DIR" ]]; then
        step "删除安装目录: $CUDA_INSTALL_DIR"
        run_cmd "rm -rf $CUDA_INSTALL_DIR"
    else
        step "安装目录不存在，跳过删除"
    fi
    
    # 清理残留文件
    step "清理残留文件"
    find / -name "*cuda-${CUDA_MAJOR_MINOR}*" 2>/dev/null | while read -r file; do
        if [[ -e "$file" ]]; then
            step "删除残留文件: $file"
            rm -rf "$file" 2>/dev/null || warning "无法删除: $file"
        fi
    done
    
    step "CUDA ${CUDA_VERSION} 卸载完成"
}



step "开始 CUDA ${CUDA_VERSION} ${ACTION} 流程"
case "$ACTION" in
    install) install_cuda ;;
    uninstall) uninstall_cuda ;;
    *) error "未知操作: $ACTION" ;;
esac

step "CUDA ${CUDA_VERSION} ${ACTION} 流程完成"