#!/bin/bash
# autopkgtest podman backend with support for NVIDIA GPUs in rootless containers
#
# This is just a thin wrapper around autopkgtest-virt-podman. It adds all the
# arguments necessary for using NVIDIA GPUs in the container.
#
# Author: Christian Kastner <ckk@kvr.at>
# License: MIT

function usage() {
    cat >&2 <<-EOF

	autopkgtest podman backend with support for NVIDIA GPUs

	This is a simple wrapper around autopkgtest-virt-podman that sets up the
	correct options to pass in devices and map necessary GIDs. All options on
	the command line are passed on directly to autopkgtest-virt-podman, so the
	reader is referred to its man page autopkgtest-virt-podman(1).

	Synopsis:
	  $0 -h

	  $0 autopkgtest [...] -- podman+cuda [options] image [-- extra podman-run args...]

	Examples:

	  # Create an image first

	  \$ gpuisol-podman-create -m http://10.1.2.3:9999/debian cuda cuda/debian:unstable

	  # Run autopkgtests for src:foo, using packages from the Archive

	  \$ autopkgtest -B foo -- podman+cuda cuda/debian:unstable

	EOF
}

[ "${1:-}" = "-h" ] && usage && exit 0

userNAME="$(whoami)"
renderGID="$(getent group render | cut -d: -f3)"
# By policy
videoGID=44

# Sanity checks
if [ -z "$renderGID" ]; then
    cat >&2 <<-EOF
	Group 'render' does not exist on this system. Are you sure that you are on
	the right system? This group should have been autmatically created by the
	udev package."
	EOF
    exit 1
elif ! groups "$userNAME" | grep -q '\brender\b'; then
    echo "'$userNAME' is not in group 'render'." >&2
    exit 1
elif ! groups "$userNAME" | grep -q '\bvideo\b'; then
    echo "'$userNAME' is not in group 'video'." >&2
    exit 1
elif [ "$(cat /proc/sys/kernel/unprivileged_userns_clone)" != "1" ]; then
    echo "unprivileged_userns_clone not enabled." >&2
    exit 1
elif ! [ -c /dev/nvidiactl ]; then
    echo "Device /dev/nvidiactl does not exist - is the 'nvidia' module loaded?" >&2
    exit 1
elif ! grep -q "$userNAME:$renderGID:1" /etc/subgid; then
    echo "No subgid mapping for group 'render'. Run gpuisol-podman-setup" >&2
    exit 1
elif ! grep -q "$userNAME:$videoGID:1" /etc/subgid; then
    echo "No subgid mapping for group 'video'. Run gpuisol-podman-setup" >&2
    exit 1
elif ! grep -q -E "$userNAME:[0-9]{6,}:6553[4-6]" /etc/subgid; then
    echo "No large subgid mapping for '$(whoami)'. Run gpuisol-podman-setup" >&2
    exit 1
fi
[ -c /dev/nvidia-uvm ] || nvidia-modprobe -c 0 -u

maxsubGID="$(sed -nr "s/$userNAME:[0-9]{6,}:(6553[4-6])/\1/p" /etc/subgid)"
maxsubGID=$((maxsubGID - renderGID - 1))
exec autopkgtest-virt-podman "$@" \
    --device=/dev/dri \
    --device=/dev/nvidiactl \
    --device=/dev/nvidia-modeset \
    --device=/dev/nvidia-uvm \
    --device=/dev/nvidia-uvm-tools \
    --device=/dev/nvidia0 \
    --gidmap=0:0:1 \
    --gidmap=44:1:1 \
    --gidmap="$renderGID":2:1 \
    --gidmap=1:3:43 \
    --gidmap=45:46:$((renderGID - videoGID - 1)) \
    --gidmap=$((renderGID + 1)):$((renderGID + 2)):$maxsubGID
