327 lines
8.2 KiB
Bash
Executable File
327 lines
8.2 KiB
Bash
Executable File
#!/bin/sh
|
|
# bee-nvidia-recover — drain NVIDIA clients, then reset a GPU or reload drivers.
|
|
|
|
set -u
|
|
|
|
log() {
|
|
echo "[bee-nvidia-recover] $*"
|
|
}
|
|
|
|
log_blocker() {
|
|
echo "[bee-nvidia-recover] blocker: $*"
|
|
}
|
|
|
|
usage() {
|
|
cat <<'EOF'
|
|
usage:
|
|
bee-nvidia-recover restart-drivers
|
|
bee-nvidia-recover reset-gpu <index>
|
|
EOF
|
|
}
|
|
|
|
unit_exists() {
|
|
systemctl cat "$1" >/dev/null 2>&1
|
|
}
|
|
|
|
unit_is_active() {
|
|
systemctl is-active --quiet "$1" 2>/dev/null
|
|
}
|
|
|
|
stop_unit_if_active() {
|
|
unit="$1"
|
|
if unit_is_active "$unit"; then
|
|
log "stopping $unit"
|
|
systemctl stop "$unit"
|
|
return 0
|
|
fi
|
|
return 1
|
|
}
|
|
|
|
start_unit_if_marked() {
|
|
unit="$1"
|
|
marker="$2"
|
|
if [ "$marker" = "1" ] && unit_exists "$unit"; then
|
|
log "starting $unit"
|
|
systemctl start "$unit"
|
|
fi
|
|
}
|
|
|
|
wait_for_process_exit() {
|
|
name="$1"
|
|
tries=0
|
|
while pgrep -x "$name" >/dev/null 2>&1; do
|
|
tries=$((tries + 1))
|
|
if [ "$tries" -ge 15 ]; then
|
|
log "WARN: $name is still running after stop request"
|
|
return 1
|
|
fi
|
|
sleep 1
|
|
done
|
|
return 0
|
|
}
|
|
|
|
log_pid_details() {
|
|
pid="$1"
|
|
line=$(ps -p "$pid" -o pid=,comm=,args= 2>/dev/null | sed 's/^[[:space:]]*//')
|
|
if [ -n "$line" ]; then
|
|
log_blocker "$line"
|
|
else
|
|
log_blocker "pid $pid"
|
|
fi
|
|
}
|
|
|
|
collect_gpu_compute_pids() {
|
|
index="$1"
|
|
if ! command -v nvidia-smi >/dev/null 2>&1; then
|
|
return 0
|
|
fi
|
|
nvidia-smi --id="$index" \
|
|
--query-compute-apps=pid \
|
|
--format=csv,noheader,nounits 2>/dev/null \
|
|
| sed 's/^[[:space:]]*//;s/[[:space:]]*$//' \
|
|
| grep -E '^[0-9]+$' || true
|
|
}
|
|
|
|
collect_gpu_device_pids() {
|
|
index="$1"
|
|
dev="/dev/nvidia$index"
|
|
[ -e "$dev" ] || return 0
|
|
if command -v fuser >/dev/null 2>&1; then
|
|
fuser "$dev" 2>/dev/null \
|
|
| tr ' ' '\n' \
|
|
| sed 's/[^0-9].*$//' \
|
|
| grep -E '^[0-9]+$' || true
|
|
fi
|
|
}
|
|
|
|
collect_gpu_holder_pids() {
|
|
index="$1"
|
|
{
|
|
collect_gpu_compute_pids "$index"
|
|
collect_gpu_device_pids "$index"
|
|
} | awk 'NF' | sort -u
|
|
}
|
|
|
|
kill_pid_list() {
|
|
pids="$1"
|
|
[ -n "$pids" ] || return 0
|
|
|
|
for pid in $pids; do
|
|
log_pid_details "$pid"
|
|
done
|
|
log "terminating GPU holder PIDs: $(echo "$pids" | tr '\n' ' ' | sed 's/[[:space:]]*$//')"
|
|
for pid in $pids; do
|
|
kill -TERM "$pid" >/dev/null 2>&1 || true
|
|
done
|
|
sleep 1
|
|
for pid in $pids; do
|
|
if kill -0 "$pid" >/dev/null 2>&1; then
|
|
log "forcing GPU holder PID $pid to exit"
|
|
kill -KILL "$pid" >/dev/null 2>&1 || true
|
|
fi
|
|
done
|
|
}
|
|
|
|
gpu_has_display_holders() {
|
|
index="$1"
|
|
holders=$(collect_gpu_device_pids "$index")
|
|
[ -n "$holders" ] || return 1
|
|
for pid in $holders; do
|
|
comm=$(ps -p "$pid" -o comm= 2>/dev/null | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
|
|
case "$comm" in
|
|
Xorg|Xwayland|X|gnome-shell)
|
|
return 0
|
|
;;
|
|
esac
|
|
done
|
|
return 1
|
|
}
|
|
|
|
stop_nv_hostengine_if_running() {
|
|
if pgrep -x nv-hostengine >/dev/null 2>&1; then
|
|
pgrep -af "^nv-hostengine$" 2>/dev/null | while IFS= read -r line; do
|
|
[ -n "$line" ] || continue
|
|
log_blocker "$line"
|
|
done
|
|
log "stopping nv-hostengine"
|
|
pkill -TERM -x nv-hostengine >/dev/null 2>&1 || true
|
|
wait_for_process_exit nv-hostengine || pkill -KILL -x nv-hostengine >/dev/null 2>&1 || true
|
|
hostengine_was_active=1
|
|
return 0
|
|
fi
|
|
return 1
|
|
}
|
|
|
|
stop_fabricmanager_if_active() {
|
|
if unit_exists nvidia-fabricmanager.service && stop_unit_if_active nvidia-fabricmanager.service; then
|
|
log_blocker "service nvidia-fabricmanager.service"
|
|
fabric_was_active=1
|
|
return 0
|
|
fi
|
|
return 1
|
|
}
|
|
|
|
stop_display_stack_if_active() {
|
|
stopped=1
|
|
for unit in display-manager.service lightdm.service; do
|
|
if unit_exists "$unit" && stop_unit_if_active "$unit"; then
|
|
log_blocker "service $unit"
|
|
display_was_active=1
|
|
stopped=0
|
|
fi
|
|
done
|
|
return "$stopped"
|
|
}
|
|
|
|
try_gpu_reset() {
|
|
index="$1"
|
|
log "resetting GPU $index"
|
|
nvidia-smi -r -i "$index"
|
|
}
|
|
|
|
drain_gpu_clients() {
|
|
display_was_active=0
|
|
fabric_was_active=0
|
|
hostengine_was_active=0
|
|
|
|
if pgrep -x nv-hostengine >/dev/null 2>&1; then
|
|
pgrep -af "^nv-hostengine$" 2>/dev/null | while IFS= read -r line; do
|
|
[ -n "$line" ] || continue
|
|
log_blocker "$line"
|
|
done
|
|
log "stopping nv-hostengine"
|
|
pkill -TERM -x nv-hostengine >/dev/null 2>&1 || true
|
|
wait_for_process_exit nv-hostengine || pkill -KILL -x nv-hostengine >/dev/null 2>&1 || true
|
|
hostengine_was_active=1
|
|
fi
|
|
|
|
if unit_exists nvidia-fabricmanager.service && stop_unit_if_active nvidia-fabricmanager.service; then
|
|
log_blocker "service nvidia-fabricmanager.service"
|
|
fabric_was_active=1
|
|
fi
|
|
|
|
for unit in display-manager.service lightdm.service; do
|
|
if unit_exists "$unit" && stop_unit_if_active "$unit"; then
|
|
log_blocker "service $unit"
|
|
display_was_active=1
|
|
fi
|
|
done
|
|
|
|
for dev in /dev/nvidia[0-9]*; do
|
|
[ -e "$dev" ] || continue
|
|
holders=$(collect_gpu_device_pids "${dev#/dev/nvidia}")
|
|
kill_pid_list "$holders"
|
|
done
|
|
}
|
|
|
|
restore_gpu_clients() {
|
|
if command -v nvidia-smi >/dev/null 2>&1; then
|
|
if nvidia-smi -pm 1 >/dev/null 2>&1; then
|
|
log "enabled NVIDIA persistence mode"
|
|
else
|
|
log "WARN: failed to enable NVIDIA persistence mode"
|
|
fi
|
|
fi
|
|
|
|
if [ "${hostengine_was_active:-0}" = "1" ] && command -v nv-hostengine >/dev/null 2>&1 && ! pgrep -x nv-hostengine >/dev/null 2>&1; then
|
|
log "starting nv-hostengine"
|
|
nv-hostengine
|
|
fi
|
|
|
|
start_unit_if_marked nvidia-fabricmanager.service "${fabric_was_active:-0}"
|
|
start_unit_if_marked display-manager.service "${display_was_active:-0}"
|
|
if [ "${display_was_active:-0}" = "1" ] && unit_exists lightdm.service && ! unit_is_active lightdm.service; then
|
|
start_unit_if_marked lightdm.service "1"
|
|
fi
|
|
}
|
|
|
|
restart_drivers() {
|
|
drain_gpu_clients
|
|
for mod in nvidia_uvm nvidia_drm nvidia_modeset nvidia; do
|
|
if lsmod | awk '{print $1}' | grep -qx "$mod"; then
|
|
log "unloading module $mod"
|
|
rmmod "$mod"
|
|
fi
|
|
done
|
|
rm -f /dev/nvidiactl /dev/nvidia-uvm /dev/nvidia-uvm-tools /dev/nvidia[0-9]* 2>/dev/null || true
|
|
log "reloading NVIDIA driver stack"
|
|
/usr/local/bin/bee-nvidia-load
|
|
restore_gpu_clients
|
|
}
|
|
|
|
reset_gpu() {
|
|
index="$1"
|
|
display_was_active=0
|
|
fabric_was_active=0
|
|
hostengine_was_active=0
|
|
|
|
holders=$(collect_gpu_holder_pids "$index")
|
|
if [ -n "$holders" ]; then
|
|
kill_pid_list "$holders"
|
|
fi
|
|
if try_gpu_reset "$index"; then
|
|
restore_gpu_clients
|
|
return 0
|
|
fi
|
|
|
|
stop_nv_hostengine_if_running || true
|
|
holders=$(collect_gpu_holder_pids "$index")
|
|
if [ -n "$holders" ]; then
|
|
kill_pid_list "$holders"
|
|
fi
|
|
if try_gpu_reset "$index"; then
|
|
restore_gpu_clients
|
|
return 0
|
|
fi
|
|
|
|
stop_fabricmanager_if_active || true
|
|
holders=$(collect_gpu_holder_pids "$index")
|
|
if [ -n "$holders" ]; then
|
|
kill_pid_list "$holders"
|
|
fi
|
|
if try_gpu_reset "$index"; then
|
|
restore_gpu_clients
|
|
return 0
|
|
fi
|
|
|
|
if gpu_has_display_holders "$index"; then
|
|
stop_display_stack_if_active || true
|
|
holders=$(collect_gpu_holder_pids "$index")
|
|
if [ -n "$holders" ]; then
|
|
kill_pid_list "$holders"
|
|
fi
|
|
if try_gpu_reset "$index"; then
|
|
restore_gpu_clients
|
|
return 0
|
|
fi
|
|
fi
|
|
|
|
holders=$(collect_gpu_holder_pids "$index")
|
|
if [ -n "$holders" ]; then
|
|
log "GPU $index still has holders after targeted drain"
|
|
kill_pid_list "$holders"
|
|
fi
|
|
try_gpu_reset "$index"
|
|
rc=$?
|
|
restore_gpu_clients
|
|
return "$rc"
|
|
}
|
|
|
|
cmd="${1:-}"
|
|
case "$cmd" in
|
|
restart-drivers)
|
|
restart_drivers
|
|
;;
|
|
reset-gpu)
|
|
if [ "$#" -ne 2 ]; then
|
|
usage >&2
|
|
exit 2
|
|
fi
|
|
reset_gpu "$2"
|
|
;;
|
|
*)
|
|
usage >&2
|
|
exit 2
|
|
;;
|
|
esac
|