#!/bin/bash

# this script is based on NVIDIA's documentation and scripts
# it create a drop-in configuration file for the containerd service
# Something like this:
# $ cat /etc/systemd/system/containerd.service.d/containerd-override.conf
# [Service]
# DeviceAllow=/dev/nvidia0
# DeviceAllow=/dev/nvidia1
# DeviceAllow=char-nvidia-caps
# DeviceAllow=/dev/nvidiactl
# DeviceAllow=/dev/nvidia-modeset
# DeviceAllow=/dev/nvidia-uvm
# DeviceAllow=/dev/nvidia-uvm-tools

set -eu

create_containerd_override()
{
    # If nvidia-smi is not found raise err and exit
    if ! which nvidia-smi > /dev/null 2>&1; then
        echo "nvidia-smi not found. Returning error."
        return 1
    fi

    MINOR_NUM_3D_GPU="/tmp/NV_3D_GPU_MINOR_NUMBERS.txt"
    MINOR_NUM_ALL_GPU="/tmp/NV_ALL_GPU_MINOR_NUMBERS.txt"
    rm -f ${MINOR_NUM_3D_GPU} > /dev/null 2>&1
    rm -f ${MINOR_NUM_ALL_GPU} > /dev/null 2>&1

    # Retrieve all 3D controller class NVIDIA GPUs
    NV_3D_GPU_BDF_LIST=$(lspci | grep "3D controller" | grep -i NVIDIA | cut -d' ' -f1)
    for BDF in ${NV_3D_GPU_BDF_LIST}; do
        MINOR_NUMBER=$(nvidia-smi -q -i "$BDF" | grep "Minor Number" | awk '{print $4}')
        echo "$MINOR_NUMBER" >> ${MINOR_NUM_3D_GPU}
    done

    if [[ -f ${MINOR_NUM_3D_GPU} && ! -s ${MINOR_NUM_3D_GPU} ]]; then
        rm -rf ${MINOR_NUM_3D_GPU} ${MINOR_NUM_ALL_GPU} || echo "Failed to remove temporary files"
        return 1
    fi

    # Retrive Minor Numbers of all Nvidia GPUs
    nvidia-smi -q | grep Minor | awk '{print $4}' > ${MINOR_NUM_ALL_GPU}

    # Get Minor Numbers of the Nvidia GPU which are NOT of Class 3D controller
    OUT=$(comm -13  <(sort ${MINOR_NUM_3D_GPU}) <(sort ${MINOR_NUM_ALL_GPU}))

    CONF_DIR="/etc/systemd/system/containerd.service.d"
    CONF_FILE="$CONF_DIR/containerd-override.conf"

    # Remove if there exist any containerd-override.conf
    rm -f $CONF_FILE > /dev/null 2>&1
    mkdir -p $CONF_DIR > /dev/null 2>&1
    touch $CONF_FILE
    echo "[Service]" > $CONF_FILE

    # Add all /dev/nvidia* as DeviceAllow
    DEV_NVIDIA=$(find /dev/ -name \* -type c | grep nvidia | grep -v nvidia-caps | sort)
    for DEV in ${DEV_NVIDIA}; do
        sed -i "$ aDeviceAllow=${DEV}" $CONF_FILE
    done

    # Add all nvidia-caps devices
    sed -i "$ aDeviceAllow=char-nvidia-caps" $CONF_FILE

    # Remove GPUs which are NOT of Class '3D Controller'
    for NUMBER in ${OUT}; do
        sed -i "/\bnvidia${NUMBER}\b/ d" $CONF_FILE
    done

    echo "Successfully configured nvidia docker to only use Compute GPUs"

    systemctl daemon-reload

    rm -rf ${MINOR_NUM_3D_GPU} ${MINOR_NUM_ALL_GPU} || echo "Failed to remove temporary files"
    return 0
}

### MAIN ###
if [[ $EUID -ne 0 ]]; then
    echo "Requires sudo to run"
    exit 1
fi

create_containerd_override
