316 lines
9.9 KiB
Bash
316 lines
9.9 KiB
Bash
#!/bin/bash
|
|
# ==============================================================================
|
|
# WSL GPU Setup Script for Local LLM Agent
|
|
# ==============================================================================
|
|
# This script configures NVIDIA GPU support in WSL Ubuntu-24.04 for:
|
|
# - Docker GPU acceleration
|
|
# - vLLM backend
|
|
# - CUDA toolkit
|
|
#
|
|
# Prerequisites:
|
|
# - Windows 11 with WSL2
|
|
# - NVIDIA GPU with recent drivers (>= 525.xx)
|
|
# - Ubuntu-24.04 WSL distribution
|
|
#
|
|
# Usage:
|
|
# ./scripts/setup-wsl-gpu.sh
|
|
#
|
|
# Author: ISEM Development Team
|
|
# Version: 1.0.0
|
|
# ==============================================================================
|
|
|
|
set -euo pipefail
|
|
|
|
# Colors for output
|
|
RED='\033[0;31m'
|
|
GREEN='\033[0;32m'
|
|
YELLOW='\033[1;33m'
|
|
BLUE='\033[0;34m'
|
|
NC='\033[0m' # No Color
|
|
|
|
# Logging functions
|
|
log_info() {
|
|
echo -e "${BLUE}[INFO]${NC} $1"
|
|
}
|
|
|
|
log_success() {
|
|
echo -e "${GREEN}[SUCCESS]${NC} $1"
|
|
}
|
|
|
|
log_warning() {
|
|
echo -e "${YELLOW}[WARNING]${NC} $1"
|
|
}
|
|
|
|
log_error() {
|
|
echo -e "${RED}[ERROR]${NC} $1"
|
|
}
|
|
|
|
# Check if running in WSL
|
|
check_wsl() {
|
|
log_info "Checking WSL environment..."
|
|
|
|
if [[ ! -f /proc/version ]] || ! grep -qi microsoft /proc/version; then
|
|
log_error "This script must be run inside WSL (Windows Subsystem for Linux)"
|
|
exit 1
|
|
fi
|
|
|
|
log_success "Running in WSL environment"
|
|
}
|
|
|
|
# Check Ubuntu version
|
|
check_ubuntu_version() {
|
|
log_info "Checking Ubuntu version..."
|
|
|
|
if [[ -f /etc/os-release ]]; then
|
|
. /etc/os-release
|
|
if [[ "$ID" == "ubuntu" ]]; then
|
|
log_success "Ubuntu $VERSION_ID detected"
|
|
else
|
|
log_warning "Expected Ubuntu, found $ID. Proceeding anyway..."
|
|
fi
|
|
else
|
|
log_warning "Could not detect OS version"
|
|
fi
|
|
}
|
|
|
|
# Check for existing NVIDIA driver (should be provided by Windows)
|
|
check_nvidia_driver() {
|
|
log_info "Checking NVIDIA driver..."
|
|
|
|
if command -v nvidia-smi &> /dev/null; then
|
|
DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1)
|
|
log_success "NVIDIA driver version: $DRIVER_VERSION"
|
|
|
|
# Show GPU info
|
|
log_info "GPU Information:"
|
|
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
|
|
else
|
|
log_error "NVIDIA driver not detected!"
|
|
log_info "In WSL2, the NVIDIA driver is provided by Windows."
|
|
log_info "Please ensure:"
|
|
log_info " 1. NVIDIA Game Ready or Studio driver >= 525.xx is installed on Windows"
|
|
log_info " 2. WSL is updated: wsl --update"
|
|
log_info " 3. Restart WSL: wsl --shutdown"
|
|
exit 1
|
|
fi
|
|
}
|
|
|
|
# Install CUDA Toolkit
|
|
install_cuda_toolkit() {
|
|
log_info "Installing CUDA Toolkit 12.6..."
|
|
|
|
# Check if already installed
|
|
if command -v nvcc &> /dev/null; then
|
|
CUDA_VERSION=$(nvcc --version | grep release | awk '{print $6}' | cut -c2-)
|
|
log_info "CUDA Toolkit $CUDA_VERSION already installed"
|
|
|
|
read -p "Do you want to reinstall? (y/N) " -n 1 -r
|
|
echo
|
|
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
|
|
log_info "Skipping CUDA Toolkit installation"
|
|
return
|
|
fi
|
|
fi
|
|
|
|
# Add NVIDIA CUDA repository
|
|
log_info "Adding NVIDIA CUDA repository..."
|
|
|
|
# Download and install keyring
|
|
wget -q https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb
|
|
sudo dpkg -i cuda-keyring_1.1-1_all.deb
|
|
rm cuda-keyring_1.1-1_all.deb
|
|
|
|
# Update package list
|
|
sudo apt-get update
|
|
|
|
# Install CUDA Toolkit 12.6
|
|
log_info "Installing CUDA Toolkit 12.6 (this may take a while)..."
|
|
sudo apt-get install -y cuda-toolkit-12-6
|
|
|
|
# Add CUDA to PATH
|
|
if ! grep -q "cuda-12.6" ~/.bashrc; then
|
|
echo '' >> ~/.bashrc
|
|
echo '# CUDA 12.6' >> ~/.bashrc
|
|
echo 'export PATH=/usr/local/cuda-12.6/bin:$PATH' >> ~/.bashrc
|
|
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
|
|
fi
|
|
|
|
# Export for current session
|
|
export PATH=/usr/local/cuda-12.6/bin:$PATH
|
|
export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64:$LD_LIBRARY_PATH
|
|
|
|
log_success "CUDA Toolkit 12.6 installed successfully"
|
|
}
|
|
|
|
# Install Docker (if not already installed)
|
|
install_docker() {
|
|
log_info "Checking Docker installation..."
|
|
|
|
if command -v docker &> /dev/null; then
|
|
DOCKER_VERSION=$(docker --version | awk '{print $3}' | cut -d',' -f1)
|
|
log_success "Docker $DOCKER_VERSION already installed"
|
|
else
|
|
log_info "Installing Docker..."
|
|
|
|
# Install prerequisites
|
|
sudo apt-get update
|
|
sudo apt-get install -y ca-certificates curl gnupg
|
|
|
|
# Add Docker's official GPG key
|
|
sudo install -m 0755 -d /etc/apt/keyrings
|
|
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
|
|
sudo chmod a+r /etc/apt/keyrings/docker.gpg
|
|
|
|
# Add repository
|
|
echo \
|
|
"deb [arch="$(dpkg --print-architecture)" signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
|
|
"$(. /etc/os-release && echo "$VERSION_CODENAME")" stable" | \
|
|
sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
|
|
|
|
# Install Docker
|
|
sudo apt-get update
|
|
sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
|
|
|
|
# Add current user to docker group
|
|
sudo usermod -aG docker $USER
|
|
|
|
log_success "Docker installed successfully"
|
|
log_warning "You may need to log out and log back in for docker group to take effect"
|
|
fi
|
|
}
|
|
|
|
# Install NVIDIA Container Toolkit
|
|
install_nvidia_container_toolkit() {
|
|
log_info "Installing NVIDIA Container Toolkit..."
|
|
|
|
# Check if already installed
|
|
if command -v nvidia-ctk &> /dev/null; then
|
|
log_info "NVIDIA Container Toolkit already installed"
|
|
|
|
read -p "Do you want to reinstall? (y/N) " -n 1 -r
|
|
echo
|
|
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
|
|
log_info "Skipping NVIDIA Container Toolkit installation"
|
|
return
|
|
fi
|
|
fi
|
|
|
|
# Add NVIDIA Container Toolkit repository
|
|
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
|
|
|
|
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' | \
|
|
sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
|
|
|
|
# Install
|
|
sudo apt-get update
|
|
sudo apt-get install -y nvidia-container-toolkit
|
|
|
|
# Configure Docker to use NVIDIA runtime
|
|
sudo nvidia-ctk runtime configure --runtime=docker
|
|
|
|
# Restart Docker
|
|
sudo systemctl restart docker
|
|
|
|
log_success "NVIDIA Container Toolkit installed and configured"
|
|
}
|
|
|
|
# Verify GPU access in Docker
|
|
verify_docker_gpu() {
|
|
log_info "Verifying GPU access in Docker..."
|
|
|
|
# Test with nvidia-smi in container
|
|
if docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi &> /dev/null; then
|
|
log_success "GPU is accessible from Docker containers!"
|
|
|
|
# Show GPU info from container
|
|
log_info "GPU info from container:"
|
|
docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv
|
|
else
|
|
log_error "GPU is NOT accessible from Docker containers"
|
|
log_info "Try the following:"
|
|
log_info " 1. Restart Docker: sudo systemctl restart docker"
|
|
log_info " 2. Restart WSL: wsl --shutdown (from Windows)"
|
|
log_info " 3. Ensure Windows NVIDIA driver is >= 525.xx"
|
|
exit 1
|
|
fi
|
|
}
|
|
|
|
# Test vLLM container (optional)
|
|
test_vllm() {
|
|
log_info "Testing vLLM container (optional)..."
|
|
|
|
read -p "Do you want to test vLLM container? This will download ~10GB. (y/N) " -n 1 -r
|
|
echo
|
|
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
|
|
log_info "Skipping vLLM test"
|
|
return
|
|
fi
|
|
|
|
log_info "Pulling vLLM image..."
|
|
docker pull vllm/vllm-openai:latest
|
|
|
|
log_info "Testing vLLM startup (will exit after verification)..."
|
|
# Just verify it starts, don't load a model
|
|
timeout 30 docker run --rm --gpus all vllm/vllm-openai:latest --help && \
|
|
log_success "vLLM container works!" || \
|
|
log_warning "vLLM test timed out (this is OK, container works)"
|
|
}
|
|
|
|
# Print summary
|
|
print_summary() {
|
|
echo ""
|
|
echo "============================================================"
|
|
echo -e "${GREEN}WSL GPU Setup Complete!${NC}"
|
|
echo "============================================================"
|
|
echo ""
|
|
echo "Installed components:"
|
|
echo " - CUDA Toolkit 12.6"
|
|
echo " - Docker with GPU support"
|
|
echo " - NVIDIA Container Toolkit"
|
|
echo ""
|
|
echo "Quick verification commands:"
|
|
echo " nvidia-smi # Check GPU from WSL"
|
|
echo " docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi"
|
|
echo ""
|
|
echo "Next steps:"
|
|
echo " 1. Restart your terminal or run: source ~/.bashrc"
|
|
echo " 2. Start the vLLM stack:"
|
|
echo " cd /path/to/local-llm-agent"
|
|
echo " docker-compose -f docker-compose.vllm.yml up -d"
|
|
echo ""
|
|
echo "Documentation: docs/70-onboarding/WSL-GPU-SETUP.md"
|
|
echo "============================================================"
|
|
}
|
|
|
|
# Main execution
|
|
main() {
|
|
echo "============================================================"
|
|
echo "WSL GPU Setup Script for Local LLM Agent"
|
|
echo "============================================================"
|
|
echo ""
|
|
|
|
check_wsl
|
|
check_ubuntu_version
|
|
check_nvidia_driver
|
|
|
|
echo ""
|
|
read -p "Continue with installation? (Y/n) " -n 1 -r
|
|
echo
|
|
if [[ $REPLY =~ ^[Nn]$ ]]; then
|
|
log_info "Installation cancelled"
|
|
exit 0
|
|
fi
|
|
|
|
install_cuda_toolkit
|
|
install_docker
|
|
install_nvidia_container_toolkit
|
|
verify_docker_gpu
|
|
test_vllm
|
|
|
|
print_summary
|
|
}
|
|
|
|
# Run main
|
|
main "$@"
|