Skip to content

Capture prior OS state before tuning, and restore it exactly on revert #24

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rammsguns merged 3 commits into from
Aug 12, 2026
Merged

Capture prior OS state before tuning, and restore it exactly on revert #24

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
73c62fc
Capture the prior state before tuning, and restore it on revert
rammsguns Aug 12, 2026
04b7f67
Mark the cross-file state variables as intentionally external
rammsguns Aug 12, 2026
574d85f
Merge remote-tracking branch 'origin/main' into os-tune-rollback
rammsguns Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
271 changes: 213 additions & 58 deletions 10-os-tune.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,41 @@
#!/usr/bin/env bash
# OS + GPU level tuning for LLM inference. Idempotent. Needs sudo.
# Everything here is reversible; see 19-os-revert.sh
# OS + GPU level tuning for LLM inference. Idempotent, transactional, and
# reversible against the state this machine was ACTUALLY in -- not against a
# set of defaults someone assumed.
#
# Every setting's prior value is captured to a root-owned state file before it
# is changed, one setting at a time, so a crash halfway through still leaves
# 19-os-revert.sh enough to undo exactly what happened. Files under /etc that
# existed before llm-rig are backed up byte for byte and never destroyed. See
# lib/ostune.sh for the ownership model.
#
# Usage:
# ./10-os-tune.sh # capture, then tune. Needs sudo.
# ./10-os-tune.sh --dry-run # print every intended mutation. No sudo.
#
# Undo with ./19-os-revert.sh
set -uo pipefail
source "$(dirname "$0")/lib/detect.sh"
detect_hw
RIG_SRC_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$RIG_SRC_DIR/lib/detect.sh"
source "$RIG_SRC_DIR/lib/ostune.sh"

DRY=0
while (( $# )); do
case "$1" in
--dry-run|--plan) DRY=1 ;;
-h|--help) sed -n '2,18p' "${BASH_SOURCE[0]}"; exit 0 ;;
*) die "unknown argument: $1" ;;
esac
shift
done

detect_hw
c_info "Tuning for $GPU_NAME / ${RAM_GB}GB RAM / $PHYS_CORES cores"

# --- 1. GPU persistence mode ------------------------------------------------
# Without this the driver tears down GPU state between processes, adding
# seconds of latency to the first request after an idle period.
c_info "GPU persistence mode"
sudo nvidia-smi -pm 1 >/dev/null && c_ok "persistence mode on"
# --- what we intend to set --------------------------------------------------
# Computed before anything is touched, so --dry-run and the real run agree by
# construction rather than by two lists being kept in step by hand.

# --- 2. Power + clocks ------------------------------------------------------
# MEASURED on 2x RTX A4000 (see 70-thermal-sweep.sh), heat-soaked:
#
# W pp16384 t/s/W tg128 peakT sustClk
Expand All @@ -34,38 +56,20 @@ sudo nvidia-smi -pm 1 >/dev/null && c_ok "persistence mode on"
#
# Override with POWER_PCT=85 if your chassis actually has thermal headroom.
POWER_PCT="${POWER_PCT:-100}"
MAXW=$(nvidia-smi --query-gpu=power.max_limit --format=csv,noheader,nounits | head -1 | cut -d. -f1)
MINW=$(nvidia-smi --query-gpu=power.min_limit --format=csv,noheader,nounits | head -1 | cut -d. -f1)
TGT=$(( MAXW * POWER_PCT / 100 )); (( TGT < MINW )) && TGT=$MINW
c_info "Power limit: max ${MAXW}W -> setting ${TGT}W (${POWER_PCT}%)"
sudo nvidia-smi -pl "$TGT" >/dev/null 2>&1 && c_ok "power limit ${TGT}W" \
|| c_warn "could not set power limit (locked VBIOS?) -- harmless, skipping"

