setup.sh

#!/usr/bin/env bash
# shellcheck disable=SC1091

set -e
SCRIPT_DIR=$(readlink -f "$(dirname "$0")")
CLR_GREEN="\033[32m"
CLR_RED="\033[31m"
CLR_RESET="\033[0m"
INSTALL_NVIDIA=true
DOWNLOAD_ARTIFACTS=false

#### Functions ####
print_help() {
    echo "Setup runtime environment for Autoware Open AD Kit"
    echo "Usage: setup.sh [OPTIONS]"
    echo "Options:"
    echo "  --help                  Display this help message"
    echo "  -h                      Display this help message"
    echo "  --no-nvidia             Skip installation of NVIDIA container toolkit"
    echo "  --download-artifacts    Download Autoware artifacts"
    echo ""
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --help|-h)
                print_help
                exit 0
                ;;
            --no-nvidia)
                INSTALL_NVIDIA=false
                shift
                ;;
            --download-artifacts)
                DOWNLOAD_ARTIFACTS=true
                shift
                ;;
            *)
                echo "Unknown option: $1"
                print_help
                exit 1
                ;;
        esac
    done
}

install_nvidia_container_toolkit() {
    echo "Installing NVIDIA Container Toolkit..."

    # Remove any pre-existing NVIDIA container toolkit repo configuration
    sudo rm -f /etc/apt/sources.list.d/nvidia-container-toolkit*.list
    sudo rm -f /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

    # Install prerequisites for GPG key handling
    sudo apt-get update && sudo apt-get install -y --no-install-recommends \
        ca-certificates curl gnupg2

    # Add NVIDIA container toolkit GPG key (dearmored per official NVIDIA docs)
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
        sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

    # Add NVIDIA container toolkit repository using official repo list
    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 > /dev/null

    # Install NVIDIA Container Toolkit
    sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit

    # Add NVIDIA runtime support to docker engine
    sudo nvidia-ctk runtime configure --runtime=docker

    # Restart docker daemon
    sudo systemctl restart docker

    echo -e "${CLR_GREEN}NVIDIA Container Toolkit installed successfully!${CLR_RESET}"
}

install_docker() {
    echo "Installing Docker..."

    if command -v docker &> /dev/null; then
        echo "Docker is already installed. Skipping installation."
        return
    fi

    # Remove old docker packages
    for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do
        sudo apt-get remove -y "$pkg" 2>/dev/null || true
    done

    # Install ca-certificates curl
    sudo apt-get update && sudo apt-get install -y ca-certificates curl
    sudo install -m 0755 -d /etc/apt/keyrings
    sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
    sudo chmod a+r /etc/apt/keyrings/docker.asc

    # Add Docker's official GPG key
    echo \
    "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \
    $(. /etc/os-release && echo "${UBUNTU_CODENAME:-$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 user to docker group
    sudo groupadd docker 2>/dev/null || true
    sudo usermod -aG docker "$USER"

    echo -e "${CLR_GREEN}Docker installed successfully!${CLR_RESET}"
    echo -e "${CLR_GREEN}Please log out and log back in for Docker group changes to take effect.${CLR_RESET}"
}

download_autoware_artifacts() {
    echo "Downloading Autoware artifacts..."

    # Remove apt installed ansible (In Ubuntu 22.04, ansible the version is old)
    sudo apt-get purge ansible

    # Install pipx
    sudo apt-get -y update
    sudo apt-get -y install pipx

    # Add pipx to the system PATH
    python3 -m pipx ensurepath

    # Install ansible
    pipx install --include-deps --force "ansible==6.*"

    # Clone Autoware
    git clone https://github.com/autowarefoundation/autoware.git ~/autoware

    # Get the user home directory from the sudo user
    USER_HOME=$(getent passwd "$SUDO_USER" | cut -d: -f6)

    # Install Autoware artifacts
    cd ~/autoware # The root directory of the cloned repository
    ansible-galaxy collection install -f -r "ansible-galaxy-requirements.yaml"
    ansible-playbook autoware.dev_env.download_artifacts -e "data_dir=$USER_HOME/autoware_data"
    chown -R "$SUDO_USER:$SUDO_USER" "$USER_HOME/autoware_data"

    # Remove the cloned Autoware directory
    rm -rf ~/autoware

    echo -e "${CLR_GREEN}Autoware artifacts downloaded successfully!${CLR_RESET}"
}

#### Main ####
parse_args "$@"

if ! sudo -n true 2>/dev/null; then
    echo -e "${CLR_RED}This script requires sudo privileges. Please run with a user that has sudo access.${CLR_RESET}"
    exit 1
fi

[ "$DOWNLOAD_ARTIFACTS" = true ] && download_autoware_artifacts && exit 0

install_docker
[ "$INSTALL_NVIDIA" = true ] && install_nvidia_container_toolkit

echo -e "${CLR_GREEN}Setup completed!${CLR_RESET}"