ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge origin/main into image-generation

Browse source 
Resolves the app-sidebar conflict: main added Hub and Projects rows
inline while this branch renders the nav from navRows in the order and
pin state set under Settings -> Appearance. Kept the data-driven
rendering, having checked both of main's additions are already
represented there - the projects row carries the same icon, label,
active check, handlers and inline New project button.

Main also replaced the sidebar's inline name field with
NewProjectDialog, which owns its own state, so the button no longer
resets a name draft: it sets the move target and opens the dialog, as
main's other call sites do.
This commit is contained in:
Daniel Han 2026-07-26 12:37:00 +00:00
commit f7d54a757f
59 changed files with 7612 additions and 529 deletions
Download patch file Download diff file Expand all files Collapse all files

46
.github/workflows/studio-backend-ci.yml vendored
View file

@ -30,6 +30,13 @@ on:
- 'unsloth/**'
- 'unsloth_cli/**'
- 'tests/**'
# The root installers: tests/sh/*.sh and tests/studio/install/* assert
# against these two files, so a change here must run the suite that
# covers it. Without them an install-only edit (the shape most AMD/ROCm
# routing fixes take) skipped Backend CI entirely.
- 'install.sh'
- 'install.ps1'
- 'scripts/**'
- 'pyproject.toml'
- '.github/workflows/studio-backend-ci.yml'
push:
@ -217,27 +224,32 @@ jobs:
tests/studio/test_xpu_spoof_pipeline.py
- name: Shell installer tests
# Subset that does not depend on a writable / pristine install.sh
# tree; test_install_host_defaults.sh checks install.ps1 layout
# which has drifted (separate followup).
# Auto-discovered rather than allowlisted. The old hardcoded list had
# silently fallen seven files behind tests/run_all.sh, including
# test_strixhalo_wsl_reroute.sh -- the only shell coverage of the ROCm
# WSL reroute -- so that suite never ran on a PR. Skips are explicit,
# each with a reason, and tests/studio/test_ci_shell_suite_coverage.py
# fails if this step stops discovering the directory or the skip list
# grows without one.
#
# Skipped:
# test_install_host_defaults.sh: asserts an install.ps1 layout that
# has drifted (separate followup).
# test_install_rollback_lifecycle.sh: already runs on both platforms
# in cross-platform-parity-ci.yml.
run: |
set -e
for s in \
tests/sh/test_get_torch_index_url.sh \
tests/sh/test_mac_intel_compat.sh \
tests/sh/test_node_decision.sh \
tests/sh/test_studio_home_node_dir.sh \
tests/sh/test_system_node_readonly.sh \
tests/sh/test_nvcc_meets_llama_minimum.sh \
tests/sh/test_resolve_cuda_archs.sh \
tests/sh/test_staged_validation_enabled.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh \
tests/sh/test_torch_flavor.sh \
tests/sh/test_with_llama_cpp_dir_flag.sh \
tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do
skip="test_install_host_defaults.sh test_install_rollback_lifecycle.sh"
found=0
for s in tests/sh/test_*.sh; do
case " $skip " in
*" $(basename "$s") "*) echo "skipping $s (see workflow comment)"; continue ;;
esac
found=$((found + 1))
echo "::group::$s"
bash "$s"
echo "::endgroup::"
done
[ "$found" -gt 0 ] || { echo "::error::no shell tests discovered under tests/sh"; exit 1; }
echo "ran $found shell installer test files"

18
install.ps1
View file

@ -1917,12 +1917,14 @@ exit 0
# (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU.
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
@{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: RX 9070 XT / 9070 GRE / 9070 / 9080)
@{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: RX 9060 XT / 9060)
@{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
@{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
@{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
@{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
@{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33)
@{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point)
@{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family
@ -2203,6 +2205,7 @@ exit 0
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
@ -2224,6 +2227,7 @@ exit 0
$torchFloorMap = @{
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
"gfx1152" = "torch>=2.11.0,<2.12.0"
}
# Companion ranges track the torch ceiling so pip resolves a consistent
# trio on AMD's per-arch index (each published independently). Mirrors
@ -2231,10 +2235,12 @@ exit 0
$torchvisionFloorMap = @{
"gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0"
"gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0"
"gfx1152" = "torchvision>=0.26.0,<0.27.0"
}
$torchaudioFloorMap = @{
"gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0"
"gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0"
"gfx1152" = "torchaudio>=2.11.0,<2.12.0"
}
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
if ($archFamily) {
@ -2264,7 +2270,7 @@ exit 0
$_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
}
# Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare.
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $TorchIndexUrl
$ROCmTorchFloor = "torch>=2.11.0,<2.12.0"

121
install.sh
View file

@ -655,6 +655,15 @@ _apt_distro_description() {
)
}
# ── Helper: can the controlling terminal actually be opened for reading? ──
# `test -r` only checks permission bits, which look fine in containers and
# systemd units where open() then fails with ENXIO. Probe with a real open.
# The subshell is required: in dash a failed redirection on the special
# builtin `:` exits the whole script.
_can_read_tty() {
( : </dev/tty ) >/dev/null 2>&1
}
# ── Helper: install packages via apt, escalating to sudo only if needed ──
# Usage: _smart_apt_install pkg1 pkg2 pkg3 ...
_smart_apt_install() {
@ -695,24 +704,63 @@ _smart_apt_install() {
echo " from your distro's official repositories (not a third-party tarball)."
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo ""
printf " Accept? [Y/n] "
if [ -r /dev/tty ]; then
read -r REPLY </dev/tty || REPLY="y"
else
REPLY="y"
fi
case "$REPLY" in
[nN]*)
if _can_read_tty; then
printf " Accept? [Y/n] "
# The device opened, so a failed read is EOF, not consent: decline,
# as the autostart prompt below does. Enter is still yes (a
# successful read of an empty line).
read -r REPLY </dev/tty || REPLY="n"
case "$REPLY" in
[nN]*)
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
;;
esac
# Mirror the headless branch: on a sudoers denial, a wrong password
# or an apt error, say what to run by hand instead of letting set -e
# abort on a bare sudo/apt message.
if sudo apt-get update -y </dev/null &&
sudo apt-get install -y $_STILL_MISSING </dev/null; then
:
else
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " Could not install these packages: $_STILL_MISSING"
echo " See the error above."
echo " Please install them first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
;;
*)
sudo apt-get update -y </dev/null
sudo apt-get install -y $_STILL_MISSING </dev/null
;;
esac
fi
else
# Nobody can answer a prompt or type a password here. -n makes sudo
# refuse rather than prompt into a closed stdin, which is how #7307
# died. Probe with the real commands: `sudo -l` answers whether they
# are *authorized*, not whether running them needs authentication.
# -k ignores any cached timestamp, so only a real NOPASSWD rule gets
# through, not someone's sudo in another shell minutes ago. Per
# sudo(8), -k alongside a command ignores the cached credentials and
# "will not update" them, so other sessions keep theirs.
echo " No terminal to confirm on; trying passwordless sudo."
if sudo -n -k apt-get update -y </dev/null &&
sudo -n -k apt-get install -y $_STILL_MISSING </dev/null; then
echo " Installed with passwordless sudo."
else
echo ""
echo " Could not install these packages: $_STILL_MISSING"
echo " Detected ${_ad_desc}."
# Either sudo refused, or apt failed on a bad repo, dpkg lock or
# network outage. sudo exits 1 on an auth/config problem and
# when the command cannot be executed, but otherwise passes the
# command's own status through, so state both causes.
echo " Either sudo needs a password here, or apt-get itself"
echo " failed; see the error above. With no terminal to"
echo " authenticate on, this cannot be done unattended."
echo " Please install them first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
fi
fi
else
echo ""
echo " sudo is not available on this system."
@ -2260,6 +2308,7 @@ _amd_arch_index_family_for_gfx() {
gfx1201|gfx1200) echo gfx120X-all ;;
gfx1151) echo gfx1151 ;;
gfx1150) echo gfx1150 ;;
gfx1152) echo gfx1152 ;;
gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;;
gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;;
gfx90a) echo gfx90a ;;
@ -2271,12 +2320,14 @@ _amd_arch_index_family_for_gfx() {
# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable).
_infer_amd_gfx_arch_from_gpu_name() {
case "$1" in
*"9070 XT"*|*9080*) echo gfx1201 ;;
*9070*|*9060*) echo gfx1200 ;;
*9070*|*9080*) echo gfx1201 ;;
*9060*) echo gfx1200 ;;
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;;
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;;
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;;
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;;
*"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) echo gfx1150 ;;
*"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1152 ;;
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) echo gfx1102 ;;
*"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) echo gfx1101 ;;
*"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) echo gfx1100 ;;
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;;
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;;
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;;
@ -2316,10 +2367,14 @@ _infer_linux_amd_gfx_arch() {
echo gfx1151
return 0
fi
if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]' /proc/cpuinfo 2>/dev/null; then
echo gfx1150
return 0
fi
if [ -n "$_gpu_evidence" ] && grep -qiE '860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
echo gfx1152
return 0
fi
if command -v lspci >/dev/null 2>&1; then
# A non-AMD controller can enumerate first (Intel/ASPEED before an AMD
# dGPU), so scan every display-class line and take the first AMD one
@ -3055,7 +3110,7 @@ if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
# whole handoff (a user-set override re-exports unchanged).
export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"
case "$_linux_inferred_gfx" in
gfx1201|gfx1200|gfx1151|gfx1150)
gfx1201|gfx1200|gfx1151|gfx1150|gfx1152)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
@ -3124,7 +3179,7 @@ fi
# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a
# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced.
case "$_torch_index_leaf" in
rocm7.2|gfx120x-all|gfx1151|gfx1150)
rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
@ -3243,7 +3298,7 @@ case "$_torch_index_leaf" in
fi
_strix_gfx=""
case "$_runtime_gfx" in
gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;;
esac
# Skip rocm7.13+ generic indexes: they already ship the fixes, so the
# arch build (rocm7.13) would be a downgrade rather than a rescue.
@ -3339,12 +3394,14 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
# gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on
# gfx1102 (bash case has no negative lookahead like the PS tables).
case "$_gpu_disp_mkt" in
*"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
*9070*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 (Navi 48)
*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 (Navi 44)
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
*"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
*"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _gpu_disp_gfx="gfx1101" ;; # RDNA 3 (Navi 32)
*"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point)
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _gpu_disp_gfx="gfx1030" ;; # RDNA 2 (Navi 21)
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _gpu_disp_gfx="gfx1032" ;; # RDNA 2 (Navi 23)
@ -4046,9 +4103,11 @@ echo ""
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then
echo ""
printf " Start Unsloth Studio now? [Y/n] "
# No readable answer (closed/EOF tty) defaults to no; Enter is still yes.
if [ -r /dev/tty ]; then
# Prompt only when something can answer: `test -r` passes on the unopenable
# /dev/tty found in containers, leaving a dangling question in the log.
if _can_read_tty; then
printf " Start Unsloth Studio now? [Y/n] "
read -r _reply </dev/tty || _reply="n"
else
_reply="n"

68
studio/backend/colab.py
View file

@ -218,11 +218,10 @@ def _colab_login_html(username: str, password: str) -> str:
Unsloth Studio Login (Colab)
</h2>
<p style="color: #333333; margin: 0 0 12px 0; font-size: 14px; font-weight: bold;">
Log in to Studio with the Cloudflare link above using these credentials. This cell
is visible only in your notebook session.
Log in as <code>{username}</code> with this password. This cell is visible only in
your notebook session.
</p>
<p style="color: #333333; margin: 0; font-size: 14px; font-family: monospace; font-weight: bold;">
Username: <code>{username}</code><br>
Password: <code>{password}</code>
</p>
</div>
@ -441,8 +440,29 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
return False
def _shareable_link_html(cloudflare_url: str) -> str:
"""Branded card for the shareable Cloudflare link, styled like the show_link banner."""
def _shareable_link_html(
cloudflare_url: str,
password: "str | None" = None,
username: "str | None" = None,
) -> str:
"""Branded card for the shareable Cloudflare link, styled like the show_link banner.
*password* renders under the link so the credential sits in the card with the button
it unlocks. The username is always the default admin, so it reads inline.
"""
login_block = ""
if password:
login_block = f"""
<p style="color: #000000; margin: 16px 0 0 0; font-size: 20px; font-weight: 800;">
Password
</p>
<p style="margin: 6px 0 0 0;"><code style="display: inline-block; font-size: 24px;
font-weight: 800; text-decoration: underline; background: #f3f3f3;
padding: 4px 10px; border-radius: 6px;">{password}</code></p>
<p style="color: #666666; margin: 6px 0 0 0; font-size: 12px;">
Log in as <code>{username}</code> with this password. Shown only in your
notebook session, and never included in the shared link.
</p>"""
return f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
@ -460,11 +480,12 @@ def _shareable_link_html(cloudflare_url: str) -> str:
Open Unsloth Studio
</a>
<p style="color: #333333; margin: 12px 0 0 0; font-size: 14px; font-weight: bold;">
This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab.
This Cloudflare HTTPS link works from any device, so you can share it with anyone.
</p>
<p style="color: #333333; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
🔗 {cloudflare_url}
</p>
🔗 <a href="{cloudflare_url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
style="color: #000000; text-decoration: underline; cursor: pointer;">{cloudflare_url}</a>
</p>{login_block}
</div>
"""
@ -555,28 +576,37 @@ def _show_and_embed(
cloudflare_url = cloudflare_url,
)
# Fold the credentials into the link card rather than a second card below it.
credentials_shown = False
if cloudflare_url:
try:
from IPython.display import HTML, display
display(HTML(_shareable_link_html(cloudflare_url)))
username, password = colab_login if colab_login else (None, None)
display(HTML(_shareable_link_html(cloudflare_url, password, username)))
credentials_shown = bool(colab_login)
except Exception as e:
logger.info(f"Could not render Cloudflare link card ({e}).")
if colab_login:
if colab_login and not credentials_shown:
try:
_show_colab_login_credentials(*colab_login)
except Exception as e:
logger.info(f"Could not render Colab login card ({e}).")
try:
show_link(
port,
_url = url,
has_cloudflare_link = bool(cloudflare_url),
cloudflare_requested = cloudflare_requested,
)
except Exception as e:
logger.info(f"Could not render Unsloth link card ({e}).")
# With a tunnel up the embed below is skipped, so the ready card would only restate
# the link card and print a proxy URL that 404s outside this tab.
skip_ready_card = _is_colab_runtime() and bool(cloudflare_url)
if not skip_ready_card:
try:
show_link(
port,
_url = url,
has_cloudflare_link = bool(cloudflare_url),
cloudflare_requested = cloudflare_requested,
)
except Exception as e:
logger.info(f"Could not render Unsloth link card ({e}).")
# On Colab with a working tunnel, skip the in-cell proxy embed (often blank).
if _is_colab_runtime() and cloudflare_url:

137
studio/backend/core/export/export.py
View file

@ -81,6 +81,82 @@ _PYTORCH_MISSING_MESSAGE = (
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
def _multi_gpu_device_map_kwargs() -> dict:
"""``device_map`` kwargs for sharding a checkpoint across every visible GPU.
unsloth's ``from_pretrained`` defaults to ``device_map="sequential"``, which stacks
the whole model on GPU0 and OOMs multi-GPU hosts whose other GPUs sit empty (#7053).
Returns ``{"device_map": "balanced"}`` only on a real multi-GPU CUDA/ROCm host
(mirroring the inference loader's ``get_device_map``), else empty so single-GPU, CPU
and MLX loads keep the loader default."""
if _IS_MLX:
return {}
try:
from utils.hardware import get_device_map, get_parent_visible_gpu_ids
visible = get_parent_visible_gpu_ids()
if len(visible) > 1:
device_map = get_device_map(visible)
elif not visible:
# UUID/MIG masks resolve to no numeric ids; get_device_map(None) falls back
# to the visible-GPU count, so a multi-GPU UUID/MIG host still shards.
device_map = get_device_map(None)
else:
return {}
if device_map == "balanced":
return {"device_map": device_map}
except Exception as exc:
logger.debug(f"multi-GPU device_map resolution failed; using loader default: {exc}")
return {}
def _is_oom_error(exc: BaseException) -> bool:
"""True for an accelerator OOM, however it is spelled.
accelerate and transformers re-raise it as a plain ``RuntimeError`` on several paths
and ROCm/XPU use their own classes, so match the message too.
"""
if torch is not None:
oom_types = tuple(
t
for t in (
getattr(torch, "OutOfMemoryError", None),
getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None),
getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None),
)
if isinstance(t, type)
)
if oom_types and isinstance(exc, oom_types):
return True
return "out of memory" in f"{type(exc).__name__}: {exc}".lower()
def _is_cpu_spill_rejection(exc: BaseException) -> bool:
"""bitsandbytes refuses a map that spills to CPU/disk with a plain ``ValueError``.
Busy secondary GPUs can make ``balanced`` spill to CPU even where the old sequential
load fit on GPU0, and that message says nothing about memory, so the retry has to
match it explicitly. See transformers ``quantizers/quantizer_bnb_4bit.py``.
"""
return "dispatched on the cpu or the disk" in str(exc).lower()
class _CpuSpillRetry(Exception):
"""A multi-GPU load that succeeded but left modules offloaded to CPU/disk."""
def _cpu_offloaded_modules(model) -> int:
"""Count the modules a load parked on CPU or disk.
Only bitsandbytes refuses such a map; a full-precision load accepts it, leaves the
parameters on meta and dies much later in safetensors with "Cannot copy out of meta
tensor". Nothing raises at load time, so inspect the map directly. PEFT re-dispatches
when attaching an adapter, so in practice this catches merged checkpoints.
"""
device_map = getattr(model, "hf_device_map", None) or {}
return sum(1 for target in device_map.values() if str(target) in ("cpu", "disk"))
def _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
@ -271,6 +347,7 @@ class ExportBackend:
load_in_4bit: bool = True,
trust_remote_code: bool = False,
hf_token: Optional[str] = None,
_device_map_override: Optional[dict] = None,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
@ -303,6 +380,14 @@ class ExportBackend:
# Skip the Hub when offline so a no-internet export uses the local cache.
local_files_only = _hf_offline()
# Shard across every visible GPU instead of stacking on GPU0 (#7053); {} on
# single-GPU/CPU/MLX. _device_map_override is the single-device retry below.
_device_map_kw = (
_multi_gpu_device_map_kwargs()
if _device_map_override is None
else _device_map_override
)
# Run the type-detection probes in the forced-offline window (else a gated
# base 404s); it covers is_vision_model's Hub reads + the transformers-5
# subprocess, and local_files_only makes detect_audio_type's requests.get skip.
@ -328,6 +413,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "whisper":
@ -343,6 +429,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "snac":
@ -355,6 +442,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "bicodec":
@ -368,6 +456,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "dac":
@ -380,6 +469,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self.is_vision:
@ -392,6 +482,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -405,8 +496,16 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
# Only when we asked for the multi-GPU map: a single-GPU host has no second
# placement to retry on, so leave its behaviour untouched.
_offloaded = _cpu_offloaded_modules(model) if _device_map_kw else 0
if _device_map_override is None and _offloaded:
del model
raise _CpuSpillRetry(f"{_offloaded} module(s) offloaded to CPU/disk")
if _IS_MLX:
# MLX doesn't use PeftModel — detect LoRA via adapter_config.json
self.is_peft = adapter_config.exists()
@ -429,11 +528,41 @@ class ExportBackend:
return True, f"Loaded {model_type} model{peft_info} successfully"
except Exception as e:
logger.error(f"Error loading checkpoint: {e}")
import traceback
# Sharding is an optimisation, never a requirement. "balanced" budgets from the
# free memory read BEFORE this process opens a CUDA context on each GPU, so when
# a training or chat job already owns the others the shard can OOM, or spill to
# CPU and be refused by bitsandbytes, where the old single-device load succeeded.
# Fall back once before giving up.
if (
_device_map_override is None
and (
isinstance(e, _CpuSpillRetry) or _is_oom_error(e) or _is_cpu_spill_rejection(e)
)
and _multi_gpu_device_map_kwargs()
):
# Retry outside this block: the live traceback pins the half-built model's
# frames, so an in-block retry inherits the exhausted device.
retry_reason = str(e)
else:
logger.error(f"Error loading checkpoint: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
logger.warning(
f"Multi-GPU export load unusable ({retry_reason}); retrying on "
f"the single-device loader default."
)
self.cleanup_memory()
return self.load_checkpoint(
checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
hf_token = hf_token,
_device_map_override = {},
)
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""

275
studio/backend/core/inference/llama_cpp.py
View file

@ -33,6 +33,7 @@ from typing import (
List,
Literal,
Mapping,
MutableMapping,
Optional,
Union,
)
@ -306,6 +307,9 @@ def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]":
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
# A transport error can arrive before the child is reapable; a request path cannot
# afford the 5s the background MTP reload spends on the same race.
_RESPAWN_REAP_GRACE_S = 1.0
def _finalize_reasoning_only_cumulative(
@ -2098,6 +2102,9 @@ class LlamaCppBackend:
# Serialises mid-session respawns so many generations hitting a killed
# server trigger at most one reload (see _respawn_if_dead).
self._respawn_lock = threading.Lock()
# Bumped by every unload. load_model clears _cancel_event, so a respawn that
# raced an unload needs a signal that survives the clear (see _respawn_if_dead).
self._unload_epoch = 0
# Set by the in-app updater while it swaps prebuilt binaries; load_model()
# rejects fast so no server starts from a half-swapped binary.
self._llama_update_in_progress = False
@ -3182,7 +3189,7 @@ class LlamaCppBackend:
@staticmethod
def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
"""True only for AMD unified-memory APUs (gfx1150/gfx1151/gfx1152), where
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it
hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the
selected GPUs, so a dGPU on a mixed host is not treated as unified-memory;
@ -3212,7 +3219,9 @@ class LlamaCppBackend:
)
arch_by_id[pid] = _arch.split(":")[0].strip().lower()
for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id):
if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}:
# gfx1152 is Krackan Point (Radeon 860M/840M), the third RDNA 3.5
# APU: same shared GPU/system-RAM pool as Strix Point/Halo.
if arch_by_id.get(_i) in {"gfx1150", "gfx1151", "gfx1152"}:
return True
except Exception:
return False
@ -3696,6 +3705,14 @@ class LlamaCppBackend:
# aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
# V cache types that llama.cpp can run WITHOUT flash attention. Only the V
# axis has the dependency: a quantized V cache (q8_0/q4_0/q4_1/q5_0/q5_1/
# iq4_nl) aborts init with "V cache quantization requires flash_attn", while
# a quantized K cache runs fine without FA. So the flash-attn-off crash-
# recovery fallback must reset a quantized V cache to f16 before it can
# launch (and leaves K alone). These three are the only non-quantized types.
_NON_QUANTIZED_KV_TYPES = frozenset({"f16", "bf16", "f32"})
# Main-model placement settings that Manual mode owns. They must not leak
# from Studio's parent environment into llama-server and silently override
# the command assembled from the current request. Draft-model placement is
@ -6149,6 +6166,21 @@ class LlamaCppBackend:
cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode)
)
@staticmethod
def _canonical_long_flag(name: str) -> str:
"""Return ``name`` with llama.cpp's long-option underscore normalization.
llama.cpp runs ``std::replace(arg.begin(), arg.end(), '_', '-')`` on any
argv token that starts with ``--`` before looking it up, so a legal
pass-through spelling like ``--cache_type_v`` parses as
``--cache-type-v``. Mirror that here so managed-flag matching sees the
same canonical name. Short flags (``-ctv``) never carry underscores and
keep their exact spelling; pass only the flag name (no attached value).
"""
if name.startswith("--"):
return name.replace("_", "-")
return name
@staticmethod
def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]:
"""Return cmd with flash attention forced off, or None when its effective
@ -6181,8 +6213,76 @@ class LlamaCppBackend:
out[i + 1] = "off"
elif explicit(i) is None: # bare flag (reads as on) -> explicit off
out[i] = f"{tok}=off"
# A quantized V cache requires flash attention in llama.cpp: the init
# aborts with "V cache quantization requires flash_attn". A quantized K
# cache has no such requirement and runs fine without FA, so it is left
# untouched -- resetting it would needlessly enlarge the K cache and can
# OOM a memory-constrained config. Studio launches with FA on, so a
# quantized --cache-type-v is legal at launch but would make THIS FA-off
# retry crash on init instead of recovering. Reset a quantized V cache --
# main and draft (the draft context shares the global --flash-attn flag,
# so its V cache aborts too) -- to f16 (the llama.cpp default);
# non-quantized types -- f16/bf16/f32 -- run fine without FA and are left
# untouched. The value is rewritten in place so the list length is
# preserved for downstream slices, matching the flash-attn flip above.
_v_cache_flags = (
"--cache-type-v",
"-ctv",
"--cache-type-v-draft",
"--spec-draft-type-v",
"-ctvd",
)
_cache_reset = False
for i, tok in enumerate(out):
# llama.cpp rewrites '_' to '-' for any argv token starting with
# '--' before matching, so a legal pass-through spelling such as
# --cache_type_v parses as --cache-type-v and still enables a
# quantized V cache. Canonicalize the flag name the same way so the
# reset recognizes the underscore aliases too; short flags (-ctv)
# and the type value are left untouched.
name = LlamaCppBackend._canonical_long_flag(tok.partition("=")[0])
if name not in _v_cache_flags:
continue
if "=" in tok:
flag, _, value = tok.partition("=")
if value.strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES:
out[i] = f"{flag}=f16"
_cache_reset = True
elif i + 1 < len(out):
if out[i + 1].strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES:
out[i + 1] = "f16"
_cache_reset = True
if _cache_reset:
logger.info(
"V cache dtype reset to f16 because flash attention was disabled "
"by the crash-recovery fallback (quantized V cache requires flash "
"attention in llama.cpp; the K cache is left untouched)."
)
return out
@staticmethod
def _drop_env_quantized_v_cache(env: MutableMapping[str, str]) -> bool:
"""Drop an inherited quantized V-cache env var (main or draft) in place
before a flash-attn-off retry, returning True if anything was removed.
The argv rewrite in ``_with_flash_attn_off`` only reaches flags on the
command line. Studio deliberately lets an env-only cache type reach the
child untouched (an asymmetric K/V env must survive), so a quantized V
cache set purely through ``LLAMA_ARG_CACHE_TYPE_V`` (or the draft
``LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V``) would still abort the FA-off retry
with "V cache quantization requires flash_attn". Dropping it lets
llama.cpp fall back to the f16 default. Only V is dropped: a quantized K
cache runs fine without flash attention, so its env var is preserved.
"""
dropped = False
for var in ("LLAMA_ARG_CACHE_TYPE_V", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"):
value = (env.get(var) or "").strip().lower()
if value and value not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES:
env.pop(var, None)
dropped = True
return dropped
@staticmethod
def _strip_mmproj_args(cmd: list[str]) -> list[str]:
"""Return cmd without the '--mmproj <path>' pair (text-only retry).
@ -8339,6 +8439,13 @@ class LlamaCppBackend:
_fa_rc,
)
self._kill_process()
# The argv rewrite can't reach an env-only quantized V
# cache; drop it so the FA-off child doesn't abort on it.
if self._drop_env_quantized_v_cache(env):
logger.info(
"Dropped inherited quantized V-cache env for the "
"--flash-attn off retry (requires flash attention)."
)
cmd = _fa_cmd
healthy = _spawn_and_wait(_fa_cmd, label = "-noflash")
@ -8384,6 +8491,13 @@ class LlamaCppBackend:
_probe_rc,
)
self._kill_process()
# The argv rewrite can't reach an env-only quantized V
# cache; drop it so the FA-off child doesn't abort on it.
if self._drop_env_quantized_v_cache(env):
logger.info(
"Dropped inherited quantized V-cache env for the "
"--flash-attn off retry (requires flash attention)."
)
cmd = _fa_cmd
healthy = (
_spawn_and_wait(_fa_cmd, label = "-noflash-mtp")
@ -9200,6 +9314,7 @@ class LlamaCppBackend:
"""Terminate the subprocess and cancel any in-flight download."""
self._cancel_event.set()
with self._lock:
self._unload_epoch += 1
self._kill_process()
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
self._model_identifier = None
@ -9999,15 +10114,18 @@ class LlamaCppBackend:
return False
if not self._mtp_runtime_fallback_active:
return False
if not self._last_load_kwargs or self._process is None:
# Read before claiming: a raise after the claim strands the flag, and nothing
# else clears it, blocking every later respawn.
kwargs = self._last_load_kwargs
proc = self._process
if not kwargs or proc is None:
return False
# Single-flight: the first failure claims the reload.
with self._mtp_runtime_fallback_lock:
if self._mtp_runtime_fallback_in_progress:
return False
self._mtp_runtime_fallback_in_progress = True
snapshot = dict(self._last_load_kwargs)
proc = self._process
snapshot = dict(kwargs)
def _recover():
try:
@ -10055,7 +10173,14 @@ class LlamaCppBackend:
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
try:
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
except RuntimeError as exc:
# Release the claim: a reload that never started would block respawn forever.
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
logger.error(f"Could not start the MTP-crash reload: {exc}")
return False
return True
def _start_mtp_crash_watchdog(self) -> None:
@ -10527,6 +10652,21 @@ class LlamaCppBackend:
finally:
_cancel_closed.set()
def _server_socket_is_open(self, timeout_s: float = 0.15) -> bool:
"""True if anything still accepts on the server port.
The listening socket dies with the process, so this tells a live server
from a dead one without waiting for the child to become reapable.
"""
port = self._port
if not port:
return False
try:
with socket.create_connection(("127.0.0.1", port), timeout = timeout_s):
return True
except OSError:
return False
def _respawn_if_dead(self) -> bool:
"""Relaunch the llama-server if its process has exited.
@ -10536,28 +10676,114 @@ class LlamaCppBackend:
recover, returning True once healthy. Serialised on ``_respawn_lock`` so
many generations hitting the dead server trigger at most one reload.
"""
# Read outside the lock so a queued caller can tell the replacement from the child
# its own error came from; otherwise each burns the grace wait below, and that
# sleep is held under the lock, so the waits serialise.
served_by = self._process
with self._respawn_lock:
proc = self._process
if proc is None:
return False
if proc.poll() is None:
# Process is alive: either a concurrent caller already respawned
# it (healthy), or this connection error wasn't a dead server.
if self._cancel_event.is_set():
# unload_model sets this before it kills, so the child can still be
# accepting. Reporting it healthy would aim the retry at a server
# that is deliberately going away.
return False
if proc is not served_by:
# Replaced while we queued: this child never served our request.
return self._healthy
kwargs = self._last_load_kwargs
if not kwargs:
return False
logger.warning(
f"llama-server for '{self._model_identifier}' exited "
f"(code {proc.returncode}); respawning to recover the session"
)
with self._lock:
self._healthy = False
if proc.poll() is None:
# Still serving, so the error was transient. Charging it the grace below
# would cost a second per caller, serialised under this lock.
if self._server_socket_is_open():
return self._healthy
# A closing server can beat its own exit status: calling it alive returns
# the stale _healthy and spends the retry on the corpse.
deadline = time.monotonic() + _RESPAWN_REAP_GRACE_S
while proc.poll() is None and time.monotonic() < deadline:
time.sleep(0.05)
if proc.poll() is None:
# Alive: either a concurrent caller already respawned it (healthy), or
# this connection error wasn't a dead server.
return self._healthy
with self._mtp_runtime_fallback_lock:
if self._mtp_runtime_fallback_in_progress:
# An MTP-free reload owns this corpse; replaying the old kwargs
# restarts the crashing config and aborts that reload.
logger.info("Respawn skipped: an MTP-free reload is already recovering.")
return False
# The RLock lets the load_model below re-enter it.
with self._serial_load_lock:
if self._process is not proc:
logger.info("Respawn skipped: a newer load is already active.")
return self._healthy
# Snapshot under _lock, the one unload_model holds, so a teardown is
# either wholly before us (flag set) or wholly after (epoch bumped).
# _serial_load_lock alone would not exclude it: unload never takes it.
with self._lock:
if self._cancel_event.is_set():
logger.info("Respawn skipped: the model was unloaded.")
return False
kwargs = dict(self._last_load_kwargs or {})
if not kwargs:
return False
epoch = self._unload_epoch
self._healthy = False
logger.warning(
f"llama-server for '{self._model_identifier}' exited "
f"(code {proc.returncode}); respawning to recover the session"
)
try:
started = bool(self.load_model(**kwargs))
except Exception as exc:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
if started and self._unload_epoch != epoch:
# An unload landed mid-reload. load_model cleared _cancel_event on
# the way in, so the epoch is the only surviving evidence; undo the
# replacement rather than leave a model the user stopped running.
logger.info("Respawn undone: the model was unloaded during the reload.")
self.unload_model()
return False
return started
@contextlib.contextmanager
def _open_chat_stream_with_respawn_retry(self, payload: dict, cancel_event):
"""Open a chat stream, respawning a dead llama-server once before streaming.
Retry only when opening the response fails: once it is open a consumer may
already have emitted content or tool events, so a replay could duplicate
output and side effects. ``base_url`` is resolved per attempt because a
respawn may pick a new port. The budget is one retry per model request, not
per chat turn, so a long tool loop never discards a completed tool.
A child dying after the accept but before the headers surfaces as
ReadError/WriteError/RemoteProtocolError rather than ConnectError, and which
one differs per OS. llama-server flushes its 200 at slot start, so that window
is an upload still in flight or a request behind busy slots; a death during
decode arrives with the response open and is not replayed. Timeouts are
excluded: the server is slow, not dead, and a replay would spend the
first-token budget twice.
"""
for attempt in range(2):
response_opened = False
try:
return bool(self.load_model(**kwargs))
except Exception as exc:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
url = f"{self.base_url}/v1/chat/completions"
with self._open_stream(url, payload, cancel_event) as opened:
response_opened = True
yield opened
return
except (httpx.NetworkError, httpx.RemoteProtocolError) as exc:
if response_opened:
raise
if self._maybe_recover_from_mtp_crash(exc):
raise RuntimeError("Lost connection to llama-server") from exc
if attempt == 0 and self._respawn_if_dead():
logger.warning(
"llama-server was unreachable; respawned it and retrying the generation"
)
continue
raise
def generate_chat_completion(
self,
@ -10855,7 +11081,6 @@ class LlamaCppBackend:
yield _ev
conversation.extend(_auto["messages"])
url = f"{self.base_url}/v1/chat/completions"
_accumulated_completion_tokens = 0
_accumulated_predicted_ms = 0.0
_accumulated_predicted_n = 0
@ -11115,7 +11340,7 @@ class LlamaCppBackend:
_text_args_name = ""
_confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions
with self._open_stream(url, payload, cancel_event) as (
with self._open_chat_stream_with_respawn_retry(payload, cancel_event) as (
response,
first_token_deadline,
):
@ -12152,7 +12377,7 @@ class LlamaCppBackend:
_stream_done = False
try:
with self._open_stream(url, stream_payload, cancel_event) as (
with self._open_chat_stream_with_respawn_retry(stream_payload, cancel_event) as (
response,
first_token_deadline,
):

22
studio/backend/core/inference/tools.py
View file

@ -49,6 +49,7 @@ from loggers import get_logger
logger = get_logger(__name__)
_EXEC_TIMEOUT = 300 # 5 minutes
_DISABLE_DNS_PINNING_ENV = "UNSLOTH_STUDIO_DISABLE_DNS_PINNING"
# Splits the UI source-map from the result; loops strip it (like __IMAGES__).
RAG_SOURCES_SENTINEL = "\n__RAG_SOURCES__:"
@ -4194,13 +4195,18 @@ def _fetch_url_raw(
budget_error = _fetch_budget_exceeded(deadline, cancel_event)
if budget_error is not None:
return budget_error, "", ""
# Pin to the validated IP (prevents DNS rebinding): rewrite URL to
# the IP, set the Host header.
cp = urlparse(current_url)
# Bracket IPv6 addresses so the netloc is valid in a URL.
ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip
ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str
pinned_url = urlunparse(cp._replace(netloc = ip_netloc))
validated_netloc = f"[{current_host}]" if ":" in current_host else current_host
if cp.port:
validated_netloc = f"{validated_netloc}:{cp.port}"
if os.environ.get(_DISABLE_DNS_PINNING_ENV) == "1":
# Enterprise proxies need the hostname in CONNECT for policy and TLS interception.
request_url = urlunparse(cp._replace(netloc = validated_netloc))
else:
# Pin to the validated IP to prevent DNS rebinding.
ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip
ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str
request_url = urlunparse(cp._replace(netloc = ip_netloc))
opener = urllib.request.build_opener(
_NoRedirect,
@ -4209,11 +4215,11 @@ def _fetch_url_raw(
headers = {
"User-Agent": ua,
"Host": current_host,
"Host": validated_netloc,
}
if extra_headers:
headers.update(extra_headers)
req = urllib.request.Request(pinned_url, headers = headers)
req = urllib.request.Request(request_url, headers = headers)
try:
# Cap the socket timeout at the time left on the overall deadline
# so a single slow hop cannot outlast the whole fetch budget.

17
studio/backend/core/training/worker.py
View file

@ -764,8 +764,8 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
- ``gcn_arch``: canonical arch string (e.g. ``"gfx1151"``) when a known
attribute is present, else ``""``.
- ``is_unified``: ``True`` for AMD APUs with a shared GPU/system-RAM pool
(gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower
``set_per_process_memory_fraction`` cap to leave OS headroom.
(gfx1150 Strix Point, gfx1151 Strix Halo, gfx1152 Krackan Point) — these
need a lower ``set_per_process_memory_fraction`` cap to leave OS headroom.
Classification priority:
1. ``props.is_integrated`` truthy (hipDeviceProp_t.integrated -- the
@ -778,6 +778,7 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
- gfx1151 Strix Halo / Gorgon Halo: ``Radeon 8065S`` (Ryzen AI
Max+ 495), ``Radeon 8060S`` (Ryzen AI MAX+
395), ``Radeon 8050S`` (cut-down SKU)
- gfx1152 Krackan Point: ``Radeon 860M``, ``Radeon 840M``
"""
gcn_arch = ""
for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"):
@ -797,9 +798,13 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
return gcn_arch, True
if gcn_arch:
return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"}
# gfx1152 is Krackan Point, the third RDNA 3.5 APU: same shared
# GPU/system-RAM pool as Strix Point (gfx1150) and Strix Halo (gfx1151).
return gcn_arch, gcn_arch in {"gfx1150", "gfx1151", "gfx1152"}
# Arch attrs absent — fall back to device-name matching.
# Arch attrs absent — fall back to device-name matching. Only reached under
# _hw.IS_ROCM, so the NVIDIA GeForce 840M cannot collide with the Krackan
# markers here.
dev_lower = (getattr(props, "name", "") or "").lower()
is_unified = (
"890m" in dev_lower
@ -807,6 +812,8 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
or "8065s" in dev_lower
or "8060s" in dev_lower
or "8050s" in dev_lower
or "860m" in dev_lower
or "840m" in dev_lower
)
return gcn_arch, is_unified
@ -2828,7 +2835,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# On ROCm, exhausting VRAM can hang the HIP driver instead of raising.
# set_per_process_memory_fraction caps the allocator so PyTorch raises
# OutOfMemoryError first (NVIDIA already has a graceful OOM path).
# Unified-memory APUs (gfx1150/gfx1151) share GPU+system RAM, so use 0.80
# Unified-memory APUs (gfx1150/gfx1151/gfx1152) share GPU+system RAM, so use 0.80
# vs 0.90 for discrete. Classify via gcnArchName, else device-name markers.
# Non-fatal: skipped if torch is not importable.
if _hw.IS_ROCM:

9
studio/backend/routes/rag.py
View file

@ -47,7 +47,14 @@ _SAFE = re.compile(r"[^A-Za-z0-9._-]+")
def _sanitize_filename(name: str) -> str:
base = os.path.basename(name or "").strip() or "document"
base = _SAFE.sub("_", base)
return base[:200]
if len(base) <= 200:
return base
# Trim the stem, not the extension: _save_upload gates on the extension, so
# a plain truncation would reject a long-named .txt as "unsupported".
stem, ext = os.path.splitext(base)
if not ext or len(ext) > 32:
return base[:200]
return stem[: 200 - len(ext)] + ext
def _save_upload(file: UploadFile) -> tuple[str, str]:

93
studio/backend/run.py
View file

@ -991,10 +991,88 @@ class _TeeStream:
except Exception:
pass
def close(self):
# We do NOT own the console stream (it is the terminal / Jupyter kernel
# stream we wrapped), so closing the tee must never take the server down.
# Flush the log copy, then forward close() to the wrapped stream
# best-effort: on Colab that stream is an ipykernel OutStream whose
# close() can raise (see _harden_console_close / ipython/ipykernel#867).
try:
self._log_fh.flush()
except Exception:
pass
try:
self._stream.close()
except Exception:
pass
def __getattr__(self, name):
return getattr(self._stream, name)
_WATCH_FD_THREAD_ATTR = "watch_fd_thread"
def _is_missing_watch_fd_thread(exc):
"""True only for ipython/ipykernel#867's missing-``watch_fd_thread`` error.
``AttributeError.name`` exists from Python 3.10; the message carries the
attribute name on every version (possibly with a "Did you mean" tail), so
check both and let every other AttributeError through.
"""
if getattr(exc, "name", None) == _WATCH_FD_THREAD_ATTR:
return True
return _WATCH_FD_THREAD_ATTR in str(exc)
def _harden_console_close(stream):
"""Stop a displaced console stream's close() from aborting Studio startup.
``_setup_server_disk_logging`` replaces ``sys.stdout``/``sys.stderr`` with a
tee. That changes the object identity of the console stream, so a third-party
logging handler that captured the ORIGINAL stream (notably Colab's ``absl``
logging handler, whose ``close()`` skips ``sys.stdout``/``sys.stderr`` but not
a stream that is no longer either) treats it as an ordinary stream and calls
``close()`` on it during logging teardown -- ``uvicorn.Config()`` ->
``logging.config.dictConfig()`` -> ``logging.shutdown()``.
A Jupyter/Colab ``ipykernel`` ``OutStream`` created with ``watchfd=False``
(the Colab default, and every in-process kernel) never gains a
``watch_fd_thread``, yet the ``OutStream.close()`` shipped in the affected
ipykernel versions joins that thread unconditionally and raises
``AttributeError: 'OutStream' object has no attribute 'watch_fd_thread'``
(ipython/ipykernel#867). That AttributeError propagates out of
``uvicorn.Config(...)`` and aborts startup ("Unsloth Studio failed to start").
Wrap the stream's ``close()`` in a transparent pass-through that swallows
ONLY that specific teardown AttributeError. A healthy close() (a real console
stream, or an OutStream with fd-watching on) runs to completion exactly as
before and any other error still propagates, so nothing changes off Colab. A
stream whose ``close`` cannot be reassigned keeps its original close().
"""
try:
_orig_close = stream.close
except Exception:
return
def _safe_close(*args, **kwargs):
try:
return _orig_close(*args, **kwargs)
except AttributeError as exc:
if not _is_missing_watch_fd_thread(exc):
# A real teardown failure; never hide it.
raise
# ipython/ipykernel#867: watchfd=False OutStream.close() joins a
# thread that was never created. Nothing to clean up; keep going.
return None
try:
stream.close = _safe_close
except (AttributeError, TypeError):
# A stream that forbids setting instance attributes; leave it as-is.
pass
def _setup_server_disk_logging():
"""Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim
faulthandler at the same file so hard crashes (access violations /
@ -1037,6 +1115,11 @@ def _setup_server_disk_logging():
# the stderr the server already captures.
os.environ.setdefault("PYTHONFAULTHANDLER", "1")
# Replacing the console streams orphans them from third-party "is this the
# live console?" checks, so guard their close() first (ipython/ipykernel#867).
_harden_console_close(sys.stdout)
_harden_console_close(sys.stderr)
sys.stdout = _TeeStream(sys.stdout, log_fh)
sys.stderr = _TeeStream(sys.stderr, log_fh)
@ -1802,6 +1885,12 @@ def _build_arg_parser():
default = None,
help = "Force server-side tools off for every request.",
)
parser.add_argument(
"--disable-dns-pinning",
action = "store_true",
help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens "
"DNS-rebinding protection; hostname and redirect validation remain enabled.",
)
parser.add_argument(
"--parallel",
"--n-parallel",
@ -1841,6 +1930,10 @@ if __name__ == "__main__":
parser.error(
"--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare"
)
if args.disable_dns_pinning:
os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1"
else:
os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0")
kwargs = dict(
host = args.host,

4
studio/backend/tests/test_amd_apu_unified_memory.py
View file

@ -2,7 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs
(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS."""
(gfx1150/gfx1151/gfx1152), never for discrete AMD, NVIDIA, CPU or macOS."""
from __future__ import annotations
@ -35,6 +35,8 @@ def _fake_torch(
[
("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped)
("6.2.0", ["gfx1150"], True), # Strix Point APU
("6.2.0", ["gfx1152"], True), # Krackan Point APU (Radeon 860M/840M)
("6.2.0", ["gfx1152:sramecc-:xnack-"], True), # same, feature flags stripped
("6.2.0", ["gfx1100"], False), # discrete RDNA3
("6.2.0", ["gfx1201"], False), # discrete RDNA4
("6.2.0", ["gfx942"], False), # MI300X (data center)

125
studio/backend/tests/test_colab_embed.py
View file

@ -336,17 +336,71 @@ def test_colab_login_html_includes_credentials():
html = colab._colab_login_html("unsloth", "alpha-beta-gamma-delta")
assert "unsloth" in html
assert "alpha-beta-gamma-delta" in html
# The username is fixed, so it reads inline rather than as its own field.
assert "Username:" not in html
def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch):
def test_shareable_link_html_embeds_password_under_the_link():
"""The credential belongs in the same card as the button it unlocks."""
html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth")
assert "share.trycloudflare.com" in html
assert "secret-pass" in html
# Username is stated inline, not as its own labelled field.
assert "Username:" not in html
assert "unsloth" in html
# The password must sit after the link, not above it.
assert html.index("share.trycloudflare.com") < html.index("secret-pass")
def test_shareable_link_html_renders_the_url_as_a_link():
"""The printed URL is an anchor, using the popup-safe open the button uses."""
html = colab._shareable_link_html("https://share.trycloudflare.com")
assert '<a href="https://share.trycloudflare.com"' in html
assert ">https://share.trycloudflare.com</a>" in html
assert html.count("window.open(this.href,'_blank')") == 2
def test_shareable_link_html_emphasises_the_password():
"""The password is the one thing to copy, so it is enlarged and underlined."""
html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth")
pw_tag = html[html.index("Password") : html.index("secret-pass")]
assert "font-size: 24px" in pw_tag
assert "text-decoration: underline" in pw_tag
def test_shareable_link_html_password_has_no_adjacent_whitespace():
"""Whitespace beside the password is selected with it on a double click."""
html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth")
before, after = html.split("secret-pass", 1)
assert before.endswith(">")
assert after.startswith("<")
# Label on its own line, so nothing shares the password's text node.
assert "Password:" not in html
# Plain selectable text: user-select overrides break double click to select.
assert "user-select" not in html
def test_shareable_link_html_omits_login_block_without_password():
html = colab._shareable_link_html("https://share.trycloudflare.com")
assert "Password" not in html
def test_show_and_embed_folds_login_into_the_cloudflare_card(monkeypatch):
"""One card, not two: the tunnel card carries the password itself."""
displayed: list[str] = []
ipython_display = SimpleNamespace(
HTML = lambda html: SimpleNamespace(html = html),
display = lambda html: displayed.append(html.html),
)
login_cards: list[tuple] = []
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
monkeypatch.setattr(
colab,
"_show_colab_login_credentials",
lambda *args: login_cards.append(args),
)
monkeypatch.setattr(
colab,
"show_link",
@ -360,9 +414,74 @@ def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch):
colab_login = ("unsloth", "secret-pass"),
)
assert len(displayed) == 2
assert len(displayed) == 1
assert "share.trycloudflare.com" in displayed[0]
assert "secret-pass" in displayed[1]
assert "secret-pass" in displayed[0]
assert login_cards == []
def test_show_and_embed_keeps_separate_login_card_without_tunnel(monkeypatch):
"""No tunnel card to fold into, so the standalone login card still renders."""
login_cards: list[tuple] = []
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
monkeypatch.setattr(
colab,
"_show_colab_login_credentials",
lambda *args: login_cards.append(args),
)
monkeypatch.setattr(
colab,
"show_link",
lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None,
)
monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True)
colab._show_and_embed(8888, colab_login = ("unsloth", "secret-pass"))
assert login_cards == [("unsloth", "secret-pass")]
def test_show_and_embed_skips_ready_card_when_tunnel_is_up(monkeypatch):
"""The ready card only restates the tunnel card and prints a proxy URL that 404s."""
calls: list[str] = []
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
monkeypatch.setattr(
colab,
"show_link",
lambda port,
*,
_url = None,
has_cloudflare_link = False,
cloudflare_requested = False: calls.append("show_link"),
)
monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True)
colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com")
assert calls == []
def test_show_and_embed_keeps_ready_card_without_tunnel(monkeypatch):
"""Without a tunnel the ready card is the only guidance, so it must stay."""
calls: list[str] = []
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
monkeypatch.setattr(
colab,
"show_link",
lambda port,
*,
_url = None,
has_cloudflare_link = False,
cloudflare_requested = False: calls.append("show_link"),
)
monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True)
colab._show_and_embed(8888)
assert calls == ["show_link"]
def test_show_and_embed_skips_iframe_on_colab_when_cloudflare_ready(monkeypatch):

242
studio/backend/tests/test_export_multi_gpu_device_map.py Normal file
View file

@ -0,0 +1,242 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Export checkpoint loading must shard across every visible GPU (#7053): the
``device_map="sequential"`` loader default stacks the whole model on GPU0 and OOMs
while the other GPUs sit empty. The loader now passes ``device_map="balanced"``, but
only on a real multi-GPU CUDA/ROCm host, so single-GPU, CPU and MLX are untouched."""
from __future__ import annotations
import contextlib
import sys
import types
from pathlib import Path
_BACKEND_DIR = Path(__file__).resolve().parent.parent
if str(_BACKEND_DIR) not in sys.path:
sys.path.insert(0, str(_BACKEND_DIR))
_TESTS_DIR = Path(__file__).resolve().parent
if str(_TESTS_DIR) not in sys.path:
sys.path.insert(0, str(_TESTS_DIR))
# Reuse the absolute-paths test's stub harness for loading core/export/export.py
# without torch/unsloth.
from test_export_absolute_paths import ( # noqa: E402
_install_export_backend_stubs,
_load_module,
)
def _export_mod(monkeypatch):
_install_export_backend_stubs(monkeypatch)
return _load_module("test_core_export_backend_device_map", "core/export/export.py", monkeypatch)
def _stub_hardware(monkeypatch, visible, device_map):
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: visible, raising = False)
monkeypatch.setattr(hw, "get_device_map", lambda ids: device_map, raising = False)
# ── _multi_gpu_device_map_kwargs ──
def test_multi_gpu_host_gets_balanced(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0, 1, 2], "balanced")
assert mod._multi_gpu_device_map_kwargs() == {"device_map": "balanced"}
def test_single_gpu_host_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0], "sequential")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_non_balanced_resolution_keeps_loader_default(monkeypatch):
# >1 visible id but a non-CUDA device resolves to "sequential": pass nothing.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0, 1], "sequential")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_uuid_mig_mask_falls_back_to_count_detection(monkeypatch):
# UUID/MIG masks resolve to NO numeric ids ([]), but get_device_map(None) still
# detects >1 GPU, so the empty list must route there, not to the loader default.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [], raising = False)
monkeypatch.setattr(
hw,
"get_device_map",
lambda ids: "balanced" if ids is None else "sequential",
raising = False,
)
assert mod._multi_gpu_device_map_kwargs() == {"device_map": "balanced"}
def test_no_visible_gpus_keeps_loader_default(monkeypatch):
# Empty mask / CPU host: get_device_map(None) resolves "sequential" -> {}.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [], raising = False)
monkeypatch.setattr(hw, "get_device_map", lambda ids: "sequential", raising = False)
assert mod._multi_gpu_device_map_kwargs() == {}
def test_mlx_host_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
# The stubs set _IS_MLX = True; even a multi-GPU view must yield no device_map.
_stub_hardware(monkeypatch, [0, 1], "balanced")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_hardware_probe_failure_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
def _boom():
raise RuntimeError("no GPUs")
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", _boom, raising = False)
assert mod._multi_gpu_device_map_kwargs() == {}
# ── load_checkpoint forwards the kwargs to from_pretrained ──
class _RecordingLoader:
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
return types.SimpleNamespace(), types.SimpleNamespace()
def _load_text_checkpoint(monkeypatch, tmp_path, device_map_kwargs):
mod = _export_mod(monkeypatch)
_RecordingLoader.calls = []
monkeypatch.setattr(mod, "FastLanguageModel", _RecordingLoader)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: device_map_kwargs)
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
assert ok, message
assert len(_RecordingLoader.calls) == 1
return _RecordingLoader.calls[0]
def test_load_checkpoint_forwards_balanced_device_map(monkeypatch, tmp_path):
kwargs = _load_text_checkpoint(monkeypatch, tmp_path, {"device_map": "balanced"})
assert kwargs["device_map"] == "balanced"
def test_load_checkpoint_omits_device_map_on_single_gpu(monkeypatch, tmp_path):
kwargs = _load_text_checkpoint(monkeypatch, tmp_path, {})
assert "device_map" not in kwargs # loader default (sequential) untouched
# ── a load that succeeds but offloads to CPU/disk ──
def test_cpu_offloaded_modules_counts_cpu_and_disk(monkeypatch):
mod = _export_mod(monkeypatch)
model = types.SimpleNamespace(hf_device_map = {"a": 0, "b": "cpu", "c": 1, "d": "disk"})
assert mod._cpu_offloaded_modules(model) == 2
def test_cpu_offloaded_modules_ignores_gpu_only_and_missing_maps(monkeypatch):
mod = _export_mod(monkeypatch)
assert mod._cpu_offloaded_modules(types.SimpleNamespace(hf_device_map = {"a": 0})) == 0
assert mod._cpu_offloaded_modules(types.SimpleNamespace(hf_device_map = None)) == 0
assert mod._cpu_offloaded_modules(types.SimpleNamespace()) == 0
class _SpillThenCleanLoader:
"""First call offloads to CPU (bf16 accepts it silently), second is clean."""
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
device_map = {"model.layers.0": 0} if len(cls.calls) > 1 else {"model.layers.0": "cpu"}
return types.SimpleNamespace(hf_device_map = device_map), types.SimpleNamespace()
def _run_spill_loader(monkeypatch, tmp_path, device_map_kwargs):
mod = _export_mod(monkeypatch)
_SpillThenCleanLoader.calls = []
monkeypatch.setattr(mod, "FastLanguageModel", _SpillThenCleanLoader)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: device_map_kwargs)
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
return ok, message, _SpillThenCleanLoader.calls
def test_successful_load_that_offloads_to_cpu_retries_single_device(monkeypatch, tmp_path):
# Nothing raises, so only hf_device_map catches it; the parameters would otherwise
# stay on meta and kill the export inside safetensors.
ok, message, calls = _run_spill_loader(monkeypatch, tmp_path, {"device_map": "balanced"})
assert ok, message
assert len(calls) == 2
assert calls[0]["device_map"] == "balanced"
assert "device_map" not in calls[1]
def test_single_gpu_offload_is_left_alone(monkeypatch, tmp_path):
# No multi-GPU map was requested, so there is nothing to retry on.
ok, message, calls = _run_spill_loader(monkeypatch, tmp_path, {})
assert ok, message
assert len(calls) == 1
def test_retry_result_is_kept_even_if_it_also_offloads(monkeypatch, tmp_path):
# The retry runs with _device_map_override set, so it must never recurse again.
mod = _export_mod(monkeypatch)
class _AlwaysSpills:
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
return types.SimpleNamespace(hf_device_map = {"a": "cpu"}), types.SimpleNamespace()
monkeypatch.setattr(mod, "FastLanguageModel", _AlwaysSpills)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: {"device_map": "balanced"})
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
assert ok, message
assert len(_AlwaysSpills.calls) == 2

138
studio/backend/tests/test_gguf_load_cache_reuse.py
View file

@ -9,10 +9,14 @@ No GPU, network, or subprocesses are required.
from __future__ import annotations
import asyncio
import importlib.util
import logging
import sys
import threading
import types as _types
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import pytest
@ -28,7 +32,12 @@ _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
# routes/inference.py binds structlog.get_logger at import time, and setdefault
# keeps a bare stub an earlier test left behind: repair it rather than rely on order.
_structlog_stub.get_logger = lambda *_args, **_kwargs: logging.getLogger("structlog_stub")
sys.modules.setdefault("structlog", _structlog_stub)
if not hasattr(sys.modules["structlog"], "get_logger"):
sys.modules["structlog"].get_logger = _structlog_stub.get_logger
try:
import httpx # noqa: F401
@ -120,6 +129,22 @@ def _fail_get_paths_info(*_args, **_kwargs):
raise AssertionError("cached reuse must return before the sizing preflight")
def _load_route_module(name: str, relative_path: str):
"""Import a route module under a private name so patches can't leak."""
spec = importlib.util.spec_from_file_location(name, Path(_BACKEND_DIR) / relative_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
async def _inline_to_thread(func, /, *args, **kwargs):
return func(*args, **kwargs)
async def _no_gguf_gpu_ids(*_args, **_kwargs):
return None
class TestLoadReusesCachedCopy:
def test_download_uses_selected_cache_for_lookup_preflight_and_write(
self, tmp_path, monkeypatch
@ -809,3 +834,116 @@ class TestLoadHubDownloadExclusion:
Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py"
).read_text()
assert "@_with_gguf_load_marker\n def load_model(" in llama_source
def _capture_hub_guard_require_mmproj(
self,
stored_extra_args,
request_extra_args = None,
):
"""Drive /load's GGUF path and return the hub guard's require_mmproj.
The guard reports a conflicting download, so the 409 is the observation
point and no llama-server ever starts.
"""
import core.inference.llama_cpp as llama_cpp_module
from fastapi import HTTPException
from models.inference import LoadRequest
route = _load_route_module(
"inference_route_module_for_inherited_extra_args_test",
"routes/inference.py",
)
captured = {}
def _fake_blocks(
repo,
variant,
*,
require_mmproj,
hf_token = None,
):
captured["repo"] = repo
captured["variant"] = variant
captured["require_mmproj"] = require_mmproj
return True
# A vision GGUF: require_mmproj is True unless the extras say --no-mmproj.
config = SimpleNamespace(
is_gguf = True,
is_lora = False,
is_vision = True,
is_audio = False,
audio_type = None,
has_audio_input = False,
gguf_hf_repo = REPO,
gguf_variant = VARIANT,
gguf_file = None,
gguf_mmproj_file = None,
identifier = REPO,
display_name = REPO,
)
# Pass-through extras the running backend recorded for the last load.
llama_backend = SimpleNamespace(
is_loaded = False,
extra_args = list(stored_extra_args),
extra_args_source = (REPO, VARIANT),
hf_variant = VARIANT,
model_identifier = REPO,
)
request = LoadRequest(
model_path = REPO,
gguf_variant = VARIANT,
llama_extra_args = request_extra_args,
)
with (
patch.object(
route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: config),
),
patch.object(route, "get_llama_cpp_backend", lambda: llama_backend),
patch.object(
route,
"get_inference_backend",
lambda: SimpleNamespace(active_model_name = None),
),
patch.object(route, "_resolve_gguf_gpu_ids_for_request", _no_gguf_gpu_ids),
patch.object(route, "_guard_chat_load_against_training", return_value = None),
patch.object(route, "_effective_load_in_4bit", return_value = False),
patch.object(route, "_hf_offline_if_dns_dead", nullcontext),
patch.object(route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(llama_cpp_module, "_hub_download_blocks_gguf_load", _fake_blocks),
):
with pytest.raises(HTTPException) as exc_info:
asyncio.run(
route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(
state = SimpleNamespace(llama_parallel_slots = 1),
),
),
current_subject = "test-user",
)
)
assert exc_info.value.status_code == 409
assert captured["repo"] == REPO
return captured["require_mmproj"]
def test_inherited_extra_args_shape_hub_guard_require_mmproj(self):
# Inheritance must resolve before the hub-download guard: an inherited
# --no-mmproj decides require_mmproj, so resolving later rejects a load
# over a download the effective arguments disable (#7251).
assert self._capture_hub_guard_require_mmproj(["--no-mmproj"]) is False
# Control: nothing to inherit, so a vision GGUF still needs its mmproj.
assert self._capture_hub_guard_require_mmproj([]) is True
# An explicit request list wins over the stored one, both ways.
assert (
self._capture_hub_guard_require_mmproj([], request_extra_args = ["--no-mmproj"]) is False
)
assert (
self._capture_hub_guard_require_mmproj(["--no-mmproj"], request_extra_args = []) is True
)

418
studio/backend/tests/test_grouped_mm_rdna4_fallback.py Normal file
View file

@ -0,0 +1,418 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Numerics + gating for the RDNA4 _grouped_mm CPU fallback (PRs #7276 / #7292).
RDNA4 (gfx1200/gfx1201) ships a null HIP `_grouped_mm` kernel on ROCm <= 7.12
(fixed in 7.13, ROCm/TheRock #5284). Training MoE models there crashes with
0xC0000005 on Windows and a plain segfault on Linux, so worker.py registers a
Python mm/bmm fallback on the CUDA dispatch key.
The fallback is silent, GPU-gated, and reimplements a matmul: if it is wrong, an
RX 9070 user does not crash, they train on quietly wrong gradients. Until now the
only coverage was `assert '_gm_lib.impl("_grouped_mm"' in source` -- the math was
never executed once, in any suite.
worker.py cannot be imported here (module-level structlog/backend imports), so
`_install_grouped_mm_cpu_fallback` is lifted out with ast and driven with a fake
`torch_mod` that forwards to real CPU torch. That also pins the op surface: the
fallback may only use the ops the fake exposes, and the registration is captured
instead of hitting a real CUDA dispatch key that CI runners do not have.
The two gates around it are exec'd straight out of the source so this file tests
the shipped expressions rather than a copy of them.
"""
import ast
import re
import textwrap
from pathlib import Path
from types import SimpleNamespace
import pytest
import torch
_WORKER_PATH = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py"
_WORKER_SOURCE = _WORKER_PATH.read_text(encoding = "utf-8")
def _load_installer():
"""exec just _install_grouped_mm_cpu_fallback out of worker.py."""
tree = ast.parse(_WORKER_SOURCE)
fn = [
n
for n in tree.body
if isinstance(n, ast.FunctionDef) and n.name == "_install_grouped_mm_cpu_fallback"
]
assert fn, "_install_grouped_mm_cpu_fallback not found in core/training/worker.py"
ns: dict = {}
exec(compile(ast.Module(body = fn, type_ignores = []), str(_WORKER_PATH), "exec"), ns)
return ns["_install_grouped_mm_cpu_fallback"]
_install_grouped_mm_cpu_fallback = _load_installer()
class _RecordingLibrary:
"""Stands in for torch.library.Library: captures the registration instead of
binding it to a CUDA dispatch key no CI runner has."""
def __init__(self, namespace, kind):
self.namespace = namespace
self.kind = kind
self.registrations = []
def impl(self, name, fn, dispatch_key):
self.registrations.append((name, fn, dispatch_key))
class _RecordingLogger:
def __init__(self):
self.info_calls = []
self.warning_calls = []
def info(self, *args, **kwargs):
self.info_calls.append(args)
def warning(self, *args, **kwargs):
self.warning_calls.append(args)
def _fake_torch():
"""Real CPU torch behind the exact op surface the fallback is allowed to use.
Anything else the fallback reaches for raises AttributeError here, which is
the point: a new dependency has to be a deliberate edit, not a silent one."""
return SimpleNamespace(
library = SimpleNamespace(Library = _RecordingLibrary),
mm = torch.mm,
bmm = torch.bmm,
matmul = torch.matmul,
cat = torch.cat,
zeros = torch.zeros,
)
@pytest.fixture
def fallback():
"""The registered _grouped_mm implementation, plus the Library it landed on."""
torch_mod = _fake_torch()
logger = _RecordingLogger()
lib = _install_grouped_mm_cpu_fallback(torch_mod, logger, "test")
assert lib.registrations, "the fallback registered nothing"
name, fn, key = lib.registrations[0]
return SimpleNamespace(fn = fn, lib = lib, logger = logger, name = name, key = key)
class TestRegistration:
"""Where the override lands. Getting the namespace or dispatch key wrong is a
silent no-op: training still crashes on the null HIP kernel."""
def test_overrides_aten_grouped_mm_on_the_cuda_key(self, fallback):
assert fallback.lib.namespace == "aten"
assert fallback.lib.kind == "IMPL"
assert fallback.name == "_grouped_mm"
# ROCm dispatches through the CUDA key; "HIP"/"PrivateUse1" would not bind.
assert fallback.key == "CUDA"
def test_registers_exactly_once(self, fallback):
assert len(fallback.lib.registrations) == 1
def test_returns_the_library_so_the_caller_can_keep_it_alive(self, fallback):
"""A dropped Library is garbage collected and the override silently
unregisters mid-run; worker.py parks it in a module global."""
assert isinstance(fallback.lib, _RecordingLibrary)
assert "_WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(" in _WORKER_SOURCE
def test_logs_the_patch_with_its_label(self, fallback):
assert fallback.logger.info_calls, "the patch must be visible in the run log"
assert "test" in fallback.logger.info_calls[0]
class TestUngroupedNumerics:
"""offs=None: plain matmul, one path per rank combination. The 3-D case is
the regression #7292 fixed -- an unconditional mm() broke MoE experts."""
def test_2d_by_2d_matches_mm(self, fallback):
a = torch.randn(6, 4)
b = torch.randn(4, 5)
torch.testing.assert_close(fallback.fn(a, b), torch.mm(a, b))
def test_3d_by_3d_matches_bmm(self, fallback):
a = torch.randn(3, 6, 4)
b = torch.randn(3, 4, 5)
torch.testing.assert_close(fallback.fn(a, b), torch.bmm(a, b))
def test_3d_by_2d_matches_matmul(self, fallback):
a = torch.randn(3, 6, 4)
b = torch.randn(4, 5)
torch.testing.assert_close(fallback.fn(a, b), torch.matmul(a, b))
def test_2d_by_3d_matches_matmul(self, fallback):
a = torch.randn(6, 4)
b = torch.randn(3, 4, 5)
torch.testing.assert_close(fallback.fn(a, b), torch.matmul(a, b))
def test_non_contiguous_inputs_are_handled(self, fallback):
"""Transposed views reach _grouped_mm constantly; every path calls
.contiguous() and this catches it if one stops."""
a = torch.randn(4, 6).t()
b = torch.randn(5, 4).t()
torch.testing.assert_close(fallback.fn(a, b), torch.mm(a, b))
class TestGroupedNumerics:
"""offs=[end-row of each group], the MoE token-routing layout."""
def test_matches_per_group_mm_with_3d_weights(self, fallback):
a = torch.randn(7, 4)
b = torch.randn(3, 4, 5)
offs = torch.tensor([2, 5, 7])
expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[1], a[5:7] @ b[2]], dim = 0)
torch.testing.assert_close(fallback.fn(a, b, offs), expected)
def test_shared_2d_weight_is_reused_for_every_group(self, fallback):
a = torch.randn(7, 4)
b = torch.randn(4, 5)
offs = torch.tensor([2, 5, 7])
torch.testing.assert_close(fallback.fn(a, b, offs), a @ b)
def test_empty_group_produces_no_rows(self, fallback):
"""An expert that routed zero tokens (offs[i] == offs[i-1]) must
contribute nothing, not a stray row."""
a = torch.randn(5, 4)
b = torch.randn(3, 4, 5)
offs = torch.tensor([2, 2, 5])
expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[2]], dim = 0)
got = fallback.fn(a, b, offs)
assert got.shape == (5, 5)
torch.testing.assert_close(got, expected)
def test_rows_past_the_last_offset_are_not_dropped(self, fallback):
"""Trailing tokens beyond offs[-1] go through the last expert; dropping
them would silently shrink the output instead of raising."""
a = torch.randn(7, 4)
b = torch.randn(3, 4, 5)
offs = torch.tensor([2, 5])
expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[1], a[5:7] @ b[-1]], dim = 0)
got = fallback.fn(a, b, offs)
assert got.shape[0] == a.shape[0]
torch.testing.assert_close(got, expected)
def test_zero_rows_returns_an_empty_result_not_an_error(self, fallback):
a = torch.randn(0, 4)
b = torch.randn(3, 4, 5)
offs = torch.tensor([], dtype = torch.int64)
got = fallback.fn(a, b, offs)
assert got.shape == (0, 5)
assert got.dtype == a.dtype
def test_offsets_may_arrive_as_a_device_tensor_of_any_int_dtype(self, fallback):
a = torch.randn(4, 4)
b = torch.randn(2, 4, 5)
expected = torch.cat([a[0:2] @ b[0], a[2:4] @ b[1]], dim = 0)
for dtype in (torch.int32, torch.int64):
torch.testing.assert_close(
fallback.fn(a, b, torch.tensor([2, 4], dtype = dtype)), expected
)
class TestBiasAndDtype:
def test_bias_is_added(self, fallback):
a = torch.randn(6, 4)
b = torch.randn(4, 5)
bias = torch.randn(5)
torch.testing.assert_close(fallback.fn(a, b, None, bias), torch.mm(a, b) + bias)
def test_bias_is_added_on_the_grouped_path_too(self, fallback):
a = torch.randn(4, 4)
b = torch.randn(2, 4, 5)
bias = torch.randn(5)
offs = torch.tensor([2, 4])
expected = torch.cat([a[0:2] @ b[0], a[2:4] @ b[1]], dim = 0) + bias
torch.testing.assert_close(fallback.fn(a, b, offs, bias), expected)
def test_out_dtype_is_honoured(self, fallback):
a = torch.randn(6, 4)
b = torch.randn(4, 5)
got = fallback.fn(a, b, None, None, torch.float64)
assert got.dtype == torch.float64
torch.testing.assert_close(got, torch.mm(a, b).to(torch.float64))
def test_promotion_from_bias_is_cast_back_to_the_input_dtype(self, fallback):
"""Without the restore, a promoted result changes the autograd dtype
downstream of every MoE layer."""
a = torch.randn(6, 4, dtype = torch.float32)
b = torch.randn(4, 5, dtype = torch.float32)
bias = torch.randn(5, dtype = torch.float64)
got = fallback.fn(a, b, None, bias)
assert got.dtype == torch.float32
def test_out_dtype_wins_over_the_input_dtype_restore(self, fallback):
a = torch.randn(6, 4, dtype = torch.float32)
b = torch.randn(4, 5, dtype = torch.float32)
bias = torch.randn(5, dtype = torch.float64)
got = fallback.fn(a, b, None, bias, torch.float64)
assert got.dtype == torch.float64
def test_bf16_inputs_stay_bf16(self, fallback):
"""The dtype training actually runs in."""
a = torch.randn(6, 4).to(torch.bfloat16)
b = torch.randn(4, 5).to(torch.bfloat16)
got = fallback.fn(a, b)
assert got.dtype == torch.bfloat16
torch.testing.assert_close(got.float(), (a.float() @ b.float()), rtol = 2e-2, atol = 2e-2)
def _exec_source_snippet(anchor: str, last_line: str, **variables):
"""Run a slice of worker.py verbatim, so the gate under test is the shipped
one and not a copy that can drift."""
start = _WORKER_SOURCE.find(anchor)
assert start != -1, f"gate snippet not found in worker.py: {anchor!r}"
start = _WORKER_SOURCE.rfind("\n", 0, start) + 1 # keep the indent for dedent()
end = _WORKER_SOURCE.find(last_line, start)
assert end != -1, f"end of gate snippet not found: {last_line!r}"
snippet = textwrap.dedent(_WORKER_SOURCE[start : end + len(last_line)])
ns = {"re": re, **variables}
exec(compile(snippet, str(_WORKER_PATH), "exec"), ns)
return ns
class TestLinuxHipVersionGate:
"""PR #7292's Linux gate. Too low a floor keeps the slow Python fallback on
fixed ROCm 7.13+; too high reintroduces the segfault on 7.12."""
_ANCHOR = '_m = re.match(r"(\\d+)\\.(\\d+)", _hip_str)'
_LAST = '_hip_lt_713 = "rocmsdk" not in _ver'
def _decide(self, hip_str, version):
ns = _exec_source_snippet(self._ANCHOR, self._LAST, _hip_str = hip_str, _ver = version.lower())
return ns["_hip_lt_713"]
@pytest.mark.parametrize(
"hip_str,version,affected",
[
("7.12.0", "2.10.0+rocm7.12.0", True), # the broken kernel
("7.6.0", "2.9.0+rocm7.6.0", True),
("6.4.0", "2.8.0+rocm6.4.0", True),
("7.13.0", "2.11.0+rocm7.13.0", False), # AMD's fix
("7.14.0", "2.11.0+rocm7.14.0", False),
("8.0.0", "2.12.0+rocm8.0.0", False),
],
)
def test_torch_version_hip_decides_when_present(self, hip_str, version, affected):
assert self._decide(hip_str, version) is affected
@pytest.mark.parametrize(
"version,affected",
[
("2.10.0+rocm7.12.0", True),
("2.11.0+rocm7.13.0", False),
("2.11.0+rocm7.14.0", False),
],
)
def test_falls_back_to_the_rocm_tag_in_torch_version(self, version, affected):
"""AMD SDK / Radeon wheels leave torch.version.hip unset."""
assert self._decide("", version) is affected
def test_unknown_version_is_assumed_affected(self):
"""Fallback is slow but correct; a missed guard is a crash."""
assert self._decide("", "2.9.0+unknown") is True
def test_rocmsdk_wheels_without_a_version_are_assumed_fixed(self):
"""rocmsdk wheels post-date the gfx120X fix."""
assert self._decide("", "2.10.0+rocmsdk20260107") is False
class TestLinuxRdna4NameMatch:
"""The name regex is the fallback when a wheel omits gcnArchName."""
def _pattern(self):
"""Read whatever pattern worker.py currently uses, not a copy of the one
it used when this test was written. Anchoring on the literal pattern text
would make a *widened* regex -- the dangerous edit, since it silently
forces the slow Python fallback onto RDNA3 users -- fail as "moved"
instead of being checked against the cases below."""
m = re.search(r"re\.search\(r\"([^\"]+)\",\s*_lin_name\)", _WORKER_SOURCE)
assert m, "could not locate the RDNA4 device-name regex in worker.py"
return m.group(1)
def test_name_is_lowercased_before_matching(self):
"""The pattern is all-lowercase, so it only works against a lowercased
name. Device names arrive mixed case ("AMD Radeon RX 9070 XT")."""
assert self._pattern() == self._pattern().lower(), "pattern is not all-lowercase"
assert re.search(
r"_lin_name\s*=\s*\(getattr\(_props,\s*\"name\",\s*\"\"\)\s*or\s*\"\"\)\.lower\(\)",
_WORKER_SOURCE,
), "worker.py must lowercase the device name before matching the RDNA4 pattern"
def test_name_match_is_only_a_fallback_when_arch_is_unknown(self):
"""gcnArchName is authoritative when present. Letting the name regex fire
alongside a known arch would misclassify any card whose marketing name
happens to look RDNA4."""
assert re.search(
r"not _lin_arch and re\.search\(r\"[^\"]+\",\s*_lin_name\)", _WORKER_SOURCE
), "the RDNA4 name regex must be guarded by `not _lin_arch`"
@pytest.mark.parametrize(
"name,is_rdna4",
[
("AMD Radeon RX 9070 XT", True),
("AMD Radeon RX 9060 XT", True),
("Radeon RX9070", True),
("AMD Radeon AI PRO R9700", True),
("AMD Radeon RX 7900 XTX", False), # RDNA3, kernel is fine
("AMD Radeon 8060S Graphics", False), # Strix Halo
("AMD Radeon RX 6800 XT", False),
("NVIDIA GeForce RTX 4090", False),
],
)
def test_matches_only_rdna4_cards(self, name, is_rdna4):
assert bool(re.search(self._pattern(), name.lower())) is is_rdna4
class TestLinuxGateStructure:
"""The block is a few hundred lines into run_training_process and can only be
checked structurally; these pin the parts a refactor would quietly drop."""
def _linux_block(self):
start = _WORKER_SOURCE.find("1f-linux")
assert start != -1, "the Linux ROCm gfx120X guard (#7292) is gone from worker.py"
end = _WORKER_SOURCE.find("1g.", start)
assert end != -1
return _WORKER_SOURCE[start:end]
def test_gated_on_linux_and_rocm(self):
block = self._linux_block()
assert 'sys.platform.startswith("linux")' in block
assert "_hw.IS_ROCM" in block, "guard must not run on NVIDIA/CPU hosts"
def test_requires_both_rdna4_and_an_affected_hip(self):
block = self._linux_block()
assert "if _rdna4 and _hip_lt_713:" in block
def test_scans_every_visible_device(self):
"""device_map="balanced" can place layers on a later card, so checking
device 0 alone misses the RDNA4 GPU."""
block = self._linux_block()
assert "for _i in range(_torch_lin.cuda.device_count()):" in block
def test_matches_both_rdna4_arch_ids(self):
block = self._linux_block()
assert '("gfx1200", "gfx1201")' in block
def test_failure_to_patch_is_non_fatal(self):
"""A broken patch attempt must not take down the whole training run."""
block = self._linux_block()
assert "except Exception" in block
assert "logger.warning" in block
def test_windows_and_linux_share_one_implementation(self):
"""Two copies of this fallback would drift; #7292 deliberately hoisted it."""
assert _WORKER_SOURCE.count("def _install_grouped_mm_cpu_fallback(") == 1
assert _WORKER_SOURCE.count("_install_grouped_mm_cpu_fallback(") >= 3 # def + win32 + linux
if __name__ == "__main__":
pytest.main([__file__, "-v"])

195
studio/backend/tests/test_llama_cpp_mmproj_fallback.py
View file

@ -250,6 +250,201 @@ class TestFlashAttnOff:
assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"]
_drop_env_v = LlamaCppBackend._drop_env_quantized_v_cache
class TestFlashAttnOffQuantizedKvCache:
"""Only the V cache requires flash attention in llama.cpp (init aborts with
"V cache quantization requires flash_attn"); a quantized K cache runs fine
without FA. Studio launches FA on, so a quantized --cache-type-v is legal at
launch but would make the FA-off crash-recovery retry crash on init. The
fallback must reset a quantized V cache (main and draft) to f16 while leaving
the K cache and non-quantized (f16/bf16/f32) types unchanged -- resetting K
would needlessly enlarge it and can OOM a memory-constrained config."""
_QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"]
_NON_QUANTIZED = ["f16", "bf16", "f32"]
@pytest.mark.parametrize("qtype", _QUANTIZED)
def test_quantized_v_reset_k_preserved(self, qtype):
cmd = [
"llama-server",
"--flash-attn",
"on",
"--cache-type-k",
qtype,
"--cache-type-v",
qtype,
]
out = _flash_off(cmd)
assert out is not None
# FA flipped off AND the V axis reset to f16; the K axis is preserved so
# the FA-off retry keeps its memory budget (quantized K is FA-independent).
assert out[out.index("--flash-attn") + 1] == "off"
assert out[out.index("--cache-type-k") + 1] == qtype
assert out[out.index("--cache-type-v") + 1] == "f16"
assert len(out) == len(cmd)
@pytest.mark.parametrize("qtype", _QUANTIZED)
def test_quantized_draft_v_reset(self, qtype):
# The draft context shares the global --flash-attn flag, so its quantized
# V cache aborts too and must be reset; the draft K cache is preserved.
for v_flag, k_flag in (
("--cache-type-v-draft", "--cache-type-k-draft"),
("--spec-draft-type-v", "--spec-draft-type-k"),
("-ctvd", "-ctkd"),
):
cmd = ["llama-server", "-fa", "on", k_flag, qtype, v_flag, qtype]
out = _flash_off(cmd)
assert out is not None
assert out[out.index(v_flag) + 1] == "f16"
assert out[out.index(k_flag) + 1] == qtype
@pytest.mark.parametrize("ntype", _NON_QUANTIZED)
def test_nonquantized_cache_left_unchanged(self, ntype):
cmd = [
"llama-server",
"--flash-attn",
"on",
"--cache-type-k",
ntype,
"--cache-type-v",
ntype,
]
out = _flash_off(cmd)
assert out is not None
# Only FA flips; the non-quantized cache type is preserved verbatim.
assert out[out.index("--flash-attn") + 1] == "off"
assert out[out.index("--cache-type-k") + 1] == ntype
assert out[out.index("--cache-type-v") + 1] == ntype
def test_equals_form_quantized_v_reset(self):
out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-v=q8_0"])
assert out == ["llama-server", "--flash-attn=off", "--cache-type-v=f16"]
def test_equals_form_quantized_k_preserved(self):
out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-k=q8_0"])
assert out == ["llama-server", "--flash-attn=off", "--cache-type-k=q8_0"]
def test_short_alias_v_reset_k_preserved(self):
out = _flash_off(["llama-server", "-fa", "on", "-ctk", "q4_0", "-ctv", "q4_0"])
assert out == ["llama-server", "-fa", "off", "-ctk", "q4_0", "-ctv", "f16"]
def test_asymmetric_cache_only_v_reset(self):
# Quantized V, non-quantized K: reset V, keep K untouched.
out = _flash_off(
[
"llama-server",
"--flash-attn",
"on",
"--cache-type-k",
"f16",
"--cache-type-v",
"q8_0",
]
)
assert out[out.index("--cache-type-k") + 1] == "f16"
assert out[out.index("--cache-type-v") + 1] == "f16"
def test_no_cache_flags_still_flips_fa(self):
out = _flash_off(["llama-server", "--flash-attn", "on", "-c", "4096"])
assert out == ["llama-server", "--flash-attn", "off", "-c", "4096"]
def test_quantized_k_only_still_flips_fa_but_keeps_k(self):
# A quantized K cache with no V flag is a valid FA-off launch; the retry
# must not touch the K cache (it would waste memory for nothing).
out = _flash_off(["llama-server", "--flash-attn", "on", "--cache-type-k", "q8_0"])
assert out == ["llama-server", "--flash-attn", "off", "--cache-type-k", "q8_0"]
def test_input_not_mutated(self):
cmd = ["llama-server", "--flash-attn", "on", "--cache-type-v", "q8_0"]
_flash_off(cmd)
assert cmd[-1] == "q8_0"
@pytest.mark.parametrize(
"flag",
["--cache_type_v", "--cache-type_v", "--cache_type-v"],
)
def test_underscore_alias_v_reset(self, flag):
# llama.cpp normalizes '_' to '-' in any '--' long option before
# matching, so a pass-through --cache_type_v enables a quantized V cache
# and must be reset by the FA-off retry too (else init aborts).
out = _flash_off(["llama-server", "--flash-attn", "on", flag, "q8_0"])
assert out is not None
assert out[out.index("--flash-attn") + 1] == "off"
# The user's flag spelling is preserved; llama.cpp normalizes it anyway.
assert out[out.index(flag) + 1] == "f16"
def test_underscore_alias_draft_v_reset(self):
out = _flash_off(["llama-server", "-fa", "on", "--spec_draft_type_v", "q4_0"])
assert out is not None
assert out[out.index("--spec_draft_type_v") + 1] == "f16"
def test_underscore_alias_equals_form_v_reset(self):
out = _flash_off(["llama-server", "--flash-attn=on", "--cache_type_v=q8_0"])
assert out == ["llama-server", "--flash-attn=off", "--cache_type_v=f16"]
def test_underscore_value_not_normalized_for_nonquantized(self):
# Only the flag name is canonicalized; a non-quantized type value is
# matched verbatim and left untouched (no spurious reset).
out = _flash_off(["llama-server", "--flash-attn", "on", "--cache_type_v", "f16"])
assert out[out.index("--cache_type_v") + 1] == "f16"
assert out[out.index("--flash-attn") + 1] == "off"
def test_short_alias_underscore_not_applied(self):
# Short flags are never underscore-normalized by llama.cpp; -ctv still
# matches and resets, and an unrelated short token is left alone.
out = _flash_off(["llama-server", "-fa", "on", "-ctv", "q8_0"])
assert out == ["llama-server", "-fa", "off", "-ctv", "f16"]
class TestDropEnvQuantizedVCache:
"""The argv rewrite can't reach a cache type set purely through the
environment (Studio deliberately lets an env-only type reach the child), so
the FA-off retry separately drops a quantized V-cache env var. Only V is
dropped: a quantized K cache is FA-independent and must survive."""
_QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"]
@pytest.mark.parametrize("qtype", _QUANTIZED)
def test_drops_quantized_main_v_env(self, qtype):
env = {"LLAMA_ARG_CACHE_TYPE_V": qtype, "PATH": "/usr/bin"}
assert _drop_env_v(env) is True
assert "LLAMA_ARG_CACHE_TYPE_V" not in env
assert env["PATH"] == "/usr/bin"
@pytest.mark.parametrize("qtype", _QUANTIZED)
def test_drops_quantized_draft_v_env(self, qtype):
env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": qtype}
assert _drop_env_v(env) is True
assert "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V" not in env
def test_preserves_quantized_k_env(self):
# A quantized K cache runs without FA, so its env must not be dropped.
env = {"LLAMA_ARG_CACHE_TYPE_K": "q8_0", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q4_0"}
assert _drop_env_v(env) is False
assert env["LLAMA_ARG_CACHE_TYPE_K"] == "q8_0"
assert env["LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K"] == "q4_0"
@pytest.mark.parametrize("ntype", ["f16", "bf16", "f32", "F16", " q8_0 "])
def test_preserves_nonquantized_v_env(self, ntype):
# Non-quantized V env values (and whitespace/case variants of them) run
# fine without FA; only a genuinely quantized value is dropped.
if ntype.strip().lower() in ("q8_0",):
env = {"LLAMA_ARG_CACHE_TYPE_V": ntype}
assert _drop_env_v(env) is True
assert "LLAMA_ARG_CACHE_TYPE_V" not in env
else:
env = {"LLAMA_ARG_CACHE_TYPE_V": ntype}
assert _drop_env_v(env) is False
assert env["LLAMA_ARG_CACHE_TYPE_V"] == ntype
def test_noop_on_empty_env(self):
env = {}
assert _drop_env_v(env) is False
assert env == {}
class TestNonProjectorDiagnostic:
"""_output_has_nonprojector_diagnostic gates the signal-only text-only retry:
a hard crash that already names OOM / a bad arch / a TP limit must surface

304
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -14,6 +14,7 @@ import contextlib
import copy
import json
import sys
import threading
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
@ -55,7 +56,12 @@ def _finish(reason: str) -> str:
)
def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
def _make_backend(
monkeypatch,
streams: list[object],
payloads: list[dict],
urls: list[str] | None = None,
):
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
@ -77,7 +83,12 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
if urls is not None:
urls.append(_url)
stream = streams.pop(0)
if isinstance(stream, BaseException):
raise stream
yield type("FakeResponse", (), {"status_code": 200, "chunks": stream})()
def fake_iter_text_cancellable(
response,
@ -88,9 +99,27 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)
monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable)
monkeypatch.setattr(backend, "_maybe_recover_from_mtp_crash", lambda *_a, **_k: False)
return backend
def _patch_successful_respawn(
monkeypatch,
backend,
port: int | None = None,
) -> list[bool]:
calls: list[bool] = []
def fake_respawn():
calls.append(True)
if port is not None:
backend._port = port
return True
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
return calls
def _tool_names(payload: dict) -> list[str]:
return [
(tool.get("function") or {}).get("name")
@ -2239,7 +2268,13 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [raising_stream()], payloads)
respawn_calls: list[bool] = []
monkeypatch.setattr(
backend,
"_respawn_if_dead",
lambda: respawn_calls.append(True) or True,
)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK")
collected: list[dict] = []
@ -2270,6 +2305,271 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
# The closing card is marked as an error, not an empty success, so the UI
# renders it as failed.
assert "Error" in (closing[0].get("result") or "")
assert respawn_calls == []
def test_connect_error_before_tool_stream_respawns_and_retries(monkeypatch):
"""A dead server before the first tool-loop response is opened is safe to retry."""
import httpx
payloads: list[dict] = []
urls: list[str] = []
backend = _make_backend(
monkeypatch,
[
httpx.ConnectError("server is down"),
[_sse({"content": "Recovered."}), _done()],
],
payloads,
urls,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend, port = 49999)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert respawn_calls == [True]
assert len(payloads) == 2
assert payloads[0] == payloads[1]
assert urls == [
"http://127.0.0.1:48847/v1/chat/completions",
"http://127.0.0.1:49999/v1/chat/completions",
]
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_connect_error_after_tool_result_recovers_both_generation_paths(monkeypatch):
"""Recover either post-tool generation path without rerunning the tool."""
import httpx
for max_tool_iterations, final_text in (
(2, "The result is 1."),
(1, "Final answer."),
):
payloads: list[dict] = []
backend = _make_backend(
monkeypatch,
[
_structured_tool_call("python", {"code": "print(1)"}, "call_once"),
httpx.ConnectError("server died between turns"),
[_sse({"content": final_text}), _done()],
],
payloads,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
tool_calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
tool_calls.append((name, arguments))
return "1"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "print one"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = max_tool_iterations,
)
)
assert respawn_calls == [True]
assert tool_calls == [("python", {"code": "print(1)"})]
assert len(payloads) == 3
assert payloads[1] == payloads[2]
assert any(e.get("type") == "content" and e.get("text") == final_text for e in events)
def test_connect_error_retry_is_bounded(monkeypatch):
"""A failed retry surfaces the error without another respawn attempt."""
import httpx
payloads: list[dict] = []
backend = _make_backend(
monkeypatch,
[
httpx.ConnectError("server is down"),
httpx.ConnectError("replacement is also down"),
],
payloads,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
assert respawn_calls == [True]
assert len(payloads) == 2
def test_pre_header_transport_errors_also_respawn(monkeypatch):
"""A child that dies during prefill already accepted the socket, so it does
not surface as ConnectError. Nothing has streamed yet, so replay is safe."""
import httpx
for exc in (
httpx.RemoteProtocolError("server disconnected without sending a response"),
httpx.ReadError("connection reset by peer"),
httpx.WriteError("broken pipe"),
):
payloads: list[dict] = []
backend = _make_backend(
monkeypatch, [exc, [_sse({"content": "Recovered."}), _done()]], payloads
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert respawn_calls == [True], type(exc).__name__
assert len(payloads) == 2, type(exc).__name__
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_a_not_yet_reaped_child_does_not_burn_the_retry(monkeypatch):
"""A closing server can beat its own exit status, so poll() briefly reports it
alive. Without a grace wait _respawn_if_dead hands back the stale _healthy and the
single retry is spent on the corpse rather than on a replacement."""
import httpx
class _Dying:
# reapable only from the 4th poll, mimicking teardown lagging the socket close
def __init__(self):
self.polls = 0
self.returncode = None
def poll(self):
self.polls += 1
if self.polls > 3:
self.returncode = -9
return -9
return None
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [], payloads)
backend._process = _Dying()
backend._healthy = True
backend._respawn_lock = threading.RLock()
backend._lock = threading.RLock()
backend._mtp_runtime_fallback_lock = threading.Lock()
backend._serial_load_lock = threading.RLock()
backend._cancel_event = threading.Event()
backend._unload_epoch = 0
backend._mtp_runtime_fallback_in_progress = False
backend._mtp_runtime_fallback_active = False
backend._last_load_kwargs = {"gguf_path": "/m.gguf"}
backend._model_identifier = "m"
dying = backend._process
loads: list[dict] = []
@contextlib.contextmanager
def dead_until_respawned(
_c,
_url,
payload,
_ce,
headers = None,
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
if backend._process is dying:
raise httpx.ReadError("connection reset while shutting down")
yield type(
"FakeResponse",
(),
{"status_code": 200, "chunks": [_sse({"content": "Recovered."}), _done()]},
)()
def fake_load(**kwargs):
loads.append(kwargs)
backend._process = type("Live", (), {"poll": lambda self: None, "returncode": None})()
backend._healthy = True
return True
monkeypatch.setattr(backend, "_stream_with_retry", dead_until_respawned)
monkeypatch.setattr(backend, "load_model", fake_load)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert len(loads) == 1
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_prefill_timeout_is_not_retried(monkeypatch):
"""A slow-but-alive server must not have its first-token budget spent twice."""
import httpx
for exc in (httpx.ReadTimeout("no first token"), httpx.PoolTimeout("pool")):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [exc], payloads)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
except httpx.TimeoutException:
raised = True
assert raised, type(exc).__name__
assert respawn_calls == [], type(exc).__name__
assert len(payloads) == 1, type(exc).__name__
def test_mtp_crash_recovery_wins_over_respawn(monkeypatch):
"""An MTP crash reloads without MTP, so never respawn the same config on top."""
import httpx
for max_tool_iterations in (2, 1):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [httpx.ConnectError("mtp crash")], payloads)
monkeypatch.setattr(backend, "_maybe_recover_from_mtp_crash", lambda *_a, **_k: True)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = max_tool_iterations,
)
)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
assert respawn_calls == []
assert len(payloads) == 1
def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch):

81
studio/backend/tests/test_rag_project_source_upload.py Normal file
View file

@ -0,0 +1,81 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Project sources upload: the path the create-project dialog drives."""
import os
import pytest
from core.rag import ingestion, store
from routes.rag import _sanitize_filename
from storage import rag_db
def _wait(job_id, timeout = 30.0):
import time
deadline = time.time() + timeout
while time.time() < deadline:
status = ingestion.get_job_status(job_id)
if status and status["status"] in ("completed", "failed"):
return status
time.sleep(0.05)
raise AssertionError("ingestion did not finish in time")
def _ingest(project_id, filename, path):
return ingestion.start_ingestion(
store.project_scope(project_id), None, None, filename, path, project_id = project_id
)
def test_project_document_persists_under_its_scope(rag_home, stub_embeddings, tmp_path):
path = tmp_path / "notes.txt"
path.write_text("alpha bravo charlie " * 50, encoding = "utf-8")
_, job_id = _ingest("P1", "notes.txt", str(path))
assert _wait(job_id)["status"] == "completed"
conn = rag_db.get_connection()
try:
docs = store.list_documents(conn, store.project_scope("P1"))
assert [d["filename"] for d in docs] == ["notes.txt"]
# Scoped: a sibling project cannot see it.
assert store.list_documents(conn, store.project_scope("P2")) == []
assert store.search_lexical(conn, store.project_scope("P1"), "bravo", 5)
finally:
conn.close()
@pytest.mark.parametrize(
"raw",
[
"x" * 300 + ".txt",
"y" * 512 + ".PDF",
"../" * 80 + "deep.md",
],
)
def test_long_filenames_keep_their_extension(raw):
# _save_upload gates on the extension, so trimming it would reject the file.
out = _sanitize_filename(raw)
assert len(out) <= 200
assert os.path.splitext(out)[1].lower() == os.path.splitext(raw)[1].lower()
@pytest.mark.parametrize(
"raw",
[
"../../etc/passwd.txt",
"..\\..\\windows\\evil.txt",
"/absolute/notes.txt",
"C:\\Users\\me\\notes.txt",
],
)
def test_sanitized_filenames_carry_no_path(raw):
out = _sanitize_filename(raw)
assert "/" not in out and "\\" not in out
@pytest.mark.parametrize("raw", ["." * 300, "noext" * 100, "a" * 100 + "." + "e" * 250])
def test_sanitizer_degrades_safely(raw):
assert 0 < len(_sanitize_filename(raw)) <= 200

7
studio/backend/tests/test_rocm_oom_guard.py
View file

@ -80,6 +80,7 @@ class TestCanonicalGcnArchName:
[
("gfx1150", True), # Strix Point
("gfx1151", True), # Strix Halo
("gfx1152", True), # Krackan Point (Radeon 860M/840M)
("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete
("gfx906", False), # MI50 — discrete server GPU
("gfx1201", False), # RX 9070 XT — discrete
@ -166,9 +167,15 @@ class TestDeviceNameFallback:
# gfx1151 Gorgon Halo (Ryzen AI Max 400 refresh)
"Radeon 8065S Graphics", # Ryzen AI Max+ 495
"AMD Radeon 8065S",
# gfx1152 Krackan Point (Ryzen AI 7 350 / AI 5 340)
"Radeon 860M",
"AMD Radeon 860M Graphics",
"Radeon 840M",
"AMD Radeon 840M Graphics",
# case variants
"RADEON 8060S GRAPHICS",
"radeon 8050s",
"RADEON 860M",
],
)
def test_unified_memory_detected(self, device_name: str) -> None:

8
studio/backend/tests/test_secure_tunnel_gate.py
View file

@ -85,6 +85,14 @@ def test_arg_parser_secure_polarity_and_not_secure_alias():
assert parser.parse_args(["--not-secure", "--secure"]).secure is True
def test_arg_parser_dns_pinning_opt_out_defaults_off():
import run
parser = run._build_arg_parser()
assert parser.parse_args([]).disable_dns_pinning is False
assert parser.parse_args(["--disable-dns-pinning"]).disable_dns_pinning is True
def test_run_server_accepts_enable_tools_kwarg():
import inspect

258
studio/backend/tests/test_server_disk_logging_outstream.py Normal file
View file

@ -0,0 +1,258 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression tests for the Colab "OutStream has no attribute 'watch_fd_thread'"
startup crash.
Field report (Colab): Unsloth Studio dies at server startup with
``❌ Unsloth Studio failed to start: 'OutStream' object has no attribute
'watch_fd_thread'``.
Root cause chain:
* Colab's ipykernel ``OutStream`` is created with ``watchfd=False``, so it
never gains a ``watch_fd_thread``; the ``OutStream.close()`` shipped in the
affected ipykernel versions joins that thread unconditionally and raises
``AttributeError`` (ipython/ipykernel#867).
* ``run._setup_server_disk_logging()`` replaces ``sys.stdout``/``sys.stderr``
with a ``_TeeStream``. That changes the console object identity, so Colab's
``absl`` logging handler -- which captured the ORIGINAL OutStream and whose
``close()`` deliberately skips ``sys.stdout``/``sys.stderr`` -- no longer
recognizes it as the live console.
* ``run_server`` builds ``uvicorn.Config(...)``, whose ``configure_logging`` ->
``logging.config.dictConfig`` -> ``logging.shutdown`` closes every existing
handler. The absl handler then calls ``OutStream.close()`` on the orphaned
stream, and the AttributeError aborts startup.
These tests reproduce the mechanism with a stand-in OutStream (Colab-identical
constructs are not importable off Colab) and assert the tee/console path used at
startup survives it.
"""
from __future__ import annotations
import io
import logging
import sys
import weakref
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import run as run_mod # noqa: E402
class _ColabOutStream(io.TextIOBase):
"""Stand-in for Colab's ipykernel OutStream built with ``watchfd=False``:
no ``watch_fd_thread`` and an unguarded ``close()`` that joins it
(ipython/ipykernel#867)."""
def __init__(self, name: str, sink: io.StringIO):
self.name = name
self._sink = sink
def write(self, s):
return self._sink.write(s)
def flush(self):
pass
def writable(self):
return True
def isatty(self):
return False
def close(self):
# Never set because watchfd=False -> AttributeError, exactly as Colab.
self.watch_fd_thread.join()
def __del__(self):
# io.TextIOBase.__del__ would call our buggy close() at GC (the harmless
# "Exception ignored" tail seen in Colab); silence it so the test is clean.
pass
class _WatchingOutStream(_ColabOutStream):
"""OutStream with fd-watching ON: ``watch_fd_thread`` exists, close() is
well behaved and must keep working unchanged."""
def __init__(self, name: str, sink: io.StringIO):
super().__init__(name, sink)
self.close_ran = False
self.watch_fd_thread = type("_T", (), {"join": lambda self: None})()
def close(self):
self.watch_fd_thread.join()
self.close_ran = True
class _AbslLikeHandler(logging.StreamHandler):
"""Mirror of ``absl.logging.PythonHandler.close()``: close the captured
stream unless it is (still) one of the user-managed console streams."""
def close(self):
try:
user_managed = (sys.stderr, sys.stdout, sys.__stderr__, sys.__stdout__)
if self.stream not in user_managed and (
not hasattr(self.stream, "isatty") or not self.stream.isatty()
):
self.stream.close()
except ValueError:
pass
super().close()
class TestHardenConsoleClose:
def test_neutralizes_watchfd_false_close(self):
stream = _ColabOutStream("stdout", io.StringIO())
with pytest.raises(AttributeError):
stream.close() # baseline: the ipykernel #867 bug is real
stream = _ColabOutStream("stdout", io.StringIO())
run_mod._harden_console_close(stream)
assert stream.close() is None # swallowed, no crash
def test_healthy_close_still_runs_fully(self):
stream = _WatchingOutStream("stdout", io.StringIO())
run_mod._harden_console_close(stream)
stream.close()
assert stream.close_ran is True
def test_only_attributeerror_is_swallowed(self):
class _Boom:
def close(self):
raise ValueError("real teardown failure")
stream = _Boom()
run_mod._harden_console_close(stream)
with pytest.raises(ValueError):
stream.close()
def test_unrelated_attributeerror_still_propagates(self):
# Only #867 is neutralized; a genuine missing attribute during teardown
# must still surface instead of looking like a clean close.
class _Console:
def close(self):
return self.not_a_real_attribute
stream = _Console()
run_mod._harden_console_close(stream)
with pytest.raises(AttributeError, match = "not_a_real_attribute"):
stream.close()
def test_swallowed_across_attributeerror_message_shapes(self):
# Python 3.12 appends a "Did you mean" tail; the match must survive it,
# and pre-3.10 AttributeErrors carry no ``name``, only the message.
class _Suggesting:
def close(self):
raise AttributeError(
"'OutStream' object has no attribute 'watch_fd_thread'. "
"Did you mean: '_watch_pipe_fd'?"
)
stream = _Suggesting()
run_mod._harden_console_close(stream)
assert stream.close() is None
def test_unsettable_close_is_left_alone(self):
# A stream whose close cannot be reassigned must not raise from hardening.
class _Frozen:
__slots__ = ()
def close(self):
return "ok"
stream = _Frozen()
run_mod._harden_console_close(stream) # must not raise
assert stream.close() == "ok"
class TestTeeStreamClose:
def test_tee_close_over_buggy_stream_never_raises(self):
console = _ColabOutStream("stdout", io.StringIO())
log = io.StringIO()
tee = run_mod._TeeStream(console, log)
tee.write("before-close")
tee.close() # must not raise despite the wrapped stream's broken close
assert log.getvalue() == "before-close"
def test_tee_close_flushes_log(self):
class _FlushCounting(io.StringIO):
def __init__(self):
super().__init__()
self.flushes = 0
def flush(self):
self.flushes += 1
super().flush()
console, log = io.StringIO(), _FlushCounting()
tee = run_mod._TeeStream(console, log)
tee.write("x")
tee.close()
assert log.flushes >= 1
class TestColabStartupRegression:
"""End-to-end: the exact trigger -- an absl-style handler closing the
orphaned OutStream during the ``logging.shutdown`` that uvicorn's
``uvicorn.Config`` -> ``dictConfig`` runs -- must not crash Studio, and the
tee must keep logging afterwards.
``logging.shutdown`` is driven over a LOCAL weakref list (identical code path
to ``logging.config._clearExistingHandlers``) so the global logging state and
pytest's own capture are untouched.
"""
def _make_console_and_handlers(self, monkeypatch):
out_sink, err_sink = io.StringIO(), io.StringIO()
out_stream = _ColabOutStream("stdout", out_sink)
err_stream = _ColabOutStream("stderr", err_sink)
monkeypatch.setattr(sys, "stdout", out_stream)
monkeypatch.setattr(sys, "stderr", err_stream)
# absl-like handlers capture the ORIGINAL OutStreams (as in Colab).
handlers = [_AbslLikeHandler(sys.stdout), _AbslLikeHandler(sys.stderr)]
return out_sink, err_sink, out_stream, err_stream, handlers
def test_baseline_reproduces_crash_without_fix(self, monkeypatch):
# Prove the test exercises the real path: swapping the console identity
# (what the tee does) makes the absl-like close hit #867.
_, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch)
try:
monkeypatch.setattr(sys, "stdout", io.StringIO())
monkeypatch.setattr(sys, "stderr", io.StringIO())
with pytest.raises(AttributeError, match = "watch_fd_thread"):
logging.shutdown([weakref.ref(h) for h in handlers])
finally:
# Neutralize so a lingering handler can't crash global teardown.
run_mod._harden_console_close(out_stream)
run_mod._harden_console_close(err_stream)
for h in handlers:
try:
h.close()
except Exception:
pass
def test_startup_survives_with_harden_and_tee(self, monkeypatch):
out_sink, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch)
# Exactly what _setup_server_disk_logging does before serving:
run_mod._harden_console_close(sys.stdout)
run_mod._harden_console_close(sys.stderr)
log_fh = io.StringIO()
monkeypatch.setattr(sys, "stdout", run_mod._TeeStream(sys.stdout, log_fh))
monkeypatch.setattr(sys, "stderr", run_mod._TeeStream(sys.stderr, log_fh))
# The close-storm uvicorn triggers via dictConfig -> logging.shutdown,
# closing the absl-like handlers over the (now orphaned) OutStreams.
logging.shutdown([weakref.ref(h) for h in handlers]) # must NOT raise
# The tee still tees to both console and disk afterwards.
print("post-startup-line")
sys.stdout.flush()
assert "post-startup-line" in out_sink.getvalue()
assert "post-startup-line" in log_fh.getvalue()

353
studio/backend/tests/test_tensor_parallel.py
View file

@ -19,6 +19,7 @@ from __future__ import annotations
import asyncio
import inspect
import socket
import sys
import threading
import time
@ -528,6 +529,358 @@ def test_runtime_recovery_is_single_flight(monkeypatch):
release.set()
def test_single_flight_claim_is_released_when_the_reload_cannot_start(monkeypatch):
# Only the reload thread's finally clears the claim, so if starting it raises the
# claim must not latch: nothing else resets it, and _respawn_if_dead then refuses
# forever, for every later model.
b = _recovery_backend()
class _NoThread:
def __init__(self, *args, **kwargs):
pass
def start(self):
raise RuntimeError("can't start new thread")
monkeypatch.setattr(llama_cpp_module.threading, "Thread", _NoThread)
assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False
assert b._mtp_runtime_fallback_in_progress is False
def test_load_kwargs_are_read_once_before_the_claim(monkeypatch):
# Gate and snapshot must share one read: reading twice lets an unload null
# _last_load_kwargs in between, so dict(None) raises after the claim and strands
# the flag with no thread alive to clear it.
b = _recovery_backend()
class _CountingKwargs: # data descriptor, so it wins over the instance dict
def __init__(self, value):
self.value = value
self.reads = 0
def __get__(self, obj, owner):
if obj is None:
return self
self.reads += 1
return self.value
def __set__(self, obj, value):
self.value = value
counter = _CountingKwargs({"model_identifier": "owner/repo"})
monkeypatch.setattr(type(b), "_last_load_kwargs", counter, raising = False)
class _UnstartedThread: # keep the reload off-thread so only sync reads count
def __init__(self, *args, **kwargs):
pass
def start(self):
pass
monkeypatch.setattr(llama_cpp_module.threading, "Thread", _UnstartedThread)
assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
assert counter.reads == 1, f"read {counter.reads} times; an unload can race the claim"
def test_respawn_defers_to_an_inflight_mtp_reload(monkeypatch):
# "Already recovering" must not read as "not an MTP crash": respawning replays the
# crashing MTP kwargs and aborts the in-flight no-MTP reload on its "newer load" check.
b = _recovery_backend()
b._mtp_runtime_fallback_in_progress = True
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == []
# Once that reload finishes, an ordinary respawn works again.
b._mtp_runtime_fallback_in_progress = False
b._process.returncode = -9 # only the respawn path logs it
assert b._respawn_if_dead() is True
assert [kw.get("speculative_type") for kw in loads] == ["auto"]
def test_respawn_does_not_wait_out_the_grace_on_a_replacement(monkeypatch):
# Callers losing the same child queue on _respawn_lock and wake holding the healthy
# REPLACEMENT. Unable to tell it from their own child, each burns the reap grace, and
# that sleep is held under the lock, so N callers cost N grace periods.
class _LiveProcess(_FakeProcess):
returncode = None
def __init__(self):
self.polls = 0
def poll(self): # never reapable, so the grace loop runs to its deadline
self.polls += 1
return None
workers = 4
b = _recovery_backend()
b._healthy = True
b._process.returncode = -9 # only the respawn path logs it
live = _LiveProcess()
loads: list[dict] = []
guard = threading.Lock()
all_in_flight = threading.Event()
# Subclass this instance, not the class: a descriptor on LlamaCppBackend would
# redirect _process for every other live backend, including atexit-registered ones.
state = {"proc": b._process, "readers": set()}
class _Tracked(type(b)):
@property
def _process(self):
"""Reports when every worker has taken its pre-lock look at the child."""
with guard:
state["readers"].add(threading.get_ident())
everyone = len(state["readers"]) >= workers
if everyone:
all_in_flight.set()
return state["proc"]
@_process.setter
def _process(self, value):
state["proc"] = value
b.__class__ = _Tracked
def _load(**kwargs):
# A real load_model takes seconds, so every caller that lost this child is in
# flight before the replacement appears; waiting reproduces that ordering. The
# timeout keeps the pre-fix build, where losers cannot read until the lock is
# free, from hanging instead of failing.
all_in_flight.wait(timeout = 2)
with guard:
loads.append(kwargs)
b._process = live
b._healthy = True # the real load_model marks the new server healthy
return True
monkeypatch.setattr(b, "load_model", _load)
results: list[bool] = []
def _respawn():
outcome = b._respawn_if_dead()
with guard:
results.append(outcome)
threads = [threading.Thread(target = _respawn) for _ in range(workers)]
started = time.monotonic()
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout = 30)
elapsed = time.monotonic() - started
assert results == [True] * workers, results
assert len(loads) == 1, f"{len(loads)} reloads, expected one"
# The grace loop is the only poll() of a live process, so any count means a queued
# caller charged the wait to a server that never failed.
assert live.polls == 0, "queued caller waited out the grace on a healthy server"
assert elapsed < llama_cpp_module._RESPAWN_REAP_GRACE_S * (workers - 1)
class _DyingChild(_FakeProcess):
"""Alive for the first polls, then reapable: what a terminate() looks like."""
def __init__(
self,
code = -15,
alive_polls = 2,
on_death = None,
):
self.polls = 0
self.returncode = None
self._code = code
self._alive_polls = alive_polls
self._on_death = on_death
def poll(self):
self.polls += 1
if self.polls <= self._alive_polls:
return None
if self.returncode is None:
self.returncode = self._code
if self._on_death is not None:
self._on_death()
return self._code
def test_respawn_does_not_resurrect_a_deliberate_unload(monkeypatch):
# unload_model() sets _cancel_event before killing, so a request that loses the
# connection can watch that deliberate exit through the grace loop and call it a
# crash, with _last_load_kwargs still populated (unload clears it after the kill).
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild()
b._cancel_event.set()
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == [], "resurrected a model the user unloaded"
def test_respawn_rechecks_the_cancel_flag_after_the_grace_wait(monkeypatch):
# The unload can also begin while we are already sleeping in the grace loop.
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(on_death = b._cancel_event.set)
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == [], "checked the cancel flag only before the wait"
def test_respawn_does_not_revert_a_newer_load(monkeypatch):
# A model switch landing while we wait must win; replaying the old kwargs would
# swap the user's new model back out.
b = _recovery_backend()
b._healthy = True
replacement = _DyingChild(alive_polls = 10**6)
b._process = _DyingChild(on_death = lambda: setattr(b, "_process", replacement))
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
b._respawn_if_dead()
assert loads == [], "replayed stale kwargs over a newer load"
assert b._process is replacement
def test_respawn_still_recovers_an_ordinary_crash(monkeypatch):
# Guard rail: none of the above may disable the recovery this path exists for.
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(code = -9)
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is True
assert len(loads) == 1
class _NeverReapable(_FakeProcess):
"""A child that stays unreapable, so only the port can tell alive from dead."""
returncode = None
def poll(self):
return None
def test_a_transient_error_against_a_live_server_costs_nothing(monkeypatch):
# The reap grace must not be charged to a server that never died: the sleep is
# held under _respawn_lock, so a full grace per caller serialises into N seconds
# of added latency on an install that is working fine.
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(16)
try:
b = _recovery_backend()
b._healthy = True
b._process = _NeverReapable()
b._port = listener.getsockname()[1]
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
started = time.monotonic()
assert b._respawn_if_dead() is True
elapsed = time.monotonic() - started
assert loads == [], "a live server must not be reloaded"
assert (
elapsed < llama_cpp_module._RESPAWN_REAP_GRACE_S / 2
), f"waited {elapsed:.2f}s on a server that is still accepting"
finally:
listener.close()
def test_a_closed_port_still_waits_for_the_child_to_be_reapable(monkeypatch):
# The other half: no listener means the server really is gone, so the grace
# still runs and the reap-race fix is preserved.
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
probe.bind(("127.0.0.1", 0))
dead_port = probe.getsockname()[1]
probe.close()
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(code = -9)
b._port = dead_port
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is True
assert len(loads) == 1
def test_socket_fast_path_honours_a_pending_unload(monkeypatch):
# unload_model() sets _cancel_event before it kills, so the child is still
# accepting when the probe runs. Reporting it healthy aims the retry at a server
# that is deliberately going away.
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(8)
try:
b = _recovery_backend()
b._healthy = True
b._process = _NeverReapable()
b._port = listener.getsockname()[1]
b._cancel_event.set()
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == []
finally:
listener.close()
def test_an_unload_landing_during_the_reload_is_undone(monkeypatch):
# The cancel check cannot live under _serial_load_lock alone: unload_model never
# takes that lock, so it can land entirely between the check and load_model and
# the captured kwargs then restart a model the user stopped. load_model clears
# _cancel_event on the way in, so _unload_epoch is the surviving evidence.
b = _recovery_backend()
b._healthy = True
b._process = _FakeProcess()
b._process.returncode = -9
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
unloads: list[int] = []
real_unload = b.unload_model
monkeypatch.setattr(b, "unload_model", lambda: unloads.append(1) or real_unload())
# The warning marks the window: after the snapshot, before the reload.
real_warning = llama_cpp_module.logger.warning
fired: list[int] = []
def racing_warning(*args, **kwargs):
if not fired:
fired.append(1)
real_unload()
return real_warning(*args, **kwargs)
monkeypatch.setattr(llama_cpp_module.logger, "warning", racing_warning)
assert b._respawn_if_dead() is False
assert unloads, "the racing unload was not honoured"
def test_socket_probe_is_false_without_a_port():
# Unloaded backends have no port; the probe must not raise, and the caller
# then falls back to the poll-based grace.
b = _recovery_backend()
b._port = None
assert b._server_socket_is_open() is False
def test_runtime_recovery_rechecks_cancel_before_reload():
# recover() must re-check the cancel flag after the death poll (load_model
# clears it), so a reload scheduled just before /unload can't resurrect it.

57
studio/backend/tests/test_web_fetch_extraction.py
View file

@ -15,6 +15,8 @@ from __future__ import annotations
import sys
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
@ -715,6 +717,61 @@ def test_fetch_url_raw_missing_content_type_reported_empty(monkeypatch):
assert content_type == ""
@pytest.mark.parametrize(
"disable_dns_pinning,expected_url",
[
(False, "https://203.0.113.7:8443/page?q=1"),
(True, "https://example.com:8443/page?q=1"),
],
)
def test_fetch_url_raw_dns_pinning_proxy_opt_out(monkeypatch, disable_dns_pinning, expected_url):
import email
import urllib.request
import core.inference.tools as tools_mod
class _FakeResp:
headers = email.message_from_string("Content-Type: text/plain\n")
def __init__(self):
self._body = b"ok"
def read(self, n = -1):
body, self._body = self._body, b""
return body
requested = []
class _FakeOpener:
def open(
self,
req,
timeout = None,
):
requested.append(req)
return _FakeResp()
resolved = []
def resolve(host, port):
resolved.append((host, port))
return True, "", "203.0.113.7"
monkeypatch.setenv("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "1" if disable_dns_pinning else "0")
monkeypatch.setattr(tools_mod, "_validate_and_resolve_host", resolve)
monkeypatch.setattr(urllib.request, "build_opener", lambda *handlers: _FakeOpener())
err, body, _content_type = tools_mod._fetch_url_raw(
"https://user:secret@example.com:8443/page?q=1"
)
assert err is None
assert body == "ok"
assert resolved == [("example.com", 8443)]
assert [req.full_url for req in requested] == [expected_url]
assert requested[0].get_header("Host") == "example.com:8443"
def test_fetch_page_text_missing_content_type_html_sniffed(monkeypatch):
# A header-less server returning an HTML body must still be converted.
def fake_fetch(

12
studio/backend/utils/llama_cpp_update.py
View file

@ -54,6 +54,8 @@ logger = structlog.get_logger(__name__)
DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp"
_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate
# install_llama_prebuilt.py EXIT_NO_SPACE: out of disk, retrying will not help.
_EXIT_NO_SPACE = 4
# Background job state. Single in-flight update at a time, guarded by _job_lock.
_JOB_IDLE = _flow.JOB_IDLE
@ -496,6 +498,16 @@ def _run_llama_phase(
+ (" Reload your model to use it." if model_was_active else "")
),
}
except _flow.InstallerExit as exc:
# Raw "installer exited 4: <log tail>" says nothing actionable in the UI.
if exc.returncode == _EXIT_NO_SPACE:
logger.warning("llama update: out of disk space")
raise RuntimeError(
"Not enough disk space to install llama.cpp. Free up space or point "
"UNSLOTH_STUDIO_HOME/TMPDIR at a larger volume, then retry."
) from exc
logger.warning("llama update: failed", error = str(exc))
raise
except Exception as exc:
logger.warning("llama update: failed", error = str(exc))
raise

109
studio/frontend/src/components/app-sidebar.tsx
View file

@ -108,7 +108,6 @@ import {
archiveChatItem,
ChatSearchDialog,
clearNewChatDraft,
createChatProject,
deleteChatProject,
deleteChatItem,
listStoredChatThreads,
@ -126,6 +125,7 @@ import {
type ProjectRecord,
type SidebarItem,
} from "@/features/chat";
import { NewProjectDialog } from "@/features/chat/components/new-project-dialog";
import {
useAppearanceCustomStore,
useSettingsDialogStore,
@ -707,8 +707,9 @@ export function AppSidebar() {
aria-label="New project"
onClick={(e) => {
e.stopPropagation();
// NewProjectDialog owns its own name field, so opening it is just the
// move target plus the open flag (same as the other call sites).
setProjectCreateMoveTarget(null);
setProjectNameDraft("");
setCreatingProject(true);
}}
className="sidebar-row-action group-hover/projects-item:opacity-100 group-hover/projects-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto group-data-[collapsible=icon]:hidden"
@ -942,7 +943,6 @@ export function AppSidebar() {
});
}, [allChatItems, pendingRename]);
const [creatingProject, setCreatingProject] = useState(false);
const [projectNameDraft, setProjectNameDraft] = useState("");
const [projectCreateMoveTarget, setProjectCreateMoveTarget] =
useState<SidebarItem | null>(null);
const renameTrimmed = renameDraft.trim();
@ -1095,28 +1095,26 @@ export function AppSidebar() {
}
}
async function commitCreateProject() {
const name = projectNameDraft.trim();
if (!name) return;
// "New project" from a chat's menu moves that chat in and stays put;
// otherwise open the project, unless a slow upload outlasted the route the
// user was on when they hit create.
async function afterCreateProject(
project: ProjectRecord,
{ stayedOnRoute }: { stayedOnRoute: boolean },
) {
const moveTarget = projectCreateMoveTarget;
setProjectCreateMoveTarget(null);
if (!moveTarget) {
if (stayedOnRoute) openProject(project.id);
return;
}
try {
const project = await createChatProject(name);
if (moveTarget) {
await moveChatItemToProject(moveTarget, project.id);
if (activeThreadId === moveTarget.id) {
useChatRuntimeStore.getState().setActiveProjectId(project.id);
}
}
setCreatingProject(false);
setProjectNameDraft("");
setProjectCreateMoveTarget(null);
if (moveTarget) {
return;
} else {
openProject(project.id);
await moveChatItemToProject(moveTarget, project.id);
if (activeThreadId === moveTarget.id) {
useChatRuntimeStore.getState().setActiveProjectId(project.id);
}
} catch (err) {
toast.error(moveTarget ? "Failed to create and move chat" : "Failed to create project", {
toast.error("Failed to move chat to the new project", {
description: err instanceof Error ? err.message : undefined,
});
}
@ -1296,7 +1294,6 @@ export function AppSidebar() {
<DropdownMenuItem
onSelect={() => {
setProjectCreateMoveTarget(item);
setProjectNameDraft("");
setCreatingProject(true);
}}
>
@ -2063,6 +2060,18 @@ export function AppSidebar() {
</button>
</SidebarMenuItem>
)}
{/* Collapsed rail has no room for the cog on the profile row, so it
sits above the avatar instead. */}
<NavItem
className="hidden group-data-[collapsible=icon]:block"
icon={Settings02Icon}
label={t("shell.navigation.settings")}
active={false}
onClick={() => {
useSettingsDialogStore.getState().openDialog();
closeMobileIfOpen();
}}
/>
<SidebarMenuItem>
<DropdownMenu>
<DropdownMenuTrigger asChild>
@ -2368,58 +2377,18 @@ export function AppSidebar() {
</DialogFooter>
</DialogContent>
</Dialog>
<Dialog
<NewProjectDialog
open={creatingProject}
onOpenChange={(open) => {
setCreatingProject(open);
if (!open) {
setProjectNameDraft("");
setProjectCreateMoveTarget(null);
}
if (!open) setProjectCreateMoveTarget(null);
}}
>
<DialogContent className="corner-squircle dialog-soft-surface sm:max-w-md">
<DialogHeader>
<DialogTitle>
{projectCreateMoveTarget ? "Move to new project" : "New project"}
</DialogTitle>
</DialogHeader>
<Input
value={projectNameDraft}
onChange={(event) => setProjectNameDraft(event.target.value)}
onKeyDown={(event) => {
if (event.key === "Enter") {
event.preventDefault();
void commitCreateProject();
}
}}
autoFocus
maxLength={120}
placeholder="Project name"
aria-label="Project name"
className="focus-visible:border-input focus-visible:ring-0"
/>
<DialogFooter className="flex-wrap gap-2 sm:justify-end">
<Button
type="button"
variant="ghost"
onClick={() => {
setCreatingProject(false);
setProjectCreateMoveTarget(null);
}}
>
Cancel
</Button>
<Button
type="button"
onClick={() => void commitCreateProject()}
disabled={!projectNameDraft.trim()}
>
Create
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
title={
projectCreateMoveTarget ? "Move to new project" : "Create project"
}
submitLabel={projectCreateMoveTarget ? "Create and move" : "Create project"}
onCreated={afterCreateProject}
/>
</>
);
}

23
studio/frontend/src/components/assistant-ui/reasoning.tsx
View file

@ -346,6 +346,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
const [manualOpen, setManualOpen] = useState(false);
const [dismissedWhileStreaming, setDismissedWhileStreaming] = useState(false);
const [retainStreamingHeight, setRetainStreamingHeight] = useState(false);
const [duration, setDuration] = useState<number>(0);
const startTimeRef = useRef<number | null>(null);
@ -361,13 +362,26 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
}
}, [isReasoningStreaming]);
// Reset dismissed flag on new stream.
// Reset per-round open state. manualOpen is sticky and regenerate reuses this
// instance, so a hand-opened block would stay pinned open and never collapse.
useEffect(() => {
if (isReasoningStreaming) {
setDismissedWhileStreaming(false);
setManualOpen(false);
}
}, [isReasoningStreaming]);
// Keep the streaming height cap until the automatic close finishes. Removing
// it on the completion frame expands long reasoning to its full height before
// the collapsible can close, which makes the entire chat jump.
useEffect(() => {
const timeout = window.setTimeout(
() => setRetainStreamingHeight(isReasoningStreaming),
isReasoningStreaming ? 0 : ANIMATION_DURATION,
);
return () => window.clearTimeout(timeout);
}, [isReasoningStreaming]);
// Open while streaming (unless dismissed), or once manually opened.
const isOpen = (isReasoningStreaming && !dismissedWhileStreaming) || manualOpen;
const variant = isOpen ? "outline" : "ghost";
@ -378,6 +392,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
if (isReasoningStreaming) {
setDismissedWhileStreaming(!open);
} else {
if (open) {
setRetainStreamingHeight(false);
}
setManualOpen(open);
}
},
@ -407,7 +424,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
aria-busy={isReasoningStreaming}
streaming={isReasoningStreaming}
>
<ReasoningText streaming={isReasoningStreaming}>
<ReasoningText
streaming={isReasoningStreaming || retainStreamingHeight}
>
{children}
</ReasoningText>
</ReasoningContent>

10
studio/frontend/src/features/chat/api-provider-logo.tsx
View file

@ -40,10 +40,10 @@ interface ApiProviderLogoProps {
title?: string;
}
/**
* Renders a registry provider's logo when its asset exists under
* `public/provider-logos/`. OpenAI's is inverted in dark mode for contrast.
*/
// Monochrome logos vanish on a dark background.
const DARK_INVERT_LOGOS = new Set(["openai", "ollama", "openrouter"]);
/** Provider logo from `public/provider-logos/`; monochrome ones invert in dark mode. */
export function ApiProviderLogo({ providerType, className, title }: ApiProviderLogoProps) {
const src = apiProviderLogoSrc(providerType);
if (!src && isCustomProviderType(providerType)) {
@ -63,7 +63,7 @@ export function ApiProviderLogo({ providerType, className, title }: ApiProviderL
aria-hidden
className={cn(
"shrink-0 object-contain",
providerType === "openai" && "dark:invert",
providerType && DARK_INVERT_LOGOS.has(providerType) && "dark:invert",
className,
)}
/>

13
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -185,6 +185,10 @@ import {
listStoredChatThreads,
} from "./utils/chat-history-storage";
import { isAssistantLocalThreadId } from "./utils/thread-ids";
import {
consumeProjectSourcesPending,
hasProjectSourcesPending,
} from "@/features/rag/components/project-source-dropzone";
const ProjectSourcesPanel = lazy(() =>
@ -998,7 +1002,14 @@ function ProjectLanding({
const active = useChatActive();
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const initialActiveThreadRef = useRef<string | null>(null);
const [projectTab, setProjectTab] = useState<"chats" | "sources">("chats");
// Land on Sources when the project was just created with dropped files.
const [projectTab, setProjectTab] = useState<"chats" | "sources">(() =>
hasProjectSourcesPending(projectId) ? "sources" : "chats",
);
// Drop the marker once committed: React may replay the initializer above.
useEffect(() => {
consumeProjectSourcesPending(projectId);
}, [projectId]);
const [pendingNewThreadId, setPendingNewThreadId] = useState<string | null>(
null,
);

141
studio/frontend/src/features/chat/components/new-project-dialog.tsx
View file

@ -2,7 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useNavigate } from "@tanstack/react-router";
import { useState } from "react";
import { useEffect, useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import {
@ -12,31 +12,92 @@ import {
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import {
ProjectSourceDropzone,
type StagedSource,
uploadStagedSources,
} from "@/features/rag/components/project-source-dropzone";
import { toast } from "@/lib/toast";
import { Folder02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { createChatProject } from "../hooks/use-chat-projects";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
import type { ProjectRecord } from "../types";
// Create-project dialog usable from the composer + menu. Creating opens the new
// project straight away rather than dropping the user on the projects list.
function currentRoute(): string {
if (typeof window === "undefined") return "";
return window.location.pathname + window.location.search;
}
// Create-project dialog for the composer, sidebar, and projects page. Creating
// opens the new project; `onCreated` overrides that for callers with their own
// follow-up (the sidebar's "move this chat to a new project").
export function NewProjectDialog({
open,
onOpenChange,
title = "Create project",
submitLabel = "Create project",
onCreated,
}: {
open: boolean;
onOpenChange: (open: boolean) => void;
title?: string;
submitLabel?: string;
onCreated?: (
project: ProjectRecord,
context: { stayedOnRoute: boolean },
) => void | Promise<void>;
}) {
const navigate = useNavigate();
const [name, setName] = useState("");
const [staged, setStaged] = useState<StagedSource[]>([]);
const [busy, setBusy] = useState(false);
// Uploads outlive this component, so a slow one must not yank the user to the
// new project after they have navigated away.
const mounted = useRef(true);
useEffect(() => {
// Set on setup, not just cleared on cleanup: StrictMode replays
// setup/cleanup/setup, which would otherwise leave this false forever.
mounted.current = true;
return () => {
mounted.current = false;
};
}, []);
function reset() {
setName("");
setStaged([]);
}
// Every close path routes through here: callers keep this mounted, so a draft
// left behind would resurface (and upload) on the next project.
function close() {
if (busy) return;
reset();
onOpenChange(false);
}
async function commitCreate() {
const trimmed = name.trim();
if (!trimmed) return;
if (!trimmed || busy) return;
setBusy(true);
// Sidebar callers keep this mounted across routes, so unmounting alone
// cannot tell whether the user has moved on during a slow upload.
const origin = currentRoute();
try {
const project = await createChatProject(trimmed);
// Upload before closing so the Sources panel lists them on first fetch.
await uploadStagedSources(project.id, staged);
if (!mounted.current) return;
const stayedOnRoute = currentRoute() === origin;
onOpenChange(false);
setName("");
reset();
if (onCreated) {
await onCreated(project, { stayedOnRoute });
return;
}
if (!stayedOnRoute) return;
const runtime = useChatRuntimeStore.getState();
runtime.setActiveThreadId(null);
runtime.setActiveProjectId(project.id);
@ -45,6 +106,8 @@ export function NewProjectDialog({
toast.error("Failed to create project", {
description: err instanceof Error ? err.message : undefined,
});
} finally {
setBusy(false);
}
}
@ -52,43 +115,59 @@ export function NewProjectDialog({
<Dialog
open={open}
onOpenChange={(next) => {
if (!next) setName("");
onOpenChange(next);
if (next) {
onOpenChange(true);
return;
}
close();
}}
>
<DialogContent className="corner-squircle dialog-soft-surface sm:max-w-md">
<DialogContent className="corner-squircle dialog-soft-surface gap-5 sm:max-w-lg">
<DialogHeader>
<DialogTitle>New project</DialogTitle>
<DialogTitle className="text-ui-21">{title}</DialogTitle>
</DialogHeader>
<Input
value={name}
onChange={(e) => setName(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter") {
e.preventDefault();
void commitCreate();
}
}}
autoFocus={true}
maxLength={120}
placeholder="Project name"
aria-label="Project name"
className="focus-visible:border-input focus-visible:ring-0"
{/* Name field: folder glyph in its own cell, divided from the input. */}
<div className="flex items-stretch overflow-hidden rounded-[16px] border border-border bg-background transition-colors focus-within:border-ring has-[input:disabled]:opacity-50 dark:border-transparent dark:bg-white/[0.06]">
<span className="flex w-9 shrink-0 items-center justify-center text-muted-foreground">
<HugeiconsIcon
icon={Folder02Icon}
strokeWidth={1.75}
className="size-5"
/>
</span>
<span aria-hidden="true" className="my-3 w-px bg-border" />
<input
value={name}
onChange={(e) => setName(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter") {
e.preventDefault();
void commitCreate();
}
}}
autoFocus={true}
disabled={busy}
maxLength={120}
placeholder="Project name"
aria-label="Project name"
className="min-w-0 flex-1 bg-transparent py-4 pr-4 pl-2.5 text-base outline-none placeholder:text-muted-foreground disabled:cursor-not-allowed"
/>
</div>
<ProjectSourceDropzone
staged={staged}
onChange={setStaged}
disabled={busy}
/>
<DialogFooter className="flex-wrap gap-2 sm:justify-end">
<Button
type="button"
variant="ghost"
onClick={() => onOpenChange(false)}
>
<Button type="button" variant="ghost" disabled={busy} onClick={close}>
Cancel
</Button>
<Button
type="button"
onClick={() => void commitCreate()}
disabled={!name.trim()}
disabled={!name.trim() || busy}
>
Create
{busy ? "Creating…" : submitLabel}
</Button>
</DialogFooter>
</DialogContent>

70
studio/frontend/src/features/chat/projects-page.tsx
View file

@ -34,7 +34,6 @@ import { isTauri } from "@/lib/api-base";
import { isDownloadCancelled, pickNativeChatImport } from "@/lib/native-files";
import { toast } from "@/lib/toast";
import {
createChatProject,
deleteChatProject,
renameChatProject,
useChatProjects,
@ -42,6 +41,7 @@ import {
usePinnedProjectsStore,
type ProjectRecord,
} from "@/features/chat";
import { NewProjectDialog } from "./components/new-project-dialog";
import {
Delete02Icon,
Download01Icon,
@ -124,7 +124,6 @@ export function ProjectsPage() {
);
const [creating, setCreating] = useState(false);
const [nameDraft, setNameDraft] = useState("");
const [renaming, setRenaming] = useState<ProjectRecord | null>(null);
const [renameDraft, setRenameDraft] = useState("");
const [deleting, setDeleting] = useState<ProjectRecord | null>(null);
@ -258,21 +257,6 @@ export function ProjectsPage() {
navigate({ to: "/chat", search: { project: projectId } });
}
async function commitCreate() {
const name = nameDraft.trim();
if (!name) return;
try {
const project = await createChatProject(name);
setCreating(false);
setNameDraft("");
openProject(project.id);
} catch (err) {
toast.error("Failed to create project", {
description: err instanceof Error ? err.message : undefined,
});
}
}
async function commitRename() {
const target = renaming;
const name = renameDraft.trim();
@ -469,14 +453,7 @@ export function ProjectsPage() {
</DropdownMenuSub>
</DropdownMenuContent>
</DropdownMenu>
<Button
onClick={() => {
setNameDraft("");
setCreating(true);
}}
>
New project
</Button>
<Button onClick={() => setCreating(true)}>New project</Button>
</div>
</div>
@ -511,10 +488,7 @@ export function ProjectsPage() {
<Button
variant="outline"
className="mt-2 border-none bg-background shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-none"
onClick={() => {
setNameDraft("");
setCreating(true);
}}
onClick={() => setCreating(true)}
>
<HugeiconsIcon icon={FolderAddIcon} strokeWidth={1.75} className="size-icon" />
Create your first project
@ -674,42 +648,8 @@ export function ProjectsPage() {
</>
)}
{/* Create project */}
<Dialog
open={creating}
onOpenChange={(open) => {
if (!open) setCreating(false);
}}
>
<DialogContent className="corner-squircle dialog-soft-surface sm:max-w-md">
<DialogHeader>
<DialogTitle>New project</DialogTitle>
</DialogHeader>
<Input
value={nameDraft}
onChange={(e) => setNameDraft(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter") {
e.preventDefault();
void commitCreate();
}
}}
autoFocus
maxLength={120}
placeholder="Project name"
aria-label="Project name"
className="focus-visible:border-input focus-visible:ring-0"
/>
<DialogFooter className="flex-wrap gap-2 sm:justify-end">
<Button type="button" variant="ghost" onClick={() => setCreating(false)}>
Cancel
</Button>
<Button type="button" onClick={() => void commitCreate()} disabled={!nameDraft.trim()}>
Create
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
{/* Create project (name + drag-and-drop sources) */}
<NewProjectDialog open={creating} onOpenChange={setCreating} />
{/* Rename project */}
<Dialog

304
studio/frontend/src/features/rag/components/project-source-dropzone.tsx Normal file
View file

@ -0,0 +1,304 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { toast } from "@/lib/toast";
import { cn } from "@/lib/utils";
import { File02Icon, FolderAddIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { XIcon } from "lucide-react";
import { useCallback, useRef, useState } from "react";
import {
invalidateProjectSources,
uploadProjectDocument,
} from "../api/rag-api";
import { RAG_UPLOAD_ACCEPT } from "../types/rag";
import { resolveVisionOverrides } from "./vision-overrides";
/** A file picked before the project exists, held until create commits. */
export interface StagedSource {
id: string;
file: File;
}
// Client-side dedup key; backend dedups authoritatively by content hash.
function fileSignature(file: File): string {
return `${file.name}|${file.size}|${file.lastModified}`;
}
function formatSize(bytes: number): string {
if (!Number.isFinite(bytes) || bytes <= 0) return "";
const units = ["B", "KB", "MB", "GB"];
let value = bytes;
let unit = 0;
while (value >= 1024 && unit < units.length - 1) {
value /= 1024;
unit += 1;
}
const shown =
value >= 10 || unit === 0
? String(Math.round(value))
: value.toFixed(1).replace(/\.0$/, "");
return `${shown} ${units[unit]}`;
}
const ACCEPTED_EXTS = new Set(
RAG_UPLOAD_ACCEPT.split(",").map((ext) => ext.trim().toLowerCase()),
);
// `accept` only filters the picker, so a drop can carry anything. A folder
// arrives as an extension-less entry, which this rejects along with the types
// the backend would 400 on.
function isSupported(file: File): boolean {
const dot = file.name.lastIndexOf(".");
if (dot <= 0) return false;
return ACCEPTED_EXTS.has(file.name.slice(dot).toLowerCase());
}
/** Merge a selection into the staged list. Returns the names it would not take,
* so the caller can say so once instead of dropping them silently. */
function addStagedSources(
staged: StagedSource[],
incoming: FileList | File[],
): { next: StagedSource[]; unsupported: string[]; duplicates: string[] } {
const seen = new Set(staged.map((entry) => fileSignature(entry.file)));
const next = [...staged];
const unsupported: string[] = [];
const duplicates: string[] = [];
for (const file of Array.from(incoming)) {
if (!isSupported(file)) {
unsupported.push(file.name);
continue;
}
const signature = fileSignature(file);
if (seen.has(signature)) {
duplicates.push(file.name);
continue;
}
seen.add(signature);
next.push({
id: `staged_${Math.random().toString(36).slice(2)}`,
file,
});
}
return { next, unsupported, duplicates };
}
// Projects created with staged files, so the landing can open on Sources.
const projectsWithPendingSources = new Set<string>();
function markProjectSourcesPending(projectId: string): void {
projectsWithPendingSources.add(projectId);
}
/** Whether this project was just created with staged sources. Read-only, so it
* is safe in a render pass that React may replay. */
export function hasProjectSourcesPending(projectId: string): boolean {
return projectsWithPendingSources.has(projectId);
}
/** Drop the marker once the landing has committed. */
export function consumeProjectSourcesPending(projectId: string): void {
projectsWithPendingSources.delete(projectId);
}
/** Upload staged files to a new project. Indexing runs in the background; a
* per-file failure toasts and never blocks project creation. */
export async function uploadStagedSources(
projectId: string,
staged: StagedSource[],
): Promise<void> {
if (staged.length === 0) return;
invalidateProjectSources(projectId);
markProjectSourcesPending(projectId);
const { ocr, caption } = resolveVisionOverrides();
const documentIds = new Set<string>();
const merged: string[] = [];
for (const { file } of staged) {
try {
const result = await uploadProjectDocument(projectId, file, ocr, caption);
// Same bytes under another name: the backend hashes content, so this is
// the document already uploaded. Say so rather than imply a new source.
if (documentIds.has(result.documentId)) merged.push(file.name);
else documentIds.add(result.documentId);
} catch (error) {
toast.error(`Couldn't upload ${file.name}`, {
description: error instanceof Error ? error.message : String(error),
});
}
}
if (merged.length > 0) {
toast.info(
merged.length === 1
? `${merged[0]} matched a file already added`
: `${merged.length} files matched files already added`,
{ description: "Identical contents are stored once." },
);
}
invalidateProjectSources(projectId);
}
/** Create-project drop area: stages files until the project exists. */
export function ProjectSourceDropzone({
staged,
onChange,
disabled = false,
}: {
staged: StagedSource[];
onChange: (next: StagedSource[]) => void;
disabled?: boolean;
}) {
const inputRef = useRef<HTMLInputElement>(null);
// Count enter/leave pairs: children fire dragleave on the parent.
const dragDepth = useRef(0);
const [dragging, setDragging] = useState(false);
const addFiles = useCallback(
(files: FileList | File[]) => {
const { next, unsupported, duplicates } = addStagedSources(staged, files);
if (next.length !== staged.length) onChange(next);
if (unsupported.length > 0) {
toast.info(
unsupported.length === 1
? `Can't add ${unsupported[0]}`
: `Can't add ${unsupported.length} files`,
{ description: `Supported types: ${RAG_UPLOAD_ACCEPT}` },
);
}
// Name, size and mtime can in principle match for two different files, so
// never drop one without saying so.
if (duplicates.length > 0) {
toast.info(
duplicates.length === 1
? `${duplicates[0]} is already added`
: `${duplicates.length} files were already added`,
);
}
},
[staged, onChange],
);
const endDrag = useCallback(() => {
dragDepth.current = 0;
setDragging(false);
}, []);
return (
<div className="space-y-2.5">
<p className="text-ui-15 font-medium text-foreground">Sources</p>
{/* Panel is the drop target; the inner button owns the click so staged
rows can carry their own remove buttons. */}
<div
// preventDefault runs even while disabled: nothing else on the page
// cancels a file drop, so the browser would navigate to the file and
// kill the uploads in flight.
onDragEnter={(e) => {
e.preventDefault();
if (disabled) return;
dragDepth.current += 1;
setDragging(true);
}}
onDragOver={(e) => {
e.preventDefault();
if (disabled) return;
e.dataTransfer.dropEffect = "copy";
}}
onDragLeave={() => {
dragDepth.current = Math.max(0, dragDepth.current - 1);
if (dragDepth.current === 0) setDragging(false);
}}
onDrop={(e) => {
e.preventDefault();
if (disabled) return;
endDrag();
addFiles(Array.from(e.dataTransfer.files ?? []));
}}
className={cn(
"rounded-[22px] border border-border transition-colors dark:border-white/10",
dragging && "border-primary/60 bg-primary/5",
disabled && "opacity-60",
)}
>
<input
ref={inputRef}
type="file"
multiple={true}
accept={RAG_UPLOAD_ACCEPT}
className="hidden"
onChange={(e) => {
const files = Array.from(e.target.files ?? []);
e.target.value = "";
addFiles(files);
}}
/>
{staged.length === 0 ? (
<button
type="button"
aria-label="Add sources"
disabled={disabled}
onClick={() => inputRef.current?.click()}
className="flex w-full cursor-pointer flex-col items-center justify-center gap-3 rounded-[22px] px-6 py-12 text-center transition-colors hover:bg-muted/40 focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
>
<HugeiconsIcon
icon={FolderAddIcon}
strokeWidth={1.75}
className="size-6 text-muted-foreground"
/>
<span className="text-sm text-muted-foreground">
Add files every chat in this project can read
</span>
</button>
) : (
<div className="flex flex-col gap-1 p-2">
<ul className="max-h-52 space-y-0.5 overflow-y-auto">
{staged.map((entry) => (
<li
key={entry.id}
className="flex items-center gap-2.5 rounded-[10px] px-2.5 py-2 hover:bg-muted/50"
>
<HugeiconsIcon
icon={File02Icon}
strokeWidth={1.75}
className="size-4 shrink-0 text-muted-foreground"
/>
<span
className="min-w-0 flex-1 truncate text-ui-14 text-foreground"
title={entry.file.name}
>
{entry.file.name}
</span>
<span className="shrink-0 text-ui-11 text-muted-foreground">
{formatSize(entry.file.size)}
</span>
<button
type="button"
aria-label={`Remove ${entry.file.name}`}
disabled={disabled}
onClick={() =>
onChange(staged.filter((row) => row.id !== entry.id))
}
className="shrink-0 rounded-full text-muted-foreground transition-colors hover:text-foreground disabled:opacity-50"
>
<XIcon className="size-3.5" />
</button>
</li>
))}
</ul>
<button
type="button"
disabled={disabled}
onClick={() => inputRef.current?.click()}
className="flex items-center justify-center gap-2 rounded-[10px] py-2 text-ui-13 font-medium text-muted-foreground transition-colors hover:bg-muted/50 hover:text-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
>
<HugeiconsIcon
icon={FolderAddIcon}
strokeWidth={1.75}
className="size-4"
/>
Add files
</button>
</div>
)}
</div>
</div>
);
}

21
studio/frontend/src/features/rag/components/use-rag-documents.ts
View file

@ -1,13 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useEffect, useRef, useState } from "react";
import {
CHAT_RAG_CAPTION_KEY,
CHAT_RAG_OCR_KEY,
useChatRuntimeStore,
} from "@/features/chat";
import { toast } from "@/lib/toast";
import { useCallback, useEffect, useRef, useState } from "react";
import {
deleteDocument,
getJob,
@ -17,6 +12,7 @@ import {
uploadThreadDocument,
} from "../api/rag-api";
import type { DocumentStatus, RagDocument } from "../types/rag";
import { resolveVisionOverrides } from "./vision-overrides";
export interface TrackedDocument extends RagDocument {
progress?: number | null;
@ -263,18 +259,7 @@ export function useRagDocuments(
tempId: string,
) => {
try {
// Send vision-pass overrides only after the user has explicitly set them;
// otherwise backend env defaults own the ingest policy.
const state = useChatRuntimeStore.getState();
const hasLocal = (key: string) =>
typeof window !== "undefined" &&
window.localStorage.getItem(key) !== null;
const ocr = hasLocal(CHAT_RAG_OCR_KEY)
? state.ragOcrScanned
: undefined;
const caption = hasLocal(CHAT_RAG_CAPTION_KEY)
? state.ragCaptionFigures
: undefined;
const { ocr, caption } = resolveVisionOverrides();
const result =
activeScope.type === "kb"
? await uploadKnowledgeBaseDocument(

35
studio/frontend/src/features/rag/components/vision-overrides.ts Normal file
View file

@ -0,0 +1,35 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
CHAT_RAG_CAPTION_KEY,
CHAT_RAG_OCR_KEY,
useChatRuntimeStore,
} from "@/features/chat";
function hasLocal(key: string): boolean {
if (typeof window === "undefined") return false;
try {
return window.localStorage.getItem(key) !== null;
} catch {
// Storage can be blocked outright (sandboxed context). These overrides are
// optional, so fall back to the backend defaults rather than failing the
// upload that asked for them.
return false;
}
}
/** Ingest-time vision-pass overrides, sent only once the user has set them;
* otherwise backend env defaults own the policy. Shared by every upload path. */
export function resolveVisionOverrides(): {
ocr: boolean | undefined;
caption: boolean | undefined;
} {
const state = useChatRuntimeStore.getState();
return {
ocr: hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined,
caption: hasLocal(CHAT_RAG_CAPTION_KEY)
? state.ragCaptionFigures
: undefined,
};
}

28
studio/frontend/src/features/settings/settings-dialog.tsx
View file

@ -12,6 +12,7 @@ import { type TranslationKey, useT } from "@/i18n";
import { cn } from "@/lib/utils";
import { MicIcon } from "@/lib/mic-icon";
import {
BotIcon,
Cancel01Icon,
CloudIcon,
CpuIcon,
@ -40,6 +41,7 @@ import {
useSettingsDialogStore,
} from "./stores/settings-dialog-store";
import { AboutTab } from "./tabs/about-tab";
import { AgentsTab } from "./tabs/agents-tab";
import { ApiKeysTab } from "./tabs/api-keys-tab";
import { AppearanceTab } from "./tabs/appearance-tab";
import { ChatTab } from "./tabs/chat-tab";
@ -71,13 +73,11 @@ const TABS: TabDef[] = [
id: "resources",
labelKey: "settings.tabs.resources",
icon: CpuIcon,
badgeKey: "common.new",
},
{
id: "chat",
labelKey: "settings.tabs.chat",
icon: Message01Icon,
badgeKey: "common.new",
},
{
id: "api-keys",
@ -89,6 +89,12 @@ const TABS: TabDef[] = [
labelKey: "settings.tabs.connections",
icon: CloudIcon,
},
{
id: "agents",
labelKey: "settings.tabs.agents",
icon: BotIcon,
badgeKey: "common.new",
},
{
id: "voice",
labelKey: "settings.tabs.voice",
@ -124,6 +130,8 @@ function renderTab(tab: SettingsTab) {
return <DataTab />;
case "api-keys":
return <ApiKeysTab />;
case "agents":
return <AgentsTab />;
case "about":
return <AboutTab />;
}
@ -222,6 +230,7 @@ export function SettingsDialog() {
connections: null,
data: null,
"api-keys": null,
agents: null,
about: null,
});
@ -249,9 +258,10 @@ export function SettingsDialog() {
}
}}
className={cn(
// Cap at 880px but shrink to the viewport so it doesn't clip on
// iPad-portrait widths where a fixed width overflows.
"settings-surface !max-w-[min(880px,calc(100vw-2rem))] h-[560px] w-[min(880px,calc(100vw-2rem))] p-0 overflow-hidden",
// Cap at 960px but shrink to the viewport so it doesn't clip on
// iPad-portrait widths where a fixed width overflows. Height caps
// the same way so short viewports don't get a clipped dialog.
"settings-surface !max-w-[min(960px,calc(100vw-2rem))] h-[min(680px,calc(100dvh-2rem))] w-[min(960px,calc(100vw-2rem))] p-0 overflow-hidden",
// Soft shadow, no outline ring. Pin --radius to the light value so
// corner rounding matches in dark mode.
"shadow-border rounded-xl ring-0 [--radius:1.1rem]",
@ -266,7 +276,9 @@ export function SettingsDialog() {
</DialogDescription>
{/* Keep tab content from expanding the dialog grid. */}
<div className="flex h-full min-h-0 min-w-0 w-full max-sm:flex-col">
<aside className="font-heading flex w-[248px] shrink-0 flex-col border-r border-sidebar-border bg-muted/20 p-2 dark:border-r-0 max-sm:w-full max-sm:border-r-0 max-sm:border-b max-sm:border-sidebar-border">
{/* Match the app shell: tabs on the sidebar fill, content on the
page fill, so both track the active palette. */}
<aside className="font-heading flex w-[248px] shrink-0 flex-col border-r border-sidebar-border bg-sidebar text-sidebar-foreground p-2 dark:border-r-0 max-sm:w-full max-sm:border-r-0 max-sm:border-b max-sm:border-sidebar-border">
<div className="relative mx-1 mt-3 mb-2 max-sm:hidden">
<HugeiconsIcon
icon={Search01Icon}
@ -327,7 +339,7 @@ export function SettingsDialog() {
key={entry}
type="button"
onClick={() => openResult(tab.id, entry)}
className="flex h-[30px] items-center rounded-full pl-10 pr-2.5 text-left text-ui-14 text-foreground transition-colors hover:bg-accent hover:text-accent-foreground"
className="flex h-[30px] items-center rounded-full pl-10 pr-2.5 text-left text-ui-14 text-sidebar-foreground transition-colors hover:bg-accent hover:text-accent-foreground"
>
<span className="min-w-0 truncate">{entry}</span>
</button>
@ -411,7 +423,7 @@ export function SettingsDialog() {
</nav>
</aside>
<main className="relative flex min-h-0 min-w-0 flex-1 flex-col">
<main className="relative flex min-h-0 min-w-0 flex-1 flex-col bg-background">
<button
type="button"
onClick={closeDialog}

13
studio/frontend/src/features/settings/settings-search.ts
View file

@ -104,6 +104,19 @@ export const SETTINGS_SEARCH_INDEX: Record<SettingsTab, TranslationKey[]> = {
"settings.apiKeys.description",
"settings.apiKeys.accessTokens",
],
agents: [
// Heading and intro carry the searched terms ("unsloth start", agent names); titles do not.
"settings.agents.title",
"settings.agents.description",
"settings.agents.intro",
"settings.agents.quickstart.title",
"settings.agents.supportedAgents.title",
"settings.agents.models.title",
"settings.agents.options.title",
"settings.agents.remote.title",
"settings.agents.passthrough.title",
"settings.agents.dryRun.title",
],
connections: [],
voice: [
"settings.voice.dictation.sectionTitle",

2
studio/frontend/src/features/settings/stores/settings-dialog-store.ts
View file

@ -13,6 +13,7 @@ export type SettingsTab =
| "connections"
| "data"
| "api-keys"
| "agents"
| "about";
export type SettingsScrollTarget = "about-updates";
@ -69,6 +70,7 @@ function loadInitialTab(): SettingsTab {
"connections",
"data",
"api-keys",
"agents",
"about",
];
return valid.includes(stored as SettingsTab)

466
studio/frontend/src/features/settings/tabs/agents-tab.tsx Normal file
View file

@ -0,0 +1,466 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getClientPlatform } from "@/components/tauri/window-titlebar";
import { fetchDeviceType, usePlatformStore } from "@/config/env";
import { useT } from "@/i18n";
import type { TranslationKey } from "@/i18n";
import { getApiBase, isTauri } from "@/lib/api-base";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { Tick02Icon } from "@/lib/tick-icon";
import { cn } from "@/lib/utils";
import {
ArrowUpRight01Icon,
Book03Icon,
Copy01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useEffect, useRef, useState } from "react";
import { useChatRuntimeStore } from "@/features/chat";
import { ApiProviderLogo } from "../../chat/api-provider-logo";
import { type CodingAgentsInfo, loadCodingAgents } from "../api/coding-agents";
import {
buildAgentCommand,
isLoopbackHost,
normalizeHost,
} from "../components/agent-command";
import { SettingsSection } from "../components/settings-section";
const DOCS_URL = "https://unsloth.ai/docs/integrations/unsloth-start";
function isLoopbackBase(base: string): boolean {
try {
return isLoopbackHost(normalizeHost(new URL(base).hostname));
} catch {
return false;
}
}
// Desktop-only: a browser loopback URL may be an SSH/port forward to another host.
function canUseLocalAgentDetection(base: string): boolean {
return isTauri && isLoopbackBase(base);
}
// One timeout, reset on re-click and cleared on unmount, so the tick never leaks.
function useCopyButton(text: string) {
const [copied, setCopied] = useState(false);
const timeoutRef = useRef<number | null>(null);
useEffect(
() => () => {
if (timeoutRef.current !== null) window.clearTimeout(timeoutRef.current);
},
[],
);
const copy = async () => {
if (!(await copyToClipboard(text))) return;
setCopied(true);
if (timeoutRef.current !== null) window.clearTimeout(timeoutRef.current);
timeoutRef.current = window.setTimeout(() => {
setCopied(false);
timeoutRef.current = null;
}, 1600);
};
return { copied, copy };
}
// Ids match the backend detection list; agents without an official `logo` asset get a monogram.
// Names are untranslated, so `settings.agents.intro` lists them all to keep them searchable.
const SUPPORTED_AGENTS: {
id: string;
name: string;
logo?: string;
color?: string;
mark?: string;
}[] = [
{ id: "claude", name: "Claude Code", logo: "anthropic" },
{ id: "codex", name: "OpenAI Codex", logo: "openai" },
{ id: "hermes", name: "Hermes", color: "#8B5CF6", mark: "He" },
{ id: "openclaw", name: "OpenClaw", color: "#F59E0B", mark: "Ol" },
{ id: "opencode", name: "OpenCode", color: "#3B82F6", mark: "Oc" },
{ id: "pi", name: "Pi", color: "#EC4899", mark: "Pi" },
];
/** Official brand logo when available, else a brand-colored monogram tile. */
function AgentIcon({
logo,
color,
mark,
}: {
logo?: string;
color?: string;
mark?: string;
}) {
if (logo) {
return (
<span className="flex size-7 shrink-0 items-center justify-center overflow-hidden rounded-md">
<ApiProviderLogo providerType={logo} className="size-7 rounded-md" />
</span>
);
}
return (
<span
aria-hidden={true}
style={{ backgroundColor: color }}
className="flex size-7 shrink-0 items-center justify-center rounded-md font-heading text-ui-11 font-semibold text-white"
>
{mark}
</span>
);
}
function InlineCommand({ command }: { command: string }) {
const t = useT();
const { copied, copy } = useCopyButton(command);
return (
<>
<button
type="button"
onClick={copy}
title={copied ? t("settings.agents.copied") : t("settings.agents.copy")}
aria-label={`${
copied ? t("settings.agents.copied") : t("settings.agents.copy")
}: ${command}`}
className="inline-flex min-w-0 max-w-full items-center gap-2 rounded-md border border-border bg-muted/40 py-1.5 pl-2.5 pr-2 font-mono text-xs text-foreground transition-colors hover:bg-accent hover:text-accent-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring dark:bg-white/[0.04]"
>
{/* Truncate: a remote base makes the command long enough to push the icon out. */}
<span className="truncate whitespace-nowrap">{command}</span>
<HugeiconsIcon
icon={copied ? Tick02Icon : Copy01Icon}
strokeWidth={2}
className={cn(
"size-3.5 shrink-0",
copied ? "text-control-accent" : "text-muted-foreground",
)}
/>
</button>
<span className="sr-only" role="status" aria-live="polite">
{copied ? t("settings.agents.copied") : ""}
</span>
</>
);
}
// Flag tokens are literal; only the descriptions are localized.
const OPTION_ROWS: { flag: string; descKey: TranslationKey }[] = [
{ flag: "--model, -m", descKey: "settings.agents.options.model" },
{
flag: "--context-length",
descKey: "settings.agents.options.contextLength",
},
{ flag: "--gguf-variant", descKey: "settings.agents.options.ggufVariant" },
{
flag: "--load-in-4bit / --no-load-in-4bit",
descKey: "settings.agents.options.loadIn4bit",
},
{
flag: "--tensor-parallel / --no-tensor-parallel",
descKey: "settings.agents.options.tensorParallel",
},
{ flag: "--serve / --no-serve", descKey: "settings.agents.options.serve" },
{
flag: "--launch / --no-launch",
descKey: "settings.agents.options.launch",
},
{
flag: "--persist / --no-persist",
descKey: "settings.agents.options.persist",
},
{ flag: "--api-key", descKey: "settings.agents.options.apiKey" },
{ flag: "--yolo", descKey: "settings.agents.options.yolo" },
];
const QUICKSTART_AGENT = "claude";
// Flags only: agentCommand supplies the prefix so every example targets the Studio
// this tab shows. Kept single line so the copy pastes as-is.
const MODEL_SUFFIX_FLAGS =
"--model unsloth/gemma-4-E2B-it-GGUF:UD-Q4_K_XL --context-length 32768";
const MODEL_VARIANT_FLAGS =
"--model unsloth/gemma-4-E2B-it-GGUF --gguf-variant UD-Q4_K_XL --context-length 32768";
const REMOTE_CMD_UNIX = `export UNSLOTH_STUDIO_URL=https://studio.example.com
export UNSLOTH_API_KEY=sk-unsloth-...
unsloth start claude`;
// PowerShell uses $env: assignments; export is POSIX-only.
const REMOTE_CMD_WINDOWS = `$env:UNSLOTH_STUDIO_URL = "https://studio.example.com"
$env:UNSLOTH_API_KEY = "sk-unsloth-..."
unsloth start claude`;
// Independent alternatives, each with its own copy button (not one script).
const PASSTHROUGH_EXAMPLES = [
{ agent: "claude", flags: "--continue" },
{ agent: "codex", flags: "--persist resume --last" },
];
const DRY_RUN_FLAGS = "--no-launch";
function CommandBlock({ command }: { command: string }) {
const t = useT();
const { copied, copy } = useCopyButton(command);
return (
<div className="group relative">
<pre className="hover-scrollbar overflow-x-auto rounded-lg border border-border bg-muted/40 py-3 pl-3.5 pr-11 text-xs leading-relaxed text-foreground dark:bg-white/[0.04]">
<code className="font-mono whitespace-pre">{command}</code>
</pre>
<button
type="button"
onClick={copy}
aria-label={
copied ? t("settings.agents.copied") : t("settings.agents.copy")
}
className="absolute top-2 right-2 flex size-7 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
>
<HugeiconsIcon
icon={copied ? Tick02Icon : Copy01Icon}
className={cn("size-3.5", copied && "text-control-accent")}
strokeWidth={2}
/>
</button>
<span className="sr-only" role="status" aria-live="polite">
{copied ? t("settings.agents.copied") : ""}
</span>
</div>
);
}
export function AgentsTab() {
const t = useT();
const serverUrl = usePlatformStore((s) => s.serverUrl);
const deviceType = usePlatformStore((s) => s.deviceType);
const [info, setInfo] = useState<CodingAgentsInfo | null>(null);
const origin = typeof window !== "undefined" ? window.location.origin : "";
const localDetection = canUseLocalAgentDetection(serverUrl ?? origin);
// The remote snippet runs on the client, so use the client platform, not deviceType.
// Anchor the match: a bare includes("win") would also match "darwin".
const [isWindowsClient] = useState(() => {
const p = getClientPlatform();
return p.startsWith("win") || p.includes("windows");
});
useEffect(() => {
void fetchDeviceType({ force: true });
}, []);
// A remote backend's PATH says nothing about the machine running the copied command.
useEffect(() => {
if (!localDetection) return;
let cancelled = false;
loadCodingAgents()
.then((next) => {
if (!cancelled) setInfo(next);
})
.catch(() => {
// Best-effort; the tab still works without PATH detection.
});
return () => {
cancelled = true;
};
}, [localDetection]);
// Derive visibility from localDetection instead of clearing info in the effect.
const visibleInfo = localDetection ? info : null;
const detected = new Set(visibleInfo?.detected ?? []);
const remoteCommand = isWindowsClient ? REMOTE_CMD_WINDOWS : REMOTE_CMD_UNIX;
// `codex` needs a GGUF model (unsloth_cli's _require_gguf_for_codex exits otherwise), so flag
// its row instead of offering a failing command. Same three signals the API usage panel uses.
const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
const activeNativePathToken = useChatRuntimeStore(
(s) => s.activeNativePathToken,
);
const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
const isGguf =
activeGgufVariant != null ||
activeNativePathToken != null ||
ggufContextLength != null;
// Build from the reachable base: a bare `unsloth start` only probes 127.0.0.1:8888, but the
// desktop falls back across 8888-8908 and Studio may be remote. The browser must use its own
// origin, since /api/health reports the backend's localhost (the user's, behind a tunnel);
// the desktop has no window origin and falls back to getApiBase() while serverUrl loads.
// No --api-key: the CLI caches an explicit key per base, so a placeholder would overwrite a
// working saved one. Omitting it replays the saved key; the remote section covers first setup.
const commandBase = isTauri ? (serverUrl ?? getApiBase()) : origin;
// The command runs wherever the CLI is. For a loopback base that is this Studio's
// own host, so use deviceType, which reports wsl where the browser would claim
// Windows and emit $env: syntax bash rejects. A remote base is reached from the
// viewer's machine instead, so only the client platform describes that shell.
const commandOs =
(isLoopbackBase(commandBase) ? deviceType === "windows" : isWindowsClient)
? "windows"
: "unix";
const agentCommand = (agentId: string) =>
buildAgentCommand(commandBase, null, commandOs, agentId);
const example = (agentId: string, flags: string) =>
`${agentCommand(agentId)} ${flags}`;
return (
<div className="flex min-w-0 max-w-full flex-col gap-6">
{/* data-settings-label lets indexed settings search scroll to these. */}
<header className="flex min-w-0 flex-col gap-1">
<h1
data-settings-label={t("settings.agents.title")}
className="text-xl font-semibold font-heading"
>
{t("settings.agents.title")}
</h1>
<p
data-settings-label={t("settings.agents.description")}
className="text-xs text-muted-foreground leading-relaxed"
>
{t("settings.agents.description")}
</p>
</header>
<p
data-settings-label={t("settings.agents.intro")}
className="text-sm text-muted-foreground leading-relaxed"
>
<code className="rounded bg-muted px-1 py-0.5 font-mono text-[0.85em] text-foreground dark:bg-white/[0.08]">
unsloth start
</code>{" "}
{t("settings.agents.intro")}
</p>
<a
href={DOCS_URL}
target="_blank"
rel="noopener noreferrer"
className="inline-flex w-fit items-center gap-1.5 text-xs font-medium text-muted-foreground hover:text-foreground"
>
<HugeiconsIcon icon={Book03Icon} className="size-3.5" />
{t("settings.agents.readDocs")}
<HugeiconsIcon icon={ArrowUpRight01Icon} className="size-3" />
</a>
<SettingsSection
title={t("settings.agents.quickstart.title")}
description={t("settings.agents.quickstart.description")}
>
<div className="pt-2">
<CommandBlock command={agentCommand(QUICKSTART_AGENT)} />
</div>
</SettingsSection>
<SettingsSection
title={t("settings.agents.supportedAgents.title")}
description={t("settings.agents.supportedAgents.description")}
>
<div className="mt-1 flex flex-col divide-y divide-border/60">
{SUPPORTED_AGENTS.map((agent) => (
<div
key={agent.id}
className="flex flex-wrap items-center justify-between gap-x-4 gap-y-2 py-2.5"
>
<div className="flex min-w-0 items-center gap-3">
<AgentIcon
logo={agent.logo}
color={agent.color}
mark={agent.mark}
/>
<span className="truncate text-sm font-medium text-foreground">
{agent.name}
</span>
{detected.has(agent.id) ? (
<span className="shrink-0 rounded-full bg-control-accent/10 px-2 py-1 text-ui-10 leading-none font-semibold text-control-accent">
{t("settings.agents.quickstart.installed")}
</span>
) : null}
{agent.id === "codex" && !isGguf ? (
<span className="shrink-0 rounded-full bg-muted px-2 py-1 text-ui-10 leading-none font-semibold text-muted-foreground">
{t("settings.agents.supportedAgents.requiresGguf")}
</span>
) : null}
</div>
<InlineCommand command={agentCommand(agent.id)} />
</div>
))}
</div>
{visibleInfo !== null && detected.size === 0 ? (
<p className="pt-3 text-xs text-muted-foreground">
{t("settings.agents.quickstart.noneDetected")}
</p>
) : null}
</SettingsSection>
<SettingsSection
title={t("settings.agents.models.title")}
description={t("settings.agents.models.description")}
>
<div className="flex flex-col gap-3 pt-2">
<div className="flex flex-col gap-1.5">
<span className="text-xs font-medium text-foreground">
{t("settings.agents.models.suffixLabel")}
</span>
<CommandBlock command={example("codex", MODEL_SUFFIX_FLAGS)} />
</div>
<div className="flex flex-col gap-1.5">
<span className="text-xs font-medium text-foreground">
{t("settings.agents.models.variantLabel")}
</span>
<CommandBlock command={example("codex", MODEL_VARIANT_FLAGS)} />
</div>
</div>
</SettingsSection>
<SettingsSection
title={t("settings.agents.options.title")}
description={t("settings.agents.options.description")}
>
<div className="mt-1 flex flex-col divide-y divide-border/60">
{OPTION_ROWS.map((row) => (
<div
key={row.flag}
className="grid grid-cols-[minmax(0,11rem)_1fr] items-start gap-x-5 gap-y-1 py-2.5 max-sm:grid-cols-1"
>
<code className="min-w-0 break-words font-mono text-xs font-medium text-foreground">
{row.flag}
</code>
<span className="text-xs leading-relaxed text-muted-foreground">
{t(row.descKey)}
</span>
</div>
))}
</div>
</SettingsSection>
<SettingsSection
title={t("settings.agents.remote.title")}
description={t("settings.agents.remote.description")}
>
<div className="pt-2">
<CommandBlock command={remoteCommand} />
</div>
</SettingsSection>
<SettingsSection
title={t("settings.agents.passthrough.title")}
description={t("settings.agents.passthrough.description")}
>
<div className="flex flex-col gap-2 pt-2">
{PASSTHROUGH_EXAMPLES.map(({ agent, flags }) => (
<CommandBlock key={flags} command={example(agent, flags)} />
))}
</div>
</SettingsSection>
<SettingsSection
title={t("settings.agents.dryRun.title")}
description={t("settings.agents.dryRun.description")}
>
<div className="pt-2">
<CommandBlock command={example("claude", DRY_RUN_FLAGS)} />
</div>
</SettingsSection>
</div>
);
}

62
studio/frontend/src/i18n/locales/en.ts
View file

@ -107,6 +107,7 @@ export const en = {
connections: "Connections",
data: "Data",
apiKeys: "API",
agents: "Agents",
about: "About",
},
voice: {
@ -594,6 +595,67 @@ export const en = {
unknown: "Unknown",
},
},
agents: {
title: "Agents (unsloth start)",
description:
"Connect coding agents like Claude Code and Codex to a model running locally in Unsloth.",
intro:
"connects Claude Code, Codex, Hermes, OpenClaw, OpenCode, Pi and other agents to a model served locally by Unsloth, fully offline on your own hardware. It runs a OpenAI-compatible server for the agent and never touches your agent's config files.",
readDocs: "Read the docs",
copy: "Copy",
copied: "Copied",
quickstart: {
title: "Quickstart",
description:
"Launch an agent against the model currently loaded in Studio. Load a model first, then swap claude for any supported agent below.",
noneDetected: "No supported agent CLIs were found on your PATH.",
installed: "Installed",
},
supportedAgents: {
title: "Supported agents",
description: "Each agent launches with its own command:",
requiresGguf: "Needs a GGUF model",
},
models: {
title: "Choosing a model",
description:
"Pass --model to pick a model and quantization, and --context-length to set the window. Use a quantization suffix, or an explicit --gguf-variant flag.",
suffixLabel: "With a quantization suffix",
variantLabel: "With an explicit variant flag",
},
options: {
title: "Common options",
description:
"Unsloth flags are parsed first; anything it doesn't recognize is passed straight through to the agent.",
model:
"Select a model. Without --model, unsloth start uses the model currently loaded in Studio and errors if none is loaded.",
contextLength:
"Set the requested context length (alias: --max-seq-length).",
ggufVariant: "Choose the GGUF quantization variant.",
loadIn4bit: "Toggle 4-bit loading for Hugging Face models.",
tensorParallel: "Toggle tensor-parallel across multiple GPUs.",
serve: "Enable or disable the automatic local server.",
launch: "Launch the agent, or just print the command and environment.",
persist: "Keep Unsloth-managed agent storage between runs.",
apiKey: "Provide your Unsloth API key (or set UNSLOTH_API_KEY).",
yolo: "Skip approval prompts. Use only in trusted environments.",
},
remote: {
title: "Connect to a remote Studio",
description:
"Point unsloth start at a Studio running elsewhere by setting these before launching (or pass --api-key directly):",
},
passthrough: {
title: "Passing agent arguments",
description:
"Arguments after the Unsloth flags are forwarded to the agent itself, so native commands like resume still work:",
},
dryRun: {
title: "Preview without launching",
description:
"Add --no-launch to print the environment and command instead of launching the agent. If --model is set, the model may still be resolved and loaded.",
},
},
chat: {
title: "Chat",
description: "Customize how chat behaves on this device.",

162
studio/install_llama_prebuilt.py
View file

@ -63,6 +63,7 @@ EXIT_SUCCESS = 0
EXIT_FALLBACK = 2
EXIT_ERROR = 1
EXIT_BUSY = 3
EXIT_NO_SPACE = 4
# DiskPart-prompt suppression. RunAsInvoker does NOT stop amd-smi's runtime
# elevation (its manifest is asInvoker), so this is just harmless belt-and-
@ -3674,7 +3675,8 @@ def hydrate_source_tree(
break
except Exception as exc:
last_exc = exc
if index == len(source_urls) - 1:
# A full disk fails every mirror; stop so a later 404 cannot mask it.
if _environment_fatal_reason(exc) or index == len(source_urls) - 1:
raise
log(f"source tree download failed from {source_url}: {exc}")
if not downloaded:
@ -6000,6 +6002,14 @@ def validate_prebuilt_attempts(
)
raise ExistingInstallSatisfied(attempt, tried_fallback)
# Advisory: a few GB free usually fits, and rejecting here would also skip
# the source-build fallback.
if index == 0:
low_disk = _low_disk_warning(install_dir)
if low_disk is not None:
log(low_disk)
_log_disk_space_help()
staging_dir = create_install_staging_dir(install_dir)
quantized_path = work_dir / f"stories260K-q4-{index}.gguf"
if quantized_path.exists():
@ -6028,7 +6038,9 @@ def validate_prebuilt_attempts(
attempt_error = PrebuiltFallback(
f"candidate attempt failed before activation for {attempt.name}: {exc}"
)
if index == len(attempt_list) - 1:
if _environment_fatal_reason(exc) or index == len(attempt_list) - 1:
if attempt_error is exc:
raise
raise attempt_error from exc
log(
"selected CUDA bundle failed before activation; trying next prebuilt fallback "
@ -6149,6 +6161,132 @@ def diffusion_visual_server_backfill_needed(
return True
def _causal_chain(exc: BaseException) -> Iterable[BaseException]:
seen: set[int] = set()
current: BaseException | None = exc
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
# `raise X from None` sets __suppress_context__: the earlier exception is
# unrelated, so following __context__ anyway would misreport the cause.
if current.__cause__ is not None:
current = current.__cause__
elif current.__suppress_context__:
current = None
else:
current = current.__context__
# ERROR_HANDLE_DISK_FULL / ERROR_DISK_FULL. CPython's PC/errmap.h maps 112 to
# ENOSPC but has no case for 39, which arrives as EINVAL, so check winerror too.
_WINDOWS_DISK_FULL = (39, 112)
# A quota (NFS/XFS/container) leaves blocks this user cannot have, so the bigger
# source build is just as doomed; named apart from ENOSPC so df does not mislead.
# Guarded: the MSVC CRT has no EDQUOT, so on Windows CPython aliases it to the
# Winsock WSAEDQUOT (10069), which no file write raises.
_DISK_FULL_ERRNOS = {errno.ENOSPC: "no space left on device"}
if hasattr(errno, "EDQUOT"):
_DISK_FULL_ERRNOS[errno.EDQUOT] = "disk quota exceeded"
def _winerror_of(exc: OSError) -> Any:
"""exc.winerror, defensively. Not getattr(exc, ..., None): urllib's HTTPError
is an OSError that proxies unknown attributes to a wrapped file object and
raises KeyError (not AttributeError) on 3.9, which getattr will not swallow.
A 404 from a mirror must not crash the classifier."""
try:
return exc.winerror
except Exception:
return None
def _out_of_space_reason(exc: BaseException) -> str | None:
"""Why `exc` means the install cannot fit, or None if it means something else."""
if isinstance(exc, OSError):
reason = _DISK_FULL_ERRNOS.get(exc.errno)
if reason is not None:
return reason
if _winerror_of(exc) in _WINDOWS_DISK_FULL:
return "no space left on device"
# shutil.copytree stringifies each per-file OSError and raises Error(errors)
# outside the except block, so errno and the chain are gone and only text
# survives. OSError.__str__ returns early on winerror, so Windows reads
# "[WinError 112]" and never "[Errno 28]": match both, brackets included so
# WinError 112 does not match WinError 1120.
if isinstance(exc, shutil.Error):
text = str(exc)
for code, reason in _DISK_FULL_ERRNOS.items():
if f"[Errno {code}]" in text:
return reason
if any(f"[WinError {code}]" in text for code in _WINDOWS_DISK_FULL):
return "no space left on device"
return None
def _environment_fatal_reason(exc: BaseException) -> str | None:
for cause in _causal_chain(exc):
reason = _out_of_space_reason(cause)
if reason is not None:
return reason
return None
def _log_disk_space_help() -> None:
log(
"free up space or point TMPDIR and UNSLOTH_STUDIO_HOME at a larger "
"volume (e.g. /workspace), then re-run"
)
@contextmanager
def scratch_dir(prefix: str) -> Iterator[Path]:
"""Temp dir whose cleanup never raises: an rmtree failure on the way out would
replace the in-flight exception and lose EXIT_NO_SPACE. Not
TemporaryDirectory(ignore_cleanup_errors = True), which is 3.10+ (setup.sh
still runs this helper under the host python, and we support 3.9)."""
path = Path(tempfile.mkdtemp(prefix = prefix))
try:
yield path
finally:
shutil.rmtree(path, ignore_errors = True)
def _first_existing_ancestor(path: Path) -> Path:
current = path
while current != current.parent and not current.exists():
current = current.parent
return current
def _low_disk_warning(install_dir: Path, *, advised_gb: float = 5.0) -> str | None:
"""Advisory only, never fatal. A prebuilt install peaks well under 1 GB (the
largest published bundle is 0.77 GB, macOS is 0.01 GB), so a fixed threshold
cannot decide whether this host has room -- a real ENOSPC decides that. The
number here is the headroom a source-build fallback would want."""
advised = int(advised_gb * (1024**3))
targets = {
"build/download scratch (TMPDIR)": Path(tempfile.gettempdir()),
"llama.cpp install dir": _first_existing_ancestor(install_dir),
}
for label, path in targets.items():
try:
free = shutil.disk_usage(path).free
except OSError:
continue
if free < advised:
return (
f"low disk space for llama.cpp: {label} at {path} has "
f"{free / (1024**3):.1f} GB free (~{advised_gb:.0f} GB recommended)"
)
return None
def _fail_no_space(reason: str) -> None:
log(reason)
_log_disk_space_help()
raise SystemExit(EXIT_NO_SPACE)
def install_prebuilt(
install_dir: Path,
llama_tag: str,
@ -6217,8 +6355,7 @@ def install_prebuilt(
# recorded so the updater re-asserts it (#7213).
sync_marker_force_cpu(install_dir, persist_force_cpu)
return
with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp:
work_dir = Path(tmp)
with scratch_dir("unsloth-llama-prebuilt-") as work_dir:
probe_path = work_dir / "stories260K.gguf"
download_validation_model(probe_path, validation_model_cache_path(install_dir))
release_count = len(release_plans)
@ -6263,6 +6400,8 @@ def install_prebuilt(
except ExistingInstallSatisfied:
return
except PrebuiltFallback as exc:
if _environment_fatal_reason(exc):
raise
if release_index == release_count - 1:
raise
log(
@ -6296,6 +6435,11 @@ def install_prebuilt(
log(f"prebuilt busy reason: {exc}")
raise SystemExit(EXIT_BUSY) from exc
except PrebuiltFallback as exc:
fatal = _environment_fatal_reason(exc)
if fatal:
log(f"prebuilt install failed: {fatal}")
_log_disk_space_help()
raise SystemExit(EXIT_NO_SPACE) from exc
log("prebuilt install path failed; falling back to source build")
log(f"prebuilt fallback reason: {exc}")
report = collect_system_report(host, choice, install_dir)
@ -6466,6 +6610,10 @@ def main() -> int:
install_kind = args.install_kind,
)
except PrebuiltFallback as exc:
# A full disk is not a bad build: the CPU source rebuild needs more space.
fatal = _environment_fatal_reason(exc)
if fatal:
_fail_no_space(f"install validation failed: {fatal}")
print(str(exc), file = sys.stderr)
raise SystemExit(EXIT_FALLBACK) from exc
return EXIT_SUCCESS
@ -6595,9 +6743,15 @@ if __name__ == "__main__":
# Expected when the published repo (e.g. ggml-org/llama.cpp) has no
# prebuilt manifest. Exit quietly with EXIT_FALLBACK so the caller
# falls back to source build without a noisy "fatal helper error".
fatal = _environment_fatal_reason(exc)
if fatal:
_fail_no_space(f"prebuilt install failed: {fatal}")
log(textwrap.shorten(str(exc), width = 400, placeholder = "..."))
raise SystemExit(EXIT_FALLBACK)
except Exception as exc:
fatal = _environment_fatal_reason(exc)
if fatal:
_fail_no_space(f"prebuilt install failed: {fatal}")
message = textwrap.shorten(str(exc), width = 400, placeholder = "...")
log(f"fatal helper error: {message}")
raise SystemExit(EXIT_ERROR)

38
studio/install_python_stack.py
View file

@ -96,7 +96,9 @@ def _strix_needs_amd_arch_index(ver: tuple[int, int]) -> bool:
# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug).
# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare.
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset(
{"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"}
)
# pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't
# floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1.
@ -124,6 +126,7 @@ _WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1151": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1150": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1152": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
}
_PYTORCH_WHL_BASE = (
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
@ -369,6 +372,7 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
"gfx1200": "gfx120X-all", # RDNA 4
"gfx1151": "gfx1151",
"gfx1150": "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1152": "gfx1152", # RDNA 3.5 (Krackan Point)
"gfx1103": "gfx110X-all",
"gfx1102": "gfx110X-all", # RDNA 3
"gfx1101": "gfx110X-all",
@ -738,19 +742,18 @@ def _detect_windows_gfx_arch() -> str | None:
# prebuilts / AMD Windows torch indexes support; unknown names return None
# (callers then fall back cleanly to CPU).
_WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [
(r"9070 XT|9080", "gfx1201"), # RDNA 4 (Radeon RX 9070 XT / 9080)
(r"9070|9060", "gfx1200"), # RDNA 4 (Radeon RX 9070 / 9060)
(r"9070|9080", "gfx1201"), # RDNA 4 (Navi 48: Radeon RX 9070 XT / 9070 GRE / 9070 / 9080)
(r"9060", "gfx1200"), # RDNA 4 (Navi 44: Radeon RX 9060 XT / 9060)
# RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
(r"8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"),
# RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
(
r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]"
r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33",
"gfx1150",
),
# RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
(r"890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]", "gfx1150"),
# RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
(r"860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", "gfx1152"),
# RDNA 3 desktop / workstation (Navi 31)
(r"RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700", "gfx1100"),
(r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710", "gfx1102"), # Navi 33
(r"RX 7900|PRO W7900|PRO W7800", "gfx1100"),
(r"RX 7800|RX 7700(?!S)|PRO W7700|PRO V710", "gfx1101"), # Navi 32
(r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500", "gfx1102"), # Navi 33
# RDNA 3 iGPU (Phoenix / Hawk Point)
(r"780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme", "gfx1103"),
(r"RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900", "gfx1030"), # Navi 21
@ -777,13 +780,12 @@ def _linux_amd_gfx_from_cpuinfo() -> "str | None":
return None
if re.search(r"Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo", text, re.IGNORECASE):
return "gfx1151"
if re.search(
r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]"
r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33",
text,
re.IGNORECASE,
):
if re.search(r"890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]", text, re.IGNORECASE):
return "gfx1150"
if re.search(
r"860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", text, re.IGNORECASE
):
return "gfx1152"
return None
@ -1896,7 +1898,7 @@ def _ensure_rocm_torch() -> None:
# An explicit ROCm pin is authoritative: never auto-reroute it.
if _strix_needs_amd_arch_index(ver) and _explicit_rocm_torch_index_url() is None:
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150"}
_strix_gfx = {"gfx1151", "gfx1150", "gfx1152"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
if _detected_strix:
# Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first);

34
studio/setup.ps1
View file

@ -484,7 +484,7 @@ function Redact-InstallOutput {
# the install-spec path below and the other installers; other leaves ship <2.11 and stay default.
function Test-RocmGfx211Leaf {
param([string]$Leaf)
return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf
return @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $Leaf
}
# rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer
@ -1496,12 +1496,14 @@ if (-not $HasNvidiaSmi) {
# (gfx120X/110X/1151/1150/103X); unknown names fall back cleanly to CPU.
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (Radeon RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (Radeon RX 9070 / 9060)
@{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: Radeon RX 9070 XT / 9070 GRE / 9070 / 9080)
@{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: Radeon RX 9060 XT / 9060)
@{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
@{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
@{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
@{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31)
@{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33)
@{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point)
@{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family
@ -2773,7 +2775,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
# for those or a correct CPU venv rebuilds every update.
$_rocmWheelArches = @(
"gfx1201", "gfx1200", # RDNA 4
"gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1151", "gfx1150", "gfx1152", # RDNA 3.5 (Strix Halo/Point, Krackan Point)
"gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3
"gfx1036", "gfx1035", "gfx1034", "gfx1033", "gfx1032", "gfx1031", "gfx1030", # RDNA 2 (RX 6000)
"gfx90a", "gfx908" # MI200 / MI100
@ -3064,6 +3066,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
@ -3078,6 +3081,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
$torchFloorMap = @{
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
"gfx1152" = "torch>=2.11.0,<2.12.0"
}
# Companion ranges for torchvision/torchaudio -- must stay in sync with the
# torch ceiling so pip can always find a consistent trio on AMD's per-arch
@ -3089,10 +3093,12 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
$torchvisionFloorMap = @{
"gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0"
"gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0"
"gfx1152" = "torchvision>=0.26.0,<0.27.0"
}
$torchaudioFloorMap = @{
"gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0"
"gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0"
"gfx1152" = "torchaudio>=2.11.0,<2.12.0"
}
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
$ROCmTorchSpec = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { "torch" }
@ -3104,7 +3110,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
# GPU arch detected but not in the supported wheel map — warn explicitly
# so the user knows why they are getting CPU PyTorch instead of ROCm.
substep "[WARN] AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow"
substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx1030-1036 (RDNA 2), gfx90a, gfx908" "Yellow"
substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151/1152 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx1030-1036 (RDNA 2), gfx90a, gfx908" "Yellow"
} else {
# HIP SDK present ($HasROCm=true via amd-smi) but gcnArchName was not
# readable — warn rather than silently falling back to CPU PyTorch.
@ -3757,6 +3763,18 @@ if ($LocalLlamaCppLinked) {
}
substep "Close Unsloth or other llama.cpp users and retry" "Yellow"
exit 3
} elseif ($prebuiltExit -eq 4) {
step "llama.cpp" "not enough disk space to install llama.cpp" "Yellow"
Write-LlamaFailureLog -Output $prebuiltOutput
substep "Free up disk or move UNSLOTH_STUDIO_HOME/TEMP to a larger volume, then re-run" "Yellow"
$PreservedLlamaServerFound = $false
foreach ($_cand in @(
(Join-Path $LlamaCppDir "llama-server.exe"),
(Join-Path $LlamaCppDir "build\bin\llama-server.exe"),
(Join-Path $LlamaCppDir "build\bin\Release\llama-server.exe"))) {
if (Test-Path -LiteralPath $_cand) { $PreservedLlamaServerFound = $true; break }
}
if (-not $PreservedLlamaServerFound) { $script:LlamaCppDegraded = $true }
} else {
step "llama.cpp" "prebuilt install failed (continuing)" "Yellow"
Write-LlamaFailureLog -Output $prebuiltOutput

97
studio/setup.sh
View file

@ -67,6 +67,15 @@ fi
step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
substep() { printf " %-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
# ── Helper: can the controlling terminal actually be opened for reading? ──
# `test -r` only checks permission bits, which look fine in containers and
# systemd units where open() then fails with ENXIO. Probe with a real open.
# Mirrors install.sh's _can_read_tty; defined here too because setup.sh runs
# as its own process (install.sh invokes it, it does not source it).
_can_read_tty() {
( : </dev/tty ) >/dev/null 2>&1
}
_is_verbose() {
[ "${UNSLOTH_VERBOSE:-0}" = "1" ]
}
@ -1167,12 +1176,14 @@ elif [ "$_setup_amd_detected" = true ]; then
# gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on
# gfx1102 (bash case has no negative lookahead like the PS tables).
case "$_setup_mkt" in
*"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4
*9070*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 (Navi 48)
*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 (Navi 44)
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _setup_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
*"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
*"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _setup_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _setup_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _setup_gfx="gfx1101" ;; # RDNA 3 (Navi 32)
*"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point)
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _setup_gfx="gfx1030" ;; # RDNA 2 (Navi 21)
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _setup_gfx="gfx1032" ;; # RDNA 2 (Navi 23)
@ -1222,6 +1233,7 @@ LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
_NEED_LLAMA_SOURCE_BUILD=false
_LLAMA_CPP_DEGRADED=false
_LLAMA_CPP_NO_SPACE=false
_LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}"
_HOST_SYSTEM="$(uname -s 2>/dev/null || true)"
@ -1449,6 +1461,13 @@ else
fi
substep "close Unsloth or other llama.cpp users and retry"
exit 3
elif [ "$_PREBUILT_STATUS" -eq 4 ]; then
step "llama.cpp" "not enough disk space to install llama.cpp" "$C_WARN"
print_llama_error_log "$_PREBUILT_LOG"
rm -f "$_PREBUILT_LOG"
substep "free up disk or move UNSLOTH_STUDIO_HOME/TMPDIR to a larger volume, then re-run"
_LLAMA_CPP_NO_SPACE=true
_has_local_llama_server "$LLAMA_CPP_DIR" || _LLAMA_CPP_DEGRADED=true
else
step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
print_llama_error_log "$_PREBUILT_LOG"
@ -1500,25 +1519,46 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && grep -qi microsoft /proc/version 2>
step "gguf deps" "installed"
elif command -v sudo >/dev/null 2>&1; then
step "gguf deps" "sudo required for: $_STILL_MISSING" "$C_WARN"
printf " %-15s" ""
printf "accept? [Y/n] "
if [ -r /dev/tty ]; then
read -r REPLY </dev/tty || REPLY="y"
if _can_read_tty; then
printf " %-15s" ""
printf "accept? [Y/n] "
# The device opened, so a failed read is EOF, not consent: decline.
read -r REPLY </dev/tty || REPLY="n"
case "$REPLY" in
[nN]*)
substep "skipped -- run manually:"
substep "sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
;;
*)
# Degrade like the no-sudo branch below rather than letting
# set -e abort setup on a bare apt error: missing GGUF build
# deps are recoverable, not fatal.
if sudo apt-get update -y </dev/null &&
sudo apt-get install -y $_STILL_MISSING </dev/null; then
step "gguf deps" "installed"
else
step "gguf deps" "install failed -- run manually:" "$C_WARN"
substep "sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
fi
;;
esac
else
REPLY="y"
fi
case "$REPLY" in
[nN]*)
substep "skipped -- run manually:"
substep "sudo apt-get install -y $_STILL_MISSING"
# Nobody can answer a prompt or type a password here, so -n makes
# sudo refuse rather than prompt into a closed stdin, and -k ignores
# any cached timestamp so only a real NOPASSWD rule gets through.
# Same treatment as install.sh's _smart_apt_install. This is the WSL
# GGUF-export case noted above, where sudo does want a password.
if sudo -n -k apt-get update -y </dev/null &&
sudo -n -k apt-get install -y $_STILL_MISSING </dev/null; then
step "gguf deps" "installed (non-interactive sudo)"
else
step "gguf deps" "needs sudo, no terminal -- run manually:" "$C_WARN"
substep "sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
;;
*)
sudo apt-get update -y
sudo apt-get install -y $_STILL_MISSING
step "gguf deps" "installed"
;;
esac
fi
fi
else
step "gguf deps" "missing (no sudo) -- install manually:" "$C_WARN"
substep "apt-get install -y $_STILL_MISSING"
@ -1944,7 +1984,14 @@ else
--validate-install "$_BUILD_TMP"
)
[ -n "$_SMOKE_KIND" ] && _SMOKE_CMD+=(--install-kind "$_SMOKE_KIND")
if ! run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}"; then
_SMOKE_RC=0
run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}" || _SMOKE_RC=$?
# Exit 4 is a full disk, not a bad build: the CPU rebuild needs even
# more space, so keep what we already have.
if [ "$_SMOKE_RC" -eq 4 ]; then
substep "not enough disk space to validate the $_FB_LABEL build; keeping it" "$C_WARN"
_LLAMA_CPP_NO_SPACE=true
elif [ "$_SMOKE_RC" -ne 0 ]; then
substep "$_FB_LABEL source build failed smoke test; retrying CPU build..." "$C_WARN"
_TRY_METAL_CPU_FALLBACK=false
rm -rf "$_BUILD_TMP/build"
@ -2001,8 +2048,10 @@ fi # end _SKIP_GGUF_BUILD check
# An arm64 Linux GPU host source-builds for the GPU above. If that produced no
# binary, install the fork's arm64 CPU prebuilt (app-<tag>-linux-arm64-cpu.tar.gz)
# instead of leaving the host without llama.cpp. --cpu-fallback drops the GPU
# attributes so the CPU bundle is selected rather than re-attempting CUDA.
# attributes so the CPU bundle is selected rather than re-attempting CUDA. Skipped
# on a full disk: the retry fails the same way and buries the hint.
if [ "$_LLAMA_CPP_DEGRADED" = true ] \
&& [ "$_LLAMA_CPP_NO_SPACE" != true ] \
&& [ "$_HOST_SYSTEM" = "Linux" ] \
&& { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; }; then
substep "GPU source build unavailable; trying arm64 CPU prebuilt..."

3
tests/_zoo_rocm_spoof.py
View file

@ -25,6 +25,7 @@ _PROFILES: dict[str, tuple[str, tuple[int, int], str]] = {
"gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"),
"gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"),
"gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU
"gfx1152": ("AMD Radeon 860M", (11, 5), "7.2.1"),
"gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"),
"gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4
"gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"),
@ -73,7 +74,7 @@ def apply(gfx: str = "gfx1100", device_count: int = 1) -> None:
_p.total_memory = 16 * 1024**3
_p.multi_processor_count = 40
_p.warp_size = 32 # RDNA wavefront (CDNA is 64)
_p.is_integrated = gfx in ("gfx1150", "gfx1151")
_p.is_integrated = gfx in ("gfx1150", "gfx1151", "gfx1152")
_p.is_multi_gpu_board = False
torch.cuda.get_device_properties = lambda *a, **k: _p

190
tests/fast_inference/test_fast_inference.py Normal file
View file

@ -0,0 +1,190 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
# ruff: noqa
"""GRPO smoke test for the ``fast_inference=True`` vLLM rollout path.
Exercises the vLLM LoRA activation path (`WorkerLoRAManager`) that regressed on
vLLM >= 0.25.0 (unsloth#7283): the stacked `WeightsMapper` collapsed q/k/v and
gate/up LoRA weights onto one key, crashing adapter activation with
`IndexError`. All seven attention and MLP projections are LoRA targets so both
the fused `qkv_proj` and `gate_up_proj` families are covered.
Kept deliberately tiny so it finishes in well under a minute: a 0.6B model,
`enforce_eager`, no torch.compile, three short training steps, and short
prompts/completions. Seeded, so the asserted metrics are reproducible.
Run directly (`python tests/fast_inference/test_fast_inference.py`) or via
pytest; it skips automatically when no CUDA device is present.
"""
import math
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).parents[2]
sys.path.insert(0, str(REPO_ROOT))
import pytest
import torch
from tests.utils import header_footer_context
MODEL_NAME = "unsloth/Qwen3-0.6B"
MAX_SEQ_LENGTH = 256
LORA_RANK = 8
NUM_GENERATIONS = 2
MAX_PROMPT_LENGTH = 64
MAX_COMPLETION_LENGTH = 16
# >1 so the updated LoRA adapter is re-synced into vLLM on every step, not just
# loaded once; that repeat sync is the path that regressed.
MAX_STEPS = 3
GPU_MEMORY_UTILIZATION = 0.3
COMPILATION_CONFIG = 0
# Pins torch's global RNG (via the Trainer's set_seed), which the colocated vLLM
# sampler draws from, so the rollout and every metric below is reproducible.
SEED = 42
# Loose sanity bounds, not fitted values: they catch divergence and degenerate
# rollouts while staying valid across GPUs, models and vLLM versions.
MAX_CHARS_PER_TOKEN = 20
MAX_GRAD_NORM = 1e3
MAX_KL = 1.0
# All attention + MLP projections, so both fused vLLM LoRA families (qkv_proj and
# gate_up_proj) are exercised -- the >= 0.25.0 collision hit both.
TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
SYSTEM_PROMPT = "Respond concisely."
QUESTIONS = ["What is the capital of France?", "What is 2 + 2?"]
PROMPTS = [
[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": q}]
for q in QUESTIONS
]
def length_reward_func(completions, **kwargs) -> list[float]:
"""Reward longer completions. The fractional tie-break keeps rewards distinct
even if the model samples equal-length completions, so GRPO advantages are
never all-zero and the step stays meaningful on any vLLM/GPU combination."""
n = len(completions)
return [float(len(c[0]["content"])) + i / (n + 1) for i, c in enumerate(completions)]
def _metric(metrics, *names):
"""First present key; TRL spells some metrics differently across versions."""
for name in names:
if name in metrics:
return metrics[name]
return None
@pytest.mark.skipif(not torch.cuda.is_available(), reason = "fast_inference needs a CUDA GPU + vLLM")
def test_fast_inference():
# Import here, not at module load: importing unsloth probes for an
# accelerator and errors on CPU-only machines, so deferring keeps pytest
# collection and the skip path import-free. Unsloth must precede TRL.
from unsloth import FastLanguageModel
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer
with header_footer_context("Load model (fast_inference=True)"):
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = MODEL_NAME,
max_seq_length = MAX_SEQ_LENGTH,
load_in_4bit = False,
fast_inference = True,
max_lora_rank = LORA_RANK,
gpu_memory_utilization = GPU_MEMORY_UTILIZATION,
enforce_eager = True, # skip CUDA graph capture for fast startup
compilation_config = COMPILATION_CONFIG,
)
assert hasattr(model, "vllm_engine"), "fast_inference=True did not attach a vLLM engine"
model = FastLanguageModel.get_peft_model(
model,
r = LORA_RANK,
target_modules = TARGET_MODULES,
lora_alpha = LORA_RANK,
use_gradient_checkpointing = False,
random_state = SEED,
)
dataset = Dataset.from_dict({"prompt": PROMPTS})
with header_footer_context("GRPO config and trainer"):
training_args = GRPOConfig(
learning_rate = 5e-6,
per_device_train_batch_size = NUM_GENERATIONS,
gradient_accumulation_steps = 1,
num_generations = NUM_GENERATIONS,
max_prompt_length = MAX_PROMPT_LENGTH,
max_completion_length = MAX_COMPLETION_LENGTH,
max_steps = MAX_STEPS,
logging_steps = 1,
report_to = "none",
seed = SEED,
)
trainer = GRPOTrainer(
model = model,
processing_class = tokenizer,
reward_funcs = [length_reward_func],
args = training_args,
train_dataset = dataset,
)
# The trainer must actually route rollouts through vLLM, otherwise it would
# fall back to HF generation and never exercise WorkerLoRAManager.
assert trainer.args.use_vllm, "GRPO is not configured to use vLLM"
assert getattr(trainer, "llm", None) is not None, "GRPO did not bind a vLLM engine"
with header_footer_context("GRPO train (vLLM LoRA rollout)"):
trainer_stats = trainer.train()
assert trainer_stats is not None, "trainer.train() returned None"
assert trainer_stats.global_step == MAX_STEPS, "GRPO ran the wrong number of steps"
assert math.isfinite(trainer_stats.training_loss), "training loss is not finite"
# Without these, a rollout that silently produced nothing, or an update that
# diverged to NaN, would still pass the wiring assertions above.
steps = [log for log in trainer.state.log_history if "loss" in log]
assert len(steps) == MAX_STEPS, f"expected {MAX_STEPS} logged steps, got {len(steps)}"
# Every reward is a completion's character count, so this bounds reward and
# its spread without hard-coding model-specific values.
max_reward = MAX_COMPLETION_LENGTH * MAX_CHARS_PER_TOKEN
for i, step in enumerate(steps, start = 1):
loss = step["loss"]
grad_norm = step.get("grad_norm")
reward = step.get("reward")
zero_std = step.get("frac_reward_zero_std")
kl = step.get("kl")
# Key names differ across the supported TRL range, so accept either.
length = _metric(step, "completion_length", "completions/mean_length")
reward_std = _metric(step, "reward_std", "rewards/std")
assert math.isfinite(loss), f"step {i}: loss not finite ({loss})"
assert grad_norm is not None, f"step {i}: no grad_norm logged"
assert math.isfinite(grad_norm), f"step {i}: grad_norm not finite ({grad_norm})"
# Sign check only: a step can legitimately be near zero (0.004 observed),
# so any tighter lower bound would be flaky.
assert 0.0 < grad_norm < MAX_GRAD_NORM, f"step {i}: grad_norm {grad_norm}"
assert length is not None, f"step {i}: no completion length logged"
assert 0.0 < length <= MAX_COMPLETION_LENGTH, f"step {i}: empty rollout ({length})"
assert reward is not None, f"step {i}: no reward logged"
assert 0.0 < reward <= max_reward, f"step {i}: reward {reward} out of range"
assert reward_std is not None, f"step {i}: no reward_std logged"
assert 0.0 < reward_std <= max_reward, f"step {i}: no reward spread ({reward_std})"
assert zero_std in (None, 0.0), f"step {i}: {zero_std} of groups had no spread"
assert kl is None or math.isfinite(kl), f"step {i}: kl not finite ({kl})"
assert kl is None or abs(kl) < MAX_KL, f"step {i}: kl diverged ({kl})"
print("fast_inference GRPO rollout completed:", trainer_stats)
if __name__ == "__main__":
if torch.cuda.is_available():
test_fast_inference()
else:
print("Skipping fast_inference test: needs a CUDA GPU + vLLM")

68
tests/python/test_cross_platform_parity.py
View file

@ -249,22 +249,42 @@ class TestTorchIndexOverrideParity:
class TestGfx211AllowlistParity:
"""The gfx per-arch 2.11-floor leaves (gfx120X-all / gfx1151 / gfx1150) must be the
SAME set in every installer and its stale/mismatch check. When they diverged, a
pinned gfx110X-all / gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update."""
"""The gfx per-arch 2.11-floor leaves must be the SAME set in every installer
and its stale/mismatch check. When they diverged, a pinned gfx110X-all /
gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update.
EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"}
Each test extracts the set each installer actually holds and compares it
against EXPECTED, rather than matching one hardcoded ordering. Order and
spacing are free; membership is not. The earlier literal-string form had to
be edited in four places whenever a leaf was added, which is how adding
gfx1152 (Krackan Point) turned this class red without any installer
actually disagreeing with another."""
EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"}
@staticmethod
def _leaves(blob: str) -> set[str]:
"""The gfx leaves named in an allowlist literal, quoting-agnostic."""
return set(re.findall(r"gfx[0-9a-z-]+", blob.lower()))
def test_install_sh_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8").lower()
# install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150).
m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text)
# install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx...|gfx...).
m = re.search(r"^\s*(rocm7\.2\|[a-z0-9|.\-]*)\)", text, re.MULTILINE)
assert m, "install.sh gfx-2.11 allowlist case not found / changed"
assert self._leaves(m.group(1)) == self.EXPECTED, (
f"install.sh gfx-2.11 allowlist is {sorted(self._leaves(m.group(1)))}, "
f"expected {sorted(self.EXPECTED)}"
)
def test_install_ps1_allowlist(self):
text = INSTALL_PS1.read_text(encoding = "utf-8").lower()
m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text)
m = re.search(r"\$_pingfx211\s*=\s*@\(([^)]*)\)", text)
assert m, "install.ps1 $_pinGfx211 allowlist not found / changed"
assert self._leaves(m.group(1)) == self.EXPECTED, (
f"install.ps1 $_pinGfx211 is {sorted(self._leaves(m.group(1)))}, "
f"expected {sorted(self.EXPECTED)}"
)
def test_setup_ps1_defines_single_allowlist_helper(self):
# setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and reuse it, so
@ -273,9 +293,12 @@ class TestGfx211AllowlistParity:
assert (
"function Test-RocmGfx211Leaf" in text
), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
assert re.search(
r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()
), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
m = re.search(r"function test-rocmgfx211leaf[\s\S]{0,400}?@\(([^)]*)\)", text.lower())
assert m, "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
assert self._leaves(m.group(1)) == self.EXPECTED, (
f"Test-RocmGfx211Leaf holds {sorted(self._leaves(m.group(1)))}, "
f"expected {sorted(self.EXPECTED)}"
)
assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, (
"setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not "
"re-hardcode the allowlist (they must not diverge)"
@ -283,9 +306,12 @@ class TestGfx211AllowlistParity:
def test_stack_py_allowlist(self):
text = STACK_PY.read_text(encoding = "utf-8").lower()
assert (
'"gfx120x-all", "gfx1151", "gfx1150"' in text
), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
m = re.search(r"_rocm_gfx_torch211_leaves[^=]*=\s*frozenset\(\s*\{([^}]*)\}", text)
assert m, "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
assert self._leaves(m.group(1)) == self.EXPECTED, (
f"_ROCM_GFX_TORCH211_LEAVES is {sorted(self._leaves(m.group(1)))}, "
f"expected {sorted(self.EXPECTED)}"
)
class TestCudaLeafDigitParity:
@ -351,15 +377,21 @@ class TestCudaLeafDigitParity:
class TestKnown211SetParity:
"""The KNOWN-2.11 rocm/gfx set must be identical across all four installers:
exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}.
exactly {rocm7.2} plus TestGfx211AllowlistParity.EXPECTED.
rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively."""
def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves.
assert re.search(
r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text
), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
# The 2.11 floor case matches exactly rocm7.2 + the gfx allowlist, in
# any order: it is the same set as TestGfx211AllowlistParity.EXPECTED,
# asserted here so the rocm-version half cannot drift on its own.
m = re.search(r"^\s*(rocm7\.2\|[a-zA-Z0-9|.\-]*)\)", text, re.MULTILINE)
assert m, "install.sh 2.11 floor case (rocm7.2|gfx...) not found / changed"
alternatives = set(m.group(1).lower().split("|"))
assert alternatives == {"rocm7.2"} | TestGfx211AllowlistParity.EXPECTED, (
f"install.sh 2.11 floor is {sorted(alternatives)}, expected "
f"{sorted({'rocm7.2'} | TestGfx211AllowlistParity.EXPECTED)}"
)
# No speculative rocm7.3 anywhere.
assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"

33
tests/run_all.sh
View file

@ -7,18 +7,27 @@ set -e
TESTS_DIR="$(cd "$(dirname "$0")" && pwd)"
echo "=== Bash tests ==="
sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
sh "$TESTS_DIR/sh/test_torch_constraint.sh"
sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh"
sh "$TESTS_DIR/sh/test_staged_validation_enabled.sh"
sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh"
sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh"
sh "$TESTS_DIR/sh/test_torch_flavor.sh"
sh "$TESTS_DIR/sh/test_redact_install_output.sh"
sh "$TESTS_DIR/sh/test_install_uv_override_space.sh"
sh "$TESTS_DIR/sh/test_install_rollback_lifecycle.sh"
# Discovered, not listed: a hand-maintained list drifts (this one had fallen
# eight files behind sh/, and the Backend CI copy of it had fallen seven).
# Backend CI discovers the same directory and skips the same file, plus
# test_install_rollback_lifecycle.sh which cross-platform-parity-ci.yml already
# runs on both platforms. tests/studio/test_ci_shell_suite_coverage.py fails if
# either side stops discovering, or skips something undocumented.
# test_install_host_defaults.sh: asserts an install.ps1 layout that has
# drifted (separate followup).
SH_SKIP="test_install_host_defaults.sh"
for _t in "$TESTS_DIR"/sh/test_*.sh; do
case " $SH_SKIP " in
*" $(basename "$_t") "*) echo "skipping $(basename "$_t")"; continue ;;
esac
# bash, not sh: every file under sh/ declares a bash shebang, and three of
# them fail on bashisms under dash, which is /bin/sh on Debian and Ubuntu
# (test_apt_distro_prompt, test_studio_home_node_dir, and
# test_with_llama_cpp_dir_link_behavior). The old hand-written list happened
# to name only dash-clean files, so discovering the directory is what
# exposed this. Backend CI already invokes them with bash.
bash "$_t"
done
echo ""
echo "=== Python tests ==="

178
tests/sh/test_apt_distro_prompt.sh
View file

@ -89,6 +89,184 @@ assert_contains "mentions apt-get" "$_smart" 'sudo apt-get'
assert_contains "mentions official repos" "$_smart" "official repositories"
assert_contains "rejects tarball worry" "$_smart" "not a third-party tarball"
# ── No-TTY sudo escalation (#7307 Problem 7) ────────────────────────
# The old code assumed consent when /dev/tty was unreadable, then ran sudo with
# stdin closed, so a password-requiring host died on a raw sudo error. Drive the
# real function with /dev/tty rewritten to a fixture, the same trick used for
# /etc/os-release above, so every TTY state is reachable hermetically.
echo "=== _smart_apt_install no-TTY escalation ==="
# Closest portable stand-in for the /dev/tty inside containers and systemd
# units: the mode bits satisfy `test -r`, but open() fails with ENXIO. Callers
# must verify the shape before relying on it.
make_unopenable() {
python3 -c 'import socket,sys; socket.socket(socket.AF_UNIX).bind(sys.argv[1])' \
"$1" 2>/dev/null
}
# $1 tty: "tty" | "notty" | "unopenable"
# $2 sudo: "nopasswd" | "needspasswd" | "aptneedspasswd" | "cached" | "absent"
run_smart() {
_tty_mode="$1"; _sudo_mode="$2"
_d=$(mktemp -d -p "$_TMP_ROOT")
case "$_tty_mode" in
tty) printf 'y\n' > "$_d/tty" ;;
# Opens fine but reads EOF straight away (drained/half-closed
# terminal): openable is not the same as answerable.
eof) : > "$_d/tty" ;;
unopenable) make_unopenable "$_d/tty" ;;
esac
_f=$(mktemp -p "$_TMP_ROOT")
sed -n -e '/^_can_read_tty()/,/^}/p' \
-e '/^_smart_apt_install()/,/^}/p' "$INSTALL_SH" \
| sed -e "s#/dev/tty#$_d/tty#g" > "$_f"
(
TAURI_MODE=false
_apt_distro_description() { echo "TestOS 1.0 (debian-like)"; }
_is_pkg_installed() { return 1; } # nothing ever installs
apt-get() { return 1; } # unprivileged attempt fails
command() {
if [ "$1" = -v ] && [ "$2" = sudo ]; then
[ "$_sudo_mode" != absent ]; return $?
fi
builtin command "$@"
}
# Models real sudo: -n refuses (exit 1, nothing runs) when a password
# would be needed. -k ignores any cached timestamp for this invocation
# (sudo(8)), so only a real NOPASSWD rule counts as passwordless.
sudo() {
_noninteractive=false
_ignore_cache=false
while :; do
case "$1" in
-n) _noninteractive=true; shift ;;
-k) _ignore_cache=true; shift ;;
*) break ;;
esac
done
if [ "$_noninteractive" = true ]; then
case "$_sudo_mode" in
nopasswd) ;;
# A valid timestamp from an earlier, unrelated sudo. Without
# -k this looks passwordless; with -k it must not.
cached) [ "$_ignore_cache" = true ] && return 1 ;;
# Authorized for everything, NOPASSWD only on trivial
# commands: `sudo -l` says yes while execution still needs
# a password. Authorization is not the question to ask.
aptneedspasswd)
case " $* " in
*" apt-get "*) return 1 ;;
esac
;;
*) return 1 ;;
esac
fi
# Sudoers refuses the command outright, with or without -n.
if [ "$_sudo_mode" = denied ]; then
echo "sudo: user is not allowed to execute that" >&2
return 1
fi
echo "SUDO_RAN: $*"
}
# shellcheck disable=SC1090
. "$_f"
_smart_apt_install cmake 2>&1
echo "EXIT:$?"
) || true
}
_out=$(run_smart notty needspasswd)
assert_contains "no tty + password sudo: says it cannot run unattended" \
"$_out" "cannot be done unattended"
assert_contains "no tty + password sudo: gives the manual command" \
"$_out" "sudo apt-get update -y && sudo apt-get install -y cmake"
assert_contains "no tty + password sudo: names the distro" \
"$_out" "TestOS 1.0 (debian-like)"
case "$_out" in
*SUDO_RAN*) echo " FAIL: no tty + password sudo must not run apt-get as root"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: no tty + password sudo runs nothing as root"; PASS=$((PASS + 1)) ;;
esac
case "$_out" in
*"Accept? [Y/n]"*) echo " FAIL: must not print an unanswerable prompt"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: no dangling Accept? prompt without a tty"; PASS=$((PASS + 1)) ;;
esac
# Passwordless sudo is the one case where unattended escalation is legitimate.
_out=$(run_smart notty nopasswd)
assert_contains "no tty + passwordless sudo: still installs" "$_out" "SUDO_RAN: apt-get install -y cmake"
assert_contains "no tty + passwordless sudo: says why it proceeded" \
"$_out" "passwordless sudo"
# A readable tty must behave exactly as before: prompt, then honour the answer.
_out=$(run_smart tty needspasswd)
assert_contains "tty present: still prompts" "$_out" "Accept? [Y/n]"
assert_contains "tty present: accepts and installs" "$_out" "SUDO_RAN: apt-get install -y cmake"
# Consent given at a real tty, but the elevated apt-get fails anyway (sudoers
# denial, wrong password, apt error). The interactive branch must say what to
# run by hand, like the headless branch does, not die on the bare sudo error.
_out=$(run_smart tty denied)
assert_contains "tty + denied sudo: gives the manual command" \
"$_out" "sudo apt-get update -y && sudo apt-get install -y cmake"
# No sudo at all keeps its own message.
_out=$(run_smart notty absent)
assert_contains "no sudo binary: unchanged message" "$_out" "sudo is not available on this system"
# A /dev/tty that passes `test -r` but cannot be opened counts as no tty.
# Only assert where the platform can actually produce that shape.
_probe=$(mktemp -d -p "$_TMP_ROOT")
if make_unopenable "$_probe/tty" && [ -r "$_probe/tty" ] && ! ( : <"$_probe/tty" ) 2>/dev/null; then
_out=$(run_smart unopenable needspasswd)
assert_contains "unopenable tty: treated as no tty" "$_out" "cannot be done unattended"
case "$_out" in
*"Accept? [Y/n]"*) echo " FAIL: unopenable tty must not print a prompt"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: unopenable tty prints no prompt"; PASS=$((PASS + 1)) ;;
esac
else
echo " SKIP: this platform cannot fake a readable-but-unopenable /dev/tty"
fi
# A tty that opens but yields EOF must decline: a failed read is nobody
# answering, and calling that "yes" escalates through the branch that does
# have a terminal.
_out=$(run_smart eof needspasswd)
assert_contains "eof tty: declines instead of escalating" \
"$_out" "Please install these packages first"
case "$_out" in
*SUDO_RAN*) echo " FAIL: eof tty must not escalate"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: eof tty runs nothing as root"; PASS=$((PASS + 1)) ;;
esac
# A cached timestamp from an earlier, unrelated sudo must not count as
# passwordless: nobody answered this run's prompt and the apt-get rule still
# carries PASSWD. Asserts the -k is present and effective.
_out=$(run_smart notty cached)
assert_contains "cached credentials: says it cannot run unattended" \
"$_out" "cannot be done unattended"
case "$_out" in
*SUDO_RAN*) echo " FAIL: a cached timestamp must not authorise unattended install"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: cached credentials run nothing as root"; PASS=$((PASS + 1)) ;;
esac
# The failure message must not blame a password when apt itself failed: sudo
# passes the command's own exit status through when the command runs.
assert_contains "failure message does not blame a password exclusively" \
"$_out" "or apt-get itself"
# Authorized for apt-get but not NOPASSWD on it. Both `sudo -n true` and
# `sudo -n -l -- apt-get ...` read this as unattended, since list mode answers
# authorization, not authentication. Only running it with -n is truthful.
_out=$(run_smart notty aptneedspasswd)
assert_contains "apt-get needs a password: says it cannot run unattended" \
"$_out" "cannot be done unattended"
case "$_out" in
*SUDO_RAN*) echo " FAIL: apt-get needing a password must not run as root"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: apt-get needing a password runs nothing as root"; PASS=$((PASS + 1)) ;;
esac
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ]

391
tests/studio/install/test_llama_prebuilt_no_space.py Normal file
View file

@ -0,0 +1,391 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Out-of-disk handling in the llama.cpp prebuilt installer: ENOSPC classification
through exception chains, EXIT_NO_SPACE, and the advisory low-disk warning. Offline."""
from __future__ import annotations
import errno
import importlib.util
import shutil
import sys
import urllib.error
from pathlib import Path
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
assert SPEC is not None and SPEC.loader is not None
INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT)
M = INSTALL_LLAMA_PREBUILT
PrebuiltFallback = M.PrebuiltFallback
AssetChoice = M.AssetChoice
ApprovedReleaseChecksums = M.ApprovedReleaseChecksums
GB = 1024**3
def linux_host() -> "M.HostInfo":
return M.HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
)
def choice(name: str, tag: str = "release-2") -> "M.AssetChoice":
return AssetChoice(
repo = "unslothai/llama.cpp",
tag = tag,
name = name,
url = f"https://example.com/{name}",
source_label = "published",
install_kind = "linux-cpu",
)
def checksums(release_tag: str, llama_tag: str) -> "M.ApprovedReleaseChecksums":
return ApprovedReleaseChecksums(
repo = "unslothai/llama.cpp",
release_tag = release_tag,
upstream_tag = llama_tag,
source_commit = None,
artifacts = {},
)
def plan(llama_tag: str, release_tag: str, attempts) -> "M.InstallReleasePlan":
return M.InstallReleasePlan(
requested_tag = "latest",
llama_tag = llama_tag,
release_tag = release_tag,
attempts = attempts,
approved_checksums = checksums(release_tag, llama_tag),
)
def fake_disk_usage(free_bytes: int):
def _usage(path):
return shutil._ntuple_diskusage(100 * GB, 100 * GB - free_bytes, free_bytes)
return _usage
def install_harness(monkeypatch: pytest.MonkeyPatch, plans, *, free_bytes: int) -> list[str]:
"""Wire install_prebuilt down to a fake per-candidate validation. Returns the
list of candidate names the run actually reached."""
monkeypatch.setattr(M, "detect_host", lambda: linux_host())
monkeypatch.setattr(
M,
"resolve_simple_install_release_plans",
lambda llama_tag, host, published_repo, published_release_tag: ("latest", plans),
)
monkeypatch.setattr(
M, "download_validation_model", lambda probe_path, cache_path: probe_path.write_bytes(b"p")
)
monkeypatch.setattr(M.shutil, "disk_usage", fake_disk_usage(free_bytes))
monkeypatch.setattr(M, "existing_install_matches_plan", lambda *args, **kwargs: False)
monkeypatch.setattr(M, "existing_install_matches_choice", lambda *args, **kwargs: False)
monkeypatch.setattr(M, "activate_install_tree", lambda *args, **kwargs: None)
monkeypatch.setattr(M, "ensure_converter_scripts", lambda *args, **kwargs: None)
monkeypatch.setattr(M, "ensure_diffusion_visual_server", lambda *args, **kwargs: None)
monkeypatch.setattr(M, "collect_system_report", lambda *args, **kwargs: "report")
reached: list[str] = []
monkeypatch.setattr(
M, "validate_prebuilt_choice", lambda attempt, *a, **k: reached.append(attempt.name)
)
return reached
# ── the low-disk check is advisory, never fatal ──
def test_low_disk_warning_reports_the_starved_volume(tmp_path, monkeypatch):
monkeypatch.setattr(M.shutil, "disk_usage", fake_disk_usage(1 * GB))
reason = M._low_disk_warning(tmp_path / "llama.cpp")
assert reason is not None and "low disk space for llama.cpp" in reason
def test_low_disk_warning_silent_when_roomy(tmp_path, monkeypatch):
monkeypatch.setattr(M.shutil, "disk_usage", fake_disk_usage(50 * GB))
assert M._low_disk_warning(tmp_path / "llama.cpp") is None
def test_low_disk_warning_ignores_unstatable_paths(tmp_path, monkeypatch):
def _boom(path):
raise OSError(errno.EACCES, "permission denied")
monkeypatch.setattr(M.shutil, "disk_usage", _boom)
assert M._low_disk_warning(tmp_path / "llama.cpp") is None
def test_low_disk_does_not_block_an_install_that_fits(tmp_path, monkeypatch, capsys):
"""A 15 MB CPU bundle installs fine on a host with 3 GB free; the fixed
threshold must warn rather than reject it (and skip the source fallback)."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
only = plan("b10079", "release-2", [choice("app-b10079-linux-x64-cpu.tar.gz")])
reached = install_harness(monkeypatch, [only], free_bytes = 3 * GB)
M.install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
assert reached == ["app-b10079-linux-x64-cpu.tar.gz"]
captured = capsys.readouterr()
assert "low disk space for llama.cpp" in captured.out + captured.err
# ── ENOSPC classification ──
def test_classifies_direct_and_chained_enospc():
assert M._environment_fatal_reason(OSError(errno.ENOSPC, "No space left on device"))
for wrap in ("cause", "context"):
try:
try:
raise OSError(errno.ENOSPC, "No space left on device")
except OSError as inner:
if wrap == "cause":
raise PrebuiltFallback("download failed") from inner
raise PrebuiltFallback("download failed")
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer), wrap
def test_ignores_unrelated_errors_and_cycles():
assert M._environment_fatal_reason(OSError(errno.EACCES, "denied")) is None
first, second = PrebuiltFallback("a"), PrebuiltFallback("b")
first.__cause__, second.__cause__ = second, first
assert M._environment_fatal_reason(first) is None
def test_suppressed_context_is_not_treated_as_disk_full():
"""`raise ... from None` means the earlier ENOSPC is unrelated."""
try:
try:
raise OSError(errno.ENOSPC, "No space left on device")
except OSError:
raise PrebuiltFallback("checksum mismatch") from None
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer) is None
def test_windows_disk_full_winerrors_are_classified():
"""CPython maps ERROR_DISK_FULL (112) to ENOSPC but has no case for
ERROR_HANDLE_DISK_FULL (39), which arrives as EINVAL."""
for winerror, code in ((112, errno.ENOSPC), (39, errno.EINVAL)):
exc = OSError(code, "The disk is full")
exc.winerror = winerror
assert M._environment_fatal_reason(exc), winerror
other = OSError(errno.EACCES, "sharing violation")
other.winerror = 32
assert M._environment_fatal_reason(other) is None
def test_http_errors_in_the_chain_do_not_crash_the_classifier():
"""HTTPError is an OSError that proxies unknown attributes to a wrapped file
and raises KeyError, not AttributeError, on 3.9."""
err = urllib.error.HTTPError("https://example.com/a", 404, "Not Found", {}, None)
assert M._environment_fatal_reason(err) is None
try:
try:
raise err
except urllib.error.HTTPError as inner:
raise PrebuiltFallback("mirror failed") from inner
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer) is None
@pytest.mark.skipif(not hasattr(errno, "EDQUOT"), reason = "EDQUOT is POSIX only")
def test_quota_exhaustion_counts_as_out_of_space():
"""A quota'd home has free blocks this user cannot have, so the larger source
build is just as doomed. Reported as a quota so df does not mislead."""
assert M._environment_fatal_reason(OSError(errno.EDQUOT, "Disk quota exceeded")) == (
"disk quota exceeded"
)
try:
try:
raise OSError(errno.EDQUOT, "Disk quota exceeded")
except OSError as inner:
raise PrebuiltFallback("bundle download failed") from inner
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer) == "disk quota exceeded"
def test_a_bare_oserror_never_matches():
"""errno is None on a bare OSError, so it must not collide with a code."""
assert M._environment_fatal_reason(OSError()) is None
assert M._environment_fatal_reason(shutil.Error("copy failed")) is None
def test_flattened_markers_are_not_matched_as_prefixes():
"""Bare "WinError 112" would also match WinError 1120; the brackets pin it."""
assert (
M._environment_fatal_reason(
shutil.Error("[('a', 'b', '[WinError 1120] a serial write completed')]")
)
is None
)
assert (
M._environment_fatal_reason(
shutil.Error(f"[('a', 'b', '[Errno {errno.ENOSPC}0] not a real code')]")
)
is None
)
def test_windows_flattened_disk_full_text_is_classified():
"""copytree stringifies the per-file OSError, and on Windows str(OSError)
prints [WinError 112] and never [Errno 28] (confirmed on a real NTFS volume)."""
flattened = (
"[('D:\\\\a\\\\src\\\\big.bin', 'T:\\\\dst\\\\big.bin', "
"'[WinError 112] There is not enough space on the disk')]"
)
assert M._environment_fatal_reason(shutil.Error(flattened))
assert M._environment_fatal_reason(
shutil.Error("[('a', 'b', '[WinError 39] The disk is full')]")
)
assert (
M._environment_fatal_reason(shutil.Error("[('a', 'b', '[WinError 32] sharing violation')]"))
is None
)
def test_validate_install_mode_exits_no_space(tmp_path, monkeypatch):
"""setup.sh reacts to a failed staged validation by deleting the finished GPU
build and starting a CPU rebuild, which needs more of the space that ran out."""
def boom(*args, **kwargs):
try:
raise OSError(errno.ENOSPC, "No space left on device")
except OSError as inner:
raise PrebuiltFallback("validation model unavailable") from inner
monkeypatch.setattr(M, "validate_existing_install", boom)
monkeypatch.setattr(
sys, "argv", ["install_llama_prebuilt.py", "--validate-install", str(tmp_path)]
)
with pytest.raises(SystemExit) as caught:
M.main()
assert caught.value.code == M.EXIT_NO_SPACE
def test_validate_install_mode_still_falls_back_on_ordinary_failure(tmp_path, monkeypatch):
monkeypatch.setattr(
M,
"validate_existing_install",
lambda *a, **k: (_ for _ in ()).throw(PrebuiltFallback("llama-server crashed")),
)
monkeypatch.setattr(
sys, "argv", ["install_llama_prebuilt.py", "--validate-install", str(tmp_path)]
)
with pytest.raises(SystemExit) as caught:
M.main()
assert caught.value.code == M.EXIT_FALLBACK
def test_classifies_enospc_hidden_in_a_shutil_error(tmp_path):
"""copytree stringifies the per-file OSError, so errno and the chain are gone."""
src = tmp_path / "src" / "sub"
src.mkdir(parents = True)
(src / "f").write_text("x", encoding = "utf-8")
def boom(*args, **kwargs):
raise OSError(errno.ENOSPC, "No space left on device")
with pytest.raises(shutil.Error) as caught:
shutil.copytree(tmp_path / "src", tmp_path / "dst", copy_function = boom)
assert caught.value.errno is None
assert M._environment_fatal_reason(caught.value)
def test_source_tree_enospc_is_not_masked_by_a_later_mirror_error(tmp_path, monkeypatch):
"""A full disk fails every mirror, so the first ENOSPC must win over a 404."""
calls: list[str] = []
def fake_download(
url,
path,
*,
expected_sha256 = None,
label = None,
):
calls.append(url)
if len(calls) == 1:
raise OSError(errno.ENOSPC, "No space left on device")
raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
monkeypatch.setattr(M, "download_file_verified", fake_download)
with pytest.raises(PrebuiltFallback) as caught:
M.hydrate_source_tree(
"deadbeef",
tmp_path / "install",
tmp_path,
source_repo = "unslothai/llama.cpp",
expected_sha256 = None,
exact_source = True,
asset_url = "https://example.com/llama.cpp-source.tar.gz",
)
assert len(calls) == 1, f"stopped after the first ENOSPC, tried: {calls}"
assert M._environment_fatal_reason(caught.value)
# ── exit codes ──
def test_enospc_exits_no_space_without_trying_older_releases(tmp_path, monkeypatch):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
newer = plan("b9002", "release-2", [choice("app-b9002-linux-x64-cpu.tar.gz")])
older = plan("b9001", "release-1", [choice("app-b9001-linux-x64-cpu.tar.gz", "release-1")])
reached = install_harness(monkeypatch, [newer, older], free_bytes = 50 * GB)
def enospc(attempt, *args, **kwargs):
reached.append(attempt.name)
raise OSError(errno.ENOSPC, "No space left on device")
monkeypatch.setattr(M, "validate_prebuilt_choice", enospc)
with pytest.raises(SystemExit) as caught:
M.install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
assert caught.value.code == M.EXIT_NO_SPACE
assert reached == ["app-b9002-linux-x64-cpu.tar.gz"]
def test_ordinary_failure_still_exits_fallback(tmp_path, monkeypatch):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
only = plan("b9002", "release-2", [choice("app-b9002-linux-x64-cpu.tar.gz")])
install_harness(monkeypatch, [only], free_bytes = 50 * GB)
monkeypatch.setattr(
M,
"validate_prebuilt_choice",
lambda *a, **k: (_ for _ in ()).throw(PrebuiltFallback("checksum mismatch")),
)
with pytest.raises(SystemExit) as caught:
M.install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
assert caught.value.code == M.EXIT_FALLBACK

651
tests/studio/install/test_rocm_arch_table_parity.py Normal file
View file

@ -0,0 +1,651 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Drift guards for the AMD gfx tables that are duplicated across the installers.
The same three tables are hand-copied into up to seven places each:
gfx -> AMD index family install.sh (_amd_arch_index_family_for_gfx)
install.ps1 ($archFamilyMap)
studio/setup.ps1 ($archFamilyMap)
studio/install_python_stack.py (_GFX_TO_AMD_INDEX_ARCH)
GPU name -> gfx install.sh (_infer_amd_gfx_arch_from_gpu_name)
install.sh (case "$_gpu_disp_mkt", detection banner + env tip)
studio/setup.sh (case "$_setup_mkt")
install.ps1 ($nameArchTable)
studio/setup.ps1 ($nameArchTable)
studio/install_python_stack.py (_WIN_GPU_NAME_ARCH_TABLE)
tests/_zoo_rocm_spoof.py (_PROFILES, inverted gfx -> name)
torch>=2.11 pin allowlist install.sh (case "$_torch_index_leaf")
install.ps1 ($_pinGfx211)
studio/setup.ps1 (Test-RocmPinLeaf211)
Every copy carries a "kept in sync with" comment and nothing enforced it, which is
how the routing family of bugs kept recurring: #7264 / #7280 (Strix left on the
generic rocm7.2 index), #7293 / #7300 (fixed in one installer at a time) and #7277
(RDNA2 gfx1030-1036 added to install.ps1 / setup.ps1 / install_python_stack.py --
install.sh had to follow separately). Half-applied edits are invisible until an AMD
user on the missed path gets CPU-only PyTorch.
These tests parse each copy out of its source file and compare them, so a table
edited in one place fails CI naming the file that was missed.
Counting the copies by hand is itself unreliable -- the in-code "kept in sync
with" comments claimed four when there were seven -- so TestNoUnregisteredArchTable
below rediscovers them by scanning the repo instead of trusting this list.
"""
import ast
import fnmatch
import importlib.util
import re
import sys
from pathlib import Path
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
_INSTALL_SH = PACKAGE_ROOT / "install.sh"
_INSTALL_PS1 = PACKAGE_ROOT / "install.ps1"
_SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
_STACK_PY = PACKAGE_ROOT / "studio" / "install_python_stack.py"
_SPOOF_PY = PACKAGE_ROOT / "tests" / "_zoo_rocm_spoof.py"
def _load_stack_module():
spec = importlib.util.spec_from_file_location("studio_install_python_stack_parity", _STACK_PY)
assert spec is not None and spec.loader is not None
mod = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = mod
spec.loader.exec_module(mod)
return mod
stack_mod = _load_stack_module()
# ── Source extraction helpers ────────────────────────────────────────────────
def _sh_function_body(source: str, name: str) -> str:
"""Return a POSIX-shell function body by brace matching (same idea as
_extract_sh_function_body in test_rocm_support.py, kept local so this file
stands alone)."""
needle = f"{name}() {{"
start = source.find(needle)
assert start != -1, f"{name}() not found"
depth = 0
i = start + len(needle) - 1
while i < len(source):
if source[i] == "{":
depth += 1
elif source[i] == "}":
depth -= 1
if depth == 0:
return source[start : i + 1]
i += 1
raise AssertionError(f"unterminated {name}()")
def _sh_case_block(source: str, subject: str) -> str:
"""Return the body of `case <subject> in ... esac` (first match)."""
start = source.find(f"case {subject} in")
assert start != -1, f"case {subject} in ... not found"
end = source.find("esac", start)
assert end != -1, f"unterminated case {subject}"
return source[start:end]
def _ps_block(source: str, header: str, open_ch: str, close_ch: str) -> str:
"""Return the balanced `header <open> ... <close>` block from a PowerShell file."""
start = source.find(header)
assert start != -1, f"{header} not found"
i = source.find(open_ch, start)
assert i != -1
depth = 0
while i < len(source):
if source[i] == open_ch:
depth += 1
elif source[i] == close_ch:
depth -= 1
if depth == 0:
return source[start : i + 1]
i += 1
raise AssertionError(f"unterminated {header}")
def _strip_sh_comment(line: str) -> str:
"""Drop a trailing `# ...` comment. Safe here: no table line contains a '#'
inside a pattern."""
return line.split("#", 1)[0]
# ── Table 1: gfx -> AMD index family ─────────────────────────────────────────
def _gfx_family_map_sh() -> dict[str, str]:
body = _sh_function_body(
_INSTALL_SH.read_text(encoding = "utf-8"), "_amd_arch_index_family_for_gfx"
)
out: dict[str, str] = {}
for line in body.splitlines():
m = re.match(r"\s*(gfx[^)]*)\)\s*echo\s+(\S+)\s*;;", _strip_sh_comment(line))
if not m:
continue
for arch in m.group(1).split("|"):
out[arch.strip()] = m.group(2).strip()
return out
def _gfx_family_map_ps(path: Path) -> dict[str, str]:
block = _ps_block(path.read_text(encoding = "utf-8"), "$archFamilyMap = @{", "{", "}")
out: dict[str, str] = {}
for line in block.splitlines():
for m in re.finditer(r'"(gfx[0-9a-z]+)"\s*=\s*"([A-Za-z0-9-]+)"', _strip_sh_comment(line)):
out[m.group(1)] = m.group(2)
return out
def _gfx_family_maps() -> dict[str, dict[str, str]]:
return {
"studio/install_python_stack.py": dict(stack_mod._GFX_TO_AMD_INDEX_ARCH),
"install.sh": _gfx_family_map_sh(),
"install.ps1": _gfx_family_map_ps(_INSTALL_PS1),
"studio/setup.ps1": _gfx_family_map_ps(_SETUP_PS1),
}
class TestGfxIndexFamilyParity:
"""All four gfx -> AMD index family maps must agree, entry for entry."""
def test_every_copy_is_non_empty(self):
for where, table in _gfx_family_maps().items():
assert (
table
), f"{where}: parsed an empty gfx -> index family map (table moved or renamed?)"
def test_all_copies_identical(self):
maps = _gfx_family_maps()
reference_name = "studio/install_python_stack.py"
reference = maps[reference_name]
for where, table in maps.items():
if where == reference_name:
continue
missing = {k: v for k, v in reference.items() if k not in table}
extra = {k: v for k, v in table.items() if k not in reference}
wrong = {
k: (v, reference[k])
for k, v in table.items()
if k in reference and v != reference[k]
}
assert (
not missing
), f"{where} is missing {sorted(missing)} (present in {reference_name})"
assert not extra, f"{where} has {sorted(extra)} that {reference_name} does not"
assert not wrong, f"{where} maps {wrong} (value, expected)"
def test_rdna2_family_present_everywhere(self):
"""#7277 added gfx1030-1036 to three files; install.sh followed later.
Pin the whole RDNA2 range so the next family lands everywhere at once."""
for where, table in _gfx_family_maps().items():
for arch in (
"gfx1030",
"gfx1031",
"gfx1032",
"gfx1033",
"gfx1034",
"gfx1035",
"gfx1036",
):
assert table.get(arch) == "gfx103X-all", f"{where}: {arch} -> {table.get(arch)!r}"
class TestSupportedWheelArchList:
"""setup.ps1's $_rocmWheelArches decides whether a detected arch gets ROCm torch
at all. An arch present in the family map but absent here silently installs
CPU-only PyTorch (the 'not in supported arch list' report from r/unsloth)."""
def test_wheel_arch_list_covers_every_mapped_arch(self):
block = _ps_block(
_SETUP_PS1.read_text(encoding = "utf-8"), "$_rocmWheelArches = @(", "(", ")"
)
listed = set(re.findall(r'"(gfx[0-9a-z]+)"', block))
assert listed, "could not parse $_rocmWheelArches"
mapped = set(stack_mod._GFX_TO_AMD_INDEX_ARCH)
assert mapped - listed == set(), (
f"studio/setup.ps1 $_rocmWheelArches is missing {sorted(mapped - listed)}: "
"those arches map to an AMD index but would still fall back to CPU torch"
)
# ── Table 2: GPU marketing name -> gfx ───────────────────────────────────────
#
# Each copy is an ordered, first-match-wins table. The shell copies use case
# globs (case-sensitive); the PowerShell copies use -match regexes
# (case-insensitive, and the only place a negative lookahead is available).
# Rather than diff the patterns -- which legitimately differ in syntax -- run
# every copy against the same real GPU names and require the same answer.
def _name_table_sh_function(source: str, name: str) -> list[tuple[list[str], str]]:
body = _sh_function_body(source, name)
rows: list[tuple[list[str], str]] = []
for line in body.splitlines():
m = re.match(r"\s*(\*.*?)\)\s*echo\s+(gfx[0-9a-z]+)\s*;;", _strip_sh_comment(line))
if m:
rows.append(([p.strip() for p in m.group(1).split("|")], m.group(2)))
return rows
def _name_table_sh_case(source: str, subject: str, var: str) -> list[tuple[list[str], str]]:
"""A bare `case ... in` table that assigns to a variable rather than echoing."""
block = _sh_case_block(source, subject)
rows: list[tuple[list[str], str]] = []
for line in block.splitlines():
m = re.match(
rf'\s*(\*.*?)\)\s*{re.escape(var)}="(gfx[0-9a-z]+)"\s*;;', _strip_sh_comment(line)
)
if m:
rows.append(([p.strip() for p in m.group(1).split("|")], m.group(2)))
return rows
def _name_table_ps(path: Path) -> list[tuple[str, str]]:
block = _ps_block(path.read_text(encoding = "utf-8"), "$nameArchTable = @(", "(", ")")
return re.findall(r'@\{\s*P\s*=\s*"([^"]+)"\s*;\s*A\s*=\s*"(gfx[0-9a-z]+)"\s*\}', block)
def _match_sh(rows: list[tuple[list[str], str]], gpu_name: str) -> str | None:
"""Evaluate a shell `case` table: first arm whose glob matches wins."""
for patterns, arch in rows:
for pattern in patterns:
# Shell case globs quote literal segments: *"RX 7900"* -> *RX 7900*
if fnmatch.fnmatchcase(gpu_name, pattern.replace('"', "")):
return arch
return None
def _match_ps(rows: list[tuple[str, str]], gpu_name: str) -> str | None:
"""Evaluate a PowerShell -match table: first arm whose regex matches wins.
-match is case-insensitive; .NET and Python agree on these patterns
(alternation plus one negative lookahead)."""
for pattern, arch in rows:
if re.search(pattern, gpu_name, re.IGNORECASE):
return arch
return None
# Real strings as amd-smi / rocm-smi / WMI report them, including the two
# ordering traps: "RX 9070 XT" must beat the bare "9070" arm, and "RX 7700S"
# must beat the "RX 7700" arm.
#
# The expectation is the *AMD pip index leaf*, not the gfx id. The leaf is what
# the tables exist to produce -- it picks the wheel -- and it is what a wrong
# answer actually costs the user. Exact gfx ids are pinned separately in
# _AMD_DOCUMENTED_ARCH, sourced from AMD rather than from these tables.
_GPU_NAME_LEAF_CASES = [
("AMD Radeon RX 9070 XT", "gfx120X-all"),
("AMD Radeon RX 9070", "gfx120X-all"),
("AMD Radeon RX 9060 XT", "gfx120X-all"),
("AMD Radeon 8060S Graphics", "gfx1151"),
("AMD Ryzen AI Max+ 395 w/ Radeon 8060S Graphics", "gfx1151"),
("AMD Radeon 890M Graphics", "gfx1150"),
("AMD Radeon 880M Graphics", "gfx1150"),
("AMD Radeon 860M Graphics", "gfx1152"),
("AMD Radeon 840M Graphics", "gfx1152"),
("AMD Ryzen AI 7 350 w/ Radeon 860M", "gfx1152"),
("AMD Radeon RX 7900 XTX", "gfx110X-all"),
("AMD Radeon RX 7800 XT", "gfx110X-all"),
("AMD Radeon PRO W7900", "gfx110X-all"),
("AMD Radeon RX 7700S", "gfx110X-all"),
("AMD Radeon RX 7600 XT", "gfx110X-all"),
("AMD Radeon 780M Graphics", "gfx110X-all"),
("AMD Radeon RX 6900 XT", "gfx103X-all"),
("AMD Radeon RX 6700 XT", "gfx103X-all"),
("AMD Radeon RX 6600 XT", "gfx103X-all"),
("AMD Radeon RX 6500 XT", "gfx103X-all"),
]
# Exact gfx ids, transcribed from AMD's ROCm compatibility matrix (the "Radeon
# GPU" list at rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html),
# NOT from the installer tables. This is the ground truth the tables are supposed
# to reproduce, so it has to come from outside them.
#
# Three of these were wrong until the commit that added this table: RX 9070
# (non-XT) said gfx1200, RX 7800 XT / 7700 XT / PRO W7700 said gfx1100, and PRO
# V710 said gfx1102. Nobody was misrouted, because each wrong id happened to
# share an index leaf with the right one, which is exactly why it went unnoticed
# through five copies of the table. The leaf assertions above cannot catch that
# class of error; only an external source can.
#
# The APU rows were added after that: Krackan Point (860M / 840M) said gfx1150
# but is gfx1152, and unlike the three above that one DID change the wheel,
# since gfx1150 and gfx1152 are separate index leaves on repo.amd.com.
_AMD_DOCUMENTED_ARCH = {
# RDNA 4 -- Navi 48 is gfx1201, Navi 44 is gfx1200.
"AMD Radeon RX 9070 XT": "gfx1201",
"AMD Radeon RX 9070 GRE": "gfx1201",
"AMD Radeon RX 9070": "gfx1201",
"AMD Radeon RX 9060 XT": "gfx1200",
"AMD Radeon RX 9060": "gfx1200",
# RDNA 3 -- Navi 31 / 32 / 33.
"AMD Radeon RX 7900 XTX": "gfx1100",
"AMD Radeon PRO W7900": "gfx1100",
"AMD Radeon PRO W7800": "gfx1100",
"AMD Radeon RX 7800 XT": "gfx1101",
"AMD Radeon RX 7700 XT": "gfx1101",
"AMD Radeon PRO W7700": "gfx1101",
"AMD Radeon PRO V710": "gfx1101",
"AMD Radeon RX 7600 XT": "gfx1102",
"AMD Radeon RX 7700S": "gfx1102",
"AMD Radeon PRO W7600": "gfx1102",
# RDNA 3.5 APUs -- Strix Point is gfx1150, Krackan Point (860M/840M) is
# gfx1152, per AMD's own lemonade GPU table (src/cpp/server/system_info.cpp).
"AMD Radeon 8060S Graphics": "gfx1151",
"AMD Radeon 890M Graphics": "gfx1150",
"AMD Radeon 880M Graphics": "gfx1150",
"AMD Radeon 860M Graphics": "gfx1152",
"AMD Radeon 840M Graphics": "gfx1152",
}
def _name_tables() -> dict[str, object]:
install_sh = _INSTALL_SH.read_text(encoding = "utf-8")
return {
"install.sh:_infer_amd_gfx_arch_from_gpu_name": _name_table_sh_function(
install_sh, "_infer_amd_gfx_arch_from_gpu_name"
),
# install.sh carries the table TWICE. The second copy drives the detection
# banner and, more importantly, the "Tip: set UNSLOTH_ROCM_GFX_ARCH=<arch>"
# line, so a wrong id there gets pasted into a user's environment where it
# becomes authoritative. Neither this copy nor the two below were in this
# parity check until the arch-id fix went looking for every place the
# table lives -- six, not four.
"install.sh:_gpu_disp_gfx": _name_table_sh_case(
install_sh, '"$_gpu_disp_mkt"', "_gpu_disp_gfx"
),
"studio/setup.sh": _name_table_sh_case(
_SETUP_SH.read_text(encoding = "utf-8"), '"$_setup_mkt"', "_setup_gfx"
),
"install.ps1": _name_table_ps(_INSTALL_PS1),
"studio/setup.ps1": _name_table_ps(_SETUP_PS1),
"studio/install_python_stack.py": list(stack_mod._WIN_GPU_NAME_ARCH_TABLE),
}
def _spoof_profiles() -> dict[str, str]:
"""gfx -> marketing name out of tests/_zoo_rocm_spoof.py::_PROFILES.
Parsed with ast rather than imported: that module spoofs torch.cuda and the
AMD identity as an import side effect, which would poison every test sharing
the process."""
tree = ast.parse(_SPOOF_PY.read_text(encoding = "utf-8"))
for node in tree.body:
target = node.target if isinstance(node, ast.AnnAssign) else None
if target is not None and getattr(target, "id", "") == "_PROFILES":
return {gfx: value[0] for gfx, value in ast.literal_eval(node.value).items()}
raise AssertionError("_PROFILES not found in tests/_zoo_rocm_spoof.py")
# The spoof fixture states the mapping backwards (gfx -> the name torch should
# report), so it is the one copy written from the hardware's point of view
# instead of the installer's. That makes it a useful independent witness: it had
# gfx1101 -> "RX 7800 XT" and gfx1201 -> "RX 9070 XT" correct while all six
# installer copies were wrong, and nothing compared the two.
#
# RX 6700 XT is a known, deliberate divergence rather than drift. AMD's
# compatibility matrix documents no consumer RX 6000 card and no gfx1031 at all
# (only "AMD Radeon PRO W6800 (gfx1030)"), the installer arm is commented
# "gfx103X family", and no code consumes the exact id -- gfx1031 appears only as
# a key in the index-family maps, never as a value any name table emits. With no
# external source to correct it against, changing shipped behaviour here would be
# guesswork, so the divergence is pinned instead of silently normalised.
_SPOOF_DIVERGENCES = {
"gfx1031": "installers group Navi 22 into the gfx1030 arm; see comment above",
}
def _resolve(where: str, rows, gpu_name: str) -> str | None:
"""Shell copies are case globs; the PowerShell and Python copies are both
ordered first-match regex tables evaluated case-insensitively, so _match_ps
models either one. `where` may be "<file>:<symbol>" for the files that carry
the table more than once."""
return (
_match_sh(rows, gpu_name)
if where.split(":")[0].endswith(".sh")
else _match_ps(rows, gpu_name)
)
class TestGpuNameArchParity:
"""All four name -> gfx tables must resolve the same GPU the same way."""
def test_every_copy_is_non_empty(self):
for where, rows in _name_tables().items():
assert rows, f"{where}: parsed an empty name -> gfx table (table moved or renamed?)"
@pytest.mark.parametrize("gpu_name", [name for name, _ in _GPU_NAME_LEAF_CASES])
def test_all_copies_return_the_same_arch(self, gpu_name):
"""The drift guard proper: no expected value, just agreement. This is what
catches a table edited in one installer and not the other three, and it
stays honest even where the shipped gfx id is itself wrong."""
answers = {where: _resolve(where, rows, gpu_name) for where, rows in _name_tables().items()}
distinct = set(answers.values())
assert len(distinct) == 1, f"{gpu_name!r} resolves inconsistently: {answers}"
assert distinct != {None}, f"{gpu_name!r} is not matched by any copy of the table"
@pytest.mark.parametrize("gpu_name,expected_leaf", _GPU_NAME_LEAF_CASES)
def test_every_copy_routes_to_the_right_wheel_index(self, gpu_name, expected_leaf):
"""What the tables are for. A wrong leaf is the user-visible failure:
CPU-only torch, or a wheel built for the wrong ISA."""
families = stack_mod._GFX_TO_AMD_INDEX_ARCH
for where, rows in _name_tables().items():
arch = _resolve(where, rows, gpu_name)
assert arch is not None, f"{where}: {gpu_name!r} matched nothing"
assert (
families.get(arch) == expected_leaf
), f"{where}: {gpu_name!r} -> {arch} -> {families.get(arch)!r}, expected {expected_leaf!r}"
@pytest.mark.parametrize("gpu_name,expected_arch", sorted(_AMD_DOCUMENTED_ARCH.items()))
def test_every_copy_matches_amds_documented_arch(self, gpu_name, expected_arch):
"""The gfx id itself, against AMD's matrix rather than against a sibling
copy of the same table. Agreement between five copies proves nothing if
all five were transcribed from the same mistake."""
for where, rows in _name_tables().items():
arch = _resolve(where, rows, gpu_name)
assert (
arch == expected_arch
), f"{where}: {gpu_name!r} -> {arch!r}, AMD documents {expected_arch!r}"
def test_unknown_name_matches_nothing_anywhere(self):
"""An unrecognised card must fall through to the CPU path in every copy,
never onto a neighbouring arm."""
for where, rows in _name_tables().items():
got = _resolve(where, rows, "NVIDIA GeForce RTX 4090")
assert got is None, f"{where}: RTX 4090 matched {got!r}"
def test_inferred_arch_always_has_an_index_family(self):
"""Every arch a name table can produce must be routable to an AMD wheel
index, else detection succeeds and the install still lands on CPU torch."""
families = stack_mod._GFX_TO_AMD_INDEX_ARCH
for where, rows in _name_tables().items():
for arch in {arch for _, arch in rows}:
assert arch in families, f"{where}: {arch} has no entry in _GFX_TO_AMD_INDEX_ARCH"
def test_every_documented_gpu_resolves_somewhere(self):
"""The reverse of the AMD check above. That one asks "do the tables get
the documented cards right"; this asks "is a documented card missing
entirely", which is a silent CPU fallback rather than a wrong id.
This cannot notice a GPU AMD shipped that nobody transcribed into
_AMD_DOCUMENTED_ARCH -- doing that honestly would mean fetching AMD's
matrix at test time, which makes the suite non-hermetic and offline
runners fail. It does catch a card added to the ground-truth list, or to
one installer, without the tables being completed."""
for gpu_name in sorted(_AMD_DOCUMENTED_ARCH):
for where, rows in _name_tables().items():
assert (
_resolve(where, rows, gpu_name) is not None
), f"{where}: {gpu_name!r} matches no arm, so this card gets CPU-only torch"
class TestSpoofFixtureParity:
"""tests/_zoo_rocm_spoof.py is the seventh copy of the name/gfx mapping and
was outside every drift guard. It is the fixture other ROCm tests build their
fake AMD host from, so if it and the installers disagree, those tests exercise
a machine that cannot exist."""
def test_spoof_profiles_parse(self):
profiles = _spoof_profiles()
assert profiles, "parsed an empty _PROFILES (renamed or restructured?)"
assert all(gfx.startswith("gfx") for gfx in profiles), profiles
def test_spoof_names_resolve_back_to_their_own_arch(self):
"""Round-trip: feed each spoofed marketing name through the installer
tables and the answer must be the gfx the spoof claims to be emulating."""
tables = _name_tables()
for gfx, gpu_name in sorted(_spoof_profiles().items()):
if gfx in _SPOOF_DIVERGENCES:
continue
for where, rows in tables.items():
got = _resolve(where, rows, gpu_name)
assert (
got == gfx
), f"{where}: spoof says {gfx} is {gpu_name!r}, installer says {got!r}"
def test_divergences_are_real_and_still_diverging(self):
"""Keeps the exception list from going stale: if the installers are
corrected later, this fails and the entry has to be removed rather than
quietly suppressing a check that now passes."""
tables = _name_tables()
profiles = _spoof_profiles()
for gfx in _SPOOF_DIVERGENCES:
assert gfx in profiles, f"{gfx} is exempted but no longer in the spoof"
answers = {_resolve(w, r, profiles[gfx]) for w, r in tables.items()}
assert answers != {gfx}, f"{gfx} now agrees everywhere; drop it from _SPOOF_DIVERGENCES"
# ── The meta-guard: find copies nobody registered ────────────────────────────
# A table line names a card and gives its arch. Matching both on one line is what
# separates a real table from the many files that merely mention a gfx id (kernel
# dispatch, OOM guards, doc comments).
_MKT_NAME = re.compile(r"(RX\s*\d{4}|PRO\s*[WV]\d{3,4}|\b90[5-8]0\b)", re.IGNORECASE)
_GFX_ID = re.compile(r"gfx1[0-2][0-9a-z]{1,2}")
# Skip dirs of third-party or generated code; scanning them is slow and any hit
# would not be ours to fix.
_SCAN_SKIP_DIRS = {".git", "node_modules", ".venv", "venv", "build", "dist", "__pycache__"}
# Every file allowed to carry a name/arch table, as a repo-relative posix path.
# Adding a copy means adding it here AND wiring it into a parity check above;
# that is the point of the guard.
_REGISTERED_TABLE_FILES = {
"install.sh",
"install.ps1",
"studio/setup.sh",
"studio/setup.ps1",
"studio/install_python_stack.py",
"tests/_zoo_rocm_spoof.py",
}
# Three or more such lines means a table. One or two means prose: the two known
# single-line hits are comments ("Verified on gfx1151 (Radeon 8060S)" in
# scripts/install_rocm_wsl_strixhalo.sh, and a parenthetical in
# studio/install_llama_prebuilt.py). Real tables score 9 to 17, so the gap is
# wide and the threshold is not load-bearing.
_TABLE_LINE_THRESHOLD = 3
def _files_carrying_a_name_arch_table() -> dict[str, int]:
found: dict[str, int] = {}
for path in PACKAGE_ROOT.rglob("*"):
if path.suffix not in {".sh", ".ps1", ".py"} or not path.is_file():
continue
rel = path.relative_to(PACKAGE_ROOT).as_posix()
if any(part in _SCAN_SKIP_DIRS for part in path.relative_to(PACKAGE_ROOT).parts):
continue
# Tests that *assert* on the tables quote card names next to gfx ids by
# nature. Fixtures like _zoo_rocm_spoof.py do not start with test_ and so
# stay in scope, which is how the seventh copy surfaced.
if path.name.startswith("test_"):
continue
try:
text = path.read_text(encoding = "utf-8", errors = "ignore")
except OSError:
continue
hits = sum(
1 for line in text.splitlines() if _MKT_NAME.search(line) and _GFX_ID.search(line)
)
if hits >= _TABLE_LINE_THRESHOLD:
found[rel] = hits
return found
class TestNoUnregisteredArchTable:
"""The failure this whole file exists for is a copy of the table that nobody
knew about. Enumerating the copies by hand is the same manual step that let
them drift, so this rediscovers them from the source tree."""
def test_scan_still_finds_the_known_copies(self):
"""Guards the guard: if the heuristic stops matching (patterns reformatted
onto multiple lines, say), it would silently find nothing and pass."""
found = _files_carrying_a_name_arch_table()
missing = _REGISTERED_TABLE_FILES - set(found)
assert not missing, f"scan no longer detects known tables in {sorted(missing)}"
def test_no_unregistered_copies(self):
found = _files_carrying_a_name_arch_table()
extra = {rel: n for rel, n in found.items() if rel not in _REGISTERED_TABLE_FILES}
assert not extra, (
f"unregistered GPU-name/arch table(s): {extra}. Wire each into "
f"_name_tables() (or the spoof check) and add it to "
f"_REGISTERED_TABLE_FILES, so drift there fails CI too."
)
# ── Table 3: the torch>=2.11 pin allowlist ───────────────────────────────────
class TestTorch211PinAllowlistParity:
"""gfx120X-all / gfx1151 / gfx1150 / gfx1152 (and rocm7.2) ship the null
_grouped_mm kernel below torch 2.11, so all three installers must raise the
same floor. A leaf missing from one copy reintroduces the crash there."""
_EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"}
def test_install_sh_pins_the_same_leaves(self):
source = _INSTALL_SH.read_text(encoding = "utf-8")
idx = source.find('case "$_torch_index_leaf" in')
assert idx != -1
arm = re.search(r"\n\s*(rocm7\.2\|[^)]*)\)", source[idx:])
assert arm, "torch 2.11 pin arm not found in install.sh"
leaves = {leaf.strip() for leaf in arm.group(1).split("|")}
assert (
self._EXPECTED <= leaves
), f"install.sh pin arm missing {sorted(self._EXPECTED - leaves)}"
assert "rocm7.2" in leaves
def test_install_ps1_pins_the_same_leaves(self):
source = _INSTALL_PS1.read_text(encoding = "utf-8")
m = re.search(r"\$_pinGfx211\s*=\s*@\(([^)]*)\)", source)
assert m, "$_pinGfx211 not found in install.ps1"
leaves = set(re.findall(r"'([^']+)'", m.group(1)))
assert leaves == self._EXPECTED, f"install.ps1 pins {sorted(leaves)}"
def test_setup_ps1_pins_the_same_leaves(self):
source = _SETUP_PS1.read_text(encoding = "utf-8")
m = re.search(r"return\s+@\(([^)]*)\)\s*-contains\s*\$Leaf", source)
assert m, "the 2.11 pin allowlist helper was not found in studio/setup.ps1"
leaves = set(re.findall(r"'([^']+)'", m.group(1)))
assert leaves == self._EXPECTED, f"studio/setup.ps1 pins {sorted(leaves)}"
if __name__ == "__main__":
pytest.main([__file__, "-v"])

438
tests/studio/install/test_rocm_native_linux_lib_dirs.py Normal file
View file

@ -0,0 +1,438 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Coverage for the native-Linux system-ROCm library prepend (PR #7233).
#7233 fixed the segfault-on-launch class of AMD reports (#7208, #7310, #6276 and
the native-Linux half of #7307): a prebuilt llama.cpp ships its own libggml-hip /
HIP runtime, and on a bare-metal ROCm box that bundled runtime can disagree with
the host amdkfd driver, so the server dies the moment a model is loaded. The fix
prepends the *system* ROCm lib dirs ahead of the bundle on LD_LIBRARY_PATH.
It landed as two hand-copied helpers, one in the installer (validation-time) and
one in the serve-time launcher:
studio/install_llama_prebuilt.py _bundled_hip_present / _native_linux_system_rocm_lib_dirs
studio/backend/core/inference/llama_cpp.py same two, "mirrors" comment only
and shipped with no tests at all: the WSL sibling helper added earlier has
TestWslSystemRocmLibDirs / TestBinaryEnvWslOrdering / TestLlamaCppRuntimeWslOrdering,
the native-Linux one has nothing. Every gate here is a false-positive risk that
would silently reorder LD_LIBRARY_PATH for users the fix was never meant to touch
(WSL, NVIDIA hosts, macOS, containers without /dev/kfd), so each gate gets a test,
and both copies are run against the same fake host and required to agree.
llama_cpp.py cannot be imported from the test suite (module-level structlog /
backend imports), so its two helpers are lifted out with ast and exec'd standalone.
"""
import ast
import importlib.util
import os
import sys
from pathlib import Path
from unittest.mock import patch
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
_LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py"
_HELPERS = ("_bundled_hip_present", "_native_linux_system_rocm_lib_dirs")
def _load_prebuilt_module():
spec = importlib.util.spec_from_file_location(
"studio_install_llama_prebuilt_native", _PREBUILT_PATH
)
assert spec is not None and spec.loader is not None
mod = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = mod
spec.loader.exec_module(mod)
return mod
def _extract_functions(path: Path, names) -> dict:
"""exec just the named top-level functions out of a module that is too
heavy to import."""
tree = ast.parse(path.read_text(encoding = "utf-8"))
wanted = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
found = {n.name for n in wanted}
assert found == set(names), f"{path.name}: missing {sorted(set(names) - found)}"
module = ast.Module(body = wanted, type_ignores = [])
ns: dict = {"os": os, "sys": sys, "Path": Path}
exec(compile(module, str(path), "exec"), ns)
return ns
prebuilt_mod = _load_prebuilt_module()
llama_ns = _extract_functions(_LLAMA_CPP_PATH, _HELPERS)
def _impls():
"""The two copies of the helper, by the file they live in."""
return {
"studio/install_llama_prebuilt.py": prebuilt_mod._native_linux_system_rocm_lib_dirs,
"studio/backend/core/inference/llama_cpp.py": llama_ns[
"_native_linux_system_rocm_lib_dirs"
],
}
def _norm(paths):
"""os.path.join emits '\\' on the Windows test host; compare POSIX-style."""
return [str(p).replace("\\", "/") for p in paths]
def _fake_exists(present):
"""os.path.exists stub over a set of POSIX paths."""
present = {p.replace("\\", "/") for p in present}
def _exists(p):
return str(p).replace("\\", "/") in present
return _exists
@pytest.fixture
def bundle_dir(tmp_path):
"""A prebuilt directory that does contain a bundled HIP runtime."""
d = tmp_path / "bundle"
d.mkdir()
(d / "libggml-hip.so").write_text("")
return d
@pytest.fixture(autouse = True)
def _clean_rocm_env(monkeypatch):
for var in ("UNSLOTH_LLAMA_NO_SYSTEM_ROCM", "HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
monkeypatch.delenv(var, raising = False)
def _call(
impl,
bundle,
present,
platform = "linux",
):
"""Run one copy of the helper against a fake host.
sys.platform is patched inside the call rather than in a fixture: pytest's own
tmp_path factory branches on it, so a session-wide patch breaks the fixture on
a Windows test host."""
with patch.object(sys, "platform", platform):
with patch("os.path.exists", _fake_exists(present)):
return _norm(impl(str(bundle)))
class TestBundledHipPresent:
"""The prepend only makes sense when the prebuilt actually bundles HIP; a
CPU or CUDA build must be left alone."""
@pytest.mark.parametrize(
"where", ["studio/install_llama_prebuilt.py", "studio/backend/core/inference/llama_cpp.py"]
)
def test_detects_versioned_and_plain_sonames(self, tmp_path, where):
impl = (
prebuilt_mod._bundled_hip_present
if where == "studio/install_llama_prebuilt.py"
else llama_ns["_bundled_hip_present"]
)
plain = tmp_path / "plain"
plain.mkdir()
(plain / "libggml-hip.so").write_text("")
versioned = tmp_path / "versioned"
versioned.mkdir()
(versioned / "libggml-hip.so.0.0.1").write_text("")
cpu_only = tmp_path / "cpu"
cpu_only.mkdir()
(cpu_only / "libggml-cpu.so").write_text("")
assert impl(str(plain)) is True, f"{where}: plain soname not detected"
assert impl(str(versioned)) is True, f"{where}: versioned soname not detected"
assert impl(str(cpu_only)) is False, f"{where}: CPU-only build treated as HIP"
assert impl("") is False, f"{where}: empty binary_dir must be falsy"
assert (
impl(str(tmp_path / "does-not-exist")) is False
), f"{where}: missing dir must be falsy"
class TestNativeLinuxGates:
"""Each gate, on both copies. A gate that stops working silently reorders
LD_LIBRARY_PATH for a platform the fix was never aimed at."""
_ROCM_LIB = "/opt/rocm/lib"
_HOST = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so"}
_run = staticmethod(_call)
def test_returns_system_rocm_lib_on_a_native_rocm_host(self, bundle_dir):
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, self._HOST) == [self._ROCM_LIB], where
def test_accepts_versioned_hsa_runtime_soname(self, bundle_dir):
present = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so.1"}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == [self._ROCM_LIB], where
@pytest.mark.parametrize("platform", ["win32", "darwin"])
def test_no_op_off_linux(self, bundle_dir, platform):
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, self._HOST, platform = platform) == [], where
def test_no_op_on_wsl(self, bundle_dir):
"""WSL has its own ordering path (plus HSA_ENABLE_DXG_DETECTION); /dev/dxg
must hand off to it, not double-prepend here."""
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, self._HOST | {"/dev/dxg"}) == [], where
def test_no_op_without_amdkfd(self, bundle_dir):
"""No /dev/kfd: NVIDIA host, CPU host, or a container without the AMD
device node. Prepending system ROCm there would be pure breakage."""
present = {"/opt/rocm/lib/libhsa-runtime64.so"}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == [], where
def test_no_op_when_prebuilt_bundles_no_hip(self, tmp_path):
cpu_bundle = tmp_path / "cpu-bundle"
cpu_bundle.mkdir()
for where, impl in _impls().items():
assert self._run(impl, cpu_bundle, self._HOST) == [], where
def test_no_op_when_system_rocm_has_no_hsa_runtime(self, bundle_dir):
"""ROCm dir exists but is not a usable runtime install."""
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, {"/dev/kfd"}) == [], where
def test_opt_out_env_wins_over_everything(self, bundle_dir, monkeypatch):
monkeypatch.setenv("UNSLOTH_LLAMA_NO_SYSTEM_ROCM", "1")
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, self._HOST) == [], where
def test_opt_out_env_only_honours_exactly_one(self, bundle_dir, monkeypatch):
"""Documented switch is =1; "0"/"" must not disable the fix."""
for value in ("0", "", "false"):
monkeypatch.setenv("UNSLOTH_LLAMA_NO_SYSTEM_ROCM", value)
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, self._HOST) == [
self._ROCM_LIB
], f"{where} ({value!r})"
class TestNativeLinuxRootResolution:
"""Which ROCm roots are searched, in what order."""
_run = staticmethod(_call)
def test_env_roots_take_precedence_over_opt_rocm(self, bundle_dir, monkeypatch):
"""A user with a side-by-side ROCm (HIP_PATH) must get theirs first: the
one matching their driver, not whatever /opt/rocm happens to be."""
monkeypatch.setenv("HIP_PATH", "/usr/local/rocm7")
present = {
"/dev/kfd",
"/usr/local/rocm7/lib/libhsa-runtime64.so",
"/opt/rocm/lib/libhsa-runtime64.so",
}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == [
"/usr/local/rocm7/lib",
"/opt/rocm/lib",
], where
def test_all_three_env_roots_are_consulted_in_order(self, bundle_dir, monkeypatch):
monkeypatch.setenv("HIP_PATH", "/a")
monkeypatch.setenv("HIP_PATH_57", "/b")
monkeypatch.setenv("ROCM_PATH", "/c")
present = {
"/dev/kfd",
"/a/lib/libhsa-runtime64.so",
"/b/lib/libhsa-runtime64.so",
"/c/lib/libhsa-runtime64.so",
}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == ["/a/lib", "/b/lib", "/c/lib"], where
def test_lib64_layout_is_found(self, bundle_dir):
"""RHEL / SUSE ROCm packages install to lib64."""
present = {"/dev/kfd", "/opt/rocm/lib64/libhsa-runtime64.so"}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == ["/opt/rocm/lib64"], where
def test_lib_precedes_lib64_when_both_exist(self, bundle_dir):
present = {
"/dev/kfd",
"/opt/rocm/lib/libhsa-runtime64.so",
"/opt/rocm/lib64/libhsa-runtime64.so",
}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == [
"/opt/rocm/lib",
"/opt/rocm/lib64",
], where
def test_duplicate_roots_are_deduped(self, bundle_dir, monkeypatch):
"""ROCM_PATH=/opt/rocm is the common setup; it must not emit the dir twice."""
monkeypatch.setenv("ROCM_PATH", "/opt/rocm")
present = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so"}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == ["/opt/rocm/lib"], where
def test_empty_env_var_is_ignored(self, bundle_dir, monkeypatch):
monkeypatch.setenv("HIP_PATH", "")
present = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so"}
for where, impl in _impls().items():
assert self._run(impl, bundle_dir, present) == ["/opt/rocm/lib"], where
class TestHelperParity:
"""The two copies carry a "mirrors ..." comment and nothing enforced it."""
@pytest.mark.parametrize("name", _HELPERS)
def test_bodies_are_identical(self, name):
a = _function_ast(_PREBUILT_PATH, name)
b = _function_ast(_LLAMA_CPP_PATH, name)
assert ast.dump(a) == ast.dump(b), (
f"{name} has drifted between install_llama_prebuilt.py and llama_cpp.py; "
"the install-time and serve-time launchers must resolve the same lib dirs"
)
def _function_ast(path: Path, name: str) -> ast.FunctionDef:
"""The function's executable body, with docstring and type annotations
stripped: llama_cpp.py quotes its annotations ('list[str]') for the
older-typing lint and documents itself as mirroring the installer. Neither is
drift; the code is."""
tree = ast.parse(path.read_text(encoding = "utf-8"))
for node in tree.body:
if isinstance(node, ast.FunctionDef) and node.name == name:
for child in ast.walk(node):
if isinstance(child, (ast.AnnAssign, ast.arg)):
child.annotation = None
elif isinstance(child, ast.FunctionDef):
child.returns = None
if (
node.body
and isinstance(node.body[0], ast.Expr)
and isinstance(node.body[0].value, ast.Constant)
and isinstance(node.body[0].value.value, str)
):
node.body = node.body[1:]
return node
raise AssertionError(f"{name} not found in {path.name}")
class TestBinaryEnvNativeOrdering:
"""install-time validation launches the binary through binary_env."""
@staticmethod
def _linux_host():
return prebuilt_mod.HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = True,
)
def test_system_rocm_precedes_bundle_dir(self, tmp_path):
binary = tmp_path / "bundle" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("")
sys_rocm = tmp_path / "sysrocm" # dedupe_existing_dirs drops missing dirs
sys_rocm.mkdir()
with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []):
with patch.object(
prebuilt_mod, "_native_linux_system_rocm_lib_dirs", return_value = [str(sys_rocm)]
):
with patch.dict(os.environ, {}, clear = True):
env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host())
ld = [str(Path(p).resolve()) for p in env["LD_LIBRARY_PATH"].split(os.pathsep)]
assert ld.index(str(sys_rocm.resolve())) < ld.index(str(binary.parent.resolve()))
def test_native_path_does_not_enable_dxg_detection(self, tmp_path):
"""HSA_ENABLE_DXG_DETECTION belongs to the WSL branch only; setting it on
bare metal changes HSA agent enumeration for every native AMD user."""
binary = tmp_path / "bundle" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("")
sys_rocm = tmp_path / "sysrocm"
sys_rocm.mkdir()
with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []):
with patch.object(
prebuilt_mod, "_native_linux_system_rocm_lib_dirs", return_value = [str(sys_rocm)]
):
with patch.dict(os.environ, {}, clear = True):
env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host())
assert "HSA_ENABLE_DXG_DETECTION" not in env
def test_helper_is_asked_about_the_binary_dir_not_the_install_dir(self, tmp_path):
"""_bundled_hip_present globs the directory it is handed; passing
install_dir would look for libggml-hip.so in the wrong place and no-op."""
binary = tmp_path / "bundle" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("")
seen = []
def _spy(binary_dir = ""):
seen.append(binary_dir)
return []
with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []):
with patch.object(prebuilt_mod, "_native_linux_system_rocm_lib_dirs", _spy):
with patch.dict(os.environ, {}, clear = True):
prebuilt_mod.binary_env(binary, tmp_path, self._linux_host())
assert seen == [str(binary.parent)]
def test_no_prepend_leaves_bundle_dir_first(self, tmp_path):
binary = tmp_path / "bundle" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("")
with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []):
with patch.object(prebuilt_mod, "_native_linux_system_rocm_lib_dirs", return_value = []):
with patch.dict(os.environ, {}, clear = True):
env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host())
assert env["LD_LIBRARY_PATH"].split(os.pathsep)[0] == str(binary.parent)
class TestLlamaCppRuntimeNativeOrdering:
"""The serve-time launcher builds LD_LIBRARY_PATH inline inside a large
function, so this half stays a source check (as the WSL sibling does)."""
def test_prepends_before_binary_dir(self):
source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8")
idx_helper = source.find("lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))")
idx_binary = source.find("lib_dirs.append(binary_dir)")
assert (
idx_helper != -1
), "serve-time launcher must call the native-Linux helper with binary_dir"
assert idx_binary != -1
assert (
idx_helper < idx_binary
), "system ROCm must be searched before the bundled HIP runtime"
def test_dxg_detection_stays_on_the_wsl_branch(self):
"""HSA_ENABLE_DXG_DETECTION must be set from the WSL helper's result only."""
source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8")
idx_wsl = source.find("lib_dirs.extend(_wsl_system_rocm_lib_dirs())")
idx_dxg = source.find('env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")', idx_wsl)
idx_native = source.find("lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))")
assert idx_wsl != -1 and idx_dxg != -1 and idx_native != -1
assert idx_wsl < idx_dxg < idx_native, (
"HSA_ENABLE_DXG_DETECTION must be decided from the WSL dirs alone, before the "
"native-Linux dirs are appended to lib_dirs"
)
if __name__ == "__main__":
pytest.main([__file__, "-v"])

16
tests/studio/install/test_rocm_rdna_routing.py
View file

@ -12,6 +12,7 @@ at import) resolves from a clean process.
from __future__ import annotations
import json
import os
import subprocess
import sys
from pathlib import Path
@ -43,6 +44,15 @@ _ARCHES = {
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
# Import bitsandbytes under the real torch first. unsloth_zoo pulls it in, and it
# picks a compute backend at import: once the spoof reports an AMD GPU, it loads
# its ROCm/CUDA ops, which a CPU-only torch cannot satisfy (no libhipblas, no
# torch._C._cuda_getCurrentRawStream) and the child dies before printing RESULT.
# Nothing here tests bitsandbytes, so let it see the honest hardware.
try:
import bitsandbytes # noqa: F401
except Exception:
pass
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
@ -60,7 +70,11 @@ print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}})
@pytest.fixture(scope = "module")
def routed():
code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES))
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True)
# get_device_type() returns "mlx" before it ever looks at torch on Darwin arm64
# with mlx installed, so the spoof would be ignored. Force the GPU path to keep
# the assertion live there instead of skipping it.
env = {**os.environ, "UNSLOTH_FORCE_GPU_PATH": "1"}
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True, env = env)
line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None)
assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
return json.loads(line[len("RESULT ") :])

25
tests/studio/install/test_rocm_support.py
View file

@ -2695,9 +2695,11 @@ class TestGfxArchNameFallback:
("AMD Radeon(TM) 890M", "gfx1150"),
("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"),
("AMD Radeon RX 9070 XT", "gfx1201"),
("AMD Radeon RX 9070", "gfx1200"),
("AMD Radeon RX 7700S", "gfx1102"), # (?!S) lookahead must not hit gfx1100
("AMD Radeon RX 7700 XT", "gfx1100"),
("AMD Radeon RX 9070", "gfx1201"), # Navi 48 like the XT, not Navi 44
("AMD Radeon RX 9060 XT", "gfx1200"), # Navi 44
("AMD Radeon RX 7700S", "gfx1102"), # (?!S) lookahead must not hit gfx1101
("AMD Radeon RX 7700 XT", "gfx1101"), # Navi 32
("AMD Radeon RX 7900 XTX", "gfx1100"), # Navi 31
("AMD Radeon(TM) 780M", "gfx1103"),
("NVIDIA GeForce RTX 4090", None),
("Microsoft Basic Display Adapter", None),
@ -3845,7 +3847,7 @@ class TestStrixRocm71Override:
assert m._infer_linux_amd_gfx_arch() == "gfx1151"
def test_install_sh_cpuinfo_inference_requires_pci_evidence(self):
"""install.sh mirror of the VM/container guard: both cpuinfo greps must be
"""install.sh mirror of the VM/container guard: every cpuinfo grep must be
gated on _gpu_evidence (AMD PCI display device via _amd_gpu_present_via_pci,
or the WSL librocdxg gate), and the gate must sit before the first grep."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
@ -3855,9 +3857,9 @@ class TestStrixRocm71Override:
infer = body.find("grep -qiE 'Ryzen AI Max")
assert pci >= 0 and infer >= 0
assert pci < infer, "the PCI evidence check must run before the cpuinfo inference"
assert (
body.count('[ -n "$_gpu_evidence" ] && grep -qiE') == 2
), "both cpuinfo greps (gfx1151 and gfx1150) must be gated on _gpu_evidence"
assert body.count("grep -qiE") == body.count(
'[ -n "$_gpu_evidence" ] && grep -qiE'
), "every cpuinfo grep (gfx1151/gfx1150/gfx1152) must be gated on _gpu_evidence"
def test_lspci_scan_covers_all_display_controllers(self):
"""The lspci fallback must scan every display-class line, not just the
@ -4211,11 +4213,14 @@ class TestStrixRocm71Override:
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
# The 2.11 constraint block must switch on $_torch_index_leaf, not the full
# $TORCH_INDEX_URL (a */gfx* match false-positives on a mirror base path). Only the
# _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150) are pushed to 2.11;
# _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150 / gfx1152) go to 2.11;
# a bare gfx* would also floor gfx110X-all/gfx90a/gfx908, left bare on purpose.
assert 'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150)' in source, (
assert (
'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152)'
in source
), (
"the torch>=2.11 constraint must match the specific gfx leaves that need "
"it (rocm7.2|gfx120x-all|gfx1151|gfx1150), not a bare gfx* or the whole URL"
"it (rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152), not a bare gfx* or the URL"
)
def test_amd_rocm_mirror_env_var_respected(self):

101
tests/studio/playwright_ui_font_scale.py
View file

@ -16,6 +16,7 @@ import os
import sys
from pathlib import Path
from playwright.sync_api import TimeoutError as PWTimeout
from playwright.sync_api import sync_playwright
sys.path.insert(0, str(Path(__file__).resolve().parent))
@ -46,6 +47,18 @@ def near(
return a is not None and b is not None and abs(a - b) <= tol
_VP = 'document.querySelector("[data-radix-select-viewport]")'
SCROLL_TOP_JS = f"() => {_VP}.scrollTop"
SCROLLABLE_JS = f"() => {{ const vp = {_VP}; return !!vp && vp.scrollHeight > vp.clientHeight; }}"
VIEWPORT_STATE_JS = f"""
() => {{
const vp = {_VP};
return vp
? {{ scrollHeight: vp.scrollHeight, clientHeight: vp.clientHeight, top: vp.scrollTop }}
: null;
}}
"""
MEASURE_JS = """
() => {
const fs = (el) => (el ? parseFloat(getComputedStyle(el).fontSize) : null);
@ -83,15 +96,22 @@ def set_input(page, label, value):
def open_appearance(page):
page.keyboard.press("Control+,")
page.wait_for_timeout(700)
if page.get_by_role("dialog").count() == 0:
page.keyboard.press("Meta+,")
page.wait_for_timeout(700)
if page.get_by_role("dialog").count() == 0:
fail("settings dialog did not open")
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Appearance").first.click()
page.wait_for_timeout(600)
# The shortcut can fire before the app has wired its key handler, so press
# each chord once behind a fixed sleep and a slow boot loses the dialog.
# Alternate them on a bounded retry, waiting on the dialog itself.
dialog = page.get_by_role("dialog")
for attempt in range(10):
page.keyboard.press("Meta+," if attempt % 2 else "Control+,")
try:
dialog.first.wait_for(state = "visible", timeout = 2_000)
break
except PWTimeout:
continue
if dialog.count() == 0:
fail("settings dialog did not open after 10 attempts")
dialog.get_by_role("button").filter(has_text = "Appearance").first.click()
# Wait for the control the caller is about to drive, not a fixed interval.
page.locator("input[aria-label='UI font size']").wait_for(state = "visible", timeout = 15_000)
def main():
@ -155,39 +175,46 @@ def main():
page.wait_for_timeout(400)
step("overflowing select scrolls its Radix viewport")
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Voice").first.click()
page.wait_for_timeout(600)
voice = page.get_by_role("dialog").get_by_role("button").filter(has_text = "Voice").first
voice.click()
page.set_viewport_size({"width": 1440, "height": 480})
page.locator("[aria-label='Dictation language']").click()
page.wait_for_timeout(700)
state = page.evaluate(
"""
() => {
const vp = document.querySelector("[data-radix-select-viewport]");
return vp
? { scrollable: vp.scrollHeight > vp.clientHeight, top: vp.scrollTop }
: null;
}
"""
)
if not state or not state["scrollable"]:
fail(f"select viewport not scrollable: {state}")
for _ in range(6):
trigger = page.locator("[aria-label='Dictation language']")
trigger.wait_for(state = "visible")
trigger.click()
viewport = page.locator("[data-radix-select-viewport]")
viewport.wait_for(state = "visible")
# Wait for the overflow itself rather than a fixed sleep: the list is
# populated asynchronously, so measuring too early reads it as short.
try:
page.wait_for_function(SCROLLABLE_JS, timeout = 10_000)
except PWTimeout:
fail(f"select viewport not scrollable: {page.evaluate(VIEWPORT_STATE_JS)}")
# Radix moves focus into the listbox after the content opens, so a fixed
# burst of presses can land on the trigger and scroll nothing. Press until
# it moves instead; a real regression still fails, just after more tries.
kb_top = 0
for _ in range(40):
page.keyboard.press("ArrowDown")
page.wait_for_timeout(100)
kb_top = page.evaluate(
"() => document.querySelector('[data-radix-select-viewport]').scrollTop"
)
kb_top = page.evaluate(SCROLL_TOP_JS)
if kb_top > 0:
break
page.wait_for_timeout(50)
if not kb_top > 0:
fail(f"keyboard did not scroll the select viewport: {kb_top}")
vp_box = page.locator("[data-radix-select-viewport]").bounding_box()
fail(f"keyboard did not scroll the select viewport after 40 presses: {kb_top}")
vp_box = viewport.bounding_box()
page.mouse.move(vp_box["x"] + vp_box["width"] / 2, vp_box["y"] + 40)
page.mouse.wheel(0, -400)
page.wait_for_timeout(300)
wheel_top = page.evaluate(
"() => document.querySelector('[data-radix-select-viewport]').scrollTop"
)
if not wheel_top < kb_top:
try:
page.wait_for_function(
"top => document.querySelector('[data-radix-select-viewport]').scrollTop < top",
arg = kb_top,
timeout = 10_000,
)
except PWTimeout:
wheel_top = page.evaluate(SCROLL_TOP_JS)
fail(f"wheel did not scroll the select viewport: {kb_top} -> {wheel_top}")
page.keyboard.press("Escape")
page.set_viewport_size({"width": 1440, "height": 900})

25
tests/studio/test_chat_response_details_ui_contract.py
View file

@ -75,6 +75,31 @@ def test_response_model_badge_is_user_configurable_and_rendered_once_per_message
assert 'className="min-w-0 flex-1"' in reasoning_src
def test_reasoning_keeps_streaming_height_cap_through_automatic_collapse():
src = REASONING_TSX.read_text()
assert "const [retainStreamingHeight, setRetainStreamingHeight]" in src
assert "setRetainStreamingHeight(false)" in src
assert "setRetainStreamingHeight(isReasoningStreaming)" in src
assert "isReasoningStreaming ? 0 : ANIMATION_DURATION" in src
assert "streaming={isReasoningStreaming || retainStreamingHeight}" in src
def test_reasoning_clears_manual_open_on_a_new_stream():
"""A hand-opened block must not stay pinned open when the stream restarts.
isOpen is `(streaming && !dismissed) || manualOpen` and manualOpen is only
settable while idle, so the new-stream reset has to clear it too.
"""
src = REASONING_TSX.read_text()
marker = "setDismissedWhileStreaming(false)"
start = src.find(marker)
assert start != -1, "new-stream reset effect is missing"
effect = src[src.rfind("useEffect(() => {", 0, start) : src.find("});", start)]
assert "setManualOpen(false)" in effect
def test_response_details_metadata_is_persisted_without_backend_schema_change():
src = ADAPTER_TS.read_text()
assert "interface ResponseDetailsMetadata" in src

197
tests/studio/test_ci_shell_suite_coverage.py Normal file
View file

@ -0,0 +1,197 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Guards that the installer test suites actually run on a PR.
Two ways coverage went missing without anyone noticing:
1. Backend CI ran a hardcoded list of tests/sh/*.sh files. New tests were added
to the directory and never to the list, so by the time this was written the
list was seven files behind -- including test_strixhalo_wsl_reroute.sh, the
only shell coverage of the ROCm WSL reroute, which had never run on a PR.
tests/run_all.sh, the local entrypoint, had drifted the other way.
2. Backend CI's path filter did not include install.sh / install.ps1, while a
large share of the suites it runs (tests/sh/*, tests/studio/install/*) assert
against exactly those two files. An install-only change -- the shape most
AMD/ROCm routing fixes take, e.g. #7277 / #7293 / #7300 -- skipped the
workflow that tests it.
Both are now discovery-based. These tests fail if either reverts to a list, if a
shell test lands somewhere the discovery cannot see it, or if a skip is added
without a reason next to it.
"""
import re
from pathlib import Path
import pytest
import yaml
REPO_ROOT = Path(__file__).resolve().parents[2]
_WORKFLOWS = REPO_ROOT / ".github" / "workflows"
_BACKEND_CI = _WORKFLOWS / "studio-backend-ci.yml"
_PARITY_CI = _WORKFLOWS / "cross-platform-parity-ci.yml"
_RUN_ALL = REPO_ROOT / "tests" / "run_all.sh"
_SH_DIR = REPO_ROOT / "tests" / "sh"
# Files deliberately not run by the auto-discovered Backend CI step. Each needs
# a reason here AND in the workflow; anything else in tests/sh must run.
_EXPECTED_CI_SKIPS = {
"test_install_host_defaults.sh": "asserts an install.ps1 layout that has drifted",
"test_install_rollback_lifecycle.sh": "runs on both platforms in cross-platform-parity-ci.yml",
}
def _backend_ci() -> dict:
return yaml.safe_load(_BACKEND_CI.read_text(encoding = "utf-8"))
def _shell_step_script() -> str:
"""The `run:` body of the shell-installer step, located by name through the
parsed YAML rather than by slicing the raw file."""
for job in _backend_ci()["jobs"].values():
for step in job.get("steps", []):
if step.get("name") == "Shell installer tests":
return step["run"]
raise AssertionError("Backend CI has no 'Shell installer tests' step")
def _shell_test_files():
files = sorted(p.name for p in _SH_DIR.glob("test_*.sh"))
assert files, "tests/sh has no test_*.sh files -- did the directory move?"
return files
def _skip_list(source: str) -> set[str]:
"""The skip= / SH_SKIP= line from a discovery loop."""
m = re.search(r"^\s*(?:skip|SH_SKIP)=\"([^\"]*)\"", source, re.MULTILINE)
assert m, "no skip list found; the discovery loop must declare one (even if empty)"
return {name for name in m.group(1).split() if name}
class TestBackendCiRunsEveryShellTest:
def test_step_discovers_the_directory_instead_of_listing_files(self):
"""Matched against the parsed step script, and on the glob rather than a
verbatim line, so reformatting the loop does not turn CI red -- only
going back to a hardcoded list does."""
script = _shell_step_script()
assert re.search(r"for\s+\w+\s+in\s+tests/sh/test_\*\.sh", script), (
"Backend CI must glob tests/sh; a hardcoded list is how the ROCm WSL "
f"suite went unrun for months. Step script was:\n{script}"
)
listed = re.findall(r"tests/sh/test_[a-z0-9_]+\.sh", script)
assert not listed, f"Backend CI still names individual shell tests: {sorted(set(listed))}"
def test_step_fails_loudly_if_discovery_finds_nothing(self):
"""A moved directory must break the build, not pass vacuously."""
assert "no shell tests discovered under tests/sh" in _shell_step_script()
def test_every_shell_test_runs_or_is_a_known_skip(self):
skips = _skip_list(_shell_step_script())
unexpected = skips - set(_EXPECTED_CI_SKIPS)
assert not unexpected, (
f"Backend CI skips {sorted(unexpected)} without a reason recorded in "
"_EXPECTED_CI_SKIPS; add one or stop skipping it"
)
# Everything else in the directory is covered by the glob.
for name in _shell_test_files():
assert name not in skips or name in _EXPECTED_CI_SKIPS, name
def test_skip_entries_are_not_stale(self):
"""A skip for a deleted file quietly widens next time a name is reused."""
existing = set(_shell_test_files())
for name in _skip_list(_shell_step_script()):
assert name in existing, f"{name} is skipped but no longer exists in tests/sh"
def test_each_skip_is_documented_in_the_workflow(self):
source = _BACKEND_CI.read_text(encoding = "utf-8")
for name in _EXPECTED_CI_SKIPS:
assert (
source.count(name) >= 2
), f"{name} is skipped in Backend CI without a comment explaining why"
def test_rollback_lifecycle_really_does_run_elsewhere(self):
"""The one skip justified by 'another workflow covers it' must be true."""
assert "tests/sh/test_install_rollback_lifecycle.sh" in _PARITY_CI.read_text(
encoding = "utf-8"
)
def test_rocm_shell_suite_is_in_scope(self):
"""The suite whose absence prompted this file: it must exist and be
picked up (i.e. not skipped)."""
assert "test_strixhalo_wsl_reroute.sh" in _shell_test_files()
assert "test_strixhalo_wsl_reroute.sh" not in _skip_list(_shell_step_script())
class TestRunAllMatchesCi:
"""tests/run_all.sh is what a contributor runs before pushing. If it and CI
disagree, one of them is lying about the state of the tree."""
def test_run_all_discovers_the_directory(self):
source = _RUN_ALL.read_text(encoding = "utf-8")
assert 'for _t in "$TESTS_DIR"/sh/test_*.sh; do' in source
def test_run_all_invokes_the_tests_with_bash(self):
"""Both runners must use the interpreter the tests declare. Every file
under tests/sh/ has a bash shebang, and on Debian/Ubuntu /bin/sh is
dash, under which three of them fail on bashisms. Running them with sh
would fail the suite locally for reasons CI never reproduces."""
source = _RUN_ALL.read_text(encoding = "utf-8")
assert 'bash "$_t"' in source, "tests/run_all.sh must run tests/sh/ with bash"
assert 'sh "$_t"' not in source.replace(
'bash "$_t"', ""
), "tests/run_all.sh still invokes a discovered test with sh"
assert 'bash "$s"' in _shell_step_script(), "Backend CI must run tests/sh/ with bash"
def test_run_all_skips_are_a_subset_of_ci_skips(self):
local = _skip_list(_RUN_ALL.read_text(encoding = "utf-8"))
unexpected = local - set(_EXPECTED_CI_SKIPS)
assert not unexpected, (
f"tests/run_all.sh skips {sorted(unexpected)} that CI still runs: a "
"contributor would see green locally and red on the PR"
)
class TestBackendCiPathFilters:
"""The workflow has to fire on the files its tests assert against."""
def _paths(self) -> set[str]:
"""Read the real trigger through the YAML parser. `on:` is a YAML 1.1
boolean, so pyyaml keys it as True."""
wf = _backend_ci()
triggers = wf.get("on", wf.get(True))
assert triggers, "Backend CI has no trigger block"
paths = triggers["pull_request"]["paths"]
assert paths, "Backend CI pull_request trigger has no paths filter"
return set(paths)
@pytest.mark.parametrize(
"path,why",
[
("install.sh", "tests/sh/* and tests/studio/install/* assert against it"),
("install.ps1", "the Windows/ROCm arch tables and pin allowlist live here"),
("studio/**", "covers studio/setup.sh, studio/setup.ps1, install_python_stack.py"),
("tests/**", "test-only changes must run the tests they touch"),
],
)
def test_trigger_covers(self, path, why):
assert path in self._paths(), f"Backend CI does not run when {path} changes ({why})"
def test_installer_change_would_trigger_the_workflow(self):
"""End to end: the exact filenames the ROCm fixes edit."""
paths = self._paths()
for changed in ("install.sh", "install.ps1"):
assert changed in paths
for changed in ("studio/setup.ps1", "studio/setup.sh", "studio/install_python_stack.py"):
assert any(
changed.startswith(pattern.rstrip("*").rstrip("/"))
for pattern in paths
if pattern.endswith("/**")
), f"nothing in the path filter matches {changed}"
if __name__ == "__main__":
pytest.main([__file__, "-v"])

7
tests/studio/test_cli_studio_defaults.py
View file

@ -68,3 +68,10 @@ def test_studio_run_host_is_loopback():
f"`unsloth studio run` --host default must be '127.0.0.1' (loopback) "
f"but got '{host_default}'."
)
def test_dns_pinning_opt_out_is_registered_safe_by_default():
source = _STUDIO_CMD_PY.read_text()
for func_name in ("studio_default", "run"):
default = _find_typer_option_default(source, func_name, "--disable-dns-pinning")
assert default is False, f"{func_name} must keep DNS pinning enabled by default"

670
tests/test_compressed_export_gpu_release.py Normal file
View file

@ -0,0 +1,670 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""The compressed (FP8/NVFP4) export must free GPU weights before its llm-compressor
subprocess loads a second copy from disk, including for accelerate-dispatched multi-GPU
shards, which the old single-device-only ``.to("cpu")`` skipped and left resident.
Pulls the release/restore helpers out of unsloth/save.py via AST (importing the module
needs torch/transformers) and exercises them with fakes.
"""
from __future__ import annotations
import ast
import gc
import sys
import types
from pathlib import Path
import pytest
_SAVE_PY = Path(__file__).resolve().parent.parent / "unsloth" / "save.py"
_WANTED = {
"_accelerate_dispatch_root",
"_snapshot_dispatch_state",
"_drop_accelerator_tied_param_cache",
"_accelerate_move_guards",
"_split_tensor_path",
"_lookup_tensor",
"_share_tensor",
"_restore_dispatch_state",
"_offload_model_for_quantize_subprocess",
"_restore_model_after_quantize_subprocess",
}
_WANTED_ASSIGNS = {
"_DISPATCH_SNAPSHOT_ATTR",
"_ACCELERATE_MOVE_GUARDS",
} # module constants the helpers close over
class _FakeLogger:
def __init__(self):
self.warnings = []
def warning_once(self, msg):
self.warnings.append(msg)
def _load_helpers(fake_torch, fake_logger):
tree = ast.parse(_SAVE_PY.read_text(encoding = "utf-8"))
keep = [
node
for node in tree.body
if (isinstance(node, ast.FunctionDef) and node.name in _WANTED)
or (
isinstance(node, ast.Assign)
and any(isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets)
)
]
n_fns = sum(1 for node in keep if isinstance(node, ast.FunctionDef))
assert n_fns == len(_WANTED), "release helpers missing from save.py"
namespace = {"torch": fake_torch, "logger": fake_logger}
exec( # noqa: S102 - loading trusted repo source
compile(ast.Module(body = keep, type_ignores = []), str(_SAVE_PY), "exec"),
namespace,
)
return namespace
def _fake_torch(cuda_available = True):
t = types.ModuleType("torch")
t.cuda = types.SimpleNamespace(is_available = lambda: cuda_available)
return t
class _FakeModel:
def __init__(
self,
device_map = None,
devices = ("cuda:0",),
quantized = False,
):
if device_map is not None:
self.hf_device_map = device_map
self._devices = [types.SimpleNamespace(device = d) for d in devices]
self.moved_to = []
self.is_loaded_in_4bit = quantized
def parameters(self):
return iter(self._devices)
def to(self, target):
self.moved_to.append(str(target))
return self
@pytest.fixture
def _fake_accelerate(monkeypatch):
calls = {"removed": [], "dispatched": [], "dispatch_kwargs": [], "hooks_added": []}
accel = types.ModuleType("accelerate")
def _dispatch(model, device_map, **kwargs):
calls["dispatched"].append((model, dict(device_map)))
calls["dispatch_kwargs"].append(kwargs)
accel.dispatch_model = _dispatch
hooks = types.ModuleType("accelerate.hooks")
hooks.remove_hook_from_submodules = lambda model: calls["removed"].append(model)
hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append((module, hook))
accel.hooks = hooks
monkeypatch.setitem(sys.modules, "accelerate", accel)
monkeypatch.setitem(sys.modules, "accelerate.hooks", hooks)
return calls
def test_dispatched_multi_gpu_model_is_released_and_redispatched(_fake_accelerate):
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 0, "model.layers.1": 1}
model = _FakeModel(device_map = device_map, devices = ("cuda:0", "cuda:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model] # hooks removed before the move
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_dispatched_move_failure_redispatches_and_returns_none(_fake_accelerate):
# If .to("cpu") raises after the hooks came off, the model must be re-dispatched,
# not left hookless and half-moved.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.1": 1}
class _MoveFails(_FakeModel):
def to(self, target):
raise RuntimeError("host RAM cannot hold the sharded model")
model = _MoveFails(device_map = device_map, devices = ("cuda:0", "cuda:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token is None # offload aborted
assert _fake_accelerate["removed"] == [model] # hooks were removed...
assert _fake_accelerate["dispatched"] == [(model, device_map)] # ...then restored
def test_single_device_move_failure_restores_and_returns_none():
ns = _load_helpers(_fake_torch(), _FakeLogger())
class _MoveFails(_FakeModel):
def __init__(self):
super().__init__(devices = ("cuda:0",))
def to(self, target):
self.moved_to.append(str(target))
if target == "cpu":
raise RuntimeError("move failed")
return self
model = _MoveFails()
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token is None
# attempted the cpu move, then restored back to the original device
assert model.moved_to == ["cpu", "cuda:0"]
def test_cpu_spilled_map_still_releases_its_gpu_shards(_fake_accelerate):
# One module spilled to CPU, but the rest is the GPU memory the reload needs, and
# the spilled weights are already in host RAM, so the move is safe.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1, "model.layers.9": "cpu"}
model = _FakeModel(device_map = device_map)
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model]
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_disk_offloaded_map_is_left_alone(_fake_accelerate):
# disk/meta entries are not on the model, so moving would materialize the whole
# checkpoint into RAM.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": 0, "model.layers.9": "disk"})
assert ns["_offload_model_for_quantize_subprocess"](model) is None
assert model.moved_to == []
assert _fake_accelerate["removed"] == []
def test_all_cpu_map_is_left_alone(_fake_accelerate):
# Nothing on an accelerator: no GPU memory to reclaim, so do not churn the hooks.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": "cpu", "model.layers.0": "cpu"})
assert ns["_offload_model_for_quantize_subprocess"](model) is None
assert model.moved_to == []
assert _fake_accelerate["removed"] == []
def test_single_device_model_keeps_plain_move():
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(devices = ("cuda:0",))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert model.moved_to == ["cpu"]
assert token is not None and token[0] == "device"
ns["_restore_model_after_quantize_subprocess"](model, token)
assert model.moved_to[-1] == "cuda:0"
def test_quantized_model_is_released_when_the_stack_allows_it():
# Studio exports load 4-bit by DEFAULT, so skipping quantized models left a shard
# on every GPU. Release them too where the move is accepted.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(devices = ("cuda:0",), quantized = True)
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token == ("device", "cuda:0")
assert model.moved_to == ["cpu"]
def test_quantized_model_that_refuses_to_move_is_left_usable():
# transformers rejects .to() for some bitsandbytes builds and raises before
# anything moves, so the old behaviour must hold: no token, nothing escaping.
ns = _load_helpers(_fake_torch(), _FakeLogger())
class _Refuses(_FakeModel):
def to(self, target):
raise ValueError("`.to` is not supported for 4-bit bitsandbytes models")
model = _Refuses(devices = ("cuda:0",), quantized = True)
assert ns["_offload_model_for_quantize_subprocess"](model) is None
def test_no_cuda_is_noop_and_restore_none_is_noop():
ns = _load_helpers(_fake_torch(cuda_available = False), _FakeLogger())
model = _FakeModel()
assert ns["_offload_model_for_quantize_subprocess"](model) is None
ns["_restore_model_after_quantize_subprocess"](model, None) # must not raise
assert model.moved_to == []
def test_restore_failure_warns_instead_of_raising(_fake_accelerate):
fake_logger = _FakeLogger()
ns = _load_helpers(_fake_torch(), fake_logger)
class _ExplodingModel(_FakeModel):
def to(self, target):
raise RuntimeError("device gone")
model = _ExplodingModel(devices = ("cuda:0",))
ns["_restore_model_after_quantize_subprocess"](model, ("device", "cuda:0"))
assert fake_logger.warnings # warned, did not raise
def test_lora_merge_budgets_per_device():
# A merged tensor W lives on the GPU of its source layer, so budget against W's
# own device, not GPU0, else a sharded model OOMs GPU1+ (#7053).
src = _SAVE_PY.read_text(encoding = "utf-8")
tree = ast.parse(src)
fn = next(
(
n
for n in ast.walk(tree)
if isinstance(n, ast.FunctionDef) and n.name == "unsloth_save_model"
),
None,
)
assert fn is not None, "unsloth_save_model not found"
body = ast.get_source_segment(src, fn)
# Budget keyed on W's device, not a hardcoded device 0 / unqualified alloc.
assert "torch.cuda.memory_allocated(W.device)" in body
assert "_device_vram_budget(W.device)" in body
assert "get_device_properties(0).total_memory * maximum_memory_usage" not in body
# ── the torchao ("portable" FP8/INT8) export shares the same release ──
def _fake_torch_xpu():
t = types.ModuleType("torch")
t.cuda = types.SimpleNamespace(is_available = lambda: False)
t.xpu = types.SimpleNamespace(is_available = lambda: True)
return t
def test_dispatched_xpu_model_is_released(_fake_accelerate):
# torchao runs on Intel GPUs too, so an XPU-dispatched shard must release exactly
# like a CUDA one.
ns = _load_helpers(_fake_torch_xpu(), _FakeLogger())
device_map = {"model.embed": "xpu:0", "model.layers.0": "xpu:1"}
model = _FakeModel(device_map = device_map, devices = ("xpu:0", "xpu:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model]
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_single_device_xpu_model_is_released():
ns = _load_helpers(_fake_torch_xpu(), _FakeLogger())
model = _FakeModel(devices = ("xpu:0",))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token == ("device", "xpu:0")
assert model.moved_to == ["cpu"]
def test_torchao_export_uses_the_shared_release():
"""The torchao path must not re-inline a single-device-only ``.to("cpu")``.
A plain move is invalid on a dispatched model, so single-device-only handling left
a multi-GPU shard resident while ``device_map="auto"`` loaded a second copy.
"""
src = _SAVE_PY.read_text(encoding = "utf-8")
torchao = src.split("def _unsloth_save_torchao(", 1)[1].split("\ndef ", 1)[0]
assert "_offload_model_for_quantize_subprocess(model)" in torchao
assert "_restore_model_after_quantize_subprocess(model" in torchao
# No hand-rolled single-device gate left behind.
assert "len(_devs) == 1" not in torchao
# ── regressions for the multi-GPU dispatch branch ──
class _Child:
"""Minimal stand-in for an nn.Module leaf, enough for the dispatch walk."""
def __init__(
self,
name = "inner",
device_map = None,
):
self._modules = {}
self.__dict__["_name"] = name
if device_map is not None:
self.hf_device_map = device_map
def named_modules(self):
yield "", self
for key, child in self._modules.items():
for sub_name, sub in child.named_modules():
yield (f"{key}.{sub_name}" if sub_name else key), sub
def get_submodule(self, target):
node = self
for part in target.split("."):
node = node._modules[part]
return node
def named_parameters(self, remove_duplicate = True):
return iter(())
def named_buffers(self, remove_duplicate = True):
return iter(())
class _PeftLikeWrapper(_Child):
"""Proxies unknown attributes to the wrapped model, like ``PeftModelForCausalLM``:
``hasattr(wrapper, "_hf_hook")`` is True while ``delattr`` fails, which is what made
the offload a silent no-op."""
def __init__(self, inner):
super().__init__(name = "wrapper")
self._modules["base_model"] = inner
self.moved_to = []
def __getattr__(self, item):
return getattr(self._modules["base_model"], item)
def to(self, target):
self.moved_to.append(str(target))
return self
def parameters(self):
return iter(self._modules["base_model"]._devices)
def test_dispatch_root_is_the_inner_model_for_a_peft_style_wrapper(_fake_accelerate):
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1}
inner = _Child(device_map = device_map)
inner._devices = [types.SimpleNamespace(device = "cuda:0")]
wrapper = _PeftLikeWrapper(inner)
assert ns["_accelerate_dispatch_root"](wrapper) is inner
token = ns["_offload_model_for_quantize_subprocess"](wrapper)
# hooks must come off the INNER module, not the proxying wrapper
assert _fake_accelerate["removed"] == [inner]
assert wrapper.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
def test_dispatch_root_falls_back_to_the_model_it_was_given():
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": 0})
assert ns["_accelerate_dispatch_root"](model) is model
def test_offload_failure_is_logged_not_swallowed():
# A bare `return None` is indistinguishable from "nothing to move".
fake_logger = _FakeLogger()
ns = _load_helpers(_fake_torch(), fake_logger)
class _Explodes(_FakeModel):
@property
def hf_device_map(self):
raise RuntimeError("boom")
assert ns["_offload_model_for_quantize_subprocess"](_Explodes()) is None
assert any("boom" in w for w in fake_logger.warnings)
def test_restore_without_a_snapshot_forwards_skip_keys(_fake_accelerate):
# dispatch_model() defaults skip_keys to None, which moves every forward kwarg to
# the executing device, wrong for tensors transformers marks device-invariant.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1}
model = _FakeModel(device_map = device_map, devices = ("cuda:0", "cuda:1"))
model._skip_keys_device_placement = ["past_key_values"]
ns["_restore_model_after_quantize_subprocess"](model, ("dispatch", device_map))
assert _fake_accelerate["dispatched"] == [(model, device_map)]
assert _fake_accelerate["dispatch_kwargs"] == [{"skip_keys": ["past_key_values"]}]
def test_snapshot_restores_a_forward_patched_after_the_dispatch(_fake_accelerate):
"""accelerate restores ``forward = _old_forward`` on removal, and ``_old_forward``
is the forward from when the hook was FIRST attached. unsloth patches forwards after
the dispatch, so a naive remove/re-add throws every fused kernel away for good."""
ns = _load_helpers(_fake_torch(), _FakeLogger())
root = _Child(device_map = {"model.embed": 0, "mlp": 1})
mlp = _Child(name = "mlp")
root._modules["mlp"] = mlp
stock_forward = lambda *a, **k: "stock" # noqa: E731
fused_forward = lambda *a, **k: "unsloth-fused" # noqa: E731
mlp._hf_hook = object()
mlp._old_forward = stock_forward # captured by accelerate at dispatch time
mlp.forward = fused_forward # installed by unsloth afterwards
snapshot = ns["_snapshot_dispatch_state"](root)
# what accelerate's removal does
del mlp.__dict__["_hf_hook"]
mlp.forward = mlp._old_forward
del mlp.__dict__["_old_forward"]
assert mlp.forward() == "stock"
ns["_restore_dispatch_state"](root, snapshot)
assert mlp.forward() == "unsloth-fused"
assert mlp.__dict__["_old_forward"] is stock_forward
def test_snapshot_reties_shared_parameters(_fake_accelerate):
"""A CPU round trip repoints every tensor, so replaying the hooks alone leaves tied
weights as independent copies: double VRAM, and updates to one never reach the other."""
import torch
root = _Child(device_map = {"embed": 0, "head": 0})
shared = torch.nn.Parameter(torch.zeros(4, 4))
for name in ("embed", "head"):
child = _Child(name = name)
child._parameters = {"weight": shared}
child._buffers = {}
root._modules[name] = child
def named(remove_duplicate = True):
seen, out = set(), []
for mod_name, mod in root._modules.items():
for attr, tensor in mod._parameters.items():
if remove_duplicate and id(tensor) in seen:
continue
seen.add(id(tensor))
out.append((f"{mod_name}.{attr}", tensor))
return iter(out)
root.named_parameters = named
ns = _load_helpers(_fake_torch(), _FakeLogger())
snapshot = ns_ties = ns["_snapshot_dispatch_state"](root)
assert ns_ties[3] == [["embed.weight", "head.weight"]]
# what the replay leaves behind before the retie step
root._modules["head"]._parameters["weight"] = torch.nn.Parameter(shared.detach().clone())
assert (
root._modules["embed"]._parameters["weight"].data_ptr()
!= root._modules["head"]._parameters["weight"].data_ptr()
)
ns["_restore_dispatch_state"](root, snapshot)
assert (
root._modules["embed"]._parameters["weight"].data_ptr()
== root._modules["head"]._parameters["weight"].data_ptr()
)
def test_meta_tensors_never_form_tie_groups(_fake_accelerate):
"""Offloaded parameters all sit on meta with storage pointer 0, so grouping by
pointer alone would collapse them into one fake tie and overwrite them all."""
import torch
root = _Child(device_map = {"a": 0, "b": "cpu", "c": "cpu"})
live = torch.nn.Parameter(torch.zeros(4, 4))
offloaded = [
torch.nn.Parameter(torch.empty(4, 4, device = "meta")),
torch.nn.Parameter(torch.empty(8, 2, device = "meta")),
]
def named(remove_duplicate = True):
return iter([("a.weight", live), ("b.weight", offloaded[0]), ("c.weight", offloaded[1])])
root.named_parameters = named
ns = _load_helpers(_fake_torch(), _FakeLogger())
_hooks, places, _attrs, ties, _grads = ns["_snapshot_dispatch_state"](root)
assert ties == [] # nothing is tied here
assert "b.weight" in places # still tracked for placement
def test_accelerate_move_guards_survive_the_replay(_fake_accelerate):
"""remove_hook_from_module also deletes the to/cuda/... guards dispatch_model
installs to stop a caller moving an offloaded model."""
ns = _load_helpers(_fake_torch(), _FakeLogger())
root = _Child(device_map = {"": 0})
guard = lambda *a, **k: "blocked" # noqa: E731
root._hf_hook = object()
root.to = guard
root.cuda = guard
snapshot = ns["_snapshot_dispatch_state"](root)
del root.__dict__["_hf_hook"], root.__dict__["to"], root.__dict__["cuda"]
ns["_restore_dispatch_state"](root, snapshot)
assert root.__dict__["to"] is guard
assert root.__dict__["cuda"] is guard
def test_gradients_survive_the_offload_round_trip():
"""init_hook rebuilds the Parameter and drops .grad, so the snapshot has to carry it."""
import torch
root = _Child(device_map = {"": 0})
weight = torch.nn.Parameter(torch.zeros(4, 4))
weight.grad = torch.full((4, 4), 3.0)
root._parameters = {"weight": weight}
root.named_parameters = lambda remove_duplicate = True: iter([("weight", weight)])
ns = _load_helpers(_fake_torch(), _FakeLogger())
snapshot = ns["_snapshot_dispatch_state"](root)
assert torch.equal(snapshot[4]["weight"], torch.full((4, 4), 3.0))
# What init_hook does: same name, fresh Parameter, no grad.
replacement = torch.nn.Parameter(torch.zeros(4, 4))
assert replacement.grad is None
root._parameters = {"weight": replacement}
ns["_restore_dispatch_state"](root, snapshot)
assert replacement.grad is not None, "the restore must put the gradient back"
assert torch.equal(replacement.grad, torch.full((4, 4), 3.0))
def test_the_other_torchao_path_also_clears_the_failed_copy():
"""Both torchao paths must drop the copy and the traceback pinning it before restoring."""
src = _SAVE_PY.read_text(encoding = "utf-8")
body = src.split("\ndef _unsloth_save_torchao(", 1)[1].split("\ndef ", 1)[0]
finally_block = body.split(" finally:", 1)[1]
assert "del quantized_model" in finally_block
assert "traceback.clear_frames" in finally_block
restore_at = finally_block.index("_restore_model_after_quantize_subprocess")
assert finally_block.index("del quantized_model") < restore_at
assert finally_block.index("traceback.clear_frames") < restore_at
def test_cpu_spill_rejection_is_retryable():
"""bitsandbytes rejects a CPU-spilled map with a ValueError that says nothing about
memory, so the single-device retry has to match it explicitly."""
import importlib.util
from pathlib import Path
export_py = (
Path(__file__).resolve().parent.parent
/ "studio"
/ "backend"
/ "core"
/ "export"
/ "export.py"
)
src = ast.parse(export_py.read_text(encoding = "utf-8"))
keep = [
n
for n in src.body
if isinstance(n, ast.FunctionDef) and n.name in {"_is_oom_error", "_is_cpu_spill_rejection"}
]
assert len(keep) == 2
namespace = {"torch": None}
exec( # noqa: S102 - loading trusted repo source
compile(ast.Module(body = keep, type_ignores = []), str(export_py), "exec"), namespace
)
bnb = ValueError(
"Some modules are dispatched on the CPU or the disk. Make sure you have enough "
"GPU RAM to fit the quantized model."
)
assert not namespace["_is_oom_error"](bnb)
assert namespace["_is_cpu_spill_rejection"](bnb)
assert namespace["_is_oom_error"](RuntimeError("CUDA out of memory. Tried to allocate 1 GiB"))
assert not namespace["_is_cpu_spill_rejection"](RuntimeError("some other failure"))
def test_torchao_releases_the_quantized_copy_in_finally():
"""If save_pretrained raises, the quantized copy must still be dropped before the
original is restored, or both are resident at once."""
src = _SAVE_PY.read_text(encoding = "utf-8")
body = src.split("def _unsloth_save_torchao_with_given_config(", 1)[1].split("\ndef ", 1)[0]
finally_block = body.split(" finally:", 1)[1]
assert "del quantized_model" in finally_block
assert "_restore_model_after_quantize_subprocess(model, model_restore)" in finally_block
# and the restore must come after the copy is dropped
assert finally_block.index("del quantized_model") < finally_block.index(
"_restore_model_after_quantize_subprocess"
)
# dropping the local is not enough: the live traceback still holds the frames
assert "traceback.clear_frames" in finally_block
assert finally_block.index("traceback.clear_frames") < finally_block.index(
"_restore_model_after_quantize_subprocess"
)
def test_a_live_traceback_pins_the_failed_copy_until_its_frames_are_cleared():
"""Why the clear_frames call above is load-bearing, on plain objects."""
import sys
import traceback
import weakref
class _Copy:
pass
def _build_and_fail(sink):
copy = _Copy() # noqa: F841 -- the point is that the frame retains it
sink.append(weakref.ref(copy))
raise RuntimeError("save_pretrained failed")
def _run(clear_frames):
# try/finally with the exception still in flight, exactly as in save.py
sink = []
alive = None
try:
try:
_build_and_fail(sink)
finally:
if clear_frames:
exc = sys.exc_info()[1]
if exc is not None:
traceback.clear_frames(exc.__traceback__)
gc.collect()
alive = sink[0]() is not None
except RuntimeError:
pass
return alive
assert _run(clear_frames = False), "expected the traceback to pin the copy"
assert not _run(clear_frames = True), "clear_frames must release it"

490
unsloth/save.py
View file

@ -48,6 +48,7 @@ import functools
from transformers.models.llama.modeling_llama import logger
from .kernels import fast_dequantize, QUANT_STATE, get_lora_parameters_bias
import subprocess
import traceback
import psutil
import re
from transformers.models.llama.modeling_llama import logger
@ -1122,7 +1123,19 @@ def unsloth_save_model(
torch_dtype
)
max_vram = int(torch.cuda.get_device_properties(0).total_memory * maximum_memory_usage)
# A merged tensor lives on the GPU of its source layer, so budget against W's own
# device, not GPU0, else a sharded model OOMs GPU1+ while only GPU0 is checked.
_max_vram_by_device = {}
def _device_vram_budget(dev):
if dev.type != "cuda":
return None
idx = dev.index if dev.index is not None else torch.cuda.current_device()
if idx not in _max_vram_by_device:
_max_vram_by_device[idx] = int(
torch.cuda.get_device_properties(idx).total_memory * maximum_memory_usage
)
return _max_vram_by_device[idx]
print("Unsloth: Saving model... This might take 5 minutes ...")
@ -1138,8 +1151,15 @@ def unsloth_save_model(
if bias is not None:
state_dict[f"model.layers.{j}.{item}.bias"] = bias
if (torch.cuda.memory_allocated() + W.nbytes) < max_vram:
# Save to GPU memory
_dev_budget = _device_vram_budget(W.device)
if (
_dev_budget is not None
and (torch.cuda.memory_allocated(W.device) + W.nbytes) < _dev_budget
):
# Fits on W's own GPU
state_dict[name] = W
elif W.device.type != "cuda":
# Already off-GPU: keeping it costs no VRAM
state_dict[name] = W
# [TODO] Saving to RAM seems to leak memory???
# elif (max_ram - W.nbytes) > 0:
@ -4524,30 +4544,61 @@ def _unsloth_save_torchao_with_given_config(
else:
kwargs = {"dtype": torch.bfloat16}
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
# Else the original stays resident on every GPU while device_map="auto" below
# loads a second copy.
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
torchao_save_directory = save_directory + "-torchao"
# TorchAO does not support safe_serialization right now 0.14.0 seems broken!
safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0")
safe_serialization = False
if push_to_hub:
quantized_model.push_to_hub(
torchao_save_directory, safe_serialization = safe_serialization, token = token
# The original stays offloaded until the quantized copy is saved AND released,
# else both are resident at once and the restore OOMs.
try:
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
tokenizer.push_to_hub(torchao_save_directory, token = token)
else:
quantized_model.save_pretrained(
torchao_save_directory, safe_serialization = safe_serialization
)
tokenizer.save_pretrained(torchao_save_directory, token = token)
torchao_save_directory = save_directory + "-torchao"
# TorchAO does not support safe_serialization right now 0.14.0 seems broken!
safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0")
safe_serialization = False
if push_to_hub:
quantized_model.push_to_hub(
torchao_save_directory, safe_serialization = safe_serialization, token = token
)
tokenizer.push_to_hub(torchao_save_directory, token = token)
else:
quantized_model.save_pretrained(
torchao_save_directory, safe_serialization = safe_serialization
)
tokenizer.save_pretrained(torchao_save_directory, token = token)
finally:
# del here, not at the end of the try: if save_pretrained raises, the copy
# would otherwise still be resident while the original is restored.
quantized_model = None
del quantized_model
# A failed save leaves a live traceback whose frames still hold the copy, so
# dropping the local alone does not free its VRAM.
_exc = sys.exc_info()[1]
if _exc is not None:
traceback.clear_frames(_exc.__traceback__)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
_restore_model_after_quantize_subprocess(model, model_restore)
# Clean up the intermediate unquantized model
if os.path.exists(save_directory):
@ -4585,6 +4636,318 @@ def _print_compressed_hw_note(scheme, out_dir):
)
_DISPATCH_SNAPSHOT_ATTR = "_unsloth_dispatch_snapshot"
def _accelerate_move_guards():
"""The instance methods dispatch_model wraps to block moving an offloaded model."""
try:
from accelerate.hooks import _accelerate_added_attributes
return tuple(_accelerate_added_attributes)
except Exception:
return ("to", "cuda", "npu", "xpu", "mlu", "sdaa", "musa")
_ACCELERATE_MOVE_GUARDS = _accelerate_move_guards()
def _accelerate_dispatch_root(model):
"""The module that really owns the accelerate dispatch.
A PEFT wrapper only proxies ``_hf_hook``, so ``delattr`` fails and
``remove_hook_from_submodules`` raises before removing anything; ``hf_device_map``
keys are relative to the inner root too. Walks real children, never ``__getattr__``.
"""
node, seen = model, set()
while id(node) not in seen:
seen.add(id(node))
if "hf_device_map" in getattr(node, "__dict__", {}):
return node
children = getattr(node, "__dict__", {}).get("_modules") or {}
nxt = next(
(
children[a]
for a in ("base_model", "model")
if hasattr(children.get(a), "named_modules")
),
None,
)
if nxt is None:
return model
node = nxt
return model
def _snapshot_dispatch_state(root):
"""Hooks, tensor placements and instance forwards, so the dispatch can be replayed.
Re-deriving it with ``dispatch_model`` is not equivalent: PEFT reparents each
targeted ``Linear`` after transformers dispatched, so accelerate hooks modules that
never had any (measured: 395 -> 1379) and the logits shift enough to reorder top-5.
"""
hooks = [
(name, mod.__dict__["_hf_hook"])
for name, mod in root.named_modules()
if "_hf_hook" in mod.__dict__
]
# remove_duplicate=False: the default hides one half of every tied pair, exactly
# the half that needs re-tying below.
named = list(root.named_parameters(remove_duplicate = False)) + list(
root.named_buffers(remove_duplicate = False)
)
places = {name: tensor.device for name, tensor in named}
# Tied weights share one storage, but the CPU round trip repoints every tensor and
# tied_params_map is keyed on the old pointer, so replaying the hooks alone gives
# independent copies: double VRAM, and updates to one no longer reach the other.
# Skip meta tensors: offloaded parameters all sit on meta with pointer 0, which
# would collapse into one fake "tied" group of differently shaped tensors, and
# each side of a tie is already its own meta placeholder so nothing is lost.
groups = {}
for name, tensor in named:
if tensor.device.type == "meta":
continue
ptr = tensor.untyped_storage().data_ptr()
if ptr:
groups.setdefault(ptr, []).append(name)
ties = [names for names in groups.values() if len(names) > 1]
# Removing a hook restores `forward = _old_forward`, captured before unsloth patched
# the module, so a remove/re-add permanently drops every fused kernel installed after
# the dispatch (measured: apply_lora_mlp_swiglu on all 28 MLPs). It also deletes the
# `to`/`cuda`/... move guards, so record those too.
attrs = ("forward", "_old_forward") + tuple(_ACCELERATE_MOVE_GUARDS)
saved_attrs = {
name: {a: mod.__dict__[a] for a in attrs if a in mod.__dict__}
for name, mod in root.named_modules()
if any(a in mod.__dict__ for a in attrs)
}
# Re-adding a hook runs init_hook -> set_module_tensor_to_device, which builds a fresh
# Parameter and so drops .grad. Snapshot the gradients and reattach them on restore.
grads = {
name: getattr(tensor, "grad", None)
for name, tensor in root.named_parameters(remove_duplicate = False)
if getattr(tensor, "grad", None) is not None
}
return hooks, places, saved_attrs, ties, grads
def _drop_accelerator_tied_param_cache(snapshot) -> None:
"""Drop the GPU tensors accelerate caches in each hook's ``tied_params_map``.
Holding the hooks across the offload pins a GPU copy of the tied embedding (0.31 GB
of 1.24 GB here). The entries are keyed on the pre-move ``data_ptr`` so they are
stale anyway, and re-attaching repopulates them.
"""
for _name, hook in snapshot[0]:
cache = getattr(hook, "tied_params_map", None)
if not cache:
continue
for ptr in list(cache):
entry = cache[ptr]
for device in list(entry):
if str(device) != "cpu":
del entry[device]
if not entry:
del cache[ptr]
def _split_tensor_path(root, full_name):
"""``("model.embed_tokens.weight")`` -> ``(the module, "weight")``."""
mod_name, _, attr = full_name.rpartition(".")
try:
return (root.get_submodule(mod_name) if mod_name else root), attr
except AttributeError:
return None, attr
def _lookup_tensor(root, full_name):
mod, attr = _split_tensor_path(root, full_name)
if mod is None:
return None
for store in ("_parameters", "_buffers"):
found = (getattr(mod, store, None) or {}).get(attr)
if found is not None:
return found
return None
def _share_tensor(root, full_name, leader) -> None:
"""Point ``full_name`` back at ``leader``, restoring a tie."""
mod, attr = _split_tensor_path(root, full_name)
if mod is None:
return
for store in ("_parameters", "_buffers"):
target = getattr(mod, store, None)
if target is None or attr not in target:
continue
current = target[attr]
if current is None or current.device != leader.device or current.shape != leader.shape:
return # not actually the same tensor; leave it alone
target[attr] = leader
return
def _restore_dispatch_state(root, snapshot) -> None:
"""Replay ``_snapshot_dispatch_state``."""
from accelerate.hooks import add_hook_to_module
hooks, places, saved_attrs, ties, grads = snapshot
for name, hook in hooks:
add_hook_to_module(root.get_submodule(name) if name else root, hook)
# Re-adding a hook rewraps whatever `_old_forward` now holds, so put the exact
# callables back, `_old_forward` first.
for name, values in saved_attrs.items():
mod = root.get_submodule(name) if name else root
for attr in ("_old_forward", "forward", *_ACCELERATE_MOVE_GUARDS):
if attr in values:
mod.__dict__[attr] = values[attr]
# init_hook only re-places tensors the hooked module owns, so anything added after
# the dispatch (the LoRA adapters) is still on CPU.
for mod_name, mod in root.named_modules():
for attr in ("_parameters", "_buffers"):
store = getattr(mod, attr, None)
if not store:
continue
for tensor_name, tensor in list(store.items()):
if tensor is None:
continue
full = f"{mod_name}.{tensor_name}" if mod_name else tensor_name
want = places.get(full)
if want is None or tensor.device == want:
continue
if getattr(tensor, "quant_state", None) is not None:
# Only bitsandbytes' own .to() moves absmax/code/state2 with the data.
mod.to(want)
else:
tensor.data = tensor.data.to(want)
# Reattach the gradients init_hook discarded, on their weight's device.
for name, grad in grads.items():
tensor = _lookup_tensor(root, name)
if tensor is not None and tensor.grad is None and tensor.shape == grad.shape:
tensor.grad = grad.to(tensor.device)
# Re-tie last, once every tensor is back on its own device.
for names in ties:
leader = _lookup_tensor(root, names[0])
if leader is None:
continue
for follower in names[1:]:
_share_tensor(root, follower, leader)
# init_hook refilled tied_params_map with the pre-retie tensors, now unreferenced
# by the model but still pinned by the map.
if ties:
_drop_accelerator_tied_param_cache(snapshot)
def _offload_model_for_quantize_subprocess(model):
"""Best-effort: move the model's weights off the GPU before the quantized export
loads its own copy from disk, so the GPUs need not hold both at once. Returns an
opaque token for ``_restore_model_after_quantize_subprocess`` (None if nothing moved).
Two shapes are handled:
* single-device CUDA/XPU model -> ``.to("cpu")``, restored with ``.to(device)``;
* accelerate-dispatched model (a multi-GPU ``device_map`` shard, e.g. the Studio
multi-GPU export load) -> hooks removed and moved to CPU, restored by replaying
the dispatch. A plain ``.to("cpu")`` is invalid here, which is why the old
single-device-only move left every GPU holding a full copy. A map spilling to
CPU is still released, but disk/meta targets are left alone: accelerate keeps
those parameters off the model, so moving would materialize the whole checkpoint.
Quantized (bnb) models are attempted too rather than skipped: Studio exports load
4-bit by DEFAULT, so skipping them left a shard on every GPU. transformers refuses
``.to()`` for some bitsandbytes builds and that refusal raises before anything moves,
so the failure path restores the model and returns None, i.e. the old behaviour.
"""
try:
_has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
if not ((torch.cuda.is_available() or _has_xpu) and hasattr(model, "parameters")):
return None
device_map = getattr(model, "hf_device_map", None)
if device_map:
targets = {str(v).lower() for v in device_map.values()}
# A cpu spill is fine to move, it is already in host RAM. disk/meta is not:
# those parameters are off the model, so .to("cpu") would materialize the
# whole checkpoint into RAM.
if not all(t.isdigit() or t.startswith(("cuda", "xpu")) or t == "cpu" for t in targets):
return None
if not any(t.isdigit() or t.startswith(("cuda", "xpu")) for t in targets):
return None # nothing on an accelerator: no GPU memory to reclaim
from accelerate.hooks import remove_hook_from_submodules
# A PEFT wrapper only proxies the hooks; they live on the inner root.
root = _accelerate_dispatch_root(model)
try:
setattr(root, _DISPATCH_SNAPSHOT_ATTR, _snapshot_dispatch_state(root))
except Exception as snap_exc:
# Restore will fall back to re-deriving from the device_map.
logger.warning_once(
f"Unsloth: could not snapshot the accelerate dispatch "
f"({type(snap_exc).__name__}: {snap_exc}); re-dispatching on restore."
)
remove_hook_from_submodules(root)
try:
model.to("cpu")
except Exception:
# The move failed after the hooks came off; re-dispatch so the model is
# left usable rather than hookless and half-moved across CPU/GPUs.
_restore_model_after_quantize_subprocess(model, ("dispatch", dict(device_map)))
return None
snapshot = getattr(root, _DISPATCH_SNAPSHOT_ATTR, None)
if snapshot is not None:
_drop_accelerator_tied_param_cache(snapshot)
return ("dispatch", dict(device_map))
devices = {str(p.device) for p in model.parameters()}
if len(devices) == 1 and next(iter(devices)).startswith(("cuda", "xpu")):
device = next(model.parameters()).device
try:
model.to("cpu")
except Exception:
_restore_model_after_quantize_subprocess(model, ("device", device))
return None
return ("device", device)
except Exception as exc:
# A silent `return None` is indistinguishable from "nothing to move", which
# hides a real bug behind a merely slower export.
logger.warning_once(
f"Unsloth: could not free the model's accelerator memory before the quantized "
f"export ({type(exc).__name__}: {exc}); continuing with the model resident."
)
return None
return None
def _restore_model_after_quantize_subprocess(model, restore_token) -> None:
"""Undo ``_offload_model_for_quantize_subprocess``; warns instead of raising."""
if restore_token is None:
return
kind, value = restore_token
try:
if kind == "dispatch":
root = _accelerate_dispatch_root(model)
snapshot = root.__dict__.pop(_DISPATCH_SNAPSHOT_ATTR, None)
if snapshot is not None:
_restore_dispatch_state(root, snapshot)
else:
from accelerate import dispatch_model
# skip_keys matters: without it accelerate moves every forward kwarg
# to the executing device, wrong for device-invariant cache tensors.
dispatch_model(
root,
device_map = value,
skip_keys = getattr(root, "_skip_keys_device_placement", None),
)
else:
model.to(value) # restore the model to its original device
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device(s) after the "
"quantized export; it may remain on CPU."
)
def _unsloth_save_compressed_tensors(
model,
save_directory: Union[str, os.PathLike],
@ -4654,7 +5017,7 @@ def _unsloth_save_compressed_tensors(
# 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and
# quantize inside an isolated temp dir instead of writing ./<repo_id> into the cwd.
repo_id, work_tmp, calib_tmp, model_dev = None, None, None, None
repo_id, work_tmp, calib_tmp, model_restore = None, None, None, None
if push_to_hub:
repo_id = os.fspath(save_directory)
work_tmp = tempfile.mkdtemp(prefix = "unsloth-compressed-")
@ -4806,23 +5169,8 @@ def _unsloth_save_compressed_tensors(
cmd += ["--variant", variant]
# Free the in-memory model's CUDA memory before the subprocess loads its own copy from
# disk, so a single GPU need not hold both at once. Best-effort and restored in finally;
# skipped for quantized or multi-device models where moving is unsafe.
try:
if (
torch.cuda.is_available()
and hasattr(model, "parameters")
and not getattr(model, "is_loaded_in_4bit", False)
and not getattr(model, "is_loaded_in_8bit", False)
and not getattr(model, "is_quantized", False)
):
_devs = {str(p.device) for p in model.parameters()}
if len(_devs) == 1 and next(iter(_devs)).startswith("cuda"):
_dev = next(model.parameters()).device
model.to("cpu")
model_dev = _dev # set only after a successful move, so finally can restore
except Exception:
model_dev = None
# disk, so the GPUs need not hold both at once. Best-effort, restored in finally.
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
@ -4893,14 +5241,7 @@ def _unsloth_save_compressed_tensors(
_print_compressed_hw_note(scheme, result)
return result
finally:
if model_dev is not None:
try:
model.to(model_dev) # restore the model to its original device
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device after compressed "
"export; it may remain on CPU."
)
_restore_model_after_quantize_subprocess(model, model_restore)
if calib_tmp is not None and os.path.isdir(calib_tmp):
shutil.rmtree(calib_tmp, ignore_errors = True)
if work_tmp is not None:
@ -4959,7 +5300,7 @@ def _unsloth_save_torchao(
# Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected
# 16-bit export written to save_directory is not overwritten or deleted; the torchao output is
# the sibling "<save_directory>-<suffix>" (or the repo id on a hub push).
repo_id, work_tmp, model_dev = None, None, None
repo_id, work_tmp, model_restore = None, None, None
work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-")
if push_to_hub:
repo_id = os.fspath(save_directory)
@ -5037,25 +5378,12 @@ def _unsloth_save_torchao(
auto_model = AutoModelForCausalLM
auto_processor = AutoProcessor if is_vlm else AutoTokenizer
# 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk.
# Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit
# resident alongside the reloaded copy and OOM a device that fit the model once.
# 3) Free the in-memory model's accelerator memory before reloading a fresh copy from
# disk, else it sits resident alongside the copy and OOMs a device that fit the
# model once. Covers CUDA and XPU (torchao runs on Intel GPUs too) plus multi-GPU
# dispatched shards, which a plain .to("cpu") cannot move.
_has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
try:
if (
(torch.cuda.is_available() or _has_xpu)
and hasattr(model, "parameters")
and not getattr(model, "is_loaded_in_4bit", False)
and not getattr(model, "is_loaded_in_8bit", False)
and not getattr(model, "is_quantized", False)
):
_devs = {str(p.device) for p in model.parameters()}
if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")):
_dev = next(model.parameters()).device
model.to("cpu")
model_dev = _dev
except Exception:
model_dev = None
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
@ -5126,14 +5454,20 @@ def _unsloth_save_torchao(
)
return result
finally:
if model_dev is not None:
try:
model.to(model_dev)
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device after torchao "
"export; it may remain on CPU."
)
# A raise pins the copy in the local and the live traceback, so free both or the
# restore below OOMs.
quantized_model = None
del quantized_model
_exc = sys.exc_info()[1]
if _exc is not None:
traceback.clear_frames(_exc.__traceback__)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
_restore_model_after_quantize_subprocess(model, model_restore)
if work_tmp is not None:
shutil.rmtree(work_tmp, ignore_errors = True)
for _ in range(3):

30
unsloth_cli/commands/studio.py
View file

@ -1285,6 +1285,12 @@ def studio_default(
help = "Force server-side tools (web search, code execution) on or off for "
"every request. Default: on for every bind, with the per-chat UI toggle honored.",
),
disable_dns_pinning: bool = typer.Option(
False,
"--disable-dns-pinning",
help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens "
"DNS-rebinding protection; hostname and redirect validation remain enabled.",
),
password: str = typer.Option(
"",
"--password",
@ -1354,6 +1360,15 @@ def studio_default(
err = True,
)
raise typer.Exit(2)
if disable_dns_pinning:
typer.echo(
"Error: --disable-dns-pinning on `unsloth studio` applies to the "
f"plain-server path only. For `unsloth studio {ctx.invoked_subcommand}`, "
f"put it after the subcommand: `unsloth studio {ctx.invoked_subcommand} "
"--disable-dns-pinning ...`",
err = True,
)
raise typer.Exit(2)
# Same for --api-only: dropping it here would silently serve the UI.
if api_only:
typer.echo(
@ -1398,6 +1413,10 @@ def studio_default(
# default (plain-server path; the `run` subcommand has its own --verbose).
if verbose:
_enable_verbose_access_logs()
if disable_dns_pinning:
os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1"
else:
os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0")
# Use the studio venv if present and not already in it. Resolve the child
# launcher BEFORE the gate: a headless gate strips the seeded
@ -1739,6 +1758,13 @@ def run(
"every request. Default: on for every bind."
),
),
disable_dns_pinning: bool = typer.Option(
False,
"--disable-dns-pinning",
rich_help_panel = _RUN_PANEL_TOOLS,
help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens "
"DNS-rebinding protection; hostname and redirect validation remain enabled.",
),
tool_call_healing: Optional[bool] = typer.Option(
None,
"--enable-tool-call-healing/--disable-tool-call-healing",
@ -1944,6 +1970,10 @@ def run(
_enable_verbose_access_logs()
if not any(a in ("--verbose", "-v", "--log-verbose") for a in extra_llama_args):
extra_llama_args.append("--log-verbose")
if disable_dns_pinning:
os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1"
else:
os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0")
# Promote legacy exact `-m`/`-hfr`/`-f` back into typer params;
# clusters stay in extras.