# --- 3. CPU governor --------------------------------------------------------
c_info "CPU governor -> performance"
if need system76-power; then
sudo system76-power profile performance >/dev/null 2>&1 && c_ok "system76 performance profile"
fi
if need cpupower; then
sudo cpupower frequency-set -g performance >/dev/null 2>&1 && c_ok "governor=performance"
else
for g in /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor; do
echo performance | sudo tee "$g" >/dev/null 2>&1 || true
done
c_ok "governor written directly"
MAXW=$(nvidia-smi --query-gpu=power.max_limit --format=csv,noheader,nounits 2>/dev/null | head -1 | cut -d. -f1)
MINW=$(nvidia-smi --query-gpu=power.min_limit --format=csv,noheader,nounits 2>/dev/null | head -1 | cut -d. -f1)
TGT=""
if [[ -n "${MAXW:-}" && -n "${MINW:-}" ]]; then
TGT=$(( MAXW * POWER_PCT / 100 )); (( TGT < MINW )) && TGT=$MINW
fi

# --- 4. Transparent huge pages ---------------------------------------------
# Model weights are mmap'd in multi-GB contiguous ranges. THP=always measurably
# reduces TLB misses during prompt processing when layers live in system RAM.
c_info "Transparent huge pages -> always"
echo always | sudo tee /sys/kernel/mm/transparent_hugepage/enabled >/dev/null
echo defer+madvise | sudo tee /sys/kernel/mm/transparent_hugepage/defrag >/dev/null 2>&1 || true
c_ok "THP=always"
WANT_GOVERNOR=performance
WANT_S76=performance
WANT_THP=always
WANT_THP_DEFRAG='defer+madvise'

# --- 5. VM / memory sysctls -------------------------------------------------
c_info "Kernel sysctls"
sudo tee /etc/sysctl.d/99-llm-inference.conf >/dev/null <<EOF
sysctl_content() {
cat <<'EOF'
# Written by llm-rig 10-os-tune.sh
# Don't page out model weights under memory pressure -- swapping a resident
# model is catastrophic for latency. 1 rather than 0 keeps the OOM killer sane.
Expand All @@ -82,21 +86,19 @@ vm.vfs_cache_pressure = 50
vm.dirty_ratio = 15
vm.dirty_background_ratio = 5
EOF
sudo sysctl --system >/dev/null 2>&1
c_ok "/etc/sysctl.d/99-llm-inference.conf applied"
}

# --- 6. mlock limits --------------------------------------------------------
# --mlock pins weights in RAM. Needs an unlimited memlock rlimit.
c_info "memlock rlimit -> unlimited"
sudo tee /etc/security/limits.d/99-llm-memlock.conf >/dev/null <<EOF
limits_content() {
cat <<'EOF'
# Written by llm-rig 10-os-tune.sh
* soft memlock unlimited
* hard memlock unlimited
EOF
c_ok "limits.d written (takes effect on next login)"
}

# --- 7. Make the tuning survive reboot -------------------------------------
c_info "Persisting GPU settings across reboot"
sudo tee /etc/systemd/system/llm-gpu-tune.service >/dev/null <<EOF
unit_content() {
cat <<EOF
# Written by llm-rig 10-os-tune.sh
[Unit]
Description=LLM GPU tuning (persistence mode + power limit)
After=nvidia-persistenced.service
Expand All @@ -106,22 +108,175 @@ Wants=nvidia-persistenced.service
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/bin/nvidia-smi -pm 1
ExecStart=-/usr/bin/nvidia-smi -pl $TGT
ExecStart=-/usr/bin/nvidia-smi -pl ${TGT:-0}

[Install]
WantedBy=multi-user.target
EOF
sudo systemctl daemon-reload
sudo systemctl enable --now llm-gpu-tune.service >/dev/null 2>&1
c_ok "llm-gpu-tune.service enabled"
}

# --- the plan ---------------------------------------------------------------
# Reading current state needs no privilege, which is what makes --dry-run
# honest: it reports the same values the real run will capture.

