260 lines
12 KiB
Bash
Executable File
260 lines
12 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
#
|
|
# Restrict /dev/nvidia* to the accounts that should have the GPU, and prove it.
|
|
#
|
|
# poc/nvidia-device-access.sh show what is true right now (no root)
|
|
# sudo poc/nvidia-device-access.sh apply install the rule (needs root)
|
|
# sudo poc/nvidia-device-access.sh revert put it back as it was (needs root)
|
|
#
|
|
# Why this exists: the driver creates /dev/nvidia* as 0666 root:root and ships no
|
|
# rule setting MODE or GROUP, so on this host *any* user can open the GPU. The
|
|
# controller runs as `umbra`. Right now the only thing preventing it from loading
|
|
# a model on top of the node's work is the scheduler choosing not to - which is a
|
|
# decision, not a boundary.
|
|
#
|
|
# What this does NOT do, and it is worth being blunt about it. This is not what
|
|
# keeps two node agents off each other's VRAM. Both would run as `umbra-poc`, both
|
|
# would be in `video`, and both would still be able to open the card. On a box
|
|
# with one GPU, placement is the control and the kernel is not; see §7.3 and
|
|
# §13.14 of the specification. What this buys is that the *other* accounts on the
|
|
# box - the controller above all - stop being able to touch it at all.
|
|
#
|
|
# The two steps must happen together. Adding the rule without the group membership
|
|
# makes `nvidia-smi` fail for the agent, and the agent's response to that is to
|
|
# report a machine with no GPU rather than to fail: a silent downgrade from "a
|
|
# node with a GPU" to "a node without one", which looks like a scheduling problem
|
|
# and is not. So `apply` does both, then checks both, and puts both back if
|
|
# either fails.
|
|
|
|
set -euo pipefail
|
|
|
|
AGENT_USER="${UMBRA_AGENT_USER:-umbra-poc}"
|
|
GPU_USER="${UMBRA_GPU_USER:-ollama}"
|
|
# The account that runs the controller, and which this change is meant to lock
|
|
# out. Named rather than inferred from `id -un`, because this script runs as root
|
|
# and root can open the card whatever the rule says.
|
|
CONTROLLER_USER="${UMBRA_CONTROLLER_USER:-umbra}"
|
|
RULE_SRC="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/72-umbra-nvidia.rules"
|
|
RULE_DST="/etc/udev/rules.d/72-umbra-nvidia.rules"
|
|
HELPER_SRC="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/nvidia-restrict-permissions"
|
|
HELPER_DST="/usr/lib/udev/nvidia-restrict-permissions"
|
|
UNIT_SRC="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/umbra-nvidia-devices.service"
|
|
UNIT_DST="/etc/systemd/system/umbra-nvidia-devices.service"
|
|
STATE_FILE="/var/lib/umbra-nvidia-access.state"
|
|
NODES=(/dev/nvidiactl /dev/nvidia0 /dev/nvidia-uvm /dev/nvidia-uvm-tools /dev/nvidia-modeset)
|
|
|
|
bold() { printf '\033[1m%s\033[0m\n' "$*"; }
|
|
head_() { printf '\n\033[1m%s\033[0m\n' "$*"; }
|
|
note() { printf ' %s\n' "$*"; }
|
|
good() { printf ' \033[32mok\033[0m %s\n' "$*"; }
|
|
bad() { printf ' \033[31mFAIL\033[0m %s\n' "$*"; }
|
|
warn() { printf ' \033[33mnote\033[0m %s\n' "$*"; }
|
|
|
|
in_group() {
|
|
id -nG "$1" 2>/dev/null | tr ' ' '\n' | grep -qx "$2"
|
|
}
|
|
|
|
# Does this account still see the GPU? Not "can it read the node" - it can read
|
|
# the node as long as the mode is 0666, which is the state we are leaving. This
|
|
# asks the only question that matters, which is whether a *measurement* works,
|
|
# because that is what the agent actually needs.
|
|
# Prints one of: yes, no, or "cannot check". The third case matters: a cached
|
|
# sudo credential expires, `sudo -n` then fails, and a two-valued check reports
|
|
# "cannot measure the GPU" - a finding about the machine, when the truth is that
|
|
# this tool could not ask. That was read as a real result during the 2026-10-01
|
|
# reboot check, and it would have sent anyone reading the output to debug a
|
|
# permission problem that did not exist.
|
|
sees_gpu() {
|
|
local user="$1"
|
|
if [[ "$user" == "$(id -un)" ]]; then
|
|
nvidia-smi -L >/dev/null 2>&1 && { printf yes; return 0; }
|
|
printf no
|
|
return 1
|
|
fi
|
|
sudo -n -u "$user" -- true >/dev/null 2>&1 || { printf "cannot check"; return 2; }
|
|
if sudo -n -u "$user" -- nvidia-smi -L >/dev/null 2>&1; then
|
|
printf yes; return 0
|
|
fi
|
|
printf no
|
|
return 1
|
|
}
|
|
|
|
do_show() {
|
|
bold "NVIDIA device access on $(hostname)"
|
|
note "rule: ${RULE_DST}$([[ -f $RULE_DST ]] && echo '' || echo ' (not installed)')"
|
|
note ""
|
|
for node in "${NODES[@]}"; do
|
|
[[ -e "$node" ]] || continue
|
|
printf ' %-26s %s\n' "$node" "$(stat -c '%A %U:%G' "$node")"
|
|
done
|
|
note ""
|
|
note "accounts and their supplementary groups:"
|
|
for user in "$AGENT_USER" "$GPU_USER" "$CONTROLLER_USER"; do
|
|
id "$user" >/dev/null 2>&1 || { printf ' %-14s (no such user)\n' "$user"; continue; }
|
|
printf ' %-14s %s\n' "$user" "$(id -nG "$user" | tr ' ' ',' | sed 's/,$//')"
|
|
done
|
|
note ""
|
|
for user in "$GPU_USER" "$AGENT_USER" "$CONTROLLER_USER"; do
|
|
id "$user" >/dev/null 2>&1 || continue
|
|
case "$(sees_gpu "$user")" in
|
|
yes) good "$user can measure the GPU" ;;
|
|
no) warn "$user CANNOT measure the GPU - if that is unexpected, add the account to the video group" ;;
|
|
*) warn "$user: could not check (no passwordless sudo; the answer is unknown, not 'no')" ;;
|
|
esac
|
|
done
|
|
if [[ -f "$RULE_DST" ]] && command -v udevadm >/dev/null 2>&1; then
|
|
note ""
|
|
note "udevadm test (does the rule match?):"
|
|
udevadm test /sys/class/nvidiactl 2>&1 | grep -i "72-umbra-nvidia" | sed 's/^/ /' | head -3 || true
|
|
fi
|
|
}
|
|
|
|
require_root() {
|
|
[[ "${EUID:-$(id -u)}" -eq 0 ]] || { bad "must run as root: sudo $0 $1"; exit 2; }
|
|
}
|
|
|
|
do_apply() {
|
|
require_root apply
|
|
command -v udevadm >/dev/null 2>&1 || { bad "udevadm is not installed"; exit 2; }
|
|
[[ -r "$RULE_SRC" ]] || { bad "no rule at $RULE_SRC"; exit 2; }
|
|
|
|
bold "Restricting the NVIDIA devices"
|
|
|
|
# Record what we are changing, so `revert` is a fact rather than a hope.
|
|
local was_in_video=no
|
|
in_group "$AGENT_USER" video && was_in_video=yes
|
|
# One line, no trailing backslash. A `\` at end of line is a line
|
|
# continuation; a `\\` is an escaped backslash, so the redirect below is
|
|
# never applied and the printf goes to the terminal instead of the file.
|
|
# That is not a hypothetical: it is what this line was.
|
|
printf 'applied %s; %s was in video: %s\n' "$(date -Is)" "$AGENT_USER" "$was_in_video" > "$STATE_FILE" 2>/dev/null || true
|
|
|
|
note "adding $AGENT_USER to the video group"
|
|
usermod -aG video "$AGENT_USER"
|
|
good "$AGENT_USER is in video"
|
|
|
|
note "installing the permission helper and the rule"
|
|
install -o root -g root -m 0755 "$HELPER_SRC" "$HELPER_DST"
|
|
install -o root -g root -m 0644 "$RULE_SRC" "$RULE_DST"
|
|
udevadm control --reload-rules
|
|
# Fire the event the rule hooks. It cannot be a trigger on /dev/nvidiactl:
|
|
# udev has no database entry for those nodes - the driver created them in the
|
|
# kernel without registering them in sysfs - so `udevadm trigger
|
|
# --name-match=nvidiactl` matches nothing and reports success. The driver's
|
|
# own bind event does reach udev, and by then the nodes exist.
|
|
udevadm trigger --action=add /sys/bus/pci/drivers/nvidia >/dev/null 2>&1 || true
|
|
# Every NVIDIA function on this host, found rather than written down. The address
|
|
# was hardcoded to 0000:01:00.0 in the first version, which is an AMD root
|
|
# complex on this machine - vendor 0x1022, not 0x10de - so the trigger fired on
|
|
# the wrong device, the RUN key never executed, and the permissions looked
|
|
# right until something reset them.
|
|
local function
|
|
for function in /sys/bus/pci/devices/*/; do
|
|
[[ "$(cat "$function/vendor" 2>/dev/null)" == "0x10de" ]] || continue
|
|
udevadm trigger --action=bind "$(readlink -f "$function")" >/dev/null 2>&1 || true
|
|
done
|
|
sleep 1
|
|
good "rule loaded and the nodes re-permissioned"
|
|
|
|
# A second path for the case the udev event does not cover: boot ordering.
|
|
# UNVERIFIED until the next reboot - nothing here can test it without one,
|
|
# and claiming otherwise would be the same class of mistake as the node-matching
|
|
# rule this file replaced.
|
|
install -o root -g root -m 0644 "$UNIT_SRC" "$UNIT_DST" 2>/dev/null && {
|
|
systemctl daemon-reload
|
|
systemctl enable umbra-nvidia-devices.service >/dev/null 2>&1 || true
|
|
good "boot unit installed and enabled (unverified until the next reboot)"
|
|
} || warn "no boot unit installed; the udev event is the only path"
|
|
|
|
# ---- the checks that decide whether this stays ----------------------
|
|
head_ "Did it work"
|
|
local problem=0
|
|
for node in "${NODES[@]}"; do
|
|
[[ -e "$node" ]] || continue
|
|
local mode group
|
|
mode="$(stat -c '%A' "$node")"; group="$(stat -c '%G' "$node")"
|
|
if [[ "$group" == "video" ]]; then
|
|
good "$(printf '%-22s %s %s' "$node" "$mode" "$group")"
|
|
else
|
|
bad "$(printf '%-22s %s %s - still not video' "$node" "$mode" "$group")"
|
|
problem=1
|
|
fi
|
|
done
|
|
for user in "$GPU_USER" "$AGENT_USER"; do
|
|
id "$user" >/dev/null 2>&1 || continue
|
|
if sees_gpu "$user"; then
|
|
good "$user can still measure the GPU"
|
|
else
|
|
bad "$user can no longer measure the GPU - this change is too tight"
|
|
problem=1
|
|
fi
|
|
done
|
|
# The account this change is *for* losing access. Tested by name, not as
|
|
# "whoever is running this script": this script runs as root, and root
|
|
# bypasses file permissions entirely, so asking root whether it can open the
|
|
# card answers yes no matter what the rule says. The first version of this
|
|
# check did exactly that and reported a working rule as a failure.
|
|
if [[ "$CONTROLLER_USER" != "$GPU_USER" && "$CONTROLLER_USER" != "$AGENT_USER" ]] &&
|
|
id "$CONTROLLER_USER" >/dev/null 2>&1; then
|
|
if sudo -n -u "$CONTROLLER_USER" -- python3 -c "open('/dev/nvidia0','rb')" >/dev/null 2>&1; then
|
|
bad "$CONTROLLER_USER can still open /dev/nvidia0 - the rule did not take"
|
|
problem=1
|
|
else
|
|
good "$CONTROLLER_USER (the controller) can no longer open the card"
|
|
fi
|
|
else
|
|
warn "no distinct controller account named; cannot check that anyone lost access"
|
|
fi
|
|
if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet ollama; then
|
|
if curl -fsS --max-time 5 http://127.0.0.1:11434/api/tags >/dev/null 2>&1; then
|
|
good "ollama is still serving models on 11434"
|
|
else
|
|
bad "ollama is running but not answering on 11434"
|
|
problem=1
|
|
fi
|
|
fi
|
|
|
|
if (( problem )); then
|
|
printf '\n'
|
|
bad "not leaving this half-applied; reverting"
|
|
do_revert
|
|
return 1
|
|
fi
|
|
printf '\n'
|
|
good "the GPU is now reachable only by: $(getent group video | cut -d: -f4 | tr ',' ' ')"
|
|
note "not by $CONTROLLER_USER, who runs the controller"
|
|
note ""
|
|
note "This restricts accounts on *this* host. It says nothing about the other"
|
|
note "three: each has its own device nodes, and the shared-memory machines do"
|
|
note "not have /dev/nvidia* at all."
|
|
}
|
|
|
|
do_revert() {
|
|
[[ "${EUID:-$(id -u)}" -eq 0 ]] || { bad "must run as root"; exit 2; }
|
|
bold "Reverting to 0666 for every user"
|
|
rm -f "$RULE_DST" "$HELPER_DST" "$UNIT_DST"
|
|
systemctl daemon-reload 2>/dev/null || true
|
|
systemctl disable umbra-nvidia-devices.service >/dev/null 2>&1 || true
|
|
udevadm control --reload-rules
|
|
for node in "${NODES[@]}"; do
|
|
[[ -e "$node" ]] || continue
|
|
chmod 0666 "$node"
|
|
done
|
|
if [[ -f "$STATE_FILE" ]] && grep -q "was in video: no" "$STATE_FILE"; then
|
|
gpasswd -d "$AGENT_USER" video >/dev/null 2>&1 || true
|
|
note "removed $AGENT_USER from video"
|
|
fi
|
|
rm -f "$STATE_FILE"
|
|
good "rule removed and the nodes are 0666 again"
|
|
}
|
|
|
|
case "${1:-show}" in
|
|
show) do_show ;;
|
|
apply) do_apply ;;
|
|
revert) do_revert ;;
|
|
*) echo "usage: $0 {show|apply|revert}" >&2; exit 2 ;;
|
|
esac
|