#!/bin/bash

# vLLM Docker一键安装脚本
# 功能:自动安装Docker, NVIDIA驱动和容器工具包，并部署vLLM服务

set -e

# 检查是否为root用户
if [ "$(id -u)" -ne 0 ]; then
    echo "请使用root用户或通过sudo运行此脚本"
    exit 1
fi

# 获取GPU数量
read -p "请输入要使用的GPU数量 (例如: 1): " GPU_COUNT
if ! [[ "$GPU_COUNT" =~ ^[1-9][0-9]*$ ]]; then
    echo "错误: 请输入有效的数字"
    exit 1
fi

# 获取模型名称
read -p "请输入要下载的模型名称 (例如: meta-llama/Llama-2-7b-chat-hf): " MODEL_NAME
if [ -z "$MODEL_NAME" ]; then
    echo "错误: 模型名称不能为空"
    exit 1
fi

# 检查NVIDIA显卡
echo "检查NVIDIA显卡..."
if ! lspci | grep -i nvidia > /dev/null; then
    echo "未检测到NVIDIA显卡!"
    exit 1
fi

# 安装基础依赖
echo "安装基础依赖..."
apt-get update
apt-get install -y curl apt-transport-https ca-certificates software-properties-common

# 安装NVIDIA驱动
echo "检查NVIDIA驱动..."
if ! nvidia-smi &> /dev/null; then
    echo "安装NVIDIA驱动..."
    add-apt-repository -y ppa:graphics-drivers/ppa
    apt-get update
    apt-get install -y ubuntu-drivers-common
    ubuntu-drivers autoinstall
    echo "NVIDIA驱动安装完成，需要重启。"
    echo "请重启后再次运行此脚本继续安装。"
    exit 0
fi

# 安装Docker
echo "检查Docker..."
if ! command -v docker &> /dev/null; then
    echo "安装Docker..."
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | apt-key add -
    add-apt-repository -y "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
    apt-get update
    apt-get install -y docker-ce docker-ce-cli containerd.io
    systemctl enable --now docker
fi

# 安装NVIDIA Container Toolkit
echo "检查NVIDIA Container Toolkit..."
if ! dpkg -s nvidia-container-toolkit &> /dev/null; then
    echo "安装NVIDIA Container Toolkit..."
    distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
    curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | apt-key add -
    curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | tee /etc/apt/sources.list.d/nvidia-docker.list
    apt-get update
    apt-get install -y nvidia-container-toolkit
    systemctl restart docker
fi

# 创建vLLM服务
echo "部署vLLM容器..."
docker run -d --gpus all \
       --shm-size=1g \
       --ulimit memlock=-1 \
       --ulimit stack=67108864 \
       -e NCCL_IGNORE_DISABLED_P2P=1 \
       -p 8000:8000 \
       --name vllm_service \
       --restart always \
       -v $HOME/.cache/huggingface:/root/.cache/huggingface \
       ghcr.io/vllm/vllm:latest \
       --model $MODEL_NAME \
       --tensor-parallel-size $GPU_COUNT \
       --host 0.0.0.0 \
       --port 8000

echo "vLLM服务已成功部署!"
echo "访问地址: http://localhost:8000"
echo "您可以使用以下命令查看日志: docker logs -f vllm_service"