ostune_build_plan() {
local idx pm pl f gov

while IFS=$'\t' read -r idx pm pl; do
[[ -n "$idx" ]] || continue
ostune_plan_line "GPU $idx persistence" "$pm" "Enabled"
[[ -n "$TGT" ]] && ostune_plan_line "GPU $idx power limit (W)" "$pl" "$TGT"
done < <(ostune_gpu_state)

while IFS=$'\t' read -r f gov; do
[[ -n "$f" ]] || continue
ostune_plan_line "governor $(basename "$(dirname "$(dirname "$f")")")" "$gov" "$WANT_GOVERNOR"
done < <(ostune_governors)

need system76-power && ostune_plan_line "system76-power profile" "$(ostune_s76_profile)" "$WANT_S76"

ostune_plan_line "THP enabled" "$(ostune_sysfs_choice "$OSTUNE_THP_ENABLED")" "$WANT_THP"
ostune_plan_line "THP defrag" "$(ostune_sysfs_choice "$OSTUNE_THP_DEFRAG")" "$WANT_THP_DEFRAG"

local p
for p in "$OSTUNE_SYSCTL_FILE" "$OSTUNE_LIMITS_FILE" "$OSTUNE_UNIT_FILE"; do
case "$(ostune_file_status "$p")" in
absent) ostune_plan_line "$p" "absent" "created by llm-rig" ;;
created) ostune_plan_line "$p" "llm-rig's" "rewritten (unchanged)" ;;
adopted) ostune_plan_line "$p" "yours, backed up" "rewritten" ;;
foreign) ostune_plan_line "$p" "YOURS" "backed up, then overwritten" ;;
*-dirty) ostune_plan_line "$p" "EDITED SINCE" "REFUSED -- tune will stop here" ;;
esac
done
}

if (( DRY )); then
echo
c_info "Planned mutations (nothing has been changed):"
ostune_build_plan | ostune_plan_render
echo
c_info "State file that would be written: $OSTUNE_STATE"
if ostune_state_exists; then
c_warn "A state file already exists -- prior values captured earlier are kept,
so a re-run cannot overwrite them with values llm-rig itself set."
fi
exit 0
fi

# --- capture, then mutate ---------------------------------------------------
# Order matters throughout: capture is always the statement BEFORE the change,
# never after and never in a batch at the end.

ostune_state_init || die "cannot create the state directory $OSTUNE_STATE_DIR"
ostune_state_put state tuned_at "$(date -u +%Y-%m-%dT%H:%M:%SZ)" >/dev/null

fail=0

# --- 1. GPU persistence mode + power ----------------------------------------
# Without persistence the driver tears down GPU state between processes, adding
# seconds of latency to the first request after an idle period.
c_info "GPU persistence mode and power limit"
while IFS=$'\t' read -r idx pm pl; do
[[ -n "$idx" ]] || continue
ostune_state_put gpu "$idx.persistence" "$pm"
ostune_state_put gpu "$idx.power_limit" "$pl"
done < <(ostune_gpu_state)

# --- 8. Report --------------------------------------------------------------
if ostune_priv nvidia-smi -pm 1 >/dev/null 2>&1; then
c_ok "persistence mode on"
else
c_warn "could not set persistence mode"
fi

if [[ -n "$TGT" ]]; then
c_info "Power limit: max ${MAXW}W -> setting ${TGT}W (${POWER_PCT}%)"
ostune_priv nvidia-smi -pl "$TGT" >/dev/null 2>&1 && c_ok "power limit ${TGT}W" \
|| c_warn "could not set power limit (locked VBIOS?) -- harmless, skipping"
else
c_warn "no GPU power limits reported -- skipping power tuning"
fi

# --- 2. CPU governor --------------------------------------------------------
c_info "CPU governor -> $WANT_GOVERNOR"
if need system76-power; then
ostune_state_put cpu s76_profile "$(ostune_s76_profile)"
ostune_priv system76-power profile "$WANT_S76" >/dev/null 2>&1 \
&& c_ok "system76 $WANT_S76 profile"
fi

while IFS=$'\t' read -r f gov; do
[[ -n "$f" ]] || continue
ostune_state_put cpu "$f" "$gov"
done < <(ostune_governors)

