#!/bin/bash

# GPU worker setup script - should run as root
# Based on fabric-worker.sh with NVIDIA driver/CUDA and HTCondor GPU support added
export DEBIAN_FRONTEND=noninteractive

apt-get update && apt-get -y upgrade
apt-get install -y linux-headers-$(uname -r)
sudo apt install -y python3-pip
sudo python3 -m pip install --break-system-packages paramiko pandas sage-data-client

apt-get install -y build-essential make zlib1g-dev librrd-dev libpcap-dev autoconf automake libarchive-dev iperf3 htop bmon vim wget pkg-config git python-dev python3-pip libtool
pip install --upgrade pip

##############################
### INSTALL NVIDIA DRIVERS ###
##############################

# Install NVIDIA driver and CUDA toolkit
apt-get install -y ubuntu-drivers-common
ubuntu-drivers autoinstall

# Install CUDA toolkit
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
dpkg -i cuda-keyring_1.1-1_all.deb
rm cuda-keyring_1.1-1_all.deb
apt-get update
apt-get install -y cuda-toolkit

# Install NVIDIA Container Toolkit (for Singularity/Docker GPU support)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    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' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
apt-get update
apt-get install -y nvidia-container-toolkit

# Configure nvidia-container-toolkit for Docker
nvidia-ctk runtime configure --runtime=docker

######################
### INSTALL CONDOR ###
######################

# Remove default single-machine config that conflicts with multi-node pool
rm -f /etc/condor/config.d/00-minicondor
rm -f /etc/condor/config.d/00-security

cat << EOF > /etc/condor/config.d/50-main.config
DAEMON_LIST = MASTER, STARTD

CONDOR_HOST = $3

USE_SHARED_PORT = TRUE

NETWORK_INTERFACE = $1

ENABLE_IPV6 = FALSE

# the nodes have shared filesystem
UID_DOMAIN = \$(CONDOR_HOST)
TRUST_UID_DOMAIN = TRUE
FILESYSTEM_DOMAIN = \$(FULL_HOSTNAME)

#--     Authentication settings
SEC_PASSWORD_FILE = /etc/condor/pool_password
SEC_DEFAULT_AUTHENTICATION = REQUIRED
SEC_DEFAULT_AUTHENTICATION_METHODS = FS,PASSWORD
SEC_READ_AUTHENTICATION = OPTIONAL
SEC_CLIENT_AUTHENTICATION = OPTIONAL
SEC_ENABLE_MATCH_PASSWORD_AUTHENTICATION = TRUE
DENY_WRITE = anonymous@*
DENY_ADMINISTRATOR = anonymous@*
DENY_DAEMON = anonymous@*
DENY_NEGOTIATOR = anonymous@*
DENY_CLIENT = anonymous@*

#--     Privacy settings
SEC_DEFAULT_ENCRYPTION = OPTIONAL
SEC_DEFAULT_INTEGRITY = REQUIRED
SEC_READ_INTEGRITY = OPTIONAL
SEC_CLIENT_INTEGRITY = OPTIONAL
SEC_READ_ENCRYPTION = OPTIONAL
SEC_CLIENT_ENCRYPTION = OPTIONAL

#-- With strong security, do not use IP based controls
ALLOW_WRITE = *
ALLOW_NEGOTIATOR = *

# dynamic slots with GPU resources
SLOT_TYPE_1 = cpus=100%,disk=100%,swap=100%,gpus=100%
SLOT_TYPE_1_PARTITIONABLE = TRUE
NUM_SLOTS = 1
NUM_SLOTS_TYPE_1 = 1

# GPU discovery
use feature : GPUs

# Advertise GPU attributes
STARTD_ATTRS = \$(STARTD_ATTRS) HasGPU
HasGPU = ifThenElse(DetectedGPUs > 0, true, false)

EOF

condor_store_cred -f /etc/condor/pool_password -p p3g@sus

systemctl enable condor
systemctl restart condor

##########################
### INSTALL SINGULARITY ##
##########################

apt-get update
apt-get install -y software-properties-common
add-apt-repository -y ppa:apptainer/ppa
apt-get update
apt-get install -y apptainer

##########################
### INSTALL DOCKER      ##
##########################
cd
apt-get install -y \
    ca-certificates \
    curl \
    gnupg \
    lsb-release

curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg

echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

apt-get update && apt-get install -y docker-ce docker-ce-cli containerd.io

#groupadd docker
usermod -aG docker condor

sudo tee /etc/docker/daemon.json > /dev/null <<'DEOF'
{
  "ipv6": true,
  "fixed-cidr-v6": "2001:db8:1::/64"
}
DEOF

sudo systemctl daemon-reload || true
sudo systemctl restart docker
sudo systemctl enable docker || true

apt-get install -y docker-compose-plugin

############################
### SETUP DEFAULT USER ####
############################
cd
usermod -a -G docker ubuntu

echo 'export GOPATH=${HOME}/go' >> /home/ubuntu/.bashrc
echo 'export PATH=/usr/local/go/bin:${PATH}:${GOPATH}/bin' >> /home/ubuntu/.bashrc
echo 'export PATH=/usr/local/cuda/bin:${PATH}' >> /home/ubuntu/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}' >> /home/ubuntu/.bashrc
