#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT

# 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

log() {
    echo "dgx-desktop-docker-gpus: $*"
}

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

    MINOR_NUM_3D_GPU=$(mktemp)
    MINOR_NUM_ALL_GPU=$(mktemp)

    # 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
        log "No 3D controller class NVIDIA GPU found."
        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

    log "Successfully configured nvidia docker to only use Compute GPUs"
    return 0
}

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

create_containerd_override
