# Description:
# This script will help the user to setup rootless docker[1] on CMLab Server.
# We follow our CMLab's rootless docker guideline[2] to build the environment:
# 
# We will add the user's username into subuid & subgid so that the docker service
# can start the containers. Then we run the script provided by docker to install
# rootless docker's service.
# 
# After checking whether the docker service can start, we should allow Docker
# can access NVIDIA Container Runtime so that our Docker container can use GPUs.
#
# Then, we should change the directory of Docker to /tmp2 so that Docker can save
# images.
#
# [1] rootless docker's documentation: https://docs.docker.com/engine/security/rootless/
# [2] CMLab's rootless docker guideline: https://hackmd.io/@snsd0805/SkWMyB1ER

echo -e -n "\033[44m[CML INFO]\033[0m "
echo "Start setting up Docker service"

# insert the username into /etc/subuid & /etc/subuid
# setup-uid's source code is on https://gitea.snsd0805.com/snsd0805/CMLAB-Docker-Builder/src/branch/master/tool/setup_uid.c
setup-uid

if [ $? -eq 0 ]; then
	echo -e -n "\033[44m[CML INFO]\033[0m "
	echo "Inserted $(whoami) into /etc/subuid & /etc/subgid"
else
	echo -e -n "\033[41m[CML ERROR]\033[0m "
	echo "failed to set subuid & subgid. Please contact us: unix_manager@cmlab.csie.ntu.edu.tw"
	exit 1
fi


# Run install script provided by docker
curl -fsSL https://get.docker.com/rootless | DOCKER_BIN=/tmp2/$(whoami)/docker/bin sh
if [ $? -eq 0 ]; then
	echo -e -n "\033[44m[CML INFO]\033[0m "
	echo "Docker installed."
else
	echo -e -n "\033[41m[CML ERROR]\033[0m "
	echo "failed to install rootless docker. Please contact us: unix_manager@cmlab.csie.ntu.edu.tw"
	exit 1
fi

echo "export PATH=/tmp2/$(whoami)/docker/bin:\$PATH" >> ~/.bashrc
echo "export DOCKER_HOST=unix:///run/user/$(id -u)/docker.sock" >> ~/.bashrc
echo -e -n "\033[44m[CML INFO]\033[0m "
echo "We have some environment variables in your ~/.bashrc"

# Link NVIDIA Container Runtime to Docker & change docker image's path
mkdir /tmp2/$(whoami)/docker-data
cat << EOF > $HOME/.config/docker/daemon.json
{
    "runtimes": {
        "nvidia": {
            "args": [],
            "path": "nvidia-container-runtime"
        }
    },
    "data-root": "/tmp2/$(whoami)/docker-data"
}
EOF
if [ $? -eq 0 ]; then
	echo -e -n "\033[44m[CML INFO]\033[0m "
	echo "Linked NVIDIA Container Runtime & changed the data path."
else
	echo -e -n "\033[41m[CML ERROR]\033[0m "
	echo "failed to link NVIDIA Container Runtime to Docker. Please contact us: unix_manager@cmlab.csie.ntu.edu.tw"
	exit 1
fi

# restart docker
echo -e -n "\033[44m[CML INFO]\033[0m "
echo "Restarting Docker..."

systemctl --user restart docker && docker pull ubuntu && docker run --rm --runtime=nvidia --gpus '"device=0,1"' ubuntu:22.04 nvidia-smi

echo -e -n "\033[44m[CML INFO]\033[0m "
echo "These are 2 GPUs on this server."