223 lines
7.4 KiB
Bash
223 lines
7.4 KiB
Bash
#!/bin/bash
|
||
set -e
|
||
|
||
# 初始化核心变量
|
||
OPERATION=""
|
||
VERSION=""
|
||
# NVIDIA Toolkit核心依赖包（卸载/安装需统一处理）
|
||
NVIDIA_PACKAGES="nvidia-container-toolkit nvidia-container-toolkit-base libnvidia-container-tools libnvidia-container1"
|
||
# 编译构建依赖包（仅安装时需要，卸载不处理）
|
||
BUILD_DEPENDS="build-essential devscripts debhelper fakeroot zlib1g-dev libnuma-dev libhwloc-dev libevent-dev"
|
||
# 提取核心版本号（如从"1.17.6-1"提取"1.17.6"，用于友好输出）
|
||
EXTRACT_VERSION=""
|
||
|
||
# 日志输出函数（统一风格）
|
||
print_step() {
|
||
echo -e "\033[1;34m[*] $1\033[0m"
|
||
}
|
||
print_success() {
|
||
echo -e "\033[1;32m[+] $1\033[0m"
|
||
}
|
||
print_error() {
|
||
echo -e "\033[1;31m[-] 错误：$1\033[0m" >&2
|
||
}
|
||
|
||
# 解析命令行参数（支持--install/--uninstall --version）
|
||
parse_args() {
|
||
while [[ $# -gt 0 ]]; do
|
||
case "$1" in
|
||
--install)
|
||
OPERATION="install"
|
||
shift
|
||
;;
|
||
--uninstall)
|
||
OPERATION="uninstall"
|
||
shift
|
||
;;
|
||
--version)
|
||
if [[ -n "$2" && "$2" != --* ]]; then
|
||
VERSION="$2"
|
||
# 提取核心版本（如1.17.6-1 → 1.17.6）
|
||
EXTRACT_VERSION=$(echo "$VERSION" | awk -F'-' '{print $1}')
|
||
shift 2
|
||
else
|
||
print_error "缺少版本号，请用 --version 指定（如 1.17.6-1）"
|
||
exit 1
|
||
fi
|
||
;;
|
||
*)
|
||
print_error "未知参数：$1"
|
||
exit 1
|
||
;;
|
||
esac
|
||
done
|
||
|
||
# 验证必要参数
|
||
if [[ -z "$OPERATION" ]]; then
|
||
print_error "请指定操作：--install（安装） 或 --uninstall（卸载）"
|
||
exit 1
|
||
fi
|
||
if [[ -z "$VERSION" ]]; then
|
||
print_error "请用 --version 指定NVIDIA Container Toolkit版本"
|
||
exit 1
|
||
fi
|
||
}
|
||
|
||
# 检查Docker是否已安装（Toolkit依赖Docker）
|
||
check_docker_exist() {
|
||
if ! docker --version >/dev/null 2>&1; then
|
||
print_error "未检测到Docker环境！NVIDIA Container Toolkit依赖Docker，请先安装Docker"
|
||
exit 1
|
||
fi
|
||
print_success "Docker环境检测正常"
|
||
}
|
||
|
||
# 添加NVIDIA官方GPG密钥与软件源（静默执行）
|
||
add_nvidia_repo() {
|
||
print_step "正在配置NVIDIA官方软件源"
|
||
# 创建密钥存储目录
|
||
mkdir -p /usr/share/keyrings >/dev/null 2>&1
|
||
|
||
# 1. 导入GPG密钥（自动覆盖现有密钥）
|
||
if ! curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg >/dev/null 2>&1; then
|
||
print_error "GPG密钥下载失败（网络异常或密钥地址失效）"
|
||
exit 1
|
||
fi
|
||
|
||
# 2. 添加软件源（并绑定密钥）
|
||
if ! curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
|
||
sed "s#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g" | \
|
||
tee /etc/apt/sources.list.d/nvidia-container-toolkit.list >/dev/null 2>&1; then
|
||
print_error "NVIDIA软件源配置失败"
|
||
exit 1
|
||
fi
|
||
|
||
# 3. 更新软件源索引
|
||
if ! apt-get update -qq >/dev/null 2>&1; then
|
||
print_error "更新软件源索引失败"
|
||
exit 1
|
||
fi
|
||
print_success "NVIDIA软件源配置完成"
|
||
}
|
||
|
||
# 安装NVIDIA Container Toolkit（指定版本+静默安装）
|
||
install_nvidia_toolkit() {
|
||
print_step "正在安装NVIDIA Container Toolkit $EXTRACT_VERSION（完整版本：$VERSION）"
|
||
|
||
# 构建安装命令：为所有核心包指定版本
|
||
install_cmd="apt-get install -qq -y"
|
||
for pkg in $NVIDIA_PACKAGES; do
|
||
install_cmd="$install_cmd $pkg=$VERSION"
|
||
done
|
||
|
||
# 执行安装（静默，仅捕错）
|
||
if ! eval $install_cmd >/dev/null 2>&1; then
|
||
print_error "Toolkit安装失败（可能版本不兼容或依赖冲突）"
|
||
exit 1
|
||
fi
|
||
|
||
# 安装编译构建依赖库
|
||
print_step "正在安装编译构建依赖库"
|
||
if ! apt-get install -qq -y $BUILD_DEPENDS >/dev/null 2>&1; then
|
||
print_error "编译依赖库（如build-essential）安装失败"
|
||
exit 1
|
||
fi
|
||
|
||
# 配置Docker启用NVIDIA Runtime并重启Docker
|
||
print_step "正在配置Docker支持NVIDIA Runtime"
|
||
if ! nvidia-ctk runtime configure --runtime=docker >/dev/null 2>&1; then
|
||
print_error "Docker NVIDIA Runtime配置失败"
|
||
exit 1
|
||
fi
|
||
if ! systemctl restart docker >/dev/null 2>&1; then
|
||
print_error "重启Docker服务失败（需手动执行：systemctl restart docker）"
|
||
exit 1
|
||
fi
|
||
|
||
# 验证安装结果
|
||
if ! nvidia-container-runtime --version >/dev/null 2>&1; then
|
||
print_error "Toolkit安装后验证失败（nvidia-container-runtime命令不可用）"
|
||
exit 1
|
||
fi
|
||
print_success "NVIDIA Container Toolkit $EXTRACT_VERSION 安装成功！"
|
||
echo -e "\033[1;33m[提示] 可执行 'nvidia-container-runtime --version' 查看详细版本信息\033[0m"
|
||
}
|
||
|
||
# 卸载NVIDIA Container Toolkit（指定版本+清理）
|
||
uninstall_nvidia_toolkit() {
|
||
print_step "正在卸载NVIDIA Container Toolkit $EXTRACT_VERSION（完整版本：$VERSION）"
|
||
|
||
# 停止Docker服务（避免卸载时文件占用）
|
||
if systemctl is-active --quiet docker; then
|
||
print_step "停止Docker服务（避免文件占用）"
|
||
if ! systemctl stop docker >/dev/null 2>&1; then
|
||
print_error "Docker服务停止失败（需手动停止后重试）"
|
||
exit 1
|
||
fi
|
||
fi
|
||
|
||
# 构建卸载命令：为所有核心包指定版本（精准卸载）
|
||
uninstall_cmd="apt-get purge -qq -y"
|
||
for pkg in $NVIDIA_PACKAGES; do
|
||
uninstall_cmd="$uninstall_cmd $pkg=$VERSION"
|
||
done
|
||
|
||
# 执行卸载（静默，仅捕错）
|
||
if ! eval $uninstall_cmd >/dev/null 2>&1; then
|
||
print_error "Toolkit卸载失败（可能包已被删除或依赖冲突）"
|
||
exit 1
|
||
fi
|
||
|
||
# 清理残留文件（软件源、密钥、缓存）
|
||
print_step "清理NVIDIA残留文件"
|
||
rm -f /etc/apt/sources.list.d/nvidia-container-toolkit.list >/dev/null 2>&1
|
||
rm -f /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg >/dev/null 2>&1
|
||
apt-get autoremove -qq -y >/dev/null 2>&1
|
||
apt-get clean -qq >/dev/null 2>&1
|
||
|
||
# 重启Docker恢复默认配置
|
||
if ! systemctl start docker >/dev/null 2>&1; then
|
||
print_error "卸载后重启Docker失败（需手动执行：systemctl start docker）"
|
||
exit 1
|
||
fi
|
||
|
||
print_success "NVIDIA Container Toolkit $EXTRACT_VERSION 卸载成功！"
|
||
}
|
||
|
||
# 主函数（流程控制）
|
||
main() {
|
||
# 1. 检查root权限
|
||
if [[ $EUID -ne 0 ]]; then
|
||
print_error "请用root用户运行脚本（sudo bash nvidia-container-runtime.sh ...）"
|
||
exit 1
|
||
fi
|
||
|
||
# 2. 解析参数
|
||
parse_args "$@"
|
||
|
||
# 3. 安装/卸载分支逻辑
|
||
case "$OPERATION" in
|
||
install)
|
||
# 安装前先检查Docker环境
|
||
check_docker_exist
|
||
# 配置NVIDIA源 → 安装Toolkit → 配置Docker
|
||
add_nvidia_repo
|
||
install_nvidia_toolkit
|
||
;;
|
||
uninstall)
|
||
# 卸载Toolkit → 清理残留 → 重启Docker
|
||
uninstall_nvidia_toolkit
|
||
;;
|
||
*)
|
||
print_error "未知操作：$OPERATION"
|
||
exit 1
|
||
;;
|
||
esac
|
||
|
||
# 正常退出（状态码0）
|
||
exit 0
|
||
}
|
||
|
||
# 启动脚本
|
||
main "$@"
|