atlas-iac/services/ai-llm/scripts/gpu_pilot_reservation.sh

59 lines
2.1 KiB
Bash
Executable File

#!/usr/bin/env bash
# Preserve enough host state to restore the temporary RTX 3080 reservation.
set -euo pipefail
state=/host/var/lib/atlas-maintenance/lan-inference-20260928
host() { chroot /host "$@"; }
umask 077
mkdir -p "$state"
case "${1:-}" in
reserve)
# Wolf must be scaled down through Flux before stopping its child containers.
for attempt in $(seq 1 120); do
host pgrep -x wolf >/dev/null || break
sleep 2
done
if host pgrep -x wolf >/dev/null; then
echo 'Wolf still running; reservation refused' >&2
exit 1
fi
if ! test -f "$state/saved"; then
host systemctl show display-manager -p Id --value > "$state/display-unit"
host systemctl is-active display-manager > "$state/display-active" || true
host systemctl is-enabled display-manager > "$state/display-enabled" || true
host docker ps --filter 'name=^/Wolf' --format '{{.ID}}' > "$state/docker-running"
touch "$state/saved"
fi
unit=$(cat "$state/display-unit")
case "$unit" in sddm.service|gdm.service|gdm3.service|lightdm.service) ;;
*) echo 'Unexpected display manager; refusing host mutation' >&2; exit 1 ;;
esac
host systemctl mask --now "$unit"
while read -r container; do
test -z "$container" || host docker stop --time 30 "$container"
done < "$state/docker-running"
touch "$state/reserved"
echo 'GPU desktop and Wolf containers stopped; restoration state saved'
;;
restore)
test -f "$state/saved"
unit=$(cat "$state/display-unit")
case "$unit" in sddm.service|gdm.service|gdm3.service|lightdm.service) ;;
*) exit 1 ;;
esac
case "$(cat "$state/display-enabled")" in
masked*) ;;
*) host systemctl unmask "$unit" ;;
esac
if test "$(cat "$state/display-active")" = active; then
host systemctl start "$unit"
fi
while read -r container; do
test -z "$container" || host docker start "$container"
done < "$state/docker-running"
rm -f "$state/reserved"
echo 'Saved display and Wolf container state restored'
;;
*) echo 'Usage: gpu_pilot_reservation.sh reserve|restore' >&2; exit 2 ;;
esac