Google_Coral_PCIe_multipass.sh

#! /usr/bin/sh

# Google Coral PCIe multi-container passthrough hookscript.
#
# This script automatically configures the host and container to allow the
# PCIe version of a Google Coral TPU to be attached to multiple Proxmox
# containers at the same time.
#
# This hookscript will work for protected and unprotected containers.
#

# You can configure a container to use this hookscript this via pct with:
#
#   pct set <vmid> -hookscript <storage>:snippets/Google_Coral_PCIe_multipass.sh
#
# eg. pct set 1000 -hookscript local:snippets/Google_Coral_PCIe_multipass.sh
#

# Copyright (c) 2024 Curtis Blumer
# License: AGPLv3
# https://github.com/CurtisBlumer/ProxmoxVE-hookscripts/LICENSE
#

friendly_name="Google Coral PCIe multi-container passthrough"

echo "\nExecuting the $friendly_name hookscript ($0)\n";

vmid=$1

phase=$2

dev_dir=dev
dev_name=apex_0
dev_major=120
dev_minor=0

host_dev_dir=/$dev_dir
host_dev_path=$host_dev_dir/$dev_name

vmid_conf_file=/etc/pve/lxc/$vmid.conf

vmid_dev_dir=/var/lib/lxc/$vmid/$dev_dir
vmid_dev_path=$vmid_dev_dir/$dev_name

ct_dev_path=$dev_dir/$dev_name

lxc_cgroup2_allow_dev="lxc.cgroup2.devices.allow: c $dev_major:$dev_minor rwm"
lxc_mount_dev="lxc.mount.entry: $vmid_dev_path $ct_dev_path none bind,optional,create=file,mode=0660"

if [ $phase = 'pre-start' ]; then
    # First phase 'pre-start' will be executed before the guest
    # is started. Exiting with a code != 0 will abort the start

    echo "Starting configuration of container $vmid for $friendly_name";

    echo "Checking configuration file for container $vmid...";

    grep_lxc_cgroup2_allow_dev=$(grep "$lxc_cgroup2_allow_dev" $vmid_conf_file);
    grep_lxc_mount_dev=$(grep "$lxc_mount_dev" $vmid_conf_file);


    if [ "$grep_lxc_cgroup2_allow_dev" = "" ] ||
        [ "$grep_lxc_mount_dev" = "" ]; then

        echo "The configuration file needs updating.\n";

        echo "Attmepting to update the configuration file...";

        if [ "$grep_lxc_cgroup2_allow_dev" = "" ]; then
            echo "$lxc_cgroup2_allow_dev" | tee -a $vmid_conf_file > /dev/null 2>&1;
        fi

        if [ "$grep_lxc_mount_dev" = "" ]; then
            echo "$lxc_mount_dev" | tee -a $vmid_conf_file > /dev/null 2>&1;
        fi

        echo "Rechecking config file for container $vmid...";

        grep_lxc_cgroup2_allow_dev=$(grep "$lxc_cgroup2_allow_dev" $vmid_conf_file);
        grep_lxc_mount_dev=$(grep "$lxc_mount_dev" $vmid_conf_file);

        if [ "$grep_lxc_cgroup2_allow_dev" = "" ] ||
            [ "$grep_lxc_mount_dev" = "" ]; then

            echo "\nYou must add the following lines to $vmid_conf_file\n";

            if
                [ "$grep_lxc_cgroup2_allow_dev" = "" ]; then
                    echo "$lxc_cgroup2_allow_dev";
            fi
            
            if [ "$grep_lxc_mount_dev" = "" ]; then
                echo "$lxc_mount_dev";
            fi

            echo ""
            ErrorExit 5 "Could not automatically update the configation file.";
        else
            echo "\n\nYOU MUST REBOOT THE CONTAINER MANUALLY FOR CHANGES TO TAKE EFFECT.\n\n";
        fi
    fi

    echo "Beginning system preparations on `hostname -f`.";

    echo "Creating special devices folder for Google Coral PCIe passthru on PVE host.";
    mkdir -p $vmid_dev_dir;

    echo "Creating special '$host_dev_path' device at '$vmid_dev_path' on PVE host.";
    mknod -m 660 $vmid_dev_path c $dev_major $dev_minor;
    chown 100000:100000 $vmid_dev_path

elif [ $phase = 'post-start' ]; then
    # Second phase 'post-start' will be executed after the guest
    # successfully started.

    echo "Changing ownership on '$host_dev_path' to 'root:root' in LXC container.";
    pct exec $vmid --keep-env=0 chown root:root $host_dev_path;

    echo "$friendly_name for container $vmid started."

elif [ $phase = 'pre-stop' ]; then
    # Third phase 'pre-stop' will be executed before stopping the guest
    # via the API. Will not be executed if the guest is stopped from
    # within e.g., with a 'poweroff'

    echo "$friendly_name for container $vmid stopping...";

elif [ $phase = 'post-stop' ]; then
    # Last phase 'post-stop' will be executed after the guest stopped.
    # This should even be executed in case the guest crashes or stopped
    # unexpectedly.

    echo "$friendly_name for container $vmid stopped. Doing cleanup.";

    echo "Removing special devices folder at '$vmid_dev_dir' on `hostname -f`";
    rm -R $vmid_dev_dir

else
    ErrorExit 1 "Got unknown phase $phase for $vmid";
fi

ErrorExit()
{
    default_msg="Exiting early."
    code=$1;
    message=$2;

    echo "";
    if [ "$message" != "" ]; then echo "$message"; fi
    echo "Exiting early.\n\n";

    exit $code;
}

exit 0;