if need cpupower; then
ostune_priv cpupower frequency-set -g "$WANT_GOVERNOR" >/dev/null 2>&1 \
&& c_ok "governor=$WANT_GOVERNOR"
else
for f in $OSTUNE_CPU_GLOB; do
[[ -f "$f" ]] || continue
printf '%s\n' "$WANT_GOVERNOR" | ostune_priv tee "$f" >/dev/null 2>&1 || true
done
c_ok "governor written directly"
fi

# --- 3. Transparent huge pages ---------------------------------------------
# Model weights are mmap'd in multi-GB contiguous ranges. THP=always measurably
# reduces TLB misses during prompt processing when layers live in system RAM.
c_info "Transparent huge pages -> $WANT_THP"
ostune_state_put thp enabled "$(ostune_sysfs_choice "$OSTUNE_THP_ENABLED")"
ostune_state_put thp defrag "$(ostune_sysfs_choice "$OSTUNE_THP_DEFRAG")"
printf '%s\n' "$WANT_THP" | ostune_priv tee "$OSTUNE_THP_ENABLED" >/dev/null 2>&1 \
&& c_ok "THP=$WANT_THP" || c_warn "could not set THP"
printf '%s\n' "$WANT_THP_DEFRAG" | ostune_priv tee "$OSTUNE_THP_DEFRAG" >/dev/null 2>&1 || true

# --- 4. VM / memory sysctls -------------------------------------------------
c_info "Kernel sysctls"
if sysctl_content | ostune_install_file "$OSTUNE_SYSCTL_FILE" 644; then
ostune_priv sysctl --system >/dev/null 2>&1
c_ok "$OSTUNE_SYSCTL_FILE applied"
else
c_err "${OSTUNE_LAST_ERROR:-could not write $OSTUNE_SYSCTL_FILE}"
fail=1
fi

# --- 5. mlock limits --------------------------------------------------------
# --mlock pins weights in RAM. Needs an unlimited memlock rlimit.
c_info "memlock rlimit -> unlimited"
if limits_content | ostune_install_file "$OSTUNE_LIMITS_FILE" 644; then
c_ok "limits.d written (takes effect on next login)"
else
c_err "${OSTUNE_LAST_ERROR:-could not write $OSTUNE_LIMITS_FILE}"
fail=1
fi

# --- 6. Make the GPU tuning survive reboot ----------------------------------
c_info "Persisting GPU settings across reboot"
if unit_content | ostune_install_file "$OSTUNE_UNIT_FILE" 644; then
ostune_state_put unit llm-gpu-tune.service enabled
ostune_priv systemctl daemon-reload
ostune_priv systemctl enable --now llm-gpu-tune.service >/dev/null 2>&1
c_ok "llm-gpu-tune.service enabled"
else
c_err "${OSTUNE_LAST_ERROR:-could not write $OSTUNE_UNIT_FILE}"
fail=1
fi

# --- 7. Report --------------------------------------------------------------
echo
c_info "Prior state captured in $OSTUNE_STATE -- 19-os-revert.sh restores exactly that"
if (( fail )); then
c_warn "Some steps were refused. Everything captured above is still revertible;
nothing that was refused was changed."
fi

c_info "Post-tune state"
nvidia-smi --query-gpu=name,persistence_mode,power.limit,temperature.gpu,clocks.max.sm \
--format=csv
echo "Governor: $(cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor)"
echo "THP: $(cat /sys/kernel/mm/transparent_hugepage/enabled)"
echo "swappiness: $(cat /proc/sys/vm/swappiness)"
--format=csv 2>/dev/null
echo "Governor: $(cat "$OSTUNE_ROOT/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor" 2>/dev/null || echo "$OSTUNE_UNKNOWN")"
echo "THP: $(cat "$OSTUNE_THP_ENABLED" 2>/dev/null || echo "$OSTUNE_UNKNOWN")"
echo "swappiness: $(cat "$OSTUNE_ROOT/proc/sys/vm/swappiness" 2>/dev/null || echo "$OSTUNE_UNKNOWN")"
echo
c_warn "Log out and back in (or reboot) for the memlock limit to apply."

exit "$fail"
Loading