diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index b8f587b63e..ec437e0c32 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -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: @@ -193,6 +200,7 @@ jobs: --ignore=tests/sh \ --ignore=tests/studio/test_hardware_dispatch_matrix.py \ --ignore=tests/studio/test_is_mlx_dispatch_gate.py \ + --ignore=tests/studio/test_xpu_spoof_pipeline.py \ --ignore=tests/vllm_compat \ --ignore=tests/version_compat \ -m 'not server and not e2e' \ @@ -205,36 +213,43 @@ jobs: env: PYTHONPATH: ${{ github.workspace }}/studio UNSLOTH_COMPILE_DISABLE: '1' - # These two files mutate hardware.py module globals at runtime - # via the spoof fixtures, which leaks state into any other test - # that imports hardware. Run them in their own pytest invocation - # so the leak does not cross file boundaries. + # These files mutate hardware.py module globals at runtime via the + # spoof fixtures (CUDA/ROCm/XPU/MLX/CPU), which leaks state into any + # other test that imports hardware. Run them in their own pytest + # invocation so the leak does not cross file boundaries. run: | python -m pytest -q --tb=short \ tests/studio/test_hardware_dispatch_matrix.py \ - tests/studio/test_is_mlx_dispatch_gate.py + tests/studio/test_is_mlx_dispatch_gate.py \ + 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_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" diff --git a/install.ps1 b/install.ps1 index 9c91d4ba16..a2aff0b69a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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" diff --git a/install.sh b/install.sh index 84c5be9742..3bc2ff4c88 100755 --- a/install.sh +++ b/install.sh @@ -625,6 +625,36 @@ _is_pkg_installed() { esac } +# ── Helper: human-readable apt distro label for the sudo package prompt (#6207) ── +# Reads /etc/os-release so the Accept? prompt can say which distro we detected and +# that packages come from that distro's official apt repos (not a tarball). +_apt_distro_description() { + # Plain ( ... ) subshell — not $() — so case/;; stays bash-3.2-safe on macOS. + # Bash 3.2 misparses case arms inside command substitution and errors on `;;`. + ( + if [ ! -r /etc/os-release ]; then + printf 'a debian-like system' + exit 0 + fi + # shellcheck disable=SC1091 + . /etc/os-release 2>/dev/null || true + if [ -n "${NAME:-}" ] && [ -n "${VERSION_ID:-}" ]; then + _ad_label="$NAME $VERSION_ID" + elif [ -n "${PRETTY_NAME:-}" ]; then + _ad_label="$PRETTY_NAME" + elif [ -n "${NAME:-}" ]; then + _ad_label="$NAME" + else + printf 'a debian-like system' + exit 0 + fi + case " ${ID:-} ${ID_LIKE:-} " in + *" debian "*|*" ubuntu "*) _ad_label="${_ad_label} (debian-like)" ;; + esac + printf '%s' "$_ad_label" + ) +} + # ── Helper: install packages via apt, escalating to sudo only if needed ── # Usage: _smart_apt_install pkg1 pkg2 pkg3 ... _smart_apt_install() { @@ -655,11 +685,14 @@ _smart_apt_install() { # Step 3: Escalate -- need elevated permissions for remaining packages if command -v sudo >/dev/null 2>&1; then + _ad_desc="$(_apt_distro_description)" echo "" echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo " WARNING: We require sudo elevated permissions to install:" echo " $_STILL_MISSING" - echo " If you accept, we'll run sudo now, and it'll prompt your password." + echo " Detected ${_ad_desc}." + echo " If you accept, we'll run sudo apt-get to install these packages" + echo " from your distro's official repositories (not a third-party tarball)." echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo "" printf " Accept? [Y/n] " @@ -2227,6 +2260,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 ;; @@ -2238,12 +2272,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 ;; @@ -2283,10 +2319,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 @@ -3022,7 +3062,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" @@ -3091,7 +3131,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" @@ -3210,7 +3250,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. @@ -3306,12 +3346,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) diff --git a/pyproject.toml b/pyproject.toml index fe0ebd13b9..0f57ecf4df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,6 +93,17 @@ huggingfacenotorch = [ "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", ] +# torchcodec backend for Gemma audio / datasets>=4 (#7225). +# Pick the audio-torch* pin matching your torch minor (see TORCH_TORCHCODEC). +audio-torch210 = [ + "torchcodec>=0.10.0,<0.11.0 ; python_version >= '3.10'", +] +audio-torch290 = [ + "torchcodec>=0.8.0,<0.10.0 ; python_version >= '3.10'", +] +audio-torch280 = [ + "torchcodec>=0.6.0,<0.8.0 ; python_version >= '3.9'", +] huggingface = [ "unsloth[huggingfacenotorch]", "unsloth_zoo>=2026.7.6", @@ -532,16 +543,19 @@ cu126-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] kaggle = [ "unsloth[huggingface]", @@ -831,16 +845,19 @@ cu126-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] flashattentiontorch260abiFALSEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", @@ -1125,7 +1142,8 @@ intelgputorch210 = [ "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=1c4b44b36a557f7381e3076fb8843366742238648441d607c8d049c6da0f8886 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] intel-gpu-torch210 = [ - "unsloth[intelgputorch210]" + "unsloth[intelgputorch210]", + "unsloth[audio-torch210]", ] intelgputorch2110 = [ "unsloth_zoo[intelgpu]", @@ -1279,6 +1297,7 @@ rocm72-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] rocm711-torch2100 = [ "unsloth[amd]", @@ -1297,6 +1316,7 @@ rocm711-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] [project.urls] diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index c1be7a63a4..7bcee47c66 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -95,8 +95,8 @@ COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-i # Source: pytorch/torchcodec compatibility matrix on its README. TORCH_TORCHCODEC: dict[str, set[str]] = { "2.10": {"0.10"}, - "2.9": {"0.7", "0.8", "0.9"}, - "2.8": {"0.6"}, + "2.9": {"0.8", "0.9"}, + "2.8": {"0.6", "0.7"}, "2.7": {"0.3", "0.4", "0.5"}, "2.6": {"0.2", "0.3"}, "2.5": {"0.1", "0.2"}, diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 65b8d2b11c..1c21f8da86 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "_comment": "scan_packages.py allowlist (reviewed). Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", "version": 1, "entries": [ { @@ -303,8 +303,8 @@ "file": "openai/_base_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b", - "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14" + "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6", + "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66" }, { "package": "openai", @@ -319,8 +319,8 @@ "file": "openai/auth/_workload.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:", - "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567" + "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()", + "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0" }, { "package": "openai", @@ -343,8 +343,8 @@ "file": "openai/resources/beta/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", - "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd", + "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f" }, { "package": "openai", @@ -359,16 +359,16 @@ "file": "openai/resources/realtime/realtime.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05", - "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89" + "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5", + "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650" }, { "package": "openai", "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", - "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" + "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac", + "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f" }, { "package": "openai", diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 44282b2255..612d739806 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -1,134 +1,145 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "view-in-github", - "colab_type": "text" - }, - "source": [ - "\"Open" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6b87de59" + }, + "source": [ + "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", + "
\n", + "\n", + "\n", + " Join Discord if you need help + ⭐ Star us on Github ⭐\n", + "
\n", + "\n", + "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", + "\n", + "### Unsloth Studio\n", + "\n", + "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", + "\n", + "\n", + "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", + "\n", + "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" + ], + "id": "6b87de59" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e4206349" + }, + "source": [ + "

" + ], + "id": "e4206349" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "27da2957" + }, + "source": [ + "### Setup: Clone repo and run setup" + ], + "id": "27da2957" + }, + { + "cell_type": "code", + "metadata": { + "id": "27e68f91" + }, + "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local", + "execution_count": null, + "outputs": [], + "id": "27e68f91" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3e1771a9" + }, + "source": [ + "### Start Unsloth Studio" + ], + "id": "3e1771a9" + }, + { + "cell_type": "code", + "metadata": { + "id": "277e431e" + }, + "source": [ + "import sys\n", + "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", + "from colab import start\n", + "\n", + "# On Colab, start() auto-opens a Cloudflare link and prints admin login credentials.\n", + "# Use the Cloudflare link above the ready card to open Studio (in-cell iframes often stay blank).\n", + "start()\n", + "\n", + "# To skip the Cloudflare tunnel and try the in-notebook proxy iframe only:\n", + "# start(cloudflare=False)" + ], + "execution_count": null, + "outputs": [], + "id": "277e431e" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f2b0c6a1" + }, + "source": [ + "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", + "\n", + "Some other resources:\n", + "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", + "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", + "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", + "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", + "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", + "\n", + "
\n", + " \n", + " \n", + " \n", + "\n", + " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", + "\n", + " This notebook is licensed AGPL-3.0\n", + "
" + ], + "id": "f2b0c6a1" + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } }, - { - "cell_type": "markdown", - "id": "6b87de59", - "metadata": { - "id": "6b87de59" - }, - "source": [ - "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", - "
\n", - "\n", - "\n", - " Join Discord if you need help + ⭐ Star us on Github ⭐\n", - "
\n", - "\n", - "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", - "\n", - "### Unsloth Studio\n", - "\n", - "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", - "\n", - "\n", - "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", - "\n", - "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" - ] - }, - { - "cell_type": "markdown", - "id": "e4206349", - "metadata": { - "id": "e4206349" - }, - "source": [ - "

" - ] - }, - { - "cell_type": "markdown", - "id": "27da2957", - "metadata": { - "id": "27da2957" - }, - "source": [ - "### Setup: Clone repo and run setup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "27e68f91", - "metadata": { - "id": "27e68f91" - }, - "outputs": [], - "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local" - }, - { - "cell_type": "markdown", - "id": "3e1771a9", - "metadata": { - "id": "3e1771a9" - }, - "source": [ - "### Start Unsloth Studio" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "277e431e", - "metadata": { - "id": "277e431e" - }, - "outputs": [], - "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)" - }, - { - "cell_type": "markdown", - "id": "f2b0c6a1", - "metadata": { - "id": "f2b0c6a1" - }, - "source": [ - "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", - "\n", - "Some other resources:\n", - "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", - "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", - "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", - "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", - "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", - "\n", - "
\n", - " \n", - " \n", - " \n", - "\n", - " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", - "\n", - " This notebook is licensed AGPL-3.0\n", - "
" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "gpuType": "T4", - "provenance": [], - "include_colab_link": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat": 4, + "nbformat_minor": 5 } \ No newline at end of file diff --git a/studio/backend/assets/configs/full_finetune.yaml b/studio/backend/assets/configs/full_finetune.yaml index e398515f61..98c45dd851 100644 --- a/studio/backend/assets/configs/full_finetune.yaml +++ b/studio/backend/assets/configs/full_finetune.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: false use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/lora_text.yaml b/studio/backend/assets/configs/lora_text.yaml index 9cb6b8c700..6c6a4d8839 100644 --- a/studio/backend/assets/configs/lora_text.yaml +++ b/studio/backend/assets/configs/lora_text.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: false use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml index 841e8ba166..e569031a31 100644 --- a/studio/backend/assets/configs/model_defaults/default.yaml +++ b/studio/backend/assets/configs/model_defaults/default.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml index f7b49c75b7..7ac1c83e04 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml index be7da0f624..4cab9e9f96 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml @@ -30,6 +30,7 @@ lora: - "query" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml index d9e49bc0d5..c1f1c2a344 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml @@ -30,6 +30,7 @@ lora: - "value" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml index c3422d399f..7828feae81 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml index 529a56a527..5a4028f15b 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml @@ -29,6 +29,7 @@ lora: - "Wqkv" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml index 734115ec41..7645d11c98 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml index 1032449e8c..b746235f1f 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml index c8e5f35841..4964fea276 100644 --- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml index 251409c29d..e5f3344356 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml index 89b1d7f938..71c61f383a 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml index e3292b5972..3fe29cd800 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml index 98fe497912..cd4e3e0c4d 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml index bda5471643..97aa10e861 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml index 18392568bd..a1b1640fa2 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml index 434ac41b46..dbf60f04d4 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml index 5f0a7b26ce..54c7dd6cd4 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml index dd5ae51ab0..119440a585 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml index e53e163a04..d08e5e9547 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml index ebe344e382..a266d7a39b 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml index fb89a07133..970cac3259 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml index 4a089992ac..5bba4ccdc0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml index ae7524b7c6..ac5c6eca22 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml index 10c1abd8a5..68c2d35644 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml index fb5c1d9dea..175f9c0f17 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml index 189e5dc6b2..4f3834e7c0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml index aa51440b6a..d6d97f7e44 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml index e2d67bcb0b..4f1f54a4e6 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml index aa436117a1..127700b53b 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml index 3f2cb84a94..2412b3accf 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml @@ -37,6 +37,7 @@ lora: - "shared_mlp.output_linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml index ab756fe764..81b59c4323 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml @@ -37,6 +37,7 @@ lora: - "shared_mlp.output_linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml index 1a7a91e56f..6110d84a6c 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml index 7c7bb8dc3e..3c7fc7f238 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml index f73b0c09b6..2b0977e435 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml index ffefb29e24..1742c04a06 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml index cd986a6da1..f33726b0dd 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml index 55dd3144c6..79b30bd758 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml index 8c9cb07fb9..4ee9a5a8ed 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml index 32441c5674..da20663688 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml index 6bba9c9633..30e4440afb 100644 --- a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml +++ b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml @@ -30,6 +30,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml index f9833ce705..9bb0a93e63 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml index 0ba857cd40..ded3607a14 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml index 3476f2dd6d..2ac72f1c88 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml index eda04d21f9..a087ced1f3 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml index bcd0d20c8c..c9811f4f06 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml index 34a033e32f..e3659d9fb0 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml index 98105eaf38..ee17efc54d 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 72b5b018e1..ef836b9b55 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml @@ -33,6 +33,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index d20751b0c7..c80fad35a8 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml @@ -38,6 +38,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 8a80282a2a..034b5bd131 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml @@ -37,6 +37,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml index a973c2d4e4..d1a226be79 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml @@ -35,6 +35,7 @@ lora: - "out_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml index b0feafbd6e..1b8df5ced9 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml index 2c44c91eab..cecab7f083 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml @@ -37,6 +37,7 @@ lora: - "out_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml index e1fbc08e4d..730be338cf 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml index 2abdfd8ac3..a70ac0bd49 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 5a3c4abb48..90ead037f6 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml @@ -38,6 +38,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml index a6ce27620f..a97c557c31 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index 050774a8cd..6855ed6a35 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -33,6 +33,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml index c574714d78..1933fed2ba 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml index e803c842b3..fda4e64158 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml index 4de3d9437d..c3910e3e5b 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml index bb75b3ce52..765ffee938 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml @@ -36,6 +36,7 @@ lora: - "gate_up_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml index c305d328c2..39b30e9cee 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml index 6cee3d0949..f97e525798 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml index 20ba81df2c..e19b94ede2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml index 9930786c24..982f54b32f 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml index 775c7ce08f..5242128004 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml index 856db0c1b3..3559b636c6 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml index 5900392547..3bc6d69afc 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml index bd54b1d015..604b86dacd 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml index 9feb6dcaae..daed4ebccb 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml index a40eace253..05eef89b88 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml index c130771c32..b4580e6d71 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml index 2fb3a95c30..2eceb7d0de 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml @@ -36,6 +36,7 @@ lora: - "gate_up_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml index 152f4ae06a..032091880c 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml index 94fe000708..e0e7f4ee3d 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml index 3c325485d2..bb463849ed 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml index 5b47c3bdd2..23e2b89dd0 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/vision_lora.yaml b/studio/backend/assets/configs/vision_lora.yaml index 063a970316..a06f971523 100644 --- a/studio/backend/assets/configs/vision_lora.yaml +++ b/studio/backend/assets/configs/vision_lora.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: true use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 1762469bcf..051d80abfe 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Colab helpers for Unsloth Studio. Uses Colab's built-in proxy. -""" +"""Colab helpers for Unsloth Studio. Uses Colab's built-in proxy.""" from pathlib import Path import sys @@ -22,11 +20,9 @@ logger = get_logger(__name__) def get_colab_url(port: int = 8888) -> str: - """ - Get the Colab proxy URL for a port. + """Get the Colab proxy URL for a port. - Retries up to 3 times, validating the result is a real HTTPS Colab URL. - Falls back to http://localhost:{port} only when all attempts fail. + Retries 3x validating a real HTTPS Colab URL; falls back to localhost on failure. """ import time as _time @@ -55,28 +51,243 @@ def get_colab_url(port: int = 8888) -> str: return fallback -def show_link(port: int = 8888, *, _url: "str | None" = None): - """Display a styled clickable link to the UI. - - *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip. - """ - from IPython.display import display, HTML - - url = _url if _url is not None else get_colab_url(port) - - # Truncated display URL; try/except so an odd URL shape still renders the link. +def _short_colab_url(url: str, port: int) -> str: + """Truncated display form of a Colab proxy URL; falls back to the full URL.""" try: port_prefix = f"{port}-" idx = url.index(port_prefix) next_dash = url.index("-", idx + len(port_prefix)) - short_url = url[: next_dash + 1] + "..." + return url[: next_dash + 1] + "..." except (ValueError, IndexError): - short_url = url + return url - # Plain-text line so the URL shows even if HTML display fails. - logger.info(f"🌐 Unsloth Studio URL: {url}") - html = f""" +def _is_colab_proxy_url(url: str, port: int) -> bool: + """True when *url* looks like a real Colab kernel proxy, not a localhost fallback.""" + return bool(url and isinstance(url, str) and url.startswith("https://") and str(port) in url) + + +def _is_colab_runtime() -> bool: + """True on a hosted Colab notebook kernel. + + Reuses the backend's main Colab detector (``/content`` + Colab env / ``google.colab``) + instead of a single env var, which is not always present on hosted runtimes. + """ + try: + from main import _IS_COLAB + return bool(_IS_COLAB) + except Exception: + return False + + +def _colab_login_credentials_path() -> Path: + from auth.storage import DB_PATH + return DB_PATH.parent / ".colab_notebook_login" + + +def _store_colab_login_credentials(username: str, password: str) -> None: + """Persist Colab admin credentials for notebook re-runs after interrupt.""" + path = _colab_login_credentials_path() + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(f"{username}\n{password}\n") + try: + import os + os.chmod(path, 0o600) + except OSError: + pass + except OSError as e: + logger.info(f"Could not persist Colab login credentials ({e}).") + + +def _load_colab_login_credentials() -> "tuple[str, str] | None": + """Return stored Colab admin credentials from a previous ``start()`` run, if any.""" + path = _colab_login_credentials_path() + try: + if not path.is_file(): + return None + lines = path.read_text().splitlines() + if len(lines) >= 2 and lines[0] and lines[1]: + return lines[0], lines[1] + except OSError as e: + logger.info(f"Could not load Colab login credentials ({e}).") + return None + + +def _clear_colab_login_credentials() -> None: + """Drop the cached Colab credentials once they no longer authenticate.""" + path = _colab_login_credentials_path() + try: + path.unlink(missing_ok = True) + except OSError as e: + logger.info(f"Could not clear Colab login credentials ({e}).") + + +def _colab_credentials_still_valid(username: str, password: str) -> bool: + """True when *password* still matches the stored admin hash. + + Guards against redisplaying a cached first-run password after the user has + changed the admin password through the app, which would print credentials + that no longer authenticate to the current Cloudflare tunnel. + """ + try: + from auth.storage import get_user_and_secret + from auth.hashing import verify_password + except Exception as e: + logger.info(f"Could not load auth to validate cached Colab credentials ({e}).") + return False + try: + row = get_user_and_secret(username) + if not row: + return False + salt, pwd_hash = row[0], row[1] + return bool(verify_password(password, salt, pwd_hash)) + except Exception as e: + logger.info(f"Could not validate cached Colab credentials ({e}).") + return False + + +def _colab_wants_cloudflare(cloudflare: "bool | None") -> bool: + """Resolve whether to open a Cloudflare tunnel. + + ``None`` auto-enables on real Colab (the in-cell proxy embed is often blank); + pass ``False`` to opt out. + """ + if cloudflare is not None: + return cloudflare + return _is_colab_runtime() + + +def _finalize_colab_admin_password() -> "tuple[str, str] | None": + """Clear the bootstrap-password gate on Colab so Cloudflare tunnels can start. + + Returns ``(username, password)`` for display in the notebook. On first run the + random admin password is finalized; on later runs (e.g. after interrupt) the + stored credentials are re-displayed so the Cloudflare link stays usable. + Anyone who can read this cell already controls the runtime. + """ + if not _is_colab_runtime(): + return None + try: + from auth.storage import ( + DEFAULT_ADMIN_USERNAME, + ensure_default_admin, + generate_bootstrap_password, + get_bootstrap_password, + requires_password_change, + update_password, + ) + except Exception as e: + logger.warning( + f"Could not load auth for Colab setup ({e}); Cloudflare link may be blocked." + ) + return None + + try: + ensure_default_admin() + username = DEFAULT_ADMIN_USERNAME + if not requires_password_change(username): + creds = _load_colab_login_credentials() + if creds is not None and _colab_credentials_still_valid(username, creds[1]): + return creds + # The admin password was changed through the app after the first run, + # so the cached copy is stale; drop it instead of printing dead credentials. + _clear_colab_login_credentials() + return None + password = get_bootstrap_password() or generate_bootstrap_password() + if not update_password(username, password): + logger.warning( + "Could not finalize Colab admin password; Cloudflare link may be blocked." + ) + return None + _store_colab_login_credentials(username, password) + return username, password + except Exception as e: + logger.warning( + f"Could not finalize Colab admin password ({e}); Cloudflare link may be blocked." + ) + return None + + +def _colab_login_html(username: str, password: str) -> str: + """Notebook card with Colab admin credentials (shown once after auto-finalize).""" + return f""" +
+

+ Unsloth Studio Login (Colab) +

+

+ Log in as {username} with this password. This cell is visible only in + your notebook session. +

+

+ Password: {password} +

+
+ """ + + +def _show_colab_login_credentials(username: str, password: str) -> None: + """Display Colab admin credentials in the notebook output.""" + from IPython.display import HTML, display + + logger.info(f"🔐 Unsloth Studio login — user: {username}") + display(HTML(_colab_login_html(username, password))) + + +def _ready_card_html( + url: str, + port: int, + *, + has_cloudflare_link: bool = False, + cloudflare_requested: bool = False, +) -> str: + """Branded ready card for the in-notebook Studio view. + + Colab ``*.prod.colab.dev`` proxy URLs are session-scoped and 404 when opened as a + top-level tab or on another device, so never ``window.open`` them. On real Colab the + Cloudflare link is the supported entry point because in-cell proxy embeds often stay blank. + """ + short_url = _short_colab_url(url, port) + if _is_colab_runtime() or _is_colab_proxy_url(url, port): + if has_cloudflare_link: + embed_note = ( + "Open Studio with the Cloudflare link above. In-cell proxy previews on " + "current Colab often stay blank, so the tunnel link is the supported path." + ) + elif cloudflare_requested: + embed_note = ( + "Could not open a Cloudflare tunnel, so Studio may be unreachable on Colab. " + "Check the logs above and re-run this cell. Pass " + '' + "cloudflare=True after fixing any tunnel errors." + ) + else: + embed_note = ( + "Colab proxy links cannot be opened in a new tab (they 404 outside this " + 'notebook). Re-run with start(cloudflare=True) for a working link.' + ) + return f""" +
+

+ + Unsloth Studio is Ready! +

+

+ {embed_note} +

+

+ {short_url} +

+
+ """ + + return f"""

""" - display(HTML(html)) + + +def show_link( + port: int = 8888, + *, + _url: "str | None" = None, + has_cloudflare_link: bool = False, + cloudflare_requested: bool = False, +): + """Display a styled ready card for the UI. + + Colab proxy URLs are informational only (no new-tab open; they 404 outside the cell); + non-proxy URLs keep a clickable open button. *_url* is an optional pre-fetched proxy + URL to avoid a second eval_js round-trip. + """ + from IPython.display import display, HTML + + url = _url if _url is not None else get_colab_url(port) + logger.info(f"🌐 Unsloth Studio URL: {url}") + display( + HTML( + _ready_card_html( + url, + port, + has_cloudflare_link = has_cloudflare_link, + cloudflare_requested = cloudflare_requested, + ) + ) + ) + + +def _warn_colab_cloudflare_missing(*, use_cloudflare: bool, cloudflare_url: "str | None") -> None: + """Log a prominent warning when Colab expected a tunnel but none was opened.""" + if not use_cloudflare or cloudflare_url or not _is_colab_runtime(): + return + logger.warning( + "Colab Cloudflare tunnel unavailable — Studio is unlikely to be reachable in this " + "notebook. Check the logs above for tunnel or auth errors, then re-run start()." + ) def _bootstrap_password_pending() -> bool: """True while the default admin still owes a bootstrap-password change. - While pending, main.py injects that password into same-origin GETs, and a public - tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin - access. Fails safe to pending if the state cannot be read. + While pending, a public tunnel GET (no Origin) reads as same-origin and gets the + injected password, so sharing the link would leak admin access. Fails safe to pending. """ try: from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME @@ -121,9 +369,8 @@ def _bootstrap_password_pending() -> bool: def start_cloudflare_tunnel(port: int) -> "str | None": """Open a shareable Cloudflare quick tunnel to localhost:*port*, or None. - run_server suppresses the tunnel on Colab by design, so we start it directly. - Refused while the bootstrap password is pending; any failure collapses to None - and the Colab proxy still works. + run_server suppresses the tunnel on Colab, so we start it directly. Refused while the + bootstrap password is pending; any failure collapses to None (Colab proxy still works). """ if _bootstrap_password_pending(): logger.warning( @@ -152,9 +399,9 @@ def start_cloudflare_tunnel(port: int) -> "str | None": def _publish_cloudflare_url(cloudflare_url: "str | None") -> None: """Publish a directly-started tunnel URL onto app.state so /api/health advertises it. - run_server only sets this when it opens the tunnel itself, which it skips on Colab, - so we set it here. Otherwise the frontend's API examples fall back to an - unreachable server_url. Best-effort. + run_server sets this only when it opens the tunnel itself (skipped on Colab), so we + set it here; otherwise the frontend's API examples fall back to an unreachable + server_url. Best-effort. """ if not cloudflare_url: return @@ -183,8 +430,7 @@ def _stop_cloudflare_tunnel() -> None: def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: """True only if Unsloth Studio (not some other app) answers /api/health on *port*. - The service-marker check stops the reuse path reusing or tunneling a foreign - process that merely serves /api/health. + The service-marker check stops the reuse path reusing or tunneling a foreign process. """ import json, urllib.request try: @@ -194,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""" +

+ Password +

+

{password}

+

+ Log in as {username} with this password. Shown only in your + notebook session, and never included in the shared link. +

""" return f"""
@@ -213,40 +480,55 @@ def _shareable_link_html(cloudflare_url: str) -> str: Open Unsloth Studio

- 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.

- 🔗 {cloudflare_url} -

+ 🔗 {cloudflare_url} +

{login_block}
""" -def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None): - """Render the Unsloth header + iframe for *port*, with a shareable-link card above - when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe.""" - url = get_colab_url(port) - logger.info(f"🌐 Unsloth Studio URL: {url}") - if cloudflare_url: - logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") +# Height for serve_kernel_port_as_iframe (~82vh on a 1080p screen, clamped). +_COLAB_IFRAME_HEIGHT = 900 + +def _embed_kernel_port_iframe(port: int) -> bool: + """Embed Studio via Colab's native kernel-port iframe helper. + + Only trusted on a real Colab runtime: colabtools can import ``google.colab`` and + queue browser-side JS without appending an iframe, so callers outside Colab must use + the HTML iframe path instead. + """ + if not _is_colab_runtime(): + return False + try: + from google.colab import output as colab_output + except ImportError: + return False + try: + colab_output.serve_kernel_port_as_iframe( + port, + height = _COLAB_IFRAME_HEIGHT, + width = "100%", + ) + return True + except Exception as e: + logger.info(f"serve_kernel_port_as_iframe failed ({e}); trying HTML iframe.") + return False + + +def _embed_html_iframe(url: str, port: int) -> bool: + """Fallback embed: raw HTML iframe when the Colab helper is unavailable.""" try: from IPython.display import HTML, display + except ImportError: + return False - iframe_id = f"unsloth-studio-{port}" - - # Truncated header URL — best-effort, falls back to full URL. - try: - port_prefix = f"{port}-" - idx = url.index(port_prefix) - next_dash = url.index("-", idx + len(port_prefix)) - short_url = url[: next_dash + 1] + "..." - except (ValueError, IndexError): - short_url = url - - if cloudflare_url: - display(HTML(_shareable_link_html(cloudflare_url))) - + short_url = _short_colab_url(url, port) + iframe_id = f"unsloth-studio-{port}" + try: display( HTML(f"""
""") ) - except Exception: - # Fallback: Colab's built-in helper. + return True + except Exception as e: + logger.info(f"HTML iframe embed failed ({e}).") + return False + + +def _show_and_embed( + port: int, + *, + cloudflare_url: "str | None" = None, + colab_login: "tuple[str, str] | None" = None, + cloudflare_requested: bool = False, +): + """Render the Unsloth ready card + iframe for *port*. + + Prefer Colab's ``serve_kernel_port_as_iframe`` on real Colab; raw HTML iframe is the + fallback. Cloudflare cards stay clickable. + """ + url = get_colab_url(port) + logger.info(f"🌐 Unsloth Studio URL: {url}") + if cloudflare_url: + logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") + + _warn_colab_cloudflare_missing( + use_cloudflare = cloudflare_requested, + 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 google.colab import output as colab_output - colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") - except ImportError: - pass + from IPython.display import HTML, display + + 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 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}).") + + # 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: + return + + # Real Colab: kernel helper needs only the port (works when eval_js failed). + if _is_colab_runtime(): + if _embed_kernel_port_iframe(port): + return + _embed_html_iframe(url, port) -def start(port: int = 8888, *, cloudflare: bool = False): +def start(port: int = 8888, *, cloudflare: "bool | None" = None): """Start Unsloth Studio in Colab and display the URL. Args: port: Port to bind/serve on. - cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any - device (default OFF). It exposes Unsloth's login page beyond Colab, so it - stays an explicit opt-in; the default shows only the in-tab proxy iframe. + cloudflare: Shareable Cloudflare HTTPS link. ``None`` (default) auto-enables on + real Colab because the in-cell proxy embed is often blank; pass ``False`` to + skip the tunnel or ``True`` to force it on other runtimes. Usage: - start() # Colab-proxy iframe only (default) - start(cloudflare=True) # also open a shareable Cloudflare link + start() # Cloudflare link on Colab (auto); proxy iframe elsewhere + start(cloudflare=False) # Colab proxy iframe only (often blank on current Colab) + start(cloudflare=True) # force Cloudflare link on any runtime """ import time logger.info("🦥 Starting Unsloth Studio...") + use_cloudflare = _colab_wants_cloudflare(cloudflare) - # Fast path: Unsloth already running (cell re-run). Re-launching would collide on - # the port, so just re-show the link and iframe. + # Fast path: already running (cell re-run); re-show link/iframe instead of rebinding the port. if _is_studio_healthy(port): logger.info(f" Unsloth is already running on port {port} — reusing existing server.") # try/finally: tear the tunnel down even if interrupted mid-start/render. try: - cf_url = start_cloudflare_tunnel(port) if cloudflare else None + colab_login = _finalize_colab_admin_password() if use_cloudflare else None + cf_url = start_cloudflare_tunnel(port) if use_cloudflare else None _publish_cloudflare_url(cf_url) - _show_and_embed(port, cloudflare_url = cf_url) + _show_and_embed( + port, + cloudflare_url = cf_url, + colab_login = colab_login, + cloudflare_requested = use_cloudflare, + ) for _ in range(10000): time.sleep(300) print("=", end = "", flush = True) @@ -313,7 +664,6 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Loading backend...") from run import run_server - # Auto-detect frontend path repo_root = Path(__file__).parent.parent frontend_path = repo_root / "frontend" / "dist" @@ -323,8 +673,7 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Starting server...") try: - # cloudflare=False: this helper owns the tunnel (Colab's own - # start(cloudflare=...) drives it), so pin it off explicitly. + # cloudflare=False: this helper owns the tunnel (via start(cloudflare=...)), so pin it off. app = run_server( host = "0.0.0.0", port = port, @@ -339,14 +688,12 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.error(f"❌ Unsloth Studio failed to start: {exc}") return - # run_server auto-increments the port if in use; read back the bound port so the - # proxy URL and iframe point at the right place. + # run_server may auto-increment the port; read back the bound port for the proxy URL/iframe. actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port logger.info(f" Server started on port {actual_port}!") - # Poll health endpoint before showing the link — avoids the race where ready_event - # fires but the process hasn't finished binding. + # Poll health before showing the link: avoids the race where ready_event fires pre-bind. import urllib.request server_ready = False @@ -365,12 +712,17 @@ def start(port: int = 8888, *, cloudflare: bool = False): ) return - # Open the tunnel now the server is healthy, publish its URL for /api/health, and - # tear it down on interrupt (try/finally) rather than orphan the process. + # Server healthy: finalize Colab auth, open the tunnel, publish URL, tear down on interrupt. try: - cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None + colab_login = _finalize_colab_admin_password() if use_cloudflare else None + cf_url = start_cloudflare_tunnel(actual_port) if use_cloudflare else None _publish_cloudflare_url(cf_url) - _show_and_embed(actual_port, cloudflare_url = cf_url) + _show_and_embed( + actual_port, + cloudflare_url = cf_url, + colab_login = colab_login, + cloudflare_requested = use_cloudflare, + ) # Keep kernel alive so the daemon server thread runs. for _ in range(10000): diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index c8be50b08b..1b24c46e65 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -1048,6 +1048,21 @@ class ExportBackend: "Use the safetensors adapter instead.", None, ) + # llama.cpp's convert_lora_to_gguf.py has no concept of DoRA's + # lora_magnitude_vector tensors: it only reads the standard + # lora_A/lora_B delta, so exporting a DoRA adapter would silently + # drop the magnitude rescaling and produce a GGUF LoRA file that + # loads fine but no longer matches the trained model. + _peft_config = getattr(self.current_model, "peft_config", {}).get("default") + if getattr(_peft_config, "use_dora", False): + return ( + False, + "GGUF LoRA export is not supported for DoRA adapters: the GGUF LoRA " + "format has no way to represent DoRA's magnitude vectors, so the " + "exported file would silently lose the DoRA behavior. Use the " + "safetensors adapter instead, or merge to a full GGUF model.", + None, + ) outtype = str(gguf_outtype).lower() if outtype not in _GGUF_LORA_OUTTYPES: return ( diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 8d262bbb0f..2f46470091 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -8,6 +8,7 @@ from unsloth.chat_templates import get_chat_template from transformers import TextIteratorStreamer, TextStreamer from peft import PeftModel, PeftModelForCausalLM +import contextlib import json import sys import torch @@ -1942,8 +1943,30 @@ class InferenceBackend: + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" ) + with torch.inference_mode(): - with torch.amp.autocast("cuda", dtype = model.dtype): + # Derive the autocast device from the loaded model, not from the + # global backend: a CPU-fallback DAC on an XPU/CUDA host must not + # open a GPU autocast context around CPU tensors. + device_type = ( + model.device.type + if hasattr(model.device, "type") + else str(model.device).split(":", 1)[0] + ) + # Clamp to autocast-supported backends so exotic devices + # (e.g. "meta" during accelerate offloaded loading) do not raise. + # MPS is autocast-supported since torch 2.3, keep it in the set. + if device_type not in ("cuda", "xpu", "mps", "cpu"): + device_type = "cpu" + # CPU and XPU autocast only accept bfloat16/float16. For a + # float32 model, skip autocast entirely to avoid raising or + # producing a warning on every generate call. + autocast_dtype_supported = model.dtype in (torch.bfloat16, torch.float16) + if device_type in ("cpu", "xpu") and not autocast_dtype_supported: + autocast_ctx = contextlib.nullcontext() + else: + autocast_ctx = torch.amp.autocast(device_type, dtype = model.dtype) + with autocast_ctx: inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fc58442a79..e54b5269c1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -33,6 +33,7 @@ from typing import ( List, Literal, Mapping, + MutableMapping, Optional, Union, ) @@ -307,6 +308,26 @@ def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]": _DEFAULT_MAX_TOKENS_FLOOR = 32768 _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min + +def _finalize_reasoning_only_cumulative( + cumulative: str, reasoning_text: str, finish_reason: Optional[str], promote_reasoning_only: bool +) -> str: + """Close a live thinking block and promote it only after a clean stop. + + Local inference streams cumulative snapshots. Replacing ``...`` with + bare reasoning at EOF makes the final snapshot shorter, so suffix-based + route consumers drop the intended fallback. Keep the snapshot append-only. + A length-truncated thought is not a final answer, so close it without + promotion and let the client surface the ``length`` terminal state. Raw + consumers that do not split reasoning from visible content can disable the + fallback to avoid returning the same reasoning twice. + """ + visible_fallback = ( + reasoning_text if promote_reasoning_only and finish_reason != "length" else "" + ) + return cumulative + "" + visible_fallback + + # Only large streamed tool payloads get an early provisional card; render_html # is exempt because it needs immediate artifact feedback. _PROVISIONAL_ARGS_MIN_CHARS = 256 @@ -2023,6 +2044,10 @@ class LlamaCppBackend: self._tensor_split: Optional[List[float]] = None # User-picked physical GPU indices (None = automatic selection). self._gpu_ids: Optional[List[int]] = None + # RAW requested GPU pin, before the fit narrowed it. self._gpu_ids records the + # EFFECTIVE (fit-narrowed) pin for /status; dedupe compares this raw value so a + # [0, 1] narrowed to [0] and re-sent as [0, 1] still matches (#7239). + self._requested_gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -2494,6 +2519,46 @@ class LlamaCppBackend: """User-picked physical GPU indices, or None for automatic selection.""" return self._gpu_ids + @property + def requested_gpu_ids(self) -> Optional[List[int]]: + """RAW requested GPU pin (before the fit narrowed it), or None for auto. + gpu_ids echoes the EFFECTIVE pin for /status.""" + return self._requested_gpu_ids + + def matches_gpu_ids(self, gpu_ids: Optional[List[int]]) -> bool: + """Whether a requested pin is already satisfied by the active runner. + + A regular GGUF load may narrow the requested placement pool to the + smallest fitting subset. Accept both the original request and the + effective status-echoed subset so either can round-trip without a + needless reload. Diffusion drives one device and keeps its existing + lowest-device normalization. + """ + if self._is_diffusion: + requested = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + return requested == (self._gpu_ids or None) + + requested = sorted(int(x) for x in gpu_ids) if gpu_ids else None + raw = self._requested_gpu_ids or None + effective = self._gpu_ids or None + return requested == raw or requested == effective + + def _record_matching_gpu_request(self, gpu_ids: Optional[List[int]]) -> None: + """Adopt the caller's explicit pool after a full already-loaded match. + + Matching an effective subset avoids a reload, but the incoming request + is still the user's latest placement intent. Record it so status and a + later reload do not restore GPUs the user just removed. + """ + if self._is_diffusion: + self._requested_gpu_ids = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + else: + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + if self._last_load_kwargs is not None: + self._last_load_kwargs["gpu_ids"] = ( + list(self._requested_gpu_ids) if self._requested_gpu_ids else None + ) + @property def n_layers(self) -> Optional[int]: """Model layer count (GGUF block_count), or None if unknown.""" @@ -2737,6 +2802,7 @@ class LlamaCppBackend: "found": False, "mtp_token": None, "supports_mtp": False, + "mtp_probe_inconclusive": True, "ngram_mod_flavor": None, "supports_ngram_mod": False, "spec_draft_n_max_flag": None, @@ -2769,6 +2835,9 @@ class LlamaCppBackend: supports_no_cache_prompt = False supports_metrics = False supports_slot_save = False + saw_spec_type = False + probe_ok = False + help_text = "" try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( @@ -2780,6 +2849,7 @@ class LlamaCppBackend: check = False, env = probe_env, ) + probe_ok = result.returncode == 0 help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented # continuation), so the "argument has been removed" description @@ -2824,17 +2894,19 @@ class LlamaCppBackend: return False return "argument has been removed" not in desc - # MTP token from the --spec-type line. - spec_line = "" - for line in help_text.splitlines(): - if "--spec-type" in line: - spec_line = line - break - # PR #22673 used draft-mtp; later renamed to mtp. - if "draft-mtp" in spec_line: - mtp_token = "draft-mtp" - elif re.search(r"[|,\[]mtp[|,\]]", spec_line): - mtp_token = "mtp" + # MTP token from the full --spec-type help block (decl + indented + # continuation). First-line-only probing missed builds putting the + # enum on the next line (#7302). Prefer draft-mtp (PR #22673) over mtp. + spec_help = blocks.get("--spec-type") or "" + if not spec_help: + # Fallback: join --spec-type lines, avoiding incidental "mtp" in --help. + spec_help = "\n".join( + line for line in help_text.splitlines() if "--spec-type" in line + ) + mtp_token = cls._mtp_token_from_spec_help(spec_help) + # Only a resolved --spec-type block confirms missing MTP; empty/crash + # leaves saw_spec_type False so supports_mtp fails open. + saw_spec_type = bool(spec_help.strip()) and "--spec-type" in spec_help # ngram-mod flag flavor. Post-rename builds advertise both new # args (real) and legacy ones (stubs); pre-rename builds only @@ -2870,11 +2942,29 @@ class LlamaCppBackend: supports_slot_save = _is_real("--slot-save-path") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") + saw_spec_type = False + probe_ok = False + help_text = "" + + help_nonempty = bool(help_text.strip()) + # Confirmed only when a successful --help lists a --spec-type block with + # mtp/draft-mtp; nonempty --help without it is a definitive pre-spec + # binary; failed/empty probes stay inconclusive (#7302). + if saw_spec_type and probe_ok: + supports_mtp = mtp_token is not None + mtp_probe_inconclusive = False + elif help_nonempty and probe_ok: + supports_mtp = False + mtp_probe_inconclusive = False + else: + supports_mtp = False + mtp_probe_inconclusive = True info = { "found": True, "mtp_token": mtp_token, - "supports_mtp": mtp_token is not None, + "supports_mtp": supports_mtp, + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": ngram_mod_flavor, "supports_ngram_mod": ngram_mod_flavor is not None, "spec_draft_n_max_flag": spec_draft_n_max_flag, @@ -2890,6 +2980,21 @@ class LlamaCppBackend: cls._capability_cache[cache_key] = info return info + @staticmethod + def _mtp_token_from_spec_help(spec_help: str) -> Optional[str]: + """Extract ``draft-mtp`` / ``mtp`` from a ``--spec-type`` help snippet. + + Prefers ``draft-mtp`` (llama.cpp PR #22673) over the later bare ``mtp`` + rename. Returns ``None`` when neither token appears as an enum value. + """ + text = spec_help or "" + if "draft-mtp" in text: + return "draft-mtp" + # Bare `mtp` enum token (`|mtp|`, `,mtp,`, ...), not a substring. + if re.search(r"(? 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; @@ -3108,7 +3213,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 @@ -3592,6 +3699,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 @@ -4581,6 +4696,14 @@ class LlamaCppBackend: LlamaCppBackend._gguf_skip_value(f, atype) return None + @classmethod + def _gguf_path_is_diffusion(cls, gguf_path: str, model_identifier: str) -> bool: + """Classify a downloaded GGUF without mutating the active backend.""" + probe = object.__new__(cls) + probe._model_identifier = model_identifier + probe._read_gguf_metadata(gguf_path) + return probe._is_diffusion + def _read_gguf_metadata(self, gguf_path: str) -> None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -5032,11 +5155,14 @@ class LlamaCppBackend: # the unload reset) so /status doesn't misreport TP and an identical # re-Apply doesn't reload against stale tensor-parallel state. self._tensor_parallel = False - # Record only the single device the runner actually uses (the lowest - # selected GPU, chosen above) -- not the whole pick. The diffusion runner - # is single-device, so echoing a multi-GPU list would misreport placement - # in /status and let a re-Apply dedup against GPUs the runner never used. + # The single-device runner records only the lowest selected GPU (chosen + # above), not the whole pick, and clears any explicit pin from a prior + # chat load; a multi-GPU list would misreport placement and mis-dedup. self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None + # The frontend prefers requested_gpu_ids when hydrating the picker. + # Diffusion uses only one device, so echo the collapsed effective pin, + # not unused members of the original request. + self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -5960,6 +6086,24 @@ class LlamaCppBackend: and ("unknown" in text or "unsupported" in text or "not supported" in text) ) + @staticmethod + def _mmproj_retry_failure_message(*, projector_confirmed: bool, detail: str) -> str: + """User-facing error when the text-only --mmproj strip retry also fails. + + Confirmed projector-format mismatches keep the historical wording. + Bare signal crashes (common on some ROCm/driver paths) must not be + reported as "Vision projector incompatible" — that misled #7302. + """ + if projector_confirmed: + return ( + "Vision projector incompatible with this llama.cpp " + "build, and the text-only retry also failed: " + detail + ) + return ( + "Vision model failed to start (llama-server crashed with " + "--mmproj), and the text-only retry also failed: " + detail + ) + @staticmethod def _output_has_nonprojector_diagnostic(output: str) -> bool: """True when the output already names a concrete non-projector cause (out @@ -6016,6 +6160,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 @@ -6048,8 +6207,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 ' pair (text-only retry). @@ -6161,6 +6388,8 @@ class LlamaCppBackend: gpu_layers: int = -1, n_cpu_moe: int = 0, tensor_split: Optional[List[float]] = None, + # Explicit GPU placement pool (issue #7164). None/[] = auto-select; + # the fitter may pin the smallest subset of this pool that fits. gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused @@ -6258,15 +6487,63 @@ class LlamaCppBackend: self._cancel_event.clear() - # ── Phase 1: kill old process (under lock, fast) ────────── - with self._lock: - self._kill_process() - # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() is_vulkan_backend = self._is_vulkan_backend(binary) + # ── Vulkan-ordinal preflight (BEFORE the Phase 1 kill) ──────── + # An explicit Vulkan pin the ggml probe never enumerated cannot be honored. + # Validate it ABOVE the kill so an invalid selection leaves the live model + # untouched: CUDA ids are range-checked at the route, but Vulkan ordinals are + # not, so a stale gpu_ids=[99] used to kill the server then 400, leaving + # nothing running (#7239). _get_gpu_memory needs only the binary (safe pre- + # download) and reuses the later fit's issubset logic. Guarded on a found + # Vulkan build + a pin so a deferred not-found stays deferred for diffusion. + if is_vulkan_backend and gpu_ids and binary: + _pf_wanted = {int(x) for x in gpu_ids} + _pf_probed = {g[0] for g in self._get_gpu_memory(binary)} + if not _pf_wanted.issubset(_pf_probed): + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {sorted(_pf_wanted)} not " + f"present. Available Vulkan devices: {sorted(_pf_probed)}." + ) + + # A remote uncached GGUF may only reveal that it needs the + # single-device diffusion runner after download. On Vulkan, an + # explicit gpu_ids request cannot be mapped from ggml ordinals to + # that runner's CUDA physical index. Download and classify the main + # file before killing the healthy server so this late rejection is + # non-destructive. The Phase 2 call below reuses this cached path. + _preflight_model_path = None + if is_vulkan_backend and gpu_ids and hf_repo: + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo + with _hf_offline_if_dns_dead(): + _preflight_model_path = self._download_gguf( + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + ) + if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier): + raise ValueError( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ) + + # ── Phase 1: kill old process (under lock, fast) ────────── + with self._lock: + self._kill_process() + # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected # sibling); for -hf loads it's None here and resolved just below. @@ -6288,7 +6565,7 @@ class LlamaCppBackend: ) hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): - model_path = self._download_gguf( + model_path = _preflight_model_path or self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, hf_token = hf_token, @@ -6338,6 +6615,18 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: + # The diffusion runner pins its child by CUDA visibility mask, so a + # ggml Vulkan ordinal cannot be honored (wrong GPU / CPU fallback). + # Route and remote-download preflights reject before teardown; keep + # this as a final defense if classification ever disagrees. + if is_vulkan_backend and gpu_ids: + raise ValueError( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ) # Not a tensor/layer GGUF: clear any preserved-fallback flag from a # prior load (this path skips the command builder that clears it). self._layer_preserves_tensor_intent = False @@ -6558,6 +6847,12 @@ class LlamaCppBackend: # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). _layer_min_gpus = 1 + # An explicit Vulkan ordinal absent from the ggml probe cannot be + # honored; flag it in the fit and reject after the try (raising inside + # would be swallowed into the --fit-on fallback). Bound before the try. + _vulkan_explicit_unmatched = False + _vulkan_requested_ids: list[int] = [] + _vulkan_available_ordinals: list[int] = [] try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -6570,6 +6865,28 @@ class LlamaCppBackend: # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] + # Restrict the fit (and thus the layer plan + pin env) to the + # selected GPUs; fail-open if none match so a stale UI choice + # can't strand the load on CPU (issue #7164). + if gpu_ids: + # A Vulkan build indexes by ggml ordinal. An explicit ordinal + # absent from the probe can't be pinned, so reject after the try + # rather than fail-open onto a device the user didn't pick. + _wanted_ids = {int(x) for x in gpu_ids} + # Reject if ANY requested ordinal is absent, not only when none + # match: [0, 99] against {0, 1} silently drops 99. Comparing the + # full requested set (before filter narrows) still lets the fitter + # pick a valid subset later -- that is narrowing, not absence. + _probed_ordinals = {g[0] for g in gpus} + if is_vulkan_backend and not _wanted_ids.issubset(_probed_ordinals): + _vulkan_explicit_unmatched = True + _vulkan_requested_ids = sorted(_wanted_ids) + _vulkan_available_ordinals = sorted(_probed_ordinals) + # Restrict the probed pool to the selection; fail-open (keep the + # full pool) if none match so a stale UI choice can't strand the + # load on CPU (issue #7164). + _sel_gpus = [g for g in gpus if g[0] in _wanted_ids] + gpus = _sel_gpus if _sel_gpus else gpus total_by_idx = {idx: total for idx, _f, total in _gpu_mem} # GPU picker: restrict every mode to the chosen devices, so # auto selection only considers them and manual mask to @@ -7396,6 +7713,17 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # An unenumerated explicit Vulkan ordinal can't be pinned; fail loudly + # instead of fitting onto an unselected device. Clear the raw selection + # the early state-publish recorded so it never leaks into gpu_ids (#7239). + if _vulkan_explicit_unmatched: + self._gpu_ids = None + self._requested_gpu_ids = None + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {_vulkan_requested_ids} not " + f"present. Available Vulkan devices: {_vulkan_available_ordinals}." + ) + # GPU picker: when no narrower subset was chosen (manual, or # a failed/file-size selection), pin the whole picked set so the # model can't spill onto an unpicked GPU. @@ -7759,11 +8087,45 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) - # Vulkan pins via --device (a cmd arg, unlike the env-based - # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's - # last-wins parsing lets a user --device override Unsloth's pick. - if is_vulkan_backend and gpu_indices is not None: - cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # Vulkan pins via --device (a cmd arg), before user extras so a user + # --device wins. Fall back to raw ids when the fit did not narrow. + _vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None) + + # Record the pin actually applied (fit-narrowed gpu_indices, else the raw + # request) for the keep-warm loop, dedupe, and /status, so an explicit + # [0, 1] narrowed to [0] records [0] and /status never echoes an ordinal + # the child never saw. Auto selection (no gpu_ids) stays None (#7239). + if is_vulkan_backend: + # Only record an EXPLICIT Vulkan pin: an auto pick still narrows + + # pins below, but recording it would misreport an explicit pin and + # make dedupe miss the loaded server; mirrors the CUDA/ROCm branch. + self._gpu_ids = ( + sorted(int(x) for x in _vulkan_pin_ids) + if (gpu_ids and _vulkan_pin_ids) + else None + ) + elif gpu_ids: + # Physical pin: the fit-selected subset when the fit ran, else the raw + # user selection so an explicit choice is honoured even when the fit + # could not size the model. + _effective_pin_ids = ( + [int(x) for x in gpu_indices] + if gpu_indices is not None + else [int(x) for x in gpu_ids] + ) + self._gpu_ids = ( + sorted(int(x) for x in _effective_pin_ids) if _effective_pin_ids else None + ) + else: + self._gpu_ids = None + + # Also record the RAW requested pin (before the fit narrowed it). Load + # dedupe compares this so a [0, 1] narrowed to [0] and re-sent as [0, 1] + # still matches, while /status keeps echoing the effective pin (#7239). + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + + if is_vulkan_backend and _vulkan_pin_ids is not None: + cmd += LlamaCppBackend._vulkan_pin_args(_vulkan_pin_ids) # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Unsloth's auto-set flags. Already @@ -7832,10 +8194,10 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so - # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device - # (above), not here. + # Pin to selected GPU(s) (issue #7164; resolved above into gpu_indices). + # On ROCm, narrowing only CUDA_VISIBLE_DEVICES leaves the AMD child + # seeing the full set, so set HIP_VISIBLE_DEVICES too. Vulkan is pinned + # via --device (above), not here. # A deliberate zero-offload load with no GPU companions runs # entirely on CPU, yet a visible CUDA device still costs the child # ~0.5 GB (context + compute scratch) that the CPU-only @@ -8071,6 +8433,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") @@ -8116,6 +8485,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") @@ -8198,23 +8574,29 @@ class LlamaCppBackend: self._kill_process() # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). + _projector_msg = self._is_projector_incompatibility(out) + _signal_mmproj_guess = self._is_signal_crash( + _crash_rc + ) and not self._output_has_nonprojector_diagnostic(out) if ( launched_with_mmproj and not self._cancel_event.is_set() - and ( - self._is_projector_incompatibility(out) - or ( - self._is_signal_crash(_crash_rc) - and not self._output_has_nonprojector_diagnostic(out) - ) - ) + and (_projector_msg or _signal_mmproj_guess) ): - logger.warning( - "llama-server could not load this model's vision " - "projector (--mmproj). The installed llama.cpp build is " - "likely too old for it. Loading text-only for this " - "session; run 'unsloth studio update' to enable vision." - ) + if _projector_msg: + logger.warning( + "llama-server could not load this model's vision " + "projector (--mmproj). The installed llama.cpp build is " + "likely too old for it. Loading text-only for this " + "session; run 'unsloth studio update' to enable vision." + ) + else: + logger.warning( + "llama-server crashed while loading this model's vision " + "projector (--mmproj). Retrying text-only for this " + "session; if this persists, run 'unsloth studio update' " + "or check GPU/driver logs." + ) cmd = self._strip_mmproj_args(_last_spawn_cmd) # This retry bypasses _spawn_and_wait, so refresh the # launched-argv snapshot itself -- the zero-offload @@ -8242,14 +8624,16 @@ class LlamaCppBackend: "(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first " "GPU) before launching Unsloth Studio." ) + _retry_detail = self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + self._model_identifier, + _retry_rc, + ) raise RuntimeError( - "Vision projector incompatible with this llama.cpp " - "build, and the text-only retry also failed: " - + self._classify_llama_start_failure( - "\n".join(self._stdout_lines[-50:]), - gguf_path, - self._model_identifier, - _retry_rc, + self._mmproj_retry_failure_message( + projector_confirmed = _projector_msg, + detail = _retry_detail, ) ) else: @@ -8479,18 +8863,29 @@ class LlamaCppBackend: caps = self.probe_server_capabilities(binary) mtp_token = caps.get("mtp_token") if caps else None if not mtp_token: - logger.warning( - "Requested MTP speculative decoding but " - "llama-server lacks --spec-type mtp/draft-mtp; " - "run `unsloth studio update`. Loading without " - "speculative decoding." - ) + inconclusive = bool(caps.get("mtp_probe_inconclusive")) if caps else True + if inconclusive: + logger.info( + "Requested MTP speculative decoding but llama-server MTP " + "capability probe was inconclusive; loading without " + "speculative decoding." + ) + else: + logger.warning( + "Requested MTP speculative decoding but " + "llama-server lacks --spec-type mtp/draft-mtp; " + "run `unsloth studio update`. Loading without " + "speculative decoding." + ) # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins # over env) so the child matches the binary-capability gate and # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. flags.append("--spec-default") self._speculative_type = "default" - self._spec_fallback_reason = "binary_no_mtp" + if inconclusive: + self._spec_fallback_reason = None + else: + self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max" @@ -8756,16 +9151,10 @@ class LlamaCppBackend: ) ): return False - # A changed GPU pick must reload (compare order-insensitively; None/[] - # both mean automatic). The diffusion runner collapses a multi-GPU pick - # to its single lowest device, so self._gpu_ids holds just that device; - # normalize the request the same way, or a multi-GPU pick that resolves - # to the same device needlessly reloads. - if self._is_diffusion: - requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None - else: - requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None - if (self._gpu_ids or None) != requested_gpu_pick: + # A changed GPU pick must reload. Regular GGUF accepts either the raw + # requested placement pool or the effective status-echoed subset; + # diffusion compares its normalized single-device pick. + if not self.matches_gpu_ids(gpu_ids): return False # Compare on the canonical requested mode. With --spec-type in @@ -8823,6 +9212,7 @@ class LlamaCppBackend: current = list(self._extra_args) if self._extra_args is not None else [] if list(extra_args) != current: return False + self._record_matching_gpu_request(gpu_ids) return True def _classify_gpu_offload( @@ -8954,12 +9344,15 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + # GPU-pin state describes the active runner only; clear it so an explicit + # pin never leaks into the next (or diffusion) runner. + self._gpu_ids = None + self._requested_gpu_ids = None self._tensor_parallel = False self._gpu_memory_mode = "auto" self._gpu_layers = -1 self._n_cpu_moe = 0 self._tensor_split = None - self._gpu_ids = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None @@ -10291,6 +10684,7 @@ class LlamaCppBackend: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, seed: Optional[int] = None, + promote_reasoning_only: bool = True, _allow_respawn_retry: bool = True, ) -> Generator[Union[str, dict], None, None]: """ @@ -10373,7 +10767,12 @@ class LlamaCppBackend: # model put its whole reply in reasoning # (e.g. Qwen3 always-think). Show it as # the main response, not a thinking block. - cumulative = reasoning_text + cumulative = _finalize_reasoning_only_cumulative( + cumulative, + reasoning_text, + _metadata_finish_reason, + promote_reasoning_only, + ) yield cumulative _stream_done = True break # exit inner while @@ -10470,6 +10869,7 @@ class LlamaCppBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, seed = seed, + promote_reasoning_only = promote_reasoning_only, _allow_respawn_retry = False, ) return @@ -10511,6 +10911,7 @@ class LlamaCppBackend: confirm_tool_calls: bool = False, bypass_permissions: bool = False, permission_mode: Optional[str] = None, + promote_reasoning_only: bool = True, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -10853,7 +11254,12 @@ class LlamaCppBackend: ), } else: - cumulative_display = reasoning_accum + cumulative_display = _finalize_reasoning_only_cumulative( + cumulative_display, + reasoning_accum, + _iter_finish_reason, + promote_reasoning_only, + ) if not _suppress_visible_output: yield { "type": "content", @@ -11317,7 +11723,12 @@ class LlamaCppBackend: if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) - cumulative_display = reasoning_accum + cumulative_display = _finalize_reasoning_only_cumulative( + cumulative_display, + reasoning_accum, + _iter_finish_reason, + promote_reasoning_only, + ) if not _suppress_visible_output: yield { "type": "content", @@ -11881,7 +12292,12 @@ class LlamaCppBackend: "text": _strip_tool_markup(cumulative, final = True), } else: - cumulative = reasoning_text + cumulative = _finalize_reasoning_only_cumulative( + cumulative, + reasoning_text, + _metadata_finish_reason, + promote_reasoning_only, + ) yield {"type": "content", "text": cumulative} _stream_done = True break # exit inner while diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 409132d605..616384386d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -27,7 +27,7 @@ import uuid from io import BytesIO from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union -from utils.hardware import prepare_gpu_selection +from utils.hardware import get_device, prepare_gpu_selection # Re-exported from the shared helper so GGUF, training, and inference share one # type; kept importable here for backwards compatibility. @@ -1012,6 +1012,8 @@ class InferenceOrchestrator: ) sub_config["resolved_gpu_ids"] = resolved_gpu_ids sub_config["gpu_selection"] = gpu_selection + # Parent-detected backend for the worker's apply_gpu_ids(). + sub_config["device_backend"] = get_device().value # Recheck the sidecar reservation BEFORE tearing the old worker down, # for REPAIRS only: an install holds this same lifecycle gate, so it diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index bc9ffe85c2..b31db3faf6 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -18,6 +18,7 @@ import queue import random import re import shlex +import shutil import ssl import subprocess import sys @@ -48,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__:" @@ -328,6 +330,7 @@ def _find_blocked_commands(command: str) -> set[str]: # Directory holding the sandbox ``sitecustomize.py`` shim (code-interpreter # path remap); placed on the sandboxed child's PYTHONPATH in _build_safe_env. _SANDBOX_SITE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sandbox_site") + # ── "Approve for me" (permission_mode="auto") safety detection ────────────── # Auto mode pauses only calls classified here as potentially unsafe. The sandbox # and hard blocks (blocklist, rlimits) still apply at run time; this gate only @@ -2491,15 +2494,124 @@ def is_potentially_unsafe_tool_call(name: str, arguments: dict) -> bool: return True +def _canon_win_path(p: str) -> str: + """Canonical form for trust comparison: realpath (expands 8.3 aliases and + resolves junctions/symlinks) + normcase/normpath.""" + return os.path.normcase(os.path.normpath(os.path.realpath(p))) + + +def _augment_native_program_roots(roots: list[str]) -> list[str]: + """Add the native Program Files sibling for any x86 root by stripping the + `` (x86)`` suffix, so a 32-bit process (whose known-folder ids map only to + the x86 root) still trusts a 64-bit Git install.""" + out = list(roots) + for root in roots: + base = root.rstrip("\\/") + if base.lower().endswith(" (x86)"): + native = base[: -len(" (x86)")] + if native and native not in out: + out.append(native) + return out + + +def _windows_program_roots() -> list[str]: + """Program Files install roots, resolved ONLY from the Windows known-folder + API (SHGetKnownFolderPath). Fails closed (returns ``[]``) if the API is + unavailable: env vars (%ProgramFiles%, even %SystemDrive%) are caller- + overrideable and could relocate the trust boundary, so we never derive a + trusted root from them. On any real Windows host shell32 is present, so + this only returns empty in a broken/non-Windows environment where the + sandbox git-PATH feature is not needed anyway (#7317). + """ + roots: list[str] = [] + try: + import ctypes + from ctypes import wintypes + + # FOLDERID_ProgramFiles, _ProgramFilesX86, _ProgramFilesX64. The X64 + # id (Win10 1703+) yields the native root even from a 32-bit process, + # where the first two both map to Program Files (x86). + folder_ids = ( + "{905e63b6-c1bf-494e-b29c-65b732d3d21a}", + "{7C5A40EF-A0FB-4BFC-874A-C0F2E0B9FA8E}", + "{6D809377-6AF0-444b-8957-A3773F02200E}", + ) + _SHGet = ctypes.windll.shell32.SHGetKnownFolderPath + _CoTaskMemFree = ctypes.windll.ole32.CoTaskMemFree + for fid in folder_ids: + guid = ctypes.create_string_buffer(16) + ctypes.windll.ole32.CLSIDFromString(wintypes.LPCWSTR(fid), ctypes.byref(guid)) + ptr = ctypes.c_wchar_p() + if _SHGet(ctypes.byref(guid), 0, None, ctypes.byref(ptr)) == 0: + if ptr.value: + roots.append(ptr.value) + _CoTaskMemFree(ptr) + except Exception: + return [] + return _augment_native_program_roots(roots) + + +def _resolve_trusted_windows_git() -> tuple[str, str]: + """Find a git launcher in a TRUSTED Program Files dir. Returns + ``(canonical_dir, ext)`` or ``("", "")``. + + ``shutil.which`` returns only the first PATH match, which may be an + untrusted user shim; scan the remaining PATH entries for a later trusted + Git so bare ``git`` still resolves (#7317). + """ + exts = [e for e in (os.environ.get("PATHEXT") or ".EXE;.CMD;.BAT;.COM").split(os.pathsep)] + candidates: list[str] = [] + primary = shutil.which("git") + if primary: + candidates.append(primary) + for entry in (os.environ.get("PATH") or "").split(os.pathsep): + entry = entry.strip().strip('"') + if not entry or not os.path.isabs(entry): + continue + for ext in exts: + cand = os.path.join(entry, "git" + ext) + if os.path.isfile(cand): + candidates.append(cand) + for git_exe in candidates: + git_dir = os.path.dirname(git_exe) + if os.path.isabs(git_dir) and _is_trusted_windows_program_dir(git_dir): + return os.path.realpath(git_dir), os.path.splitext(git_exe)[1].upper() + return "", "" + + +def _is_trusted_windows_program_dir(path: str) -> bool: + """True when ``path`` sits under a system-managed Program Files root. + + Only the Program Files roots are trusted (admin-writable only), resolved + via the known-folder API so an overridden env var cannot relocate them, + never ``%SystemRoot%`` (Git does not install there and it holds + world-writable subdirs like ``Windows\\Temp``). Per-user managers + (Scoop/Choco shims under the profile) are refused. Paths are canonicalized + so 8.3 aliases and junctions still resolve to their real root (#7317). + """ + norm = _canon_win_path(path) + for root in _windows_program_roots(): + root_norm = _canon_win_path(root) + if norm == root_norm or norm.startswith(root_norm + os.sep): + return True + return False + + def _build_safe_env(workdir: str) -> dict[str, str]: """Build a minimal, credential-free environment for sandboxed subprocesses. Whitelist-built from scratch (parent env NOT inherited): only PATH/HOME/ TMPDIR/LANG/TERM/PYTHONIOENCODING/PYTHONPATH (+VIRTUAL_ENV or Windows - SystemRoot) reach the child; all credential vars (HF_TOKEN, AWS_*, etc.) - are absent. HOME points at the sandbox workdir so SDKs can't read the + SystemRoot and a minimal PATHEXT) reach the child; all credential vars + (HF_TOKEN, AWS_*, etc.) are absent. HOME points at the sandbox workdir so SDKs can't read the operator's cached creds. PYTHONPATH carries only the sandbox sitecustomize shim directory. + + PATH starts with the Studio interpreter / venv and OS system dirs so + ``python``/``pip`` stay pinned. On Windows only, Git-for-Windows install + dirs from the host PATH are appended so bare ``git`` resolves (#7317). + User-writable host PATH entries (venv, ``node_modules/.bin``, etc.) are + never inherited — they could shadow auto-safe terminal commands. """ # Start from the running interpreter's dir so 'python'/'pip' resolve to the # same environment the Unsloth server runs in. @@ -2519,6 +2631,20 @@ def _build_safe_env(workdir: str) -> dict[str, str]: else: path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"]) + # Windows Git installs live outside System32; inherit the dir of the git + # the HOST shell resolves, but ONLY when it sits under a system install + # root (Program Files, windir). A user-writable dir (Scoop/Choco shims) + # is refused: it would let an attacker drop rg.exe/jq.exe beside git and + # have an auto-approved bare command execute it (#7317). + git_ext = "" + if sys.platform == "win32": + # Append the CANONICAL (realpath) trusted git dir, scanning past any + # untrusted user shim that sorts first on PATH; the canonical path + # cannot be retargeted via a junction after the trust check. + _trusted_git_dir, git_ext = _resolve_trusted_windows_git() + if _trusted_git_dir: + path_entries.append(_trusted_git_dir) + # Deduplicate, preserving order. deduped = list(dict.fromkeys(p for p in path_entries if p)) @@ -2538,6 +2664,15 @@ def _build_safe_env(workdir: str) -> dict[str, str]: # Windows needs SystemRoot for Python/subprocess to work. if sys.platform == "win32": env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows") + # Restrict PATHEXT so cwd .BAT/.CMD cannot hijack bare names (#7317). + pathext = ".EXE;.COM" + if git_ext and git_ext not in (".EXE", ".COM"): + # Keep the host git launcher (e.g. a .CMD shim) resolvable. + pathext += ";" + git_ext + env["PATHEXT"] = pathext + # cmd/CreateProcess search cwd before PATH for bare names; disable so + # a workdir rg.exe/git.exe cannot shadow auto-approved commands. + env["NoDefaultCurrentDirectoryInExePath"] = "1" return env @@ -4060,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, @@ -4075,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. diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 9f301ba37e..367de196f7 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -794,7 +794,7 @@ def run_inference_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) - apply_gpu_ids(config.get("resolved_gpu_ids")) + apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) model_name = config["model_name"] diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 8e419849cb..b858fe6f17 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -891,6 +891,7 @@ class UnslothTrainer: use_gradient_checkpointing: str = "unsloth", use_rslora: bool = False, use_loftq: bool = False, + use_dora: bool = False, modules_to_save: list = None, ) -> bool: """ @@ -993,6 +994,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, ) # Audio VLM models support VLM-style layer selection @@ -1023,6 +1025,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, task_type = None, ) @@ -1042,6 +1045,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, ) @@ -1067,6 +1071,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, modules_to_save = modules_to_save, ) @@ -1087,6 +1092,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, modules_to_save = modules_to_save, ) @@ -1481,6 +1487,9 @@ class UnslothTrainer: SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" SNAC_SAMPLE_RATE = 24000 + + # SNAC codec unvalidated on Intel XPU; keep the pre-PR CPU + # fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" max_length = self.max_seq_length or 2048 tokenizer = self.tokenizer @@ -1642,7 +1651,8 @@ class UnslothTrainer: del snac_model gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: @@ -1669,6 +1679,8 @@ class UnslothTrainer: import numpy as np import torchaudio.transforms as T + # Spark-TTS BiCodec unvalidated on Intel XPU; keep the pre-PR CPU + # fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" # sparktts lives in the SparkAudio/Spark-TTS GitHub repo, not the HF model @@ -1857,7 +1869,8 @@ class UnslothTrainer: del audio_tokenizer gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: @@ -1894,6 +1907,8 @@ class UnslothTrainer: from datasets import Dataset as HFDataset from utils.paths import ensure_dir, tmp_root + # OuteTTS DAC/Whisper preprocess unvalidated on Intel XPU; keep the + # pre-PR CPU fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" # Clone OuteTTS repo (same as audio_codecs._load_dac) @@ -2065,7 +2080,8 @@ class UnslothTrainer: del prompt_processor gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index bfbd11a427..8592dabbfe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -30,7 +30,7 @@ from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt -from utils.hardware import prepare_gpu_selection +from utils.hardware import get_device, prepare_gpu_selection from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -196,6 +196,7 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), "use_rslora": values.get("use_rslora", False), "use_loftq": values.get("use_loftq", False), + "use_dora": values.get("use_dora", False), "train_on_completions": values.get("train_on_completions", False), "finetune_vision_layers": values.get("finetune_vision_layers", True), "finetune_language_layers": values.get("finetune_language_layers", True), @@ -219,6 +220,9 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: config[key] = values.get(key) if config["training_type"] == "Full Finetuning": config["load_in_4bit"] = False + # The parent's detected backend: the worker's apply_gpu_ids() targets the + # right visibility env var from this, without probing torch pre-mask. + config["device_backend"] = get_device().value return config @@ -452,6 +456,7 @@ class _MLXTrainerAdapter: use_gradient_checkpointing: Union[str, bool] = "unsloth", use_rslora: bool = False, use_loftq: bool = False, + use_dora: bool = False, ) -> bool: self._peft_config = { "use_lora": bool(use_lora), @@ -462,6 +467,7 @@ class _MLXTrainerAdapter: "gradient_checkpointing": use_gradient_checkpointing, "use_rslora": bool(use_rslora), "use_loftq": bool(use_loftq), + "use_dora": bool(use_dora), "finetune_vision_layers": bool(finetune_vision_layers), "finetune_language_layers": bool(finetune_language_layers), "finetune_attention_modules": bool(finetune_attention_modules), @@ -569,6 +575,7 @@ class _MLXTrainerAdapter: "gradient_checkpointing": "unsloth", "use_rslora": False, "use_loftq": False, + "use_dora": False, "finetune_vision_layers": True, "finetune_language_layers": True, "finetune_attention_modules": True, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 5ded18ea45..baf6329dae 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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 @@ -1547,6 +1554,10 @@ def _run_mlx_training(event_queue, stop_queue, config): message = "LoftQ is not supported for MLX training yet." _send("error", error = message) raise NotImplementedError(message) + if config.get("use_dora"): + message = "DoRA is not supported for MLX training yet." + _send("error", error = message) + raise NotImplementedError(message) if config.get("is_embedding"): message = "Embedding model training is not supported for MLX training yet." _send("error", error = message) @@ -2373,7 +2384,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> env = os.getenv("ENVIRONMENT_TYPE", "production"), ) - apply_gpu_ids(config.get("resolved_gpu_ids")) + apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) model_name = config["model_name"] @@ -2824,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: @@ -3186,6 +3197,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), use_rslora = config.get("use_rslora", False), use_loftq = config.get("use_loftq", False), + use_dora = config.get("use_dora", False), ) elif use_lora: _send_status(event_queue, "Configuring LoRA adapters...") @@ -3202,6 +3214,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), use_rslora = config.get("use_rslora", False), use_loftq = config.get("use_loftq", False), + use_dora = config.get("use_dora", False), ) else: _send_status(event_queue, "Preparing model for full finetuning...") @@ -3630,6 +3643,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> use_gradient_checkpointing = gradient_checkpointing, random_state = config.get("random_seed", 3407), use_rslora = config.get("use_rslora", False), + use_dora = config.get("use_dora", False), loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if config.get("use_loftq") else None, diff --git a/studio/backend/main.py b/studio/backend/main.py index a538f935ff..5af25efa74 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -438,7 +438,11 @@ def _run_llama_cpp_startup_probes(app: FastAPI) -> None: import structlog as _structlog _log = _structlog.get_logger(__name__) - if _caps.get("found") and not _caps.get("supports_mtp"): + if ( + _caps.get("found") + and not _caps.get("supports_mtp") + and not _caps.get("mtp_probe_inconclusive") + ): _msg = ( "llama.cpp prebuilt lacks MTP support " "(--spec-type mtp/draft-mtp). Run `unsloth studio update`. " diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index c7e5ffa36b..1758efe515 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -70,11 +70,23 @@ class LoadRequest(BaseModel): cache_type_kv: Optional[str] = Field( None, - description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", + description = ( + "KV cache data type for both K and V " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32')" + ), ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", + description = ( + "GPU placement pool, for example [0, 1]. Omit or pass [] to use " + "automatic selection. CUDA/ROCm and Intel XPU values are physical " + "GPU indices; Vulkan values are ggml device ordinals. Explicit " + "physical IDs are unsupported when the parent visibility mask uses " + "non-numeric or subdevice entries, including CUDA_VISIBLE_DEVICES " + "with UUID/MIG entries and ZE_AFFINITY_MASK with subdevice tokens " + "(for example '0.0,0.1') or FLAT-hierarchy tile handles. For GGUF " + "models the fitter may pin the smallest subset of this pool that fits." + ), ) speculative_type: Optional[str] = Field( None, @@ -433,7 +445,10 @@ class LoadResponse(BaseModel): ) cache_type_kv: Optional[str] = Field( None, - description = "KV cache data type for K and V (e.g. 'f16', 'bf16', 'q8_0')", + description = ( + "KV cache data type for K and V " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32')" + ), ) chat_template: Optional[str] = Field( None, @@ -485,7 +500,14 @@ class LoadResponse(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.", + ) + requested_gpu_ids: Optional[List[int]] = Field( + None, + description = ( + "GPU placement pool requested by the user before fit-time narrowing, " + "or None for automatic selection." + ), ) @@ -586,7 +608,11 @@ class InferenceStatusResponse(BaseModel): ) cache_type_kv: Optional[str] = Field( None, - description = "KV cache quantization dtype (e.g. 'q8_0'), or None for default", + description = ( + "KV cache quantization dtype " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32'), " + "or None for default" + ), ) chat_template: Optional[str] = Field( None, description = "Model's default chat template (Jinja2 source), if any" @@ -649,7 +675,14 @@ class InferenceStatusResponse(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.", + ) + requested_gpu_ids: Optional[List[int]] = Field( + None, + description = ( + "GPU placement pool requested by the user before fit-time narrowing, " + "or None for automatic selection." + ), ) llama_cpp_supports_mtp: bool = Field( True, diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 5a75246c07..4238403e00 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -47,6 +47,14 @@ class ProviderCreate(BaseModel): None, description = "Custom base URL (overrides registry default). Omit to use the default.", ) + models: list[str] = Field( + default_factory = list, + description = "Enabled model IDs for this connection", + ) + available_models: list[str] = Field( + default_factory = list, + description = "Discovered catalog model IDs last fetched for this connection", + ) class ProviderUpdate(BaseModel): @@ -55,6 +63,11 @@ class ProviderUpdate(BaseModel): display_name: Optional[str] = Field(None, description = "New display name") base_url: Optional[str] = Field(None, description = "New base URL") is_enabled: Optional[bool] = Field(None, description = "Enable or disable this provider") + models: Optional[list[str]] = Field(None, description = "Enabled model IDs for this connection") + available_models: Optional[list[str]] = Field( + None, + description = "Discovered catalog model IDs last fetched for this connection", + ) class ProviderResponse(BaseModel): @@ -65,6 +78,14 @@ class ProviderResponse(BaseModel): display_name: str = Field(..., description = "User-chosen label") base_url: str = Field(..., description = "API base URL") is_enabled: bool = Field(True, description = "Whether this provider is enabled") + models: list[str] = Field( + default_factory = list, + description = "Enabled model IDs for this connection", + ) + available_models: list[str] = Field( + default_factory = list, + description = "Discovered catalog model IDs last fetched for this connection", + ) created_at: str = Field(..., description = "ISO 8601 creation timestamp") updated_at: str = Field(..., description = "ISO 8601 last-update timestamp") diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0f88b78f9f..0aca5da72c 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -470,6 +470,7 @@ class TrainingStartRequest(BaseModel): gradient_checkpointing: str = Field("", description = "Gradient checkpointing setting") use_rslora: bool = Field(False, description = "Use RSLoRA") use_loftq: bool = Field(False, description = "Use LoftQ") + use_dora: bool = Field(False, description = "Use DoRA") train_on_completions: bool = Field(False, description = "Train on completions only") # Vision-specific LoRA parameters @@ -496,7 +497,15 @@ class TrainingStartRequest(BaseModel): # GPU selection gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", + description = ( + "Physical GPU indices to use, for example [0, 1]. Omit or pass " + "[] to use automatic selection. Explicit gpu_ids are unsupported " + "when the parent visibility mask uses non-numeric or subdevice " + "entries -- this includes CUDA_VISIBLE_DEVICES with UUID/MIG " + "entries on NVIDIA, and ZE_AFFINITY_MASK with subdevice tokens " + "(e.g. '0.0,0.1') or FLAT-hierarchy (default) tile handles on " + "Intel XPU." + ), ) # S3 dataset source configuration @@ -537,6 +546,37 @@ class TrainingStartRequest(BaseModel): raise ValueError("Either num_epochs or max_steps must be > 0; both cannot be 0.") return self + @model_validator(mode = "after") + def _validate_lora_variant_flags(self) -> "TrainingStartRequest": + # The frontend only ever sends one of these and never under Full + # Finetuning, but a direct API/YAML/CLI caller can bypass that. Nothing + # downstream breaks (full finetune ignores them, MLX rejects use_dora/ + # use_loftq outright), but reject early here for a clear error instead + # of a silently-ignored flag. + active = [ + name + for name, enabled in ( + ("use_rslora", self.use_rslora), + ("use_loftq", self.use_loftq), + ("use_dora", self.use_dora), + ) + if enabled + ] + if len(active) > 1: + raise ValueError( + f"Only one LoRA variant may be enabled at a time; got {active}. " + "use_rslora, use_loftq, and use_dora are mutually exclusive." + ) + # getattr, not self.training_type: model_construct() (used by tests that + # validate a single field in isolation) leaves required fields unset, and + # this is a mode="after" validator so it still runs on that partial instance. + if getattr(self, "training_type", None) == "Full Finetuning" and active: + raise ValueError( + f"{active[0]} requires an adapter method (LoRA/QLoRA or " + "Continued Pretraining); it has no effect under Full Finetuning." + ) + return self + class TrainingJobResponse(BaseModel): """Immediate response when training is initiated""" diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 24b6dfb36d..6a0d49b47d 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -160,11 +160,26 @@ class ChatInferenceSettings(BaseModel): fastMode: Optional[bool] = None +class ChatPresetLoadConfig(BaseModel): + model_config = ConfigDict(extra = "forbid") + + customContextLength: Optional[int] = Field(default = None, gt = 0) + maxSeqLength: Optional[float] = None + kvCacheDtype: Optional[str] = None + speculativeType: Optional[str] = None + specDraftNMax: Optional[int] = Field(default = None, ge = 1, le = 16) + tensorParallel: Optional[bool] = None + gpuMemoryMode: Optional[Literal["manual"]] = None + gpuLayers: Optional[int] = None + nCpuMoe: Optional[int] = Field(default = None, ge = 0) + + class ChatPreset(BaseModel): model_config = ConfigDict(extra = "forbid") name: str params: ChatInferenceSettings + loadConfig: Optional[ChatPresetLoadConfig] = None class ChatSettingsPayload(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b6bbbdfb5f..9d40650a56 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1794,7 +1794,16 @@ router = APIRouter() studio_router = APIRouter() -_ARTIFACT_PREVIEW_FRAME_ANCESTORS = "'self' tauri://localhost http://tauri.localhost" +# Packaged desktop runs at tauri://localhost (macOS/Linux) or http://tauri.localhost +# (Windows WebView2); the web build is same-origin ('self'). The `tauri dev` shell, +# however, serves the frontend from the Vite dev origin (http://localhost:5173), +# so the packaged allowlist alone leaves the preview blocked in dev with an +# "ancestor violates frame-ancestors" error. This shell exposes no server resource +# (it only renders postMessage'd HTML in a no-same-origin sandbox), so also allowing +# any localhost/127.0.0.1 dev origin to frame it is safe and unblocks the dev shell. +_ARTIFACT_PREVIEW_FRAME_ANCESTORS = ( + "'self' tauri://localhost http://tauri.localhost http://localhost:* http://127.0.0.1:*" +) _ARTIFACT_PREVIEW_FRAME_STRICT_CSP = ( "default-src 'none'; " "script-src 'unsafe-inline'; " @@ -3241,15 +3250,10 @@ def _request_matches_loaded_settings( ) ): return False - # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU - # request to its single lowest device (it drives one device only), so the - # backend records just that device; compare the request the same way, or a - # multi-GPU pick that resolves to the same device needlessly reloads. - if llama_backend.is_diffusion: - _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None - else: - _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None - if _req_gpu_ids != llama_backend.gpu_ids: + # A regular GGUF may narrow the requested placement pool. Accept either the + # original request or the effective status-echoed subset; diffusion keeps + # its single-device normalization. + if not llama_backend.matches_gpu_ids(request.gpu_ids): return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement @@ -3897,15 +3901,19 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: """Classify a GGUF as diffusion, normal, or unknown before it is loaded. ``None`` is important here: a remote GGUF whose header is not cached can - still be routed to the single-GPU diffusion runner after download. Treating - that case as normal would let Manual mode skip the training guard even - though the runner ignores Manual's llama-server placement controls. + still be routed to the single-GPU diffusion runner after download. Default + placement keeps that unknown case guarded until the header is available. """ identity = " ".join( str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") ).lower() - if "diffusion" in identity: - return True + # Name-only hint, used ONLY as a pre-download fallback, scoped to the + # DiffusionGemma runner family: a bare "diffusion" substring is common in + # ordinary text-model names/paths (e.g. "stable-diffusion-prompt"), and treating + # those as diffusion falsely rejects a valid Vulkan+gpu_ids GGUF (#7239). Normalize + # non-alphanumerics so "DiffusionGemma"/"diffusion-gemma" collapse to one token. + # The local header below stays authoritative. + name_says_diffusion = "diffusiongemma" in _re.sub(r"[^a-z0-9]+", "", identity) try: main = getattr(config, "gguf_file", None) @@ -3915,23 +3923,86 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: if repo and variant: from hub.utils.gguf import resolve_local_gguf_path main = resolve_local_gguf_path(repo, variant) - if not main or not Path(main).is_file(): - return None - - probe = LlamaCppBackend() - probe._read_gguf_metadata(str(main)) - if probe.is_diffusion: - return True - # A successfully decoded architecture proves that this is a normal - # llama-server GGUF. No architecture means the lightweight probe could - # not establish the routing decision, so preserve the unknown state. - if getattr(probe, "_architecture", None): - return False - return None + if main and Path(main).is_file(): + # The local GGUF header is authoritative (same probe the loader uses), so + # it can't be fooled by a "diffusion"-flavored name/path. + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A decoded architecture proves a normal llama-server GGUF; no architecture + # means the probe was inconclusive, so fall through to the name hint below. + if getattr(probe, "_architecture", None): + return False except Exception as e: logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + + # Header unavailable (remote uncached) or inconclusive: True only for the + # DiffusionGemma name family; otherwise None keeps an unknown remote GGUF guarded + # as potentially diffusion until its header proves otherwise. + return True if name_says_diffusion else None + + +async def _resolve_gguf_gpu_ids_for_request( + config: ModelConfig, gpu_ids: Optional[List[int]] +) -> Optional[List[int]]: + """Resolve and fully validate an explicit GGUF GPU placement pool. + + CUDA and ROCm use physical IDs. Vulkan uses ggml ordinals, so its device + existence check comes from the same ggml probe used by the loader. Both + /load and /validate call this before their training guard or any teardown. + """ + if not gpu_ids: return None + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + is_vulkan = LlamaCppBackend._is_vulkan_backend() + if get_device() == DeviceType.XPU and not is_vulkan: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + + if is_vulkan and _classify_diffusion_gguf(config) is True: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ), + ) + + try: + resolved = resolve_requested_gpu_ids(gpu_ids, is_vulkan = is_vulkan) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc + + if is_vulkan and resolved: + binary = LlamaCppBackend._find_llama_server_binary() + if binary: + probed = { + gpu[0] for gpu in await asyncio.to_thread(LlamaCppBackend._get_gpu_memory, binary) + } + wanted = {int(gpu_id) for gpu_id in resolved} + if not wanted.issubset(probed): + raise HTTPException( + status_code = 400, + detail = ( + f"Requested Vulkan GPU ordinal(s) {sorted(wanted)} not " + f"present. Available Vulkan devices: {sorted(probed)}." + ), + ) + + return resolved + def _guard_chat_load_against_training( config: ModelConfig, @@ -3971,8 +4042,18 @@ def _guard_chat_load_against_training( if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: return + # Vulkan GGUF pins are ggml ordinals, not CUDA physical IDs. Detect this + # before deriving a possible diffusion fallback device so an unknown remote + # GGUF never sends its ordinal through the CUDA single-device path. + is_vulkan = False + if is_gguf: + try: + is_vulkan = LlamaCppBackend._is_vulkan_backend() + except Exception as e: + logger.warning("Could not detect Vulkan backend for chat-load guard: %s", e) + diffusion_gpu = None - if is_gguf and diffusion_kind is not False: + if is_gguf and diffusion_kind is not False and not (is_vulkan and requested_gpu_ids): # Use the same token selection as the runner: an explicit pick wins, # followed by DG_GPU, the first parent-visible token, then GPU 0. diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( @@ -3999,6 +4080,7 @@ def _guard_chat_load_against_training( max_seq_length = max_seq_length, requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, + is_vulkan = is_vulkan, required_override_gb = required_override_gb, single_device_gpu = diffusion_gpu, ) @@ -4305,6 +4387,7 @@ async def _load_model_impl( # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) ): + llama_backend._record_matching_gpu_request(request.gpu_ids) logger.info( "Model already loaded (GGUF): " f"{model_log_label} variant={request.gguf_variant or llama_backend.hf_variant}, skipping reload" @@ -4351,6 +4434,7 @@ async def _load_model_impl( n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, ) else: if ( @@ -4417,41 +4501,12 @@ async def _load_model_impl( # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # GGUF supports gpu_ids: validate the pick up front (before the training - # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects - # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts - # are rejected outright: the picker's indices are torch-xpu ordinals neither - # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin - # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. - if config.is_gguf and effective_gpu_ids is not None: - from utils.hardware import DeviceType, get_device - from utils.hardware.hardware import resolve_requested_gpu_ids - - if get_device() == DeviceType.XPU: - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported on Intel XPU. " - "Omit gpu_ids to use all devices." - ), - ) - # Same reasoning for a Vulkan-only build: --device pins ggml's own - # Vulkan ordinals, so a physical pick can land on the wrong card on - # masked or non-contiguous hosts. - if LlamaCppBackend._is_vulkan_backend(): - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported with a Vulkan " - "llama.cpp build: physical GPU ids have no defined " - "mapping to Vulkan device ordinals. Omit gpu_ids to use " - "all devices." - ), - ) - try: - resolve_requested_gpu_ids(effective_gpu_ids) - except ValueError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc + # Validate the full GGUF placement pool before the training guard so an + # invalid physical ID or Vulkan ordinal is a clean 400, not a masked VRAM + # 409. The same helper is used by /validate. + gguf_gpu_ids: Optional[List[int]] = None + if config.is_gguf: + gguf_gpu_ids = await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids) if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, @@ -4575,8 +4630,9 @@ async def _load_model_impl( gpu_layers = request.gpu_layers, n_cpu_moe = request.n_cpu_moe, tensor_split = request.tensor_split, - gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, + # Issue #7164: explicit GPU pin resolved to physical ids above. + gpu_ids = gguf_gpu_ids, ) if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server @@ -4750,6 +4806,7 @@ async def _load_model_impl( n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -5043,36 +5100,8 @@ async def validate_model( # Apply the same training coexistence policy as /load before the frontend # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is - # a clean 400) before the guard sizes the model against training VRAM. - # XPU-host picks are rejected like /load (no defined mapping from the - # picker's torch-xpu ordinals to the launcher's device spaces). - if config.is_gguf and effective_gpu_ids is not None: - from utils.hardware import DeviceType, get_device - from utils.hardware.hardware import resolve_requested_gpu_ids - - if get_device() == DeviceType.XPU: - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported on Intel XPU. " - "Omit gpu_ids to use all devices." - ), - ) - if LlamaCppBackend._is_vulkan_backend(): - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported with a Vulkan " - "llama.cpp build: physical GPU ids have no defined " - "mapping to Vulkan device ordinals. Omit gpu_ids to use " - "all devices." - ), - ) - try: - resolve_requested_gpu_ids(effective_gpu_ids) - except ValueError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc + if config.is_gguf: + await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids) effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -5815,10 +5844,15 @@ async def get_status(current_subject: str = Depends(get_current_subject)): try: _bin = type(llama_backend)._find_llama_server_binary() _caps = type(llama_backend).probe_server_capabilities(_bin) - _supports_mtp = bool(_caps.get("supports_mtp", False)) + # Fail open on inconclusive probes: False means a definitive + # "binary lacks MTP" to API consumers. + _supports_mtp = bool( + _caps.get("supports_mtp", False) + or (_caps.get("found", False) and _caps.get("mtp_probe_inconclusive", False)) + ) except Exception: _bin = None - _supports_mtp = True # fail open + _supports_mtp = False # no usable binary: MTP genuinely unavailable try: from utils.llama_cpp_freshness import check_prebuilt_freshness _freshness = check_prebuilt_freshness(_bin) @@ -5897,6 +5931,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, @@ -6061,6 +6096,7 @@ async def generate_audio( # Advertised repo id after an auto-switch load, else a clean public id, # never the absolute .gguf path. model_name = _llama_public_model_id(llama_backend) + _audio_model_id = getattr(llama_backend, "model_identifier", None) or model_name gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -6079,6 +6115,7 @@ async def generate_audio( if not model_info.get("is_audio"): raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") model_name = public_model_id(backend.active_model_name) + _audio_model_id = getattr(backend, "active_model_name", None) or model_name gen = lambda: backend.generate_audio_response( text = text, temperature = payload.temperature, @@ -6090,6 +6127,13 @@ async def generate_audio( use_adapter = payload.use_adapter, ) + # Apply per-model recommended sampling + any operator UNSLOTH_SAMPLING_* pin before + # generating, so `unsloth run --temperature` (and the other pins) and per-model + # recommendations reach audio (TTS) generation too, not just chat. The gen lambdas read + # payload.* lazily at call time, so filling here takes effect; this covers both the direct + # /audio/generate route and the chat-completions audio branches that delegate here. + _fill_recommended_sampling_openai(payload, _audio_model_id) + try: wav_bytes, sample_rate = await asyncio.to_thread(gen) except Exception as e: @@ -7381,6 +7425,51 @@ async def delete_openai_container( await client.close() +def _fill_recommended_sampling_openai(payload, model_id) -> None: + """Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to a + ChatCompletionRequest in place. + + Only the sampling fields the client did NOT explicitly send (tracked via + ``model_fields_set``) are overwritten, so a client that sets a field stays byte-identical + unless an operator pins it. Fields with neither a recommendation nor a pin keep their + existing (schema-default) value. + """ + from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES + + explicit = { + f: (getattr(payload, f) if f in payload.model_fields_set else None) + for f in SAMPLING_FIELD_NAMES + } + effective = resolve_effective_sampling(model_id, explicit) + for field, value in effective.items(): + setattr(payload, field, value) + + +# /v1/completions is proxied to llama-server verbatim; its repetition knob is "repeat_penalty", +# and every other sampling field keeps its name (mirrors _build_passthrough_payload). +_COMPLETIONS_SAMPLING_BODY_KEY = {"repetition_penalty": "repeat_penalty"} + + +def _fill_recommended_sampling_completions(body: dict, model_id) -> None: + """Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to a raw + ``/v1/completions`` body in place, so the legacy (non-chat) endpoint honors the same pins as + ``/v1/chat/completions``. + + Unlike :func:`_fill_recommended_sampling_openai`, which fills a ChatCompletionRequest whose + schema already carries per-field defaults, this body is proxied to llama-server as-is. A field + with no operator pin, client value, or per-model recommendation is therefore left untouched + (``fill_defaults = False``) so llama-server keeps its own default rather than being forced onto + this schema's value. llama-server names the repetition knob ``repeat_penalty``, so read and + write that alias for the client-sent value and any pin. + """ + from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES + + explicit = {f: body.get(_COMPLETIONS_SAMPLING_BODY_KEY.get(f, f)) for f in SAMPLING_FIELD_NAMES} + effective = resolve_effective_sampling(model_id, explicit, fill_defaults = False) + for field, value in effective.items(): + body[_COMPLETIONS_SAMPLING_BODY_KEY.get(field, field)] = value + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, @@ -7710,6 +7799,13 @@ async def openai_chat_completions( completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) + # Apply recommended sampling + operator pins to the omitted fields before generating, + # so audio-input (non-whisper) generation honors `unsloth run --temperature` and + # per-model recommendations like chat does. Whisper (ASR) ignores these fields. + _fill_recommended_sampling_openai( + payload, getattr(backend, "active_model_name", None) or model_name + ) + def audio_input_generate(): if model_info.get("audio_type") == "whisper": return backend.generate_whisper_response( @@ -7853,6 +7949,18 @@ async def openai_chat_completions( ), ) + # Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to the + # fields the client omitted, so agents and API clients get the model's tuned defaults + # unless they set the field explicitly. Placed after external-provider routing (which + # returned above) so only local llama-server / transformers requests are touched, and it + # covers both the passthrough and non-passthrough branches below since both read payload.*. + _reco_model_id = ( + getattr(llama_backend, "model_identifier", None) + if using_gguf + else getattr(backend, "active_model_name", None) + ) or model_name + _fill_recommended_sampling_openai(payload, _reco_model_id) + # ── Standard OpenAI function-calling pass-through (GGUF only) ──── # When a client (opencode / Claude Code via OpenAI compat / Cursor / # Continue / ...) sends standard OpenAI `tools` without Unsloth's @@ -10630,6 +10738,10 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if _resolved_max_tokens is not None else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) + # Apply per-model recommended sampling and any operator UNSLOTH_SAMPLING_* pin to the raw + # body so /v1/completions honors the same pins as /v1/chat/completions; it is otherwise a + # verbatim proxy that would keep llama-server's defaults for every omitted sampling field. + _fill_recommended_sampling_completions(body, getattr(llama_backend, "model_identifier", None)) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) @@ -11588,6 +11700,9 @@ async def _responses_stream( detail = "Image provided but current GGUF model does not support vision.", ) + # Streaming /v1/responses builds the passthrough body directly (bypassing + # openai_chat_completions), so apply recommended sampling here too. + _fill_recommended_sampling_openai(chat_req, getattr(llama_backend, "model_identifier", None)) body = _build_openai_passthrough_body( chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) @@ -13019,14 +13134,28 @@ async def anthropic_messages( # endpoint matches /v1/chat/completions. _has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision) - temperature = payload.temperature if payload.temperature is not None else 0.6 - top_p = payload.top_p if payload.top_p is not None else 0.95 - top_k = payload.top_k if payload.top_k is not None else 20 - min_p = payload.min_p if payload.min_p is not None else 0.01 - repetition_penalty = ( - payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 + # Fill omitted sampling fields with the per-model recommendation (or an operator + # UNSLOTH_SAMPLING_* pin); an explicit client value wins unless the operator pinned it. + # Anthropic sampling fields are Optional, so None already marks "client omitted". + from utils.inference.inference_config import resolve_effective_sampling + + _anthropic_sampling = resolve_effective_sampling( + getattr(llama_backend, "model_identifier", None) or model_name, + { + "temperature": payload.temperature, + "top_p": payload.top_p, + "top_k": payload.top_k, + "min_p": payload.min_p, + "repetition_penalty": payload.repetition_penalty, + "presence_penalty": payload.presence_penalty, + }, ) - presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0 + temperature = _anthropic_sampling["temperature"] + top_p = _anthropic_sampling["top_p"] + top_k = _anthropic_sampling["top_k"] + min_p = _anthropic_sampling["min_p"] + repetition_penalty = _anthropic_sampling["repetition_penalty"] + presence_penalty = _anthropic_sampling["presence_penalty"] stop = payload.stop_sequences or None # Translate Anthropic tool_choice to OpenAI format for llama-server. Falls @@ -13235,6 +13364,7 @@ async def anthropic_messages( disable_parallel_tool_use = _disable_parallel, bypass_permissions = bool(payload.bypass_permissions), permission_mode = getattr(payload, "permission_mode", None), + promote_reasoning_only = False, ) if payload.stream: @@ -13274,6 +13404,7 @@ async def anthropic_messages( max_tokens = payload.max_tokens, stop = stop, cancel_event = cancel_event, + promote_reasoning_only = False, ) if payload.stream: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 3c0d6ff4ba..ed83a12f48 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2692,7 +2692,10 @@ async def get_kv_cache_estimate( repo_id: str = Query(..., description = "HF repo ID or local path"), quant: str = Query(..., description = "Quantization label (e.g. Q4_K_M)"), n_ctx: int = Query(..., ge = 1, description = "Context length to size the KV cache for"), - cache_type_kv: Optional[str] = Query(None, description = "KV cache dtype (e.g. q8_0)"), + cache_type_kv: Optional[str] = Query( + None, + description = "KV cache dtype (e.g. q8_0, q4_0, q5_0, iq4_nl, f32)", + ), current_subject: str = Depends(get_current_subject), ): """Estimate KV cache + weight bytes for a downloaded GGUF at n_ctx. diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 5a55c9b0bb..4e7e53f2f0 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -47,6 +47,20 @@ logger = structlog.get_logger(__name__) router = APIRouter() +def _provider_response(row: dict) -> ProviderResponse: + return ProviderResponse( + id = row["id"], + provider_type = row["provider_type"], + display_name = row["display_name"], + base_url = row["base_url"], + is_enabled = bool(row["is_enabled"]), + models = row.get("models") or [], + available_models = row.get("available_models") or [], + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + + # ── Public key for API key encryption ───────────────────────────── @@ -89,18 +103,7 @@ async def get_pricing_snapshot(current_subject: str = Depends(get_current_subjec async def list_provider_configs(current_subject: str = Depends(get_current_subject)): """List all saved provider configurations.""" rows = providers_db.list_providers() - return [ - ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) - for row in rows - ] + return [_provider_response(row) for row in rows] @router.post("/", response_model = ProviderResponse, status_code = 201) @@ -124,18 +127,12 @@ async def create_provider_config( provider_type = payload.provider_type, display_name = payload.display_name, base_url = base_url, + models = payload.models, + available_models = payload.available_models, ) row = providers_db.get_provider(provider_id) - return ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) + return _provider_response(row) @router.put("/{provider_id}", response_model = ProviderResponse) @@ -154,20 +151,14 @@ async def update_provider_config( display_name = payload.display_name, base_url = payload.base_url, is_enabled = payload.is_enabled, + models = payload.models, + available_models = payload.available_models, ) if not updated: raise HTTPException(status_code = 400, detail = "No fields to update") row = providers_db.get_provider(provider_id) - return ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) + return _provider_response(row) @router.delete("/{provider_id}", status_code = 204) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 392a4e0d02..ae65712146 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -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]: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 5b10652fdd..8be4283415 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -332,6 +332,7 @@ async def start_training( else "unsloth", "use_rslora": request.use_rslora, "use_loftq": request.use_loftq, + "use_dora": request.use_dora, "train_on_completions": request.train_on_completions, "finetune_vision_layers": request.finetune_vision_layers, "finetune_language_layers": request.finetune_language_layers, diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index 83bce8b772..8ddda11b1e 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -135,8 +135,8 @@ def can_keep_chat_during_training( resolve_requested_gpu_ids, ) - if get_device() != DeviceType.CUDA: - return False, {"mode": "non_cuda", "reason": "non_cuda"} + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + return False, {"mode": "non_accelerator", "reason": "non_accelerator"} # Full finetuning runs in 16-bit, so ignore the 4-bit request or we under-count. effective_4bit = False if training_type == "Full Finetuning" else load_in_4bit @@ -225,6 +225,7 @@ def can_load_chat_during_training( max_seq_length: int, requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, + is_vulkan: bool = False, required_override_gb: Optional[float] = None, single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: @@ -233,11 +234,15 @@ def can_load_chat_during_training( chat model against the free VRAM that remains). Sizes/places it the same way the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an even-share per-GPU floor for device_map="balanced"; GGUF sizes from - required_override_gb over the visible pool. ``single_device_gpu`` is the - exact physical device token selected by a single-device runner. - `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA - allows the load; default-deny on any CUDA case it can't size, so a load never - OOMs training.""" + required_override_gb over the visible pool. A Vulkan GGUF selection picks by ggml + Vulkan ordinal (separate index space from CUDA ids), so its requested_gpu_ids is + NOT resolved against the CUDA set (which would raise -> invalid_gpu_ids -> bypass + the OOM check); conservatively size an N-device request against the least-free + N visible GPUs instead. + ``single_device_gpu`` is the exact physical device token selected by a + single-device runner. `load_in_4bit` must be effective (LoRA can flip 4-bit + -> 16-bit). CPU/MLX allows the load; default-deny on any CUDA/XPU case it + can't size, so a load never OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -248,8 +253,8 @@ def can_load_chat_during_training( resolve_requested_gpu_ids, ) - if get_device() != DeviceType.CUDA: - return True, {"mode": "non_cuda", "reason": "non_cuda"} + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + return True, {"mode": "non_accelerator", "reason": "non_accelerator"} est_kwargs = dict( hf_token = hf_token or None, @@ -258,6 +263,11 @@ def can_load_chat_during_training( max_seq_length = max_seq_length or 2048, ) + # A Vulkan GGUF selection uses ggml Vulkan ordinals, not CUDA physical ids; + # size it against the full visible pool (GGUF self-placement) rather than + # resolving ordinals against the CUDA parent-visible set. + vulkan_gguf = is_gguf and is_vulkan + # HF auto: reuse the loader's selector; fits iff its pick clears the margin. if not requested_gpu_ids and not is_gguf: _selected, meta = auto_select_gpu_ids(model_name, **est_kwargs) @@ -283,7 +293,9 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. - if single_device_gpu is not None: + if requested_gpu_ids and vulkan_gguf: + mode = "gguf_vulkan" + elif single_device_gpu is not None: mode = "single_device" elif is_gguf: mode = "gguf" @@ -296,7 +308,17 @@ def can_load_chat_during_training( return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if single_device_gpu is not None: + if requested_gpu_ids and vulkan_gguf: + # Vulkan ordinals cannot be mapped to CUDA physical indices. Budget + # the least-free N visible cards for an N-device request. If that + # conservative subset fits, any physical mapping of the ordinals + # fits, without collapsing a multi-GPU request to one card. + visible_free = list(free_by_index.values()) + if not visible_free: + return False, {"mode": "gguf_vulkan", "reason": "no_visible_gpus"} + n_pins = min(len(requested_gpu_ids), len(visible_free)) + free_vals = sorted(visible_free)[:n_pins] + elif single_device_gpu is not None: token = str(single_device_gpu).strip() if not token: # Empty token = a CPU-only single-device runner (e.g. a CPU @@ -324,7 +346,8 @@ def can_load_chat_during_training( return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] else: - # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. + # GGUF self-placement / auto Vulkan (no requested ids): llama.cpp picks + # the GPU(s), so any visible GPU is a candidate -> size the whole pool. free_vals = list(free_by_index.values()) if not free_vals: diff --git a/studio/backend/run.py b/studio/backend/run.py index d9569c46f6..f1fc8c6062 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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, diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py index 07165cbe70..e6f40c5030 100644 --- a/studio/backend/storage/providers_db.py +++ b/studio/backend/storage/providers_db.py @@ -6,8 +6,12 @@ Same pattern as studio_db.py (module-level functions, raw sqlite3, WAL, per-function connections). API keys are NOT stored here: they live only in the browser (localStorage) and are sent encrypted per-request. + +Enabled model selections and discovered catalog IDs are stored server-side so +remote Studio clients see the same connection state (#7281). """ +import json import logging import sqlite3 import threading @@ -22,6 +26,33 @@ _schema_lock = threading.Lock() _schema_ready = False +def _encode_models_json(models: Optional[list[str]]) -> str: + if not models: + return "[]" + return json.dumps([str(model).strip() for model in models if str(model).strip()]) + + +def _decode_models_json(raw: Optional[str]) -> list[str]: + if not raw: + return [] + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + return [] + if not isinstance(parsed, list): + return [] + return [str(model).strip() for model in parsed if str(model).strip()] + + +def _row_models(row: sqlite3.Row) -> tuple[list[str], list[str]]: + return ( + _decode_models_json(row["models_json"] if "models_json" in row.keys() else None), + _decode_models_json( + row["available_models_json"] if "available_models_json" in row.keys() else None + ), + ) + + def _ensure_schema(conn: sqlite3.Connection) -> None: """Create the llm_providers table if absent. Called once per process.""" conn.execute("PRAGMA journal_mode=WAL") @@ -38,6 +69,13 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(llm_providers)").fetchall()} + if "models_json" not in existing_cols: + conn.execute("ALTER TABLE llm_providers ADD COLUMN models_json TEXT NOT NULL DEFAULT '[]'") + if "available_models_json" not in existing_cols: + conn.execute( + "ALTER TABLE llm_providers ADD COLUMN available_models_json TEXT NOT NULL DEFAULT '[]'" + ) def get_connection() -> sqlite3.Connection: @@ -59,17 +97,37 @@ def get_connection() -> sqlite3.Connection: return conn -def create_provider(id: str, provider_type: str, display_name: str, base_url: str) -> None: +def create_provider( + id: str, + provider_type: str, + display_name: str, + base_url: str, + models: Optional[list[str]] = None, + available_models: Optional[list[str]] = None, +) -> None: """Insert a new provider configuration.""" now = datetime.now(timezone.utc).isoformat() conn = get_connection() try: conn.execute( """ - INSERT INTO llm_providers (id, provider_type, display_name, base_url, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO llm_providers ( + id, provider_type, display_name, base_url, + models_json, available_models_json, + created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, provider_type, display_name, base_url, now, now), + ( + id, + provider_type, + display_name, + base_url, + _encode_models_json(models), + _encode_models_json(available_models), + now, + now, + ), ) conn.commit() finally: @@ -81,6 +139,8 @@ def update_provider( display_name: Optional[str] = None, base_url: Optional[str] = None, is_enabled: Optional[bool] = None, + models: Optional[list[str]] = None, + available_models: Optional[list[str]] = None, ) -> bool: """Update fields on an existing provider. Returns True if a row was updated.""" updates = [] @@ -94,6 +154,12 @@ def update_provider( if is_enabled is not None: updates.append("is_enabled = ?") params.append(1 if is_enabled else 0) + if models is not None: + updates.append("models_json = ?") + params.append(_encode_models_json(models)) + if available_models is not None: + updates.append("available_models_json = ?") + params.append(_encode_models_json(available_models)) if not updates: return False updates.append("updated_at = ?") @@ -128,7 +194,13 @@ def get_provider(id: str) -> Optional[dict]: conn = get_connection() try: row = conn.execute("SELECT * FROM llm_providers WHERE id = ?", (id,)).fetchone() - return dict(row) if row else None + if not row: + return None + data = dict(row) + models, available_models = _row_models(row) + data["models"] = models + data["available_models"] = available_models + return data finally: conn.close() @@ -138,6 +210,13 @@ def list_providers() -> list[dict]: conn = get_connection() try: rows = conn.execute("SELECT * FROM llm_providers ORDER BY created_at").fetchall() - return [dict(row) for row in rows] + providers: list[dict] = [] + for row in rows: + data = dict(row) + models, available_models = _row_models(row) + data["models"] = models + data["available_models"] = available_models + providers.append(data) + return providers finally: conn.close() diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py index 9fd8260bf2..be85fd56d1 100644 --- a/studio/backend/tests/test_amd_apu_unified_memory.py +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -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) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 9ccc3f44dd..621ac9aaca 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -68,16 +68,15 @@ def _emitter_client_text(events: list[str]) -> str: def test_anthropic_emitter_closes_reasoning_only_think_block(): - # A reasoning-only reply streams X live then shrinks to bare X at EOF. - # This emitter diffs cumulative snapshots and drops the shrink, so without a - # closing pass the client text would end on an unclosed . finish() - # must balance it. + # Anthropic asks the GGUF generator not to promote reasoning into a duplicate + # visible fallback, so its final cumulative snapshot only balances the block. emitter = AnthropicStreamEmitter() events = emitter.start("msg_1", "m") events += emitter.feed({"type": "content", "text": "The capital"}) events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) - # The generator's final bare-text shrink (dropped by the cumulative diff). - events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.feed( + {"type": "content", "text": "The capital of France is Paris."} + ) events += emitter.finish() assert _emitter_client_text(events) == "The capital of France is Paris." @@ -1563,6 +1562,44 @@ class TestAnthropicMessagesToolRouting: assert entry["context_length"] == 2048 assert monitor.active_count() == 0 + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("with_tools", [False, True]) + def test_reasoning_only_output_is_not_duplicated(self, monkeypatch, stream, with_tools): + reasoning = "The capital of France is Paris." + + def _gen_plain(**kwargs): + assert kwargs["promote_reasoning_only"] is False + yield f"{reasoning}" + yield f"{reasoning}" + + def _gen_tools(**kwargs): + assert kwargs["promote_reasoning_only"] is False + yield {"type": "content", "text": f"{reasoning}"} + yield {"type": "content", "text": f"{reasoning}"} + + _mock_backend( + monkeypatch, + generate_chat_completion = _gen_plain, + generate_chat_completion_with_tools = _gen_tools, + ) + payload_fields = {"stream": stream} + if with_tools: + payload_fields.update( + { + "enable_tools": True, + "tools": [{"type": "web_search_20250305", "name": "web_search"}], + } + ) + payload = _basic_payload(**payload_fields) + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + if stream: + body = self._sse_blob(self._consume_response(response)) + assert body.count(reasoning) == 1 + else: + body = json.loads(response.body) + assert body["content"][0]["text"] == f"{reasoning}" + def test_tool_use_non_streaming_records_api_monitor_reply(self, monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_audio_sampling_fill.py b/studio/backend/tests/test_audio_sampling_fill.py new file mode 100644 index 0000000000..efea18b83e --- /dev/null +++ b/studio/backend/tests/test_audio_sampling_fill.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Audio (TTS) generation applies recommended sampling + operator pins, like chat. + +Regression guard for the fix that moved the sampling fill ahead of the audio generators: a +prior version resolved sampling only after the audio branches returned, so `unsloth run +--temperature` (UNSLOTH_SAMPLING_*) and per-model recommendations never reached audio +generation. These exercise the transformers TTS path of ``generate_audio`` (the direct +``/audio/generate`` route, which the chat-completions audio branches also delegate to). +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import ChatCompletionRequest +from utils.inference import inference_config as ic + + +class _FakeLlama: + # is_loaded False forces the transformers (non-GGUF) TTS branch in generate_audio. + is_loaded = False + _is_audio = False + + +class _FakeTransformersBackend: + def __init__(self): + self.active_model_name = "some/custom-tts" + self.models = {"some/custom-tts": {"is_audio": True}} + self.captured = {} + + def generate_audio_response(self, **kwargs): + self.captured.update(kwargs) + return (b"RIFFfake", 24000) + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + ic._recommended_sampling.cache_clear() + for field in ic.SAMPLING_FIELD_NAMES: + monkeypatch.delenv(ic._SAMPLING_FIELDS[field][0], raising = False) + yield + ic._recommended_sampling.cache_clear() + + +def _run_generate_audio( + monkeypatch, + *, + recommended = None, + temperature = None, +): + backend = _FakeTransformersBackend() + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: backend) + + async def _noop_switch(*a, **k): + return None + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _noop_switch) + + # Recommendation source == the Chat UI's .inference block. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(recommended or {})) + ic._recommended_sampling.cache_clear() + + kwargs = {"model": "some/custom-tts", "messages": [{"role": "user", "content": "hi"}]} + if temperature is not None: + kwargs["temperature"] = temperature + payload = ChatCompletionRequest(**kwargs) + + asyncio.run(inference_route.generate_audio(payload, request = None, current_subject = "t")) + return backend.captured + + +def test_audio_uses_recommended_sampling_when_omitted(monkeypatch): + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0, "top_k": 64}) + assert captured["temperature"] == 1.0 + assert captured["top_k"] == 64 + + +def test_audio_operator_pin_overrides_client(monkeypatch): + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0}, temperature = 0.2) + assert captured["temperature"] == 0.9 # operator pin wins even over an explicit client value + + +def test_audio_client_explicit_preserved(monkeypatch): + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0}, temperature = 0.2) + assert captured["temperature"] == 0.2 # explicit client value preserved over recommendation diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index a60ac700bf..896bf1a6cd 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -91,6 +91,28 @@ def test_chat_settings_payload_accepts_fast_mode_presets(): assert dumped["customPresets"][0]["params"]["fastMode"] is True +def test_chat_settings_payload_accepts_preset_load_config(): + payload = chat_history.ChatSettingsPayload.model_validate( + { + "customPresets": [ + { + "name": "GGUF preset", + "params": {"temperature": 0.7, "maxTokens": 512}, + "loadConfig": { + "customContextLength": 256, + "kvCacheDtype": "q8_0", + "tensorParallel": False, + }, + }, + ], + } + ) + + dumped = payload.model_dump(exclude_unset = True) + assert dumped["customPresets"][0]["loadConfig"]["customContextLength"] == 256 + assert dumped["customPresets"][0]["loadConfig"]["kvCacheDtype"] == "q8_0" + + def test_chat_settings_payload_accepts_nudge_tool_calls(): # extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the # frontend's persisted nudgeToolCalls needs a payload field (like diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index a5fd71b6a0..f1d973f004 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -170,6 +170,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): estimate = None, single_device_gpu = None, gpu_ids = None, + is_vulkan = False, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), @@ -185,6 +186,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): max_seq_length = 0, requested_gpu_ids = gpu_ids, is_gguf = True, + is_vulkan = is_vulkan, required_override_gb = required_override, single_device_gpu = single_device_gpu, ) @@ -234,6 +236,35 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): self.assertFalse(blocked) self.assertEqual(blocked_info["usable_gb"], 10.0) + def test_vulkan_pin_takes_precedence_over_unknown_diffusion_fallback(self): + # An uncached GGUF can carry a speculative single-device fallback while + # its explicit pin is actually a ggml Vulkan ordinal. Never interpret + # that ordinal as the same-numbered CUDA physical device. + ok, info, _ = self._run( + devices = _devices((0, 80, 0), (1, 80, 78)), + required_override = 20.0, + single_device_gpu = "0", + gpu_ids = [0], + is_vulkan = True, + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "gguf_vulkan") + self.assertEqual(info["usable_gb"], 2.0) + + def test_vulkan_multi_gpu_guard_counts_requested_devices(self): + # The ordinal mapping is unknown, so use the least-free two visible + # cards for a two-device request. Their aggregate capacity is still + # available instead of collapsing the request to one card. + ok, info, _ = self._run( + devices = _devices((0, 80, 70), (1, 80, 70), (2, 80, 0)), + required_override = 10.0, + gpu_ids = [0, 1], + is_vulkan = True, + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf_vulkan") + self.assertEqual(info["usable_gb"], 18.5) + def test_single_device_unresolved_token_sizes_against_worst_device(self): # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a # free-VRAM index. The runner still drives ONE device, so size against the @@ -295,7 +326,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): - def test_non_cuda_allows(self): + def test_non_accelerator_allows(self): with patch("utils.hardware.get_device", return_value = DeviceType.MLX): ok, info = tv.can_load_chat_during_training( model_name = "m", @@ -305,7 +336,30 @@ class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): requested_gpu_ids = None, ) self.assertTrue(ok) - self.assertEqual(info["mode"], "non_cuda") + self.assertEqual(info["mode"], "non_accelerator") + + def test_xpu_overcommit_is_refused(self): + # XPU must NOT get the blanket non-accelerator allow: an oversized + # chat model during resident training is refused, like CUDA. + with ( + patch("utils.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.auto_select_gpu_ids", + return_value = ( + None, + {"selection_mode": "auto", "required_gb": 50.0, "usable_gb": 4.0}, + ), + ), + ): + ok, info = tv.can_load_chat_during_training( + model_name = "m", + hf_token = None, + load_in_4bit = True, + max_seq_length = 0, + requested_gpu_ids = None, + ) + self.assertFalse(ok) + self.assertNotEqual(info.get("mode"), "non_accelerator") def test_no_visible_gpus_refuses(self): # GGUF with an empty device list -> no candidate GPU -> default-deny. @@ -478,58 +532,19 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_manual_known_normal_gguf_bypasses_training_estimate(self): captured = [] config = SimpleNamespace(is_gguf = True) - with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False) as classify: self._guard( config = config, captured = captured, training_active = True, decision = (False, {"reason": "must not run"}), gpu_memory_mode = "manual", + requested_gpu_ids = [1, 3], ) + classify.assert_called_once_with(config) self.assertEqual(captured, []) - def test_manual_unknown_gguf_keeps_single_device_training_guard(self): - captured = [] - config = SimpleNamespace(is_gguf = True) - with ( - patch.object(self.route, "_classify_diffusion_gguf", return_value = None), - patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), - patch.object( - self.route.LlamaCppBackend, - "_diffusion_gpu_arg", - return_value = "2", - ), - ): - self._guard( - config = config, - captured = captured, - training_active = True, - decision = (True, {"mode": "single_device"}), - gpu_memory_mode = "manual", - ) - self.assertEqual(len(captured), 1) - self.assertEqual(captured[0]["single_device_gpu"], "2") - - def test_manual_diffusion_uses_single_device_guard(self): - captured = [] - config = SimpleNamespace(is_gguf = True) - with ( - patch.object(self.route, "_classify_diffusion_gguf", return_value = True), - patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), - ): - self._guard( - config = config, - captured = captured, - training_active = True, - decision = (True, {"mode": "gguf"}), - gpu_memory_mode = "manual", - requested_gpu_ids = [3, 1], - ) - self.assertEqual(len(captured), 1) - self.assertEqual(captured[0]["single_device_gpu"], "1") - self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) - - def test_unpinned_diffusion_uses_runner_default_gpu(self): + def test_manual_diffusion_keeps_single_device_training_guard(self): captured = [] config = SimpleNamespace(is_gguf = True) with ( @@ -540,11 +555,6 @@ class TestChatLoadGuardRoute(unittest.TestCase): "_effective_gpu_count", return_value = 2, ), - patch.object( - self.route.LlamaCppBackend, - "_diffusion_gpu_arg", - return_value = "3", - ) as gpu_arg, ): self._guard( config = config, @@ -552,9 +562,11 @@ class TestChatLoadGuardRoute(unittest.TestCase): training_active = True, decision = (True, {"mode": "single_device"}), gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], ) - gpu_arg.assert_called_once_with(None, cpu_only = False) - self.assertEqual(captured[0]["single_device_gpu"], "3") + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} diff --git a/studio/backend/tests/test_colab_embed.py b/studio/backend/tests/test_colab_embed.py new file mode 100644 index 0000000000..83b2a5a82d --- /dev/null +++ b/studio/backend/tests/test_colab_embed.py @@ -0,0 +1,598 @@ +# 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 coverage for Colab iframe embedding (#7344).""" + +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import colab + + +def _mock_google_colab_modules(colab_mod): + """Mock ``google`` and ``google.colab`` for environments without Google packages.""" + google_mod = types.ModuleType("google") + google_mod.colab = colab_mod + return {"google": google_mod, "google.colab": colab_mod} + + +def test_short_colab_url_truncates_proxy_host(): + url = "https://8888-gpu-a100-s-kkb-usc1f0-9hzedjcxrlu8-f.us-central1-0.prod.colab.dev/" + assert colab._short_colab_url(url, 8888) == "https://8888-gpu-..." + + +def test_short_colab_url_falls_back_on_unexpected_shape(): + assert colab._short_colab_url("https://example.com", 8888) == "https://example.com" + + +def test_is_colab_proxy_url_requires_https_proxy(): + assert colab._is_colab_proxy_url("https://8888-test.prod.colab.dev/", 8888) is True + assert colab._is_colab_proxy_url("http://localhost:8888", 8888) is False + assert colab._is_colab_proxy_url("http://127.0.0.1:8888", 8888) is False + + +def test_ready_card_html_does_not_open_colab_proxy_in_new_tab(): + """Colab proxy hosts 404 as top-level tabs (#7349 reporter); never window.open them.""" + html = colab._ready_card_html("https://8888-test.prod.colab.dev/", 8888) + assert "window.open" not in html + assert 'href="https://8888-test.prod.colab.dev/"' not in html + assert "start(cloudflare=True)" in html + + +def test_ready_card_html_points_to_cloudflare_when_link_ready(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html( + "https://8888-test.prod.colab.dev/", + 8888, + has_cloudflare_link = True, + ) + assert "Cloudflare link above" in html + + +def test_ready_card_html_warns_when_cloudflare_tunnel_missing(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html( + "https://8888-test.prod.colab.dev/", + 8888, + cloudflare_requested = True, + ) + assert "Could not open a Cloudflare tunnel" in html + + +def test_warn_colab_cloudflare_missing_logs_on_colab_without_tunnel(monkeypatch): + warnings: list[str] = [] + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr(colab.logger, "warning", lambda msg, **kwargs: warnings.append(msg)) + colab._warn_colab_cloudflare_missing(use_cloudflare = True, cloudflare_url = None) + assert warnings + assert "Cloudflare tunnel unavailable" in warnings[0] + + +def test_warn_colab_cloudflare_missing_skips_when_tunnel_ready(monkeypatch, caplog): + import logging + + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + with caplog.at_level(logging.WARNING): + colab._warn_colab_cloudflare_missing( + use_cloudflare = True, + cloudflare_url = "https://share.trycloudflare.com", + ) + assert "Cloudflare tunnel unavailable" not in caplog.text + + +def test_is_colab_runtime_uses_backend_colab_detector(monkeypatch): + fake_main = types.ModuleType("main") + fake_main._IS_COLAB = True + monkeypatch.setitem(sys.modules, "main", fake_main) + assert colab._is_colab_runtime() is True + fake_main._IS_COLAB = False + assert colab._is_colab_runtime() is False + + +def test_ready_card_html_uses_cloudflare_hint_on_colab_runtime_localhost(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html("http://localhost:8888", 8888) + assert "window.open" not in html + assert "start(cloudflare=True)" in html + + +def test_ready_card_html_keeps_open_button_for_localhost_outside_colab(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + html = colab._ready_card_html("http://localhost:8888", 8888) + assert "window.open" in html + assert 'href="http://localhost:8888"' in html + assert "Open Unsloth Studio" in html + + +def test_embed_kernel_port_iframe_uses_colab_helper(monkeypatch): + colab_output = MagicMock() + google_colab = SimpleNamespace(output = colab_output) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + with patch.dict("sys.modules", _mock_google_colab_modules(google_colab)): + assert colab._embed_kernel_port_iframe(8888) is True + colab_output.serve_kernel_port_as_iframe.assert_called_once_with( + 8888, + height = colab._COLAB_IFRAME_HEIGHT, + width = "100%", + ) + + +def test_embed_kernel_port_iframe_returns_false_without_colab(): + with patch.dict("sys.modules", _mock_google_colab_modules(None)): + assert colab._embed_kernel_port_iframe(8888) is False + + +def test_embed_kernel_port_iframe_skips_colabtools_without_runtime(monkeypatch): + """colabtools can queue JS without appending an iframe; only trust the helper on Colab.""" + colab_output = MagicMock() + google_colab = SimpleNamespace(output = colab_output) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + with patch.dict("sys.modules", _mock_google_colab_modules(google_colab)): + assert colab._embed_kernel_port_iframe(8888) is False + colab_output.serve_kernel_port_as_iframe.assert_not_called() + + +def test_show_and_embed_prefers_kernel_port_iframe(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-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: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "kernel_iframe"] + + +def test_show_and_embed_falls_back_to_html_iframe(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: False) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append((url, port)) or True, + ) + + colab._show_and_embed(8888) + + assert calls == [("https://8888-test.prod.colab.dev/", 8888)] + + +def test_colab_wants_cloudflare_auto_enables_on_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + assert colab._colab_wants_cloudflare(None) is True + assert colab._colab_wants_cloudflare(True) is True + assert colab._colab_wants_cloudflare(False) is False + + +def test_colab_wants_cloudflare_defaults_off_outside_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + assert colab._colab_wants_cloudflare(None) is False + assert colab._colab_wants_cloudflare(True) is True + + +def test_finalize_colab_admin_password_skips_outside_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + assert colab._finalize_colab_admin_password() is None + + +def test_finalize_colab_admin_password_clears_bootstrap_gate(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr(colab, "_load_colab_login_credentials", lambda: None) + stored: list[tuple[str, str]] = [] + monkeypatch.setattr( + colab, + "_store_colab_login_credentials", + lambda username, password: stored.append((username, password)), + ) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(return_value = "alpha-beta-gamma"), + generate_bootstrap_password = MagicMock(return_value = "alpha-beta-gamma"), + requires_password_change = MagicMock(return_value = True), + update_password = MagicMock(return_value = True), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result == ("unsloth", "alpha-beta-gamma") + storage.ensure_default_admin.assert_called_once() + storage.update_password.assert_called_once_with("unsloth", "alpha-beta-gamma") + assert stored == [("unsloth", "alpha-beta-gamma")] + + +def test_start_skips_finalize_when_cloudflare_disabled(monkeypatch): + import time + + finalize_calls: list[str] = [] + monkeypatch.setattr(colab, "_is_studio_healthy", lambda port: True) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_finalize_colab_admin_password", + lambda: finalize_calls.append("finalize") or ("unsloth", "secret"), + ) + monkeypatch.setattr( + colab, "start_cloudflare_tunnel", lambda port: "https://share.trycloudflare.com" + ) + monkeypatch.setattr(colab, "_publish_cloudflare_url", lambda url: None) + monkeypatch.setattr(colab, "_show_and_embed", lambda port, **kwargs: None) + monkeypatch.setattr(colab, "_stop_cloudflare_tunnel", lambda: None) + monkeypatch.setattr(time, "sleep", lambda _: (_ for _ in ()).throw(KeyboardInterrupt)) + + colab.start(cloudflare = False) + + assert finalize_calls == [] + + +def test_finalize_colab_admin_password_redisplay_on_rerun(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_load_colab_login_credentials", + lambda: ("unsloth", "saved-pass"), + ) + monkeypatch.setattr(colab, "_colab_credentials_still_valid", lambda username, password: True) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(), + generate_bootstrap_password = MagicMock(), + requires_password_change = MagicMock(return_value = False), + update_password = MagicMock(), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result == ("unsloth", "saved-pass") + storage.update_password.assert_not_called() + + +def test_finalize_colab_admin_password_drops_stale_cached_credentials(monkeypatch): + """After an in-app password change the cached first-run password no longer + authenticates, so it must not be redisplayed (#7349 Codex review).""" + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_load_colab_login_credentials", + lambda: ("unsloth", "stale-pass"), + ) + monkeypatch.setattr(colab, "_colab_credentials_still_valid", lambda username, password: False) + cleared: list[bool] = [] + monkeypatch.setattr(colab, "_clear_colab_login_credentials", lambda: cleared.append(True)) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(), + generate_bootstrap_password = MagicMock(), + requires_password_change = MagicMock(return_value = False), + update_password = MagicMock(), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result is None + assert cleared == [True] + storage.update_password.assert_not_called() + + +def test_colab_credentials_still_valid_matches_stored_hash(monkeypatch): + from auth.hashing import hash_password + + salt, pwd_hash = hash_password("right-pass") + storage = SimpleNamespace( + get_user_and_secret = MagicMock(return_value = (salt, pwd_hash, "jwt", False)), + ) + with patch.dict("sys.modules", {"auth.storage": storage}): + assert colab._colab_credentials_still_valid("unsloth", "right-pass") is True + assert colab._colab_credentials_still_valid("unsloth", "wrong-pass") is False + + +def test_colab_credentials_still_valid_false_when_user_missing(monkeypatch): + storage = SimpleNamespace(get_user_and_secret = MagicMock(return_value = None)) + with patch.dict("sys.modules", {"auth.storage": storage}): + assert colab._colab_credentials_still_valid("unsloth", "any") is False + + +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_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 'https://share.trycloudflare.com" 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", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + with patch.dict("sys.modules", {"IPython.display": ipython_display}): + colab._show_and_embed( + 8888, + cloudflare_url = "https://share.trycloudflare.com", + colab_login = ("unsloth", "secret-pass"), + ) + + assert len(displayed) == 1 + assert "share.trycloudflare.com" in displayed[0] + 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): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-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: None, + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com") + + assert calls == [] + + +def test_show_and_embed_uses_kernel_helper_on_colab_runtime_despite_localhost(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"http://localhost:{port}") + 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: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "kernel_iframe"] + + +def test_show_and_embed_skips_kernel_helper_for_localhost_outside_colab(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"http://localhost:{port}") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + + 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: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "html_iframe"] + + +def test_show_and_embed_still_embeds_when_show_link_fails(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None: (_ for _ in ()).throw(RuntimeError("no display")), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["kernel_iframe"] diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py index b4c8f5d242..e02c33a0f1 100644 --- a/studio/backend/tests/test_file_security.py +++ b/studio/backend/tests/test_file_security.py @@ -165,6 +165,23 @@ def test_skips_local_path(): assert "local" in d.reason +def test_scans_inactive_hf_cache_snapshot_path(tmp_path): + # An inactive HF cache loads by snapshot path; the gate must recover the repo id + + # commit from models--org--repo/snapshots/ and scan that exact commit, not exempt + # it and not fall back to the default branch (an older commit may hold a dropped pickle). + snapshot = tmp_path / "models--evil--repo" / "snapshots" / "deadbeef" + snapshot.mkdir(parents = True) + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status) as model_info: + d = evaluate_file_security(str(snapshot)) + assert d.blocked is True + assert model_info.call_args.args[0] == "evil/repo" + assert model_info.call_args.kwargs["revision"] == "deadbeef" + + def test_remote_gguf_named_repo_is_still_scanned(): # Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a # poisoned pickle smuggled into it is blocked. diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 6d1fac980b..ce26147b11 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -785,13 +785,21 @@ class TestLoadHubDownloadExclusion: def test_load_marker_precedes_hub_guard_and_unload(self): source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() - gguf_branch = source[source.index("if config.is_gguf:") :] + # _load_model_impl has more than one `if config.is_gguf:`, so anchor on + # the branch that actually owns the load marker rather than the first + # one in the file, which belongs to an earlier check. + marker = source.index("enter_context(gguf_load_in_flight") + gguf_branch_start = source.rindex("if config.is_gguf:", 0, marker) + gguf_branch = source[gguf_branch_start:] # The gguf_load_in_flight marker must be entered before the hub-download # guard and the unload so a concurrent load can't race the download - # manager. The llama_extra_args inheritance that used to sit between the - # marker and the guard now runs in _guard_chat_load_against_training, ahead - # of the GGUF branch, so it is no longer a landmark inside this slice. + # manager. The llama_extra_args inheritance moved out of the branch into + # _resolve_inherited_extra_args, which must still run BEFORE it: the + # inherited value (e.g. a carried --no-mmproj) shapes the guard's + # require_mmproj. Anchor on the call form so the assertion pins the + # endpoint's call site, not the function definition. + assert source.index("= _resolve_inherited_extra_args(") < gguf_branch_start assert ( gguf_branch.index("enter_context(gguf_load_in_flight") < gguf_branch.index("_hub_download_blocks_gguf_load") diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index 19ba9e3e05..271a882b11 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -591,10 +591,23 @@ def test_load_request_accepts_gpu_ids(): @pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) def test_response_models_emit_gpu_ids(model_cls): if model_cls is LoadResponse: - obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + obj = model_cls( + status = "loaded", + model = "m", + display_name = "m", + inference = {}, + gpu_ids = [1], + requested_gpu_ids = [1, 2], + ) else: - obj = model_cls(gpu_ids = [1]) + obj = model_cls(gpu_ids = [1], requested_gpu_ids = [1, 2]) assert obj.model_dump()["gpu_ids"] == [1] + assert obj.model_dump()["requested_gpu_ids"] == [1, 2] + + +def test_gguf_load_and_status_responses_include_requested_gpu_pool(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert route_src.count("requested_gpu_ids = llama_backend.requested_gpu_ids") == 3 def test_gpu_ids_property_default_and_reset(): @@ -625,6 +638,10 @@ def _target_state_gpu_ids(backend, gpu_ids): def test_gpu_ids_reload_detection_is_order_insensitive(): backend = _loaded_backend("auto") backend._gpu_ids = [0, 1] + # A real non-narrowed load records the raw request too; the non-diffusion + # dedupe now compares that raw pin (#7239). Set it to match the effective pin + # (no narrowing) so this exercises the order-insensitive comparison. + backend._requested_gpu_ids = [0, 1] # Same set, different order -> no reload. assert _target_state_gpu_ids(backend, [1, 0]) is True # Different set -> reload. @@ -633,6 +650,26 @@ def test_gpu_ids_reload_detection_is_order_insensitive(): assert _target_state_gpu_ids(backend, None) is False +def test_gpu_ids_reload_detection_accepts_raw_and_effective_pin(): + backend = _loaded_backend("auto") + backend._requested_gpu_ids = [0, 1] + backend._gpu_ids = [0] + backend._last_load_kwargs = {"gpu_ids": [0, 1], "model_identifier": "owner/repo"} + + # The original request still matches after the fitter narrows it. + assert _target_state_gpu_ids(backend, [1, 0]) is True + assert backend.requested_gpu_ids == [0, 1] + # The status response echoes the effective pin, which must also round-trip. + # Treat the incoming subset as the latest intent so status and a future + # reload do not restore GPU 1 after the user removed it. + assert _target_state_gpu_ids(backend, [0]) is True + assert backend.requested_gpu_ids == [0] + assert backend._last_load_kwargs == {"gpu_ids": [0], "model_identifier": "owner/repo"} + # A genuinely different placement pool still reloads. + assert _target_state_gpu_ids(backend, [1]) is False + assert _target_state_gpu_ids(backend, None) is False + + def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): # The diffusion runner drives only its single lowest device, so the backend # records [lowest]. A later multi-GPU request that still resolves to that @@ -642,6 +679,7 @@ def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): backend._is_diffusion = True backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick assert _target_state_gpu_ids(backend, [3, 1]) is True + assert backend.requested_gpu_ids == [1] assert _target_state_gpu_ids(backend, [1]) is True # Lowest device changes (2, not 1) -> reload. assert _target_state_gpu_ids(backend, [3, 2]) is False @@ -649,6 +687,56 @@ def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): assert _target_state_gpu_ids(backend, None) is False +def test_remote_vulkan_diffusion_preflight_runs_before_teardown(monkeypatch): + def _mark_diffusion(probe, path): + assert path == "/cache/model.gguf" + probe._is_diffusion = True + + monkeypatch.setattr(LlamaCppBackend, "_read_gguf_metadata", _mark_diffusion) + assert LlamaCppBackend._gguf_path_is_diffusion("/cache/model.gguf", "owner/model") is True + + src = inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + preflight = src.index("_preflight_model_path = self._download_gguf(") + teardown = src.index("# ── Phase 1: kill old process") + assert preflight < teardown + assert "model_path = _preflight_model_path or self._download_gguf(" in src + + +def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch): + backend = LlamaCppBackend() + killed = [] + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)]) + monkeypatch.setattr( + backend, + "_download_gguf", + lambda **_kwargs: "/cache/diffusion.gguf", + ) + monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: True) + monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True)) + monkeypatch.setattr( + llama_cpp_module, + "_resolve_repo_id_casing", + lambda repo: repo, + ) + monkeypatch.setattr( + llama_cpp_module, + "_hf_offline_if_dns_dead", + lambda: __import__("contextlib").nullcontext(), + ) + + with pytest.raises(ValueError, match = "DiffusionGemma"): + backend.load_model( + hf_repo = "owner/model", + hf_variant = "Q4_K_M", + model_identifier = "owner/model", + gpu_ids = [0], + ) + + assert killed == [] + + def test_start_diffusion_server_resets_tensor_parallel(): # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model # phase 1 only kills the process, it skips the unload reset). Diffusion is never @@ -656,18 +744,16 @@ def test_start_diffusion_server_resets_tensor_parallel(): # diffusion re-Apply reloads against stale tensor-parallel state. src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) assert "self._tensor_parallel = False" in src + assert "self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None" in src -def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): - # The route-level reload dedupe mirrors the backend: for a loaded diffusion - # model it compares the request against the single recorded device, not the - # full requested list, or a same-device multi-GPU pick reloads needlessly. +def test_route_matches_loaded_settings_uses_shared_gpu_pin_matcher(): + # Route-level and backend race dedupe must share one normalization path so + # raw, effective, and diffusion pins cannot drift apart. route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] - guard = match_impl.index("if llama_backend.is_diffusion:") - collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") - compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") - assert guard < collapse < compare + assert "if not llama_backend.matches_gpu_ids(request.gpu_ids):" in match_impl + assert "llama_backend._record_matching_gpu_request(request.gpu_ids)" in match_impl # ── Manual tensor split: child enumeration pinned to the picker's order ────── diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d4f2fbe993..3dab7ef368 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -119,7 +119,8 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaisesRegex( - ValueError, "unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG" + ValueError, + "unsupported when CUDA_VISIBLE_DEVICES uses non-numeric or subdevice", ): resolve_requested_gpu_ids([1]) @@ -130,6 +131,26 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): ): self.assertEqual(resolve_requested_gpu_ids([]), [1, 3]) + def test_vulkan_ordinals_bypass_cuda_parent_visible_validation(self): + # Vulkan build on a CPU-only torch host: no CUDA parent-visible set and a + # zero physical count, yet a valid Vulkan ordinal must not be rejected as + # a CUDA physical id (issue #7239). + with ( + patch.dict(os.environ, {}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 0), + ): + # As a CUDA physical id, [0] is outside the empty parent-visible set. + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([0]) + # As Vulkan ordinals, [0] and [0, 1] pass through unchanged. + self.assertEqual(resolve_requested_gpu_ids([0], is_vulkan = True), [0]) + self.assertEqual(resolve_requested_gpu_ids([0, 1], is_vulkan = True), [0, 1]) + # Malformed ordinals are still rejected. + with self.assertRaisesRegex(ValueError, "duplicate GPU IDs"): + resolve_requested_gpu_ids([0, 0], is_vulkan = True) + with self.assertRaisesRegex(ValueError, "non-negative"): + resolve_requested_gpu_ids([-1], is_vulkan = True) + def test_apply_gpu_ids_only_updates_cuda_visible_devices(self): with patch.dict( os.environ, @@ -846,12 +867,177 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): class TestRouteErrors(unittest.TestCase): - def test_prepare_gpu_selection_rejects_gpu_ids_on_non_cuda_backend(self): + def test_prepare_gpu_selection_rejects_gpu_ids_on_non_accelerator_backend(self): with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): with self.assertRaises(ValueError) as exc_info: prepare_gpu_selection([0], model_name = "unsloth/test") - self.assertIn("only supported on CUDA devices", str(exc_info.exception)) + self.assertIn("only supported on CUDA and Intel XPU", str(exc_info.exception)) + + def test_inference_route_resolves_gguf_gpu_ids(self): + # GGUF gpu_ids are now supported: /load routes them through the same + # resolution as non-GGUF loads (rejecting only genuinely invalid ids with + # the resolver's actionable message) rather than a blanket "not supported" + # reject, so /validate can stay consistent with /load (#7239). + import utils.hardware.hardware as hardware_mod + + inference_route = _load_route_module( + "inference_route_module_for_gguf_gpu_ids_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + def _fake_resolve(ids, is_vulkan = False): + raise ValueError("SENTINEL requested GPUs are outside the parent-visible set") + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + # Patch both the package re-export and the defining module so the stub + # fires no matter which import path the route uses. + patch("utils.hardware.resolve_requested_gpu_ids", _fake_resolve), + patch.object(hardware_mod, "resolve_requested_gpu_ids", _fake_resolve), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route._load_model_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + + # The selection was routed through resolution (not the old blanket reject). + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("SENTINEL", exc_info.exception.detail) + self.assertNotIn("not supported for GGUF", exc_info.exception.detail) + + def test_load_rejects_unavailable_vulkan_ordinal_before_training_guard(self): + inference_route = _load_route_module( + "inference_route_module_for_vulkan_preflight_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [99]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch("utils.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(inference_route, "_classify_diffusion_gguf", return_value = None), + patch.object( + inference_route.LlamaCppBackend, + "_is_vulkan_backend", + return_value = True, + ), + patch.object( + inference_route.LlamaCppBackend, + "_find_llama_server_binary", + return_value = "/tmp/llama-server", + ), + patch.object( + inference_route.LlamaCppBackend, + "_get_gpu_memory", + return_value = [(0, 8 * 1024**3, 16 * 1024**3)], + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ) as training_guard, + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route._load_model_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("Vulkan GPU ordinal(s) [99]", exc_info.exception.detail) + training_guard.assert_not_called() + + def test_vulkan_ordinals_are_allowed_on_xpu_hosts(self): + import utils.hardware.hardware as hardware_mod + + inference_route = _load_route_module( + "inference_route_module_for_xpu_vulkan_test", + "routes/inference.py", + ) + config = SimpleNamespace(is_gguf = True) + + with ( + patch("utils.hardware.get_device", return_value = DeviceType.XPU), + patch.object( + inference_route.LlamaCppBackend, + "_is_vulkan_backend", + return_value = True, + ), + patch.object(inference_route, "_classify_diffusion_gguf", return_value = False), + patch.object(hardware_mod, "resolve_requested_gpu_ids", return_value = [0, 1]), + patch.object( + inference_route.LlamaCppBackend, + "_find_llama_server_binary", + return_value = None, + ), + ): + resolved = asyncio.run( + inference_route._resolve_gguf_gpu_ids_for_request(config, [1, 0]) + ) + + self.assertEqual(resolved, [0, 1]) def test_inference_route_validates_gpu_ids_for_gguf(self): # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still @@ -861,7 +1047,7 @@ class TestRouteErrors(unittest.TestCase): import utils.hardware.hardware as hardware_mod inference_route = _load_route_module( - "inference_route_module_for_gguf_gpu_ids_test", + "inference_route_module_for_gguf_gpu_ids_test2", "routes/inference.py", ) request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) @@ -886,6 +1072,17 @@ class TestRouteErrors(unittest.TestCase): "ModelConfig", SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), + # Patch both the package re-export and the defining module so the stub + # fires no matter which import path the route uses. + patch( + "utils.hardware.resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), patch.object( inference_route, "_guard_chat_load_against_training", @@ -893,11 +1090,6 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), - patch.object( - hardware_mod, - "resolve_requested_gpu_ids", - side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), - ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -1439,18 +1631,61 @@ class TestAutoSelectWithNoneRequired(_GpuCacheResetMixin, unittest.TestCase): self.assertEqual(metadata["selection_mode"], "fallback_all") -class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase): - def test_auto_select_returns_non_cuda_for_xpu(self): - with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): +class TestXpuSelection(_GpuCacheResetMixin, unittest.TestCase): + def test_auto_select_supports_xpu(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = (1.0, {}), + ), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "devices": [ + {"index": 0, "vram_total_gb": 8, "vram_used_gb": 1}, + ] + }, + ), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": None, + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0], + ), + ): selected, metadata = auto_select_gpu_ids("unsloth/test") - self.assertIsNone(selected) - self.assertEqual(metadata["selection_mode"], "non_cuda") + self.assertEqual(selected, [0]) + self.assertEqual(metadata["selection_mode"], "auto") - def test_prepare_gpu_selection_rejects_explicit_ids_on_xpu(self): - with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): - with self.assertRaisesRegex(ValueError, "only supported on CUDA"): - prepare_gpu_selection([0], model_name = "unsloth/test") + def test_prepare_gpu_selection_accepts_explicit_ids_on_xpu(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": "0", + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0], + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 1), + ): + selected, metadata = prepare_gpu_selection([0], model_name = "unsloth/test") + + self.assertEqual(selected, [0]) + self.assertEqual(metadata["selection_mode"], "explicit") class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase): diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py index 733933271b..ba6d057123 100644 --- a/studio/backend/tests/test_gpu_selection_sandbox.py +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -294,13 +294,13 @@ class TestAutoSelectGpuIds(unittest.TestCase): # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB self.assertEqual(len(selected), 2) - def test_non_cuda_returns_none(self): + def test_non_accelerator_returns_none(self): from utils.hardware.hardware import auto_select_gpu_ids import utils.hardware.hardware as hw with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): selected, meta = auto_select_gpu_ids("test/model") self.assertIsNone(selected) - self.assertEqual(meta["selection_mode"], "non_cuda") + self.assertEqual(meta["selection_mode"], "non_accelerator") class TestGetDeviceMap(unittest.TestCase): diff --git a/studio/backend/tests/test_grouped_mm_rdna4_fallback.py b/studio/backend/tests/test_grouped_mm_rdna4_fallback.py new file mode 100644 index 0000000000..675b9c3210 --- /dev/null +++ b/studio/backend/tests/test_grouped_mm_rdna4_fallback.py @@ -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"]) diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 049058e511..45c8bcb032 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -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 @@ -335,3 +530,24 @@ class TestRetryContract: def test_external_kill_skips_flash_attn_retry(self): # SIGKILL (-9, OOM killer) is not a program fault: no FA-off retry. assert _signal_crash(-9) is False + + +class TestMmprojRetryFailureMessage: + """#7302: bare mmproj crashes must not be reported as projector-format.""" + + def test_confirmed_projector_keeps_historical_wording(self): + msg = LlamaCppBackend._mmproj_retry_failure_message( + projector_confirmed = True, + detail = "llama-server failed to start", + ) + assert msg.startswith("Vision projector incompatible with this llama.cpp") + assert "llama-server failed to start" in msg + + def test_bare_crash_does_not_claim_projector_incompatibility(self): + msg = LlamaCppBackend._mmproj_retry_failure_message( + projector_confirmed = False, + detail = "llama-server failed to start. Check that the GGUF file is valid", + ) + assert "Vision projector incompatible" not in msg + assert "crashed with --mmproj" in msg + assert "GGUF file is valid" in msg diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 1d15647967..27c1b17a85 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -637,7 +637,9 @@ def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch def fake_run(cmd, **kwargs): captured["cmd"] = cmd captured["env"] = kwargs.get("env") - return _types.SimpleNamespace(stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "") + return _types.SimpleNamespace( + stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "", returncode = 0 + ) monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run) @@ -678,6 +680,95 @@ def test_probe_server_capabilities_reports_outdated_binary(tmp_path): assert caps["found"] is True assert caps["mtp_token"] is None assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_reads_mtp_from_multiline_help(tmp_path): + # Enum on the indented line: first-line-only probing falsely reported + # "lacks MTP" (#7302). + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--spec-type TYPE\n" + " speculative decoding type\n" + " (none,draft-simple,draft-mtp,ngram-mod)\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["mtp_token"] == "draft-mtp" + assert caps["supports_mtp"] is True + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_empty_help_fails_open(tmp_path): + # --help prints nothing: must not claim the prebuilt lacks MTP (#7302). + fake = tmp_path / "llama-server" + fake.write_text("#!/usr/bin/env bash\nexit 0\n") + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +@_NEEDS_BASH +def test_probe_server_capabilities_no_spec_type_is_definitive(tmp_path): + # Nonempty --help without --spec-type: pre-spec binary, not inconclusive. + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--gpu-layers N\n GPU layers to offload\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_failed_help_with_output_is_inconclusive(tmp_path): + fake = tmp_path / "llama-server" + fake.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "--help" ]; then\n' + " echo 'illegal instruction'\n" + " exit 1\n" + "fi\n" + ) + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +@_NEEDS_BASH +def test_probe_server_capabilities_crash_on_help_fails_open(tmp_path): + fake = tmp_path / "llama-server" + fake.write_text("#!/usr/bin/env bash\nkill -SEGV $$\n") + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +def test_mtp_token_from_spec_help_prefers_draft_mtp(): + assert ( + LlamaCppBackend._mtp_token_from_spec_help("--spec-type none,draft-mtp,mtp,ngram-mod") + == "draft-mtp" + ) + assert LlamaCppBackend._mtp_token_from_spec_help("--spec-type [none|mtp|ngram-cache]") == "mtp" + assert LlamaCppBackend._mtp_token_from_spec_help("--spec-type none,ngram-mod") is None + # No incidental substring matches. + assert LlamaCppBackend._mtp_token_from_spec_help("prompt cache") is None def test_probe_server_capabilities_handles_missing_binary(): @@ -685,6 +776,7 @@ def test_probe_server_capabilities_handles_missing_binary(): caps = LlamaCppBackend.probe_server_capabilities("/no/such/llama-server") assert caps["found"] is False assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True assert caps["supports_cache_ram"] is False assert caps["supports_ctx_checkpoints"] is False assert caps["supports_no_cache_prompt"] is False @@ -1176,12 +1268,14 @@ def _resolver_backend( *, ngram_supported = True, mtp_token = "draft-mtp", + mtp_probe_inconclusive = False, ): """Backend with a deterministic probe so the resolver is hermetic.""" fake = { "found": True, "mtp_token": mtp_token, "supports_mtp": bool(mtp_token), + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": "new" if ngram_supported else None, "supports_ngram_mod": bool(ngram_supported), "spec_draft_n_max_flag": "--spec-draft-n-max", @@ -1879,6 +1973,24 @@ def test_spec_fallback_reason_set_when_binary_lacks_mtp(monkeypatch): assert backend.spec_fallback_reason == "binary_no_mtp" +def test_spec_fallback_reason_none_when_mtp_probe_inconclusive(monkeypatch): + backend = _resolver_backend( + monkeypatch, + mtp_token = None, + mtp_probe_inconclusive = True, + ) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason is None + + def test_spec_fallback_reason_none_when_mtp_engages(monkeypatch): backend = _resolver_backend(monkeypatch) backend._build_speculative_flags( diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index e99e227d40..cf9fde7118 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -37,6 +37,24 @@ def _done() -> str: return "data: [DONE]\n" +def _finish(reason: str) -> str: + return ( + "data: " + + json.dumps( + { + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": reason, + } + ] + } + ) + + "\n" + ) + + def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): backend = LlamaCppBackend.__new__(LlamaCppBackend) backend._process = object() @@ -299,9 +317,8 @@ def test_reasoning_streams_incrementally_with_tools(monkeypatch): def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): # A reasoning-only turn (whole answer in reasoning_content, no content, no # tool) with a tool active streams the reasoning live, then resolves to the - # bare reasoning text -- identical to the no-tool generate_chat_completion - # path -- so the non-streaming drain still returns it as `content`, not an - # empty answer. + # same text on the visible channel. The final cumulative snapshot stays + # append-only so route suffix extraction cannot drop that fallback. stream = [ _sse({"reasoning_content": "The capital of France is Paris."}), _done(), @@ -321,8 +338,49 @@ def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): content_texts = [e["text"] for e in events if e["type"] == "content"] # Reasoning streamed live during BUFFERING (the fix). assert content_texts[0] == "The capital of France is Paris." - # Resolves to bare reasoning, matching the no-tool sibling. - assert content_texts[-1] == "The capital of France is Paris." + assert content_texts[-1] == ( + "The capital of France is Paris.The capital of France is Paris." + ) + + +def _assert_reasoning_only_raw_consumer_gets_one_balanced_think_block(monkeypatch, with_tools): + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + backend = _make_backend(monkeypatch, [stream], []) + + if with_tools: + items = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + promote_reasoning_only = False, + ) + ) + cumulatives = [item["text"] for item in items if item.get("type") == "content"] + else: + items = list( + backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + promote_reasoning_only = False, + ) + ) + cumulatives = [item for item in items if isinstance(item, str)] + + assert cumulatives[-1] == "The capital of France is Paris." + assert all( + current.startswith(previous) for previous, current in zip([""] + cumulatives, cumulatives) + ) + + +def test_reasoning_only_raw_consumer_without_tools_gets_one_balanced_think_block(monkeypatch): + _assert_reasoning_only_raw_consumer_gets_one_balanced_think_block(monkeypatch, False) + + +def test_reasoning_only_raw_consumer_with_tools_gets_one_balanced_think_block(monkeypatch): + _assert_reasoning_only_raw_consumer_gets_one_balanced_think_block(monkeypatch, True) def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): @@ -392,8 +450,8 @@ def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str] def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): # Parity contract: a reasoning-only reply must reach the client identically # whether tools are on or off. Both generators stream live then - # resolve to the bare reasoning text; the route's suffix-diff + extractor - # must therefore produce the same (visible, reasoning) split for both. + # append a balanced close plus visible fallback; the route's suffix-diff + + # extractor must therefore produce the same split for both. stream = [ _sse({"reasoning_content": "The capital"}), _sse({"reasoning_content": " of France is Paris."}), @@ -430,10 +488,37 @@ def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) assert tool_out == no_tool_out # Pin the shared contract so a change to either path shows up here. - _visible, reasoning = tool_out + visible, reasoning = tool_out + assert visible == "The capital of France is Paris." assert reasoning == "The capital of France is Paris." +def test_length_truncated_reasoning_stays_append_only_without_visible_promotion(monkeypatch): + stream = [ + _sse({"reasoning_content": "The proof begins by assuming finitely many primes."}), + _finish("length"), + _done(), + ] + backend = _make_backend(monkeypatch, [stream], []) + + items = list( + backend.generate_chat_completion( + messages = [{"role": "user", "content": "Prove infinitely many primes"}], + max_tokens = 16, + ) + ) + cumulatives = [item for item in items if isinstance(item, str)] + + assert all( + current.startswith(previous) for previous, current in zip([""] + cumulatives, cumulatives) + ) + assert cumulatives[-1] == ("The proof begins by assuming finitely many primes.") + visible, reasoning = _replay_route_reasoning_extractor(cumulatives) + assert visible == "" + assert reasoning == "The proof begins by assuming finitely many primes." + assert items[-1]["finish_reason"] == "length" + + def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): # _drain_silently sibling of the structured-tool close: a bare-JSON tool call # with a live reasoning prefix must also close before draining, and diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index 7ff580c36d..ccc6b5f76a 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -51,6 +51,27 @@ def _modules_json(*paths): _COMMIT = "0123456789abcdef0123456789abcdef01234567" +def _fs_case_sensitive(root): + """Whether root's filesystem is case-sensitive (Linux yes; macOS/Windows usually no). The gate + mirrors the loader, whose file lookups follow the same rule, so some cases only exist on one.""" + probe = Path(root) / "_case_probe" + probe.write_text("x") + try: + return not (Path(root) / "_CASE_PROBE").exists() + finally: + probe.unlink() + + +def _requires_case_sensitive_fs(root): + if not _fs_case_sensitive(root): + pytest.skip("requires a case-sensitive filesystem") + + +def _requires_case_insensitive_fs(root): + if _fs_case_sensitive(root): + pytest.skip("requires a case-insensitive filesystem") + + def _make_cache( root, repo_id, @@ -382,6 +403,329 @@ def test_gate_blocks_sharded_pickle(hf_cache): assert _offline_decision("org/shard").blocked is True +def test_gate_blocks_indexed_pickle_shard_in_subdirectory(hf_cache): + # from_pretrained follows weight_map paths relative to the root index, so these nested shards + # are deserialized even though they are not direct children of the load root (iterdir misses + # them). The online gate blocks index-referenced subdir pickles; the offline gate must too. + _make_cache( + hf_cache, + "org/indexed-shard", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"layer.weight": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-shard") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_pickle_shard_with_nonstandard_stem(hf_cache): + # The index tells the loader to deserialize this file, so a pickle EXTENSION is enough -- the + # shard's stem need not match the on-disk weight-name heuristic (which only guesses bare files). + _make_cache( + hf_cache, + "org/indexed-odd", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/evil-00001-of-00001.bin"}}', + "shards/evil-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-odd") + assert decision.blocked is True + assert any(u["path"] == "shards/evil-00001-of-00001.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_safetensors_index_pointing_to_pickle_shard(hf_cache): + # load_state_dict picks safetensors vs torch.load by each shard's own suffix, so a + # model.safetensors.index.json that maps a weight to a .bin shard still deserializes it. The + # index's own existence must not suppress the shard it names. + _make_cache( + hf_cache, + "org/st-index-pickle", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/st-index-pickle") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_shard_with_no_pickle_extension(hf_cache): + # Transformers torch.loads any indexed shard not ending in .safetensors, so an unconventional + # extensionless name is still a deserialization target. + _make_cache( + hf_cache, + "org/indexed-noext", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload"}}', + "shards/payload": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-noext") + assert decision.blocked is True + assert any(u["path"] == "shards/payload" for u in decision.unsafe_files) + + +_UPPER_INDEX_FILES = { + "PYTORCH_MODEL.BIN.INDEX.JSON": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", +} + + +def test_gate_blocks_uppercase_index_on_case_insensitive_fs(hf_cache): + # On a case-insensitive volume (Windows/macOS) from_pretrained opens an oddly-cased index when it + # requests the canonical lowercase name, so the loader-mirror lookup resolves it and blocks. + _requires_case_insensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + decision = _offline_decision("org/upper-index") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_allows_uppercase_index_on_case_sensitive_fs(hf_cache): + # On a case-sensitive FS from_pretrained's os.path.isfile of the canonical lowercase name misses + # the uppercase artifact and never loads its shard, so the gate must not over-block it. + _requires_case_sensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + assert _offline_decision("org/upper-index").blocked is False + + +def test_gate_blocks_indexed_shard_named_with_backslash(hf_cache): + # On POSIX a backslash is a literal filename char, so from_pretrained joins the raw weight_map + # value and deserializes a file actually named "dir\payload.bin"; the gate must probe it verbatim. + import os + + if os.sep != "/": + pytest.skip("backslash is a path separator off POSIX") + _make_cache( + hf_cache, + "org/backslash", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "dir\\\\payload.bin"}}', + "dir\\payload.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/backslash") + assert decision.blocked is True + assert any(u["path"] == "dir\\payload.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_shard_with_uppercase_safetensors_suffix(hf_cache): + # load_state_dict's endswith(".safetensors") is case-sensitive, so a shard named payload.SAFETENSORS + # falls to torch.load. The gate must classify shard suffixes case-sensitively to match it. + _make_cache( + hf_cache, + "org/upper-suffix", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload.SAFETENSORS"}}', + "shards/payload.SAFETENSORS": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-suffix") + assert decision.blocked is True + assert any(u["path"] == "shards/payload.SAFETENSORS" for u in decision.unsafe_files) + + +def test_gate_allows_stale_safetensors_index_beside_direct_safetensors(hf_cache): + # A complete direct model.safetensors is selected before either index, so a stale + # model.safetensors.index.json referencing a .bin shard never deserializes -> must not block. + _make_cache( + hf_cache, + "org/direct-plus-stale-index", + { + "model.safetensors": "tensors", + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + assert _offline_decision("org/direct-plus-stale-index").blocked is False + + +def test_gate_blocks_pytorch_index_with_uppercase_safetensors_decoy(hf_cache): + # On a case-sensitive FS, from_pretrained asks for the canonical lowercase model.safetensors, does + # not find an uppercase decoy, and selects the pytorch index instead. The decoy must not suppress. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy", + { + "MODEL.SAFETENSORS": "decoy", + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_direct_pickle_with_uppercase_safetensors_decoy(hf_cache): + # Same decoy against a direct pytorch_model.bin: the loader selects the pickle, so the uppercase + # safetensors must not suppress it on a case-sensitive FS. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy-direct", + {"MODEL.SAFETENSORS": "decoy", "pytorch_model.bin": "pickle"}, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy-direct") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_pickle_shard_in_module_subdir(hf_cache): + # A weight index inside a sentence-transformers module load root points at a nested pickle shard. + _make_cache( + hf_cache, + "org/mod-indexed", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "0_Transformer/shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/mod-indexed") + assert decision.blocked is True + assert any( + u["path"] == "0_Transformer/shards/pytorch_model-00001-of-00001.bin" + for u in decision.unsafe_files + ) + + +def test_gate_allows_indexed_pickle_shard_with_safetensors_sibling(hf_cache): + # A base model.safetensors makes the loader ignore the pickle index entirely, so it must not + # block (mirrors the direct-file safetensors-sibling suppression). + _make_cache( + hf_cache, + "org/indexed-both", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + "model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/indexed-both").blocked is False + + +def test_gate_allows_indexed_safetensors_shard_in_subdirectory(hf_cache): + # A safetensors index lists inert shards -- following it must never block (guards against a + # scanner that flags every indexed shard regardless of format). + _make_cache( + hf_cache, + "org/st-indexed", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}' + ), + "shards/model-00001-of-00001.safetensors": "tensors", + }, + ) + with _no_network(): + assert _offline_decision("org/st-indexed").blocked is False + + +def test_gate_blocks_on_index_path_traversal(hf_cache): + # A weight_map entry escaping the snapshot via ".." is abnormal/hostile -> fail closed. + _make_cache( + hf_cache, + "org/escape", + {"pytorch_model.bin.index.json": '{"weight_map": {"w": "../../../../etc/evil.bin"}}'}, + ) + with _no_network(): + assert _offline_decision("org/escape").blocked is True + + +def test_gate_allows_symlinked_sharded_safetensors(tmp_path, monkeypatch): + # Real HF caches store snapshot files as symlinks into blobs/. A resolve()-based containment + # check would escape the snapshot and false-block every sharded model; the lexical gate must not. + import hashlib + import os + + from huggingface_hub.file_download import repo_folder_name + + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = root), + ) + repo_dir = root / repo_folder_name(repo_id = "org/sym", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text(_COMMIT) + blobs = repo_dir / "blobs" + blobs.mkdir() + snapshot = repo_dir / "snapshots" / _COMMIT + (snapshot / "shards").mkdir(parents = True) + + def _blobbed(rel, content): + digest = hashlib.sha256(content.encode()).hexdigest() + (blobs / digest).write_text(content) + target = snapshot / rel + target.parent.mkdir(parents = True, exist_ok = True) + target.symlink_to(os.path.relpath(blobs / digest, target.parent)) + + _blobbed("config.json", "{}") + _blobbed( + "model.safetensors.index.json", + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}', + ) + _blobbed("shards/model-00001-of-00001.safetensors", "tensors") + with _no_network(): + assert _offline_decision("org/sym").blocked is False + + +def test_gate_allows_index_without_weight_map(hf_cache): + # An index whose top-level JSON has no dict weight_map lets the loader resolve no shards, so it + # must not crash or block on its own (only inert safetensors are cached here). + _make_cache( + hf_cache, + "org/no-wm", + {"model.safetensors.index.json": "[]", "model.safetensors": "x"}, + ) + with _no_network(): + assert _offline_decision("org/no-wm").blocked is False + + def test_gate_allows_nothing_cached(hf_cache): with _no_network(): assert _offline_decision("org/missing").blocked is False diff --git a/studio/backend/tests/test_providers_db_models.py b/studio/backend/tests/test_providers_db_models.py new file mode 100644 index 0000000000..ca9dffbd70 --- /dev/null +++ b/studio/backend/tests/test_providers_db_models.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for provider model persistence (unslothai/unsloth#7281).""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import storage.providers_db as providers_db + + +@pytest.fixture() +def isolated_providers_db(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + db_path = tmp_path / "studio.db" + monkeypatch.setattr(providers_db, "studio_db_path", lambda: db_path) + monkeypatch.setattr(providers_db, "ensure_dir", lambda _path: None) + providers_db._schema_ready = False + yield db_path + providers_db._schema_ready = False + + +def test_create_and_list_provider_models(isolated_providers_db: Path): + providers_db.create_provider( + id = "ollama1", + provider_type = "ollama", + display_name = "Home Ollama", + base_url = "http://127.0.0.1:11434", + models = ["llama3.2", "qwen2.5"], + available_models = ["llama3.2", "qwen2.5", "mistral"], + ) + + row = providers_db.get_provider("ollama1") + assert row is not None + assert row["models"] == ["llama3.2", "qwen2.5"] + assert row["available_models"] == ["llama3.2", "qwen2.5", "mistral"] + + listed = providers_db.list_providers() + assert len(listed) == 1 + assert listed[0]["models"] == ["llama3.2", "qwen2.5"] + + +def test_update_provider_models(isolated_providers_db: Path): + providers_db.create_provider( + id = "vllm1", + provider_type = "vllm", + display_name = "Remote vLLM", + base_url = "http://studio-host:8000/v1", + models = ["meta-llama/Llama-3.2-1B-Instruct"], + available_models = ["meta-llama/Llama-3.2-1B-Instruct"], + ) + + assert providers_db.update_provider( + id = "vllm1", + models = ["meta-llama/Llama-3.2-3B-Instruct"], + available_models = [ + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ], + ) + + row = providers_db.get_provider("vllm1") + assert row is not None + assert row["models"] == ["meta-llama/Llama-3.2-3B-Instruct"] + assert row["available_models"] == [ + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ] diff --git a/studio/backend/tests/test_rag_project_source_upload.py b/studio/backend/tests/test_rag_project_source_upload.py new file mode 100644 index 0000000000..fd20816b56 --- /dev/null +++ b/studio/backend/tests/test_rag_project_source_upload.py @@ -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 diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index ad46f6ee41..5cdbe4f2a5 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -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: diff --git a/studio/backend/tests/test_sampling_resolution.py b/studio/backend/tests/test_sampling_resolution.py new file mode 100644 index 0000000000..1ebbae2502 --- /dev/null +++ b/studio/backend/tests/test_sampling_resolution.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Effective sampling resolution: per-model recommendation + operator pins. + +Precedence per field: operator UNSLOTH_SAMPLING_* pin -> client explicit value -> +per-model recommendation (load_inference_config) -> static schema default. +""" + +import pytest + +from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES +from utils.inference import inference_config as ic + +_SCHEMA_DEFAULTS = { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.01, + "repetition_penalty": 1.0, + "presence_penalty": 0.0, +} + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + # The recommended lookup is lru-cached; clear it so a patched config takes effect. + ic._recommended_sampling.cache_clear() + for field in SAMPLING_FIELD_NAMES: + monkeypatch.delenv(ic._SAMPLING_FIELDS[field][0], raising = False) + yield + ic._recommended_sampling.cache_clear() + + +def _all_omitted(): + return {f: None for f in SAMPLING_FIELD_NAMES} + + +def _set_recommended(monkeypatch, mapping): + # _recommended_sampling sources from load_inference_config -- the exact block the Chat UI + # seeds from -- so patch that directly. Fields absent from `mapping` fall to schema defaults. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(mapping)) + ic._recommended_sampling.cache_clear() + + +def test_recommended_applies_when_client_omits(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 1.0 + assert eff["top_k"] == 64 + assert eff["min_p"] == 0.0 + # A field with no recommendation keeps the static schema default. + assert eff["top_p"] == 0.95 + + +def test_client_explicit_beats_recommended(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0}) + eff = resolve_effective_sampling("some/model", {**_all_omitted(), "temperature": 0.2}) + assert eff["temperature"] == 0.2 + + +def test_operator_pin_beats_client_and_recommended(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0}) + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + eff = resolve_effective_sampling("some/model", {**_all_omitted(), "temperature": 0.2}) + assert eff["temperature"] == 0.9 + + +def test_unknown_model_matches_ui_inference_block(monkeypatch): + # An unknown model gets the same values the Chat UI would seed (load_inference_config's + # default.yaml fallback: temp 0.7 / top_k -1), NOT the request schema defaults. + ui_block = { + "temperature": 0.7, + "top_p": 0.95, + "top_k": -1, + "min_p": 0.01, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, + } + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(ui_block)) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/unknown-model", _all_omitted()) + assert eff["temperature"] == 0.7 + assert eff["top_k"] == -1 + assert eff["min_p"] == 0.01 + + +def test_empty_recommendation_falls_back_to_schema_defaults(monkeypatch): + # If load_inference_config yields nothing usable, the resolver falls back to the request + # schema defaults. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: {}) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff == _SCHEMA_DEFAULTS + + +@pytest.mark.parametrize( + "model", + ["unsloth/gemma-4-E4B", "unsloth/Qwen3-4B", "unsloth/Qwen3.5-9B", "someorg/unknown-xyz"], +) +def test_recommendation_matches_ui_source(model): + # Parity guard: what the server recommends for omitted fields equals the Chat UI's source + # (load_inference_config) for every field the UI adopts (mergeBackendRecommendedInference). + ic._recommended_sampling.cache_clear() + ui = ic.load_inference_config(model) + rec = ic._recommended_sampling(model) + for f in ic._UI_RECOMMENDED_FIELDS: + cleaned = ic._clean_sampling_value(f, ui.get(f)) + if cleaned is not None: + assert rec.get(f) == cleaned, f"{model}:{f} rec={rec.get(f)} ui={ui.get(f)}" + + +def test_repetition_penalty_not_auto_recommended(monkeypatch): + # The Chat UI's mergeBackendRecommendedInference never adopts a backend repetition_penalty + # (e.g. lfm2's family value 1.05), so the server must not auto-apply one either. It stays at + # the schema default unless the client sends it or an operator pins it. + monkeypatch.setattr( + ic, "load_inference_config", lambda mid: {"temperature": 0.7, "repetition_penalty": 1.05} + ) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/lfm2-model", _all_omitted()) + assert eff["temperature"] == 0.7 # a UI-adopted field is recommended + assert eff["repetition_penalty"] == 1.0 # rep is NOT auto-recommended (matches the UI) + # An operator can still pin it explicitly. + monkeypatch.setenv("UNSLOTH_SAMPLING_REPETITION_PENALTY", "1.05") + eff2 = resolve_effective_sampling("some/lfm2-model", _all_omitted()) + assert eff2["repetition_penalty"] == 1.05 + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("0.5", 0.5), + ("abc", None), # unparseable + ("9.0", None), # above temperature max (2.0) + ("-1", None), # below temperature min (0.0) + (" ", None), # blank + ("nan", None), # NaN would pass a naive range check + ("inf", None), # non-finite + ("-inf", None), # non-finite + ], +) +def test_operator_override_parsing(monkeypatch, raw, expected): + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", raw) + assert ic._operator_sampling_override("temperature") == expected + + +def test_out_of_range_recommendation_is_dropped(monkeypatch): + # A malformed model recommendation (out of range) is ignored, so the request keeps the + # schema default rather than forwarding a bad value to llama-server. + _set_recommended(monkeypatch, {"temperature": 5.0, "top_k": 64}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 0.6 # 5.0 is outside [0, 2] -> schema default + assert eff["top_k"] == 64 # a valid recommendation is still applied + + +def test_operator_override_top_k_int_and_range(monkeypatch): + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "40") + assert ic._operator_sampling_override("top_k") == 40 + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "200") # above max 100 + assert ic._operator_sampling_override("top_k") is None + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "-1") # min allowed + assert ic._operator_sampling_override("top_k") == -1 + + +@pytest.mark.parametrize( + "field, val", + [ + ("top_k", 10**400), # oversized int on an int field: int() ok, but math.isfinite raises + ("top_k", float("nan")), # NaN reaching an int field: int(nan) raises ValueError + ("top_k", float("inf")), # inf reaching an int field: int(inf) raises OverflowError + ( + "temperature", + 10**400, + ), # oversized int on a float field: float(huge_int) raises OverflowError + ], +) +def test_clean_sampling_value_rejects_unrepresentable(field, val): + # None of these may raise; each is unusable and must be dropped to None (regression: an + # oversized value used to raise OverflowError before the range check could drop it). + assert ic._clean_sampling_value(field, val) is None + + +def test_oversized_operator_override_ignored(monkeypatch): + # A huge integer string parses via int() but overflows float(); math.isfinite would raise + # OverflowError and 500 the request. It must be ignored like any other bad override and the + # field must fall back to the schema default -- no exception. + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "9" * 400) + assert ic._operator_sampling_override("top_k") is None + _set_recommended(monkeypatch, {}) # no per-model recommendation -> schema default applies + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["top_k"] == 20 # schema default, resolved without raising + + +def test_oversized_recommendation_ignored(monkeypatch): + # A malformed per-model recommendation carrying an oversized int must not raise while + # resolving either; the field simply falls back to the schema default. + _set_recommended(monkeypatch, {"temperature": 10**400, "top_k": 64}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 0.6 # oversized -> dropped -> schema default + assert eff["top_k"] == 64 # a valid recommendation is still applied + + +def test_fill_recommended_sampling_openai_payload(monkeypatch): + from models.inference import ChatCompletionRequest + from routes.inference import _fill_recommended_sampling_openai + + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + + # Client sent only temperature; top_k / min_p were omitted. + payload = ChatCompletionRequest( + model = "m", messages = [{"role": "user", "content": "hi"}], temperature = 0.2 + ) + _fill_recommended_sampling_openai(payload, "some/model") + assert payload.temperature == 0.2 # explicit client value preserved + assert payload.top_k == 64 # recommended fills the omitted field + assert payload.min_p == 0.0 + assert payload.top_p == 0.95 # no recommendation -> schema default unchanged + + +def test_fill_recommended_sampling_openai_operator_pin_overrides_client(monkeypatch): + from models.inference import ChatCompletionRequest + from routes.inference import _fill_recommended_sampling_openai + + monkeypatch.setattr(ic, "load_model_defaults", lambda mid: {}) + monkeypatch.setattr(ic, "get_family_inference_params", lambda mid: {}) + ic._recommended_sampling.cache_clear() + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + + payload = ChatCompletionRequest( + model = "m", messages = [{"role": "user", "content": "hi"}], temperature = 0.2 + ) + _fill_recommended_sampling_openai(payload, "some/model") + assert payload.temperature == 0.9 # operator pin wins even over an explicit client value + + +def test_fill_recommended_sampling_completions_body(monkeypatch): + # /v1/completions is a raw proxy: recommendations fill omitted fields, but a field with no + # recommendation and no pin is left absent so llama-server keeps its own default (unlike the + # chat schema, which carries per-field defaults). + from routes.inference import _fill_recommended_sampling_completions + + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + + body = {"prompt": "hi", "temperature": 0.2} + _fill_recommended_sampling_completions(body, "some/model") + assert body["temperature"] == 0.2 # explicit client value preserved + assert body["top_k"] == 64 # recommendation fills the omitted field + assert body["min_p"] == 0.0 + # No recommendation and no pin -> NOT injected (llama-server keeps its default). + assert "top_p" not in body + assert "presence_penalty" not in body + assert "repeat_penalty" not in body + + +def test_fill_recommended_sampling_completions_operator_pin(monkeypatch): + # An operator pin overrides the client's raw-body value, and the repetition pin is written + # under llama-server's "repeat_penalty" key (the schema field is repetition_penalty). + from routes.inference import _fill_recommended_sampling_completions + + monkeypatch.setattr(ic, "load_inference_config", lambda mid: {}) + ic._recommended_sampling.cache_clear() + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + monkeypatch.setenv("UNSLOTH_SAMPLING_REPETITION_PENALTY", "1.2") + + body = {"prompt": "hi", "temperature": 0.2, "repeat_penalty": 1.05} + _fill_recommended_sampling_completions(body, "some/model") + assert body["temperature"] == 0.9 # operator pin wins over the client's explicit value + assert body["repeat_penalty"] == 1.2 # repetition pin lands on llama-server's key + assert "repetition_penalty" not in body # never leak the schema field name into the body diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py index 2970b1a6bb..64201477e3 100644 --- a/studio/backend/tests/test_sandbox_tools.py +++ b/studio/backend/tests/test_sandbox_tools.py @@ -297,6 +297,8 @@ class TestSandboxEnvIsolation: "PYTHONPATH", "VIRTUAL_ENV", "SystemRoot", + "PATHEXT", # Windows only; minimal list so cwd scripts cannot hijack + "NoDefaultCurrentDirectoryInExePath", # Windows only; no cwd-first lookup } extras = set(env.keys()) - allowed assert not extras, f"sandbox env added unexpected keys: {extras}" @@ -305,6 +307,220 @@ class TestSandboxEnvIsolation: assert env["PYTHONPATH"].endswith("sandbox_site") assert "leak-me" not in env["PYTHONPATH"] + def test_host_git_dir_appended_after_curated(self, monkeypatch, tmp_path): + # #7317: Windows Git lives under Program Files, not System32. Sandbox + # PATH resolves bare `git` by appending the dir of the git the HOST + # shell resolves (shutil.which), after the curated prefix. + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + git_dir = prog / "Git" / "cmd" + git_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_dir / "git.exe")) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(git_dir) in parts + # Curated prefix stays ahead of host Git so Studio python/pip win. + assert parts.index(str(git_dir)) > 0 + + def test_host_path_dirs_not_inherited(self, monkeypatch, tmp_path): + """Host PATH dirs (user-writable, git-lookalike) are never inherited; + only the resolved git dir is. No git resolved -> nothing appended.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + venv_scripts = tmp_path / "venv" / "Scripts" + venv_scripts.mkdir(parents = True) + fake_git = tmp_path / "scratch" / "Git" / "cmd" + fake_git.mkdir(parents = True) + monkeypatch.setenv( + "PATH", + os.pathsep.join([str(venv_scripts), str(fake_git), os.environ.get("PATH", "")]), + ) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: None) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(venv_scripts) not in parts + # A git-suffixed but unresolved (user-writable) dir is NOT trusted. + assert str(fake_git) not in parts + + def test_git_cmd_shim_extension_added_to_pathext(self, monkeypatch, tmp_path): + """A host git resolved as a .cmd shim under a trusted root stays + resolvable under the restricted PATHEXT (cwd lookup disabled).""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + git_dir = prog / "Git" / "cmd" + git_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_dir / "git.cmd")) + env = _build_safe_env(str(tmp_path)) + assert str(git_dir) in env["PATH"].split(os.pathsep) + assert env["PATHEXT"] == ".EXE;.COM;.CMD" + + def test_user_writable_git_dir_refused(self, monkeypatch, tmp_path): + """Git resolved from a per-user manager (Scoop shims) is NOT trusted: + an attacker could drop rg.exe beside it and hit the auto-approve gate.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr( + tools_mod, "_windows_program_roots", lambda: [str(tmp_path / "Program Files")] + ) + shim_dir = tmp_path / "users" / "alice" / "scoop" / "shims" + shim_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(shim_dir / "git.exe")) + env = _build_safe_env(str(tmp_path)) + assert str(shim_dir) not in env["PATH"].split(os.pathsep) + # No trusted git launcher -> PATHEXT stays minimal. + assert env["PATHEXT"] == ".EXE;.COM" + + def test_trust_uses_known_folder_not_env_override(self, monkeypatch, tmp_path): + """Trust is driven by the resolved Program Files roots, so a git under + an attacker-overridden %ProgramFiles% env value is still refused.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "RealProgramFiles" + (real_prog).mkdir() + evil = tmp_path / "attacker" + (evil / "Git" / "cmd").mkdir(parents = True) + # Resolver returns the genuine root; env is overridden to the evil dir. + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + monkeypatch.setenv("ProgramFiles", str(evil)) + monkeypatch.setattr( + tools_mod.shutil, "which", lambda name: str(evil / "Git" / "cmd" / "git.exe") + ) + env = _build_safe_env(str(tmp_path)) + assert str(evil / "Git" / "cmd") not in env["PATH"].split(os.pathsep) + + def test_canonical_git_dir_appended(self, monkeypatch, tmp_path): + """The PATH entry is the realpath of the trusted dir, not a junction + alias, so it cannot be retargeted after the trust check.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "Program Files" + real_git = real_prog / "Git" / "cmd" + real_git.mkdir(parents = True) + link = tmp_path / "link" + try: + link.symlink_to(real_prog, target_is_directory = True) + except (OSError, NotImplementedError): + pytest.skip("symlink unsupported in this environment") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + monkeypatch.setattr( + tools_mod.shutil, + "which", + lambda name: str(link / "Git" / "cmd" / "git.exe"), + ) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(real_git) in parts # canonical, not the `link/...` alias + + def test_windows_temp_git_dir_refused(self, monkeypatch, tmp_path): + """A git under a world-writable %SystemRoot% subdir (Windows\\Temp) is + NOT trusted, even though it sits under the Windows root.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr( + tools_mod, "_windows_program_roots", lambda: [str(tmp_path / "Program Files")] + ) + temp_git = tmp_path / "Windows" / "Temp" / "Git" / "cmd" + temp_git.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(temp_git / "git.exe")) + env = _build_safe_env(str(tmp_path)) + assert str(temp_git) not in env["PATH"].split(os.pathsep) + + def test_trusted_program_dir_matches_via_realpath(self, monkeypatch, tmp_path): + """The trust check canonicalizes paths, so a symlinked/short alias of + Program Files still matches (stand-in for 8.3 PROGRA~1 on Windows).""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "Program Files" + (real_prog / "Git" / "cmd").mkdir(parents = True) + alias = tmp_path / "PROGRA~1" + try: + alias.symlink_to(real_prog, target_is_directory = True) + except (OSError, NotImplementedError): + pytest.skip("symlink unsupported in this environment") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + git_via_alias = alias / "Git" / "cmd" / "git.exe" + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_via_alias)) + env = _build_safe_env(str(tmp_path)) + parts = [os.path.normcase(os.path.realpath(p)) for p in env["PATH"].split(os.pathsep)] + assert os.path.normcase(str(real_prog / "Git" / "cmd")) in parts + + def test_scan_past_untrusted_git_shim(self, monkeypatch, tmp_path): + """When an untrusted shim sorts first on PATH, the scan still finds a + later trusted Program Files git.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + trusted_git = prog / "Git" / "cmd" + trusted_git.mkdir(parents = True) + (trusted_git / "git.EXE").write_text("") # match PATHEXT case on this FS + shim = tmp_path / "scoop" / "shims" + shim.mkdir(parents = True) + (shim / "git.EXE").write_text("") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + # shutil.which returns the untrusted shim first. + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(shim / "git.EXE")) + monkeypatch.setenv("PATH", os.pathsep.join([str(shim), str(trusted_git)])) + monkeypatch.setenv("PATHEXT", ".EXE") + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(trusted_git) in parts + assert str(shim) not in parts + + def test_program_roots_fails_closed_without_known_folder_api(self, monkeypatch): + """When the known-folder API is unavailable, no roots are trusted: env + vars (even %SystemDrive%) are caller-overrideable, so we never derive a + trusted root from them.""" + import core.inference.tools as tools_mod + + # ctypes fails on this Linux host, so the API path raises and we fail + # closed. Any attacker override of these env vars must be irrelevant. + monkeypatch.setenv("ProgramFiles", r"D:\attacker-writable") + monkeypatch.setenv("ProgramW6432", r"D:\attacker-writable") + monkeypatch.setenv("SystemDrive", "D:") + assert tools_mod._windows_program_roots() == [] + + def test_augment_native_program_roots_derives_native_sibling(self): + """A 32-bit process only sees the x86 root; the native sibling is + derived by stripping the ` (x86)` suffix.""" + import core.inference.tools as tools_mod + + roots = tools_mod._augment_native_program_roots([r"C:\Program Files (x86)"]) + lowered = [r.lower() for r in roots] + assert r"c:\program files (x86)" in lowered + assert r"c:\program files" in lowered + + def test_no_default_current_directory_in_exe_path_set_on_windows(self, monkeypatch, tmp_path): + """cmd/CreateProcess must not search cwd for bare names in the sandbox.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: None) + env = _build_safe_env(str(tmp_path)) + assert env["NoDefaultCurrentDirectoryInExePath"] == "1" + def test_home_points_at_sandbox_workdir(self, tmp_path): from core.inference.tools import _build_safe_env diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index a8c0c2305f..b491134045 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -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 diff --git a/studio/backend/tests/test_server_disk_logging_outstream.py b/studio/backend/tests/test_server_disk_logging_outstream.py new file mode 100644 index 0000000000..0ff27666a0 --- /dev/null +++ b/studio/backend/tests/test_server_disk_logging_outstream.py @@ -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() diff --git a/studio/backend/tests/test_training_config_popover_source.py b/studio/backend/tests/test_training_config_popover_source.py index 4263b012eb..452a3a1ea8 100644 --- a/studio/backend/tests/test_training_config_popover_source.py +++ b/studio/backend/tests/test_training_config_popover_source.py @@ -105,5 +105,6 @@ def test_shared_mapper_matches_backend_config_keys(): "lora_dropout", "use_rslora", "use_loftq", + "use_dora", ): assert key in src, f"run-config mapper lost backend key {key}" diff --git a/studio/backend/tests/test_training_vram_coexistence.py b/studio/backend/tests/test_training_vram_coexistence.py index 6683cb9aaa..217caaa4fb 100644 --- a/studio/backend/tests/test_training_vram_coexistence.py +++ b/studio/backend/tests/test_training_vram_coexistence.py @@ -326,12 +326,21 @@ class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase): keep, _, _ = self._run((None, meta)) self.assertFalse(keep) - def test_unload_on_non_cuda(self): + def test_unload_on_non_accelerator(self): keep, info, auto_mock = self._run(([0], {}), device = DeviceType.CPU) self.assertFalse(keep) - self.assertEqual(info["mode"], "non_cuda") + self.assertEqual(info["mode"], "non_accelerator") auto_mock.assert_not_called() + def test_xpu_gets_sized_like_cuda(self): + # XPU is a first-class training backend: the keep-guard must size it, + # not blanket-unload it as a non-accelerator. + meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0} + keep, info, auto_mock = self._run(([0], meta), device = DeviceType.XPU) + self.assertTrue(keep) + self.assertNotEqual(info.get("mode"), "non_accelerator") + auto_mock.assert_called_once() + def test_full_finetuning_forces_16bit_in_estimate(self): meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0} _keep, _info, auto_mock = self._run( diff --git a/studio/backend/tests/test_web_fetch_extraction.py b/studio/backend/tests/test_web_fetch_extraction.py index b794ee3e81..d4c3d123c3 100644 --- a/studio/backend/tests/test_web_fetch_extraction.py +++ b/studio/backend/tests/test_web_fetch_extraction.py @@ -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( diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 62b537fbac..138238533f 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -50,6 +50,11 @@ def export_capability() -> dict: return _hardware.export_capability() +def get_torch_device_str() -> str: + """Return the torch device string ("cuda", "xpu", "cpu") for the detected hardware.""" + return _hardware.get_torch_device_str() + + __all__ = [ "DeviceType", "DEVICE", @@ -75,6 +80,7 @@ __all__ = [ "estimate_required_model_memory_gb", "auto_select_gpu_ids", "prepare_gpu_selection", + "get_torch_device_str", "safe_num_proc", "safe_thread_num_proc", "dataset_map_num_proc", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 3d312d4b01..38ebc0b6d4 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -175,18 +175,64 @@ def detect_hardware() -> DeviceType: Call once at FastAPI lifespan startup; idempotent. Detection order: - 1. CUDA (NVIDIA GPU, requires torch) - 2. MLX (Apple Silicon via MLX framework) - 3. CPU (fallback) + 1. XPU-preferred hint: only on an unambiguous "prefer XPU" signal + (CUDA hidden via ``CUDA_VISIBLE_DEVICES="" / "-1"``, + ``UNSLOTH_FORCE_XPU=1``, or CUDA unavailable) AND a non-empty + ``ZE_AFFINITY_MASK`` AND ``torch.xpu`` reports a device. A stray + inherited mask is not enough: CUDA still wins on hybrid hosts. + 2. CUDA (NVIDIA GPU, requires torch) + 3. XPU (Intel GPU, requires torch with XPU support) + 4. MLX (Apple Silicon via MLX framework) + 5. CPU (fallback) """ global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, IS_ROCM CHAT_ONLY = True # reset -- only CUDA/ROCm/XPU/MLX sets it to False CHAT_ONLY_REASON = None IS_ROCM = False - # --- CUDA / ROCm: try PyTorch --- + # --- CUDA / ROCm / XPU: try PyTorch --- if _has_torch(): import torch + + # --- Explicit-XPU hint --- + # Prefer XPU on UNSLOTH_FORCE_XPU=1, or ZE_AFFINITY_MASK set + CUDA + # hidden/unavailable. A bare mask alone is NOT enough (can leak from + # unrelated Intel tooling); torch.xpu must report a device. + ze_mask = os.environ.get("ZE_AFFINITY_MASK") + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + cuda_hidden = cvd is not None and cvd.strip() in ("", "-1") + force_xpu = os.environ.get("UNSLOTH_FORCE_XPU") == "1" + try: + cuda_unavailable = not torch.cuda.is_available() + except Exception: + cuda_unavailable = True + + prefer_xpu = force_xpu or (bool(ze_mask) and (cuda_hidden or cuda_unavailable)) + if prefer_xpu: + try: + xpu_ok = hasattr(torch, "xpu") and torch.xpu.is_available() + except Exception: + xpu_ok = False + if xpu_ok: + # Forced XPU on a hybrid host: unsloth's device_type picks + # CUDA before XPU and ignores this Studio-only env var, so + # hide CUDA or spawned workers would silently train on CUDA. + if force_xpu and not cuda_hidden and not cuda_unavailable: + os.environ["CUDA_VISIBLE_DEVICES"] = "" + DEVICE = DeviceType.XPU + CHAT_ONLY = False + CHAT_ONLY_REASON = None + device_name = torch.xpu.get_device_name(0) + if force_xpu and not ze_mask: + reason = "UNSLOTH_FORCE_XPU=1" + elif force_xpu: + reason = "UNSLOTH_FORCE_XPU=1 + ZE_AFFINITY_MASK" + else: + reason = "ZE_AFFINITY_MASK hint honoured" + print(f"Hardware detected: XPU -- {device_name} ({reason})") + return DEVICE + + # --- CUDA: NVIDIA GPU --- if torch.cuda.is_available(): DEVICE = DeviceType.CUDA CHAT_ONLY = False @@ -327,9 +373,18 @@ def clear_gpu_cache(): torch.cuda.empty_cache() torch.cuda.ipc_collect() elif device == DeviceType.XPU: - import torch - torch.xpu.synchronize() - torch.xpu.empty_cache() + # Guard synchronize/empty_cache: older torch-xpu builds may lack + # them, and an unguarded AttributeError would propagate to callers. + # torch.xpu has no ipc_collect(), so do not call it here. + try: + import torch + if hasattr(torch, "xpu"): + if hasattr(torch.xpu, "synchronize"): + torch.xpu.synchronize() + if hasattr(torch.xpu, "empty_cache"): + torch.xpu.empty_cache() + except Exception as e: + logger.debug("Failed to clear XPU cache: %s", e) elif device == DeviceType.MLX: # MLX manages memory automatically; gc.collect() above is enough. pass @@ -500,14 +555,27 @@ def get_package_versions() -> Dict[str, Optional[str]]: except PackageNotFoundError: versions[name] = None - # GPU runtime version bundled with torch + # GPU runtime versions bundled with torch (CUDA, ROCm/HIP, Intel XPU) try: import torch + versions["cuda"] = getattr(torch.version, "cuda", None) versions["rocm"] = getattr(torch.version, "hip", None) + # Isolated probe: a broken Intel runtime raising in is_available() + # must not blank the already-read cuda/rocm versions. + try: + if hasattr(torch, "xpu") and torch.xpu.is_available(): + # torch.version.xpu may be None on modern builds; fall back to + # "available" so the UI distinguishes present-but-unknown from + # "package not found". + xpu_ver = getattr(torch.version, "xpu", None) + versions["xpu"] = xpu_ver if xpu_ver is not None else "available" + except Exception: + versions["xpu"] = None except Exception: versions["cuda"] = None versions["rocm"] = None + versions["xpu"] = None return versions @@ -547,6 +615,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] if mod is None: return [] + device = get_device() # free==total is a Windows-ROCm-only quirk. _win_rocm = sys.platform == "win32" and IS_ROCM devices = [] @@ -558,11 +627,30 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] used_bytes: Optional[int] # Prefer mem_get_info (system-wide) so auto-select sees other consumers. if hasattr(mod, "mem_get_info"): - free_bytes, total_bytes = mod.mem_get_info(ordinal) - used_bytes = total_bytes - free_bytes - # free==total is the broken-API sentinel, not an idle GPU. - if _win_rocm and free_bytes == total_bytes: + try: + free_bytes, total_bytes = mod.mem_get_info(ordinal) + used_bytes = total_bytes - free_bytes + except Exception as e: + if device != DeviceType.XPU: + raise + # Arc B580 and Lunar Lake can report properties while + # rejecting free-memory queries. Preserve the usable + # device and its total memory with unknown utilization. + logger.debug( + "XPU free-memory query failed for ordinal %d: %s", + ordinal, + e, + ) used_bytes = None + else: + # free==total is the broken-API sentinel, not an idle GPU. + if _win_rocm and free_bytes == total_bytes: + used_bytes = None + elif device == DeviceType.XPU: + # XPU without mem_get_info: memory_allocated() is process-local + # and misleading for placement, so return None for the + # selector's no-telemetry fallback. + used_bytes = None else: used_bytes = mod.memory_allocated(ordinal) devices.append( @@ -571,7 +659,9 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] "visible_ordinal": ordinal, "name": props.name, "total_gb": round(total_bytes / (1024**3), 2), - "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None, + "used_gb": ( + round(used_bytes / (1024**3), 2) if used_bytes is not None else None + ), } ) except Exception as e: @@ -582,6 +672,43 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] # ========== Live GPU Utilization ========== +def _xpu_hierarchy_is_composite() -> bool: + """Return True iff Level Zero is running in COMPOSITE device hierarchy. + + COMPOSITE: numeric ``ZE_AFFINITY_MASK`` entries address root GPU IDs + (tiles use ``N.M``). FLAT (the oneAPI default; also assumed when + ``ZE_FLAT_DEVICE_HIERARCHY`` is unset): entries address tile/device + handles, so mapping them back to root GPU IDs is unsafe. Only COMPOSITE + gives stable root-ID semantics. + """ + hierarchy = (os.environ.get("ZE_FLAT_DEVICE_HIERARCHY") or "FLAT").strip().upper() + return hierarchy == "COMPOSITE" + + +def _parse_ze_mask_roots(mask: str) -> list[int]: + """Parse a ``ZE_AFFINITY_MASK`` value into an ordered list of root device IDs. + + One root ID per mask token, preserving order and duplicates so logical + ordinals map 1-to-1 to physical root IDs (e.g. ``"0.0,0.1"`` -> ``[0, 0]``, + ``"2.0,0.1,0.2"`` -> ``[2, 0, 0]``); empty list if no parseable digits. + Only meaningful in COMPOSITE hierarchy -- callers needing a stable + root-ID mapping must gate on ``_xpu_hierarchy_is_composite()``. + """ + roots: list[int] = [] + if not mask: + return roots + for token in mask.split(","): + token = token.strip() + if not token: + continue + root = token.split(".", 1)[0] + # isdecimal() (not isdigit()) rejects Unicode superscripts like + # "²"/"³", which pass isdigit() but crash int() with ValueError. + if root.isdecimal(): + roots.append(int(root)) + return roots + + def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: """Query the appropriate SMI backend (amd-smi or nvidia-smi). @@ -1504,6 +1631,13 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: for td in torch_devices: total = td["total_gb"] used = td["used_gb"] + # used=None is a deliberate "telemetry unavailable" signal + # from _torch_get_per_device_info (e.g. XPU without + # mem_get_info); propagate None instead of dividing by it. On + # CUDA/ROCm used is always an int, so this stays byte-identical. + vram_pct = ( + round((used / total) * 100, 1) if used is not None and total > 0 else None + ) devices.append( { "index": td["index"], @@ -1513,9 +1647,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: "temperature_c": None, "vram_used_gb": used, "vram_total_gb": total, - "vram_utilization_pct": round((used / total) * 100, 1) - if total > 0 and used is not None - else None, + "vram_utilization_pct": vram_pct, "power_draw_w": None, "power_limit_w": None, "power_utilization_pct": None, @@ -1583,6 +1715,82 @@ _visible_gpu_count: Optional[int] = None def _get_parent_visible_gpu_spec() -> Dict[str, Any]: + # On Intel XPU, visibility is controlled by ZE_AFFINITY_MASK (Level Zero), + # not CUDA_VISIBLE_DEVICES. + if get_device() == DeviceType.XPU: + xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK") + composite = _xpu_hierarchy_is_composite() + + if xpu_mask_raw is None: + # COMPOSITE: root GPU IDs are stable physical IDs. + if composite: + return { + "raw": None, + "numeric_ids": list(range(get_physical_gpu_count())), + "supports_explicit_gpu_ids": True, + } + # FLAT (oneAPI default): ordinals are tile/device handles, not + # physical GPU IDs. numeric_ids=None so telemetry uses relative + # ordinals; explicit selection needs ZE_FLAT_DEVICE_HIERARCHY=COMPOSITE. + return { + "raw": None, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + xpu_mask = xpu_mask_raw.strip() + if xpu_mask == "": + return { + "raw": xpu_mask, + "numeric_ids": [], + "supports_explicit_gpu_ids": True, + } + + # Subdevice syntax ("N.M") expands one root into multiple + # logical devices -- not addressable by explicit root-ID selection. + has_subdevice = any("." in token.strip() for token in xpu_mask.split(",") if token.strip()) + if has_subdevice: + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + # FLAT numeric entries are tile handles, not physical GPU IDs. Keep + # numeric_ids unresolved so every telemetry and picker consumer uses + # relative torch ordinals and cannot advertise them as pinnable roots. + if not composite: + tokens = [token.strip() for token in xpu_mask.split(",") if token.strip()] + if tokens and all(token.isdecimal() for token in tokens): + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + # COMPOSITE + pure numeric (subdevice handled above). _parse_ze_mask_roots + # maps to root GPU IDs, dropping non-decimal tokens so "*"/"GPU-uuid" -> []. + roots_with_dupes = _parse_ze_mask_roots(xpu_mask) + if not roots_with_dupes: + # Unparseable mask (e.g. "*", "GPU-uuid") -- cannot map to + # physical root IDs. + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + return { + "raw": xpu_mask, + "numeric_ids": roots_with_dupes, + "supports_explicit_gpu_ids": True, + } + # ROCm uses HIP/ROCR_VISIBLE_DEVICES on top of CUDA_VISIBLE_DEVICES; check # them first. Explicit None checks (not `or`) so "" reads as "no visible GPUs". cuda_visible = None @@ -1639,24 +1847,44 @@ def get_parent_visible_gpu_ids() -> list[int]: return list(parent_visible_ids) if parent_visible_ids is not None else [] -def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: +def resolve_requested_gpu_ids( + gpu_ids: Optional[list[int]], *, is_vulkan: bool = False +) -> list[int]: parent_visible_spec = _get_parent_visible_gpu_spec() parent_visible_ids = get_parent_visible_gpu_ids() physical_gpu_count = get_physical_gpu_count() if gpu_ids is None: - return parent_visible_ids + return [] if is_vulkan else parent_visible_ids requested_ids = list(gpu_ids) if len(requested_ids) == 0: - return parent_visible_ids + return [] if is_vulkan else parent_visible_ids + + if is_vulkan: + # A Vulkan build selects by ggml Vulkan ordinal (--device VulkanN), a separate + # index space from CUDA/ROCm ids that may be empty under CPU-only torch. The + # CUDA parent-visible / physical-count checks below do not apply; only reject + # malformed ordinals (issue #7239). + if len(set(requested_ids)) != len(requested_ids): + raise ValueError(f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed.") + negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0] + if negative_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: GPU IDs must be non-negative. " + f"Rejected IDs: {negative_ids}." + ) + return requested_ids if not parent_visible_spec["supports_explicit_gpu_ids"]: + env_var_name = ( + "ZE_AFFINITY_MASK" if get_device() == DeviceType.XPU else "CUDA_VISIBLE_DEVICES" + ) raise ValueError( f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are " - f"unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries " - f"({parent_visible_spec['raw']!r}). Omit gpu_ids to use the " - "parent-visible devices." + f"unsupported when {env_var_name} uses non-numeric or subdevice " + f"entries ({parent_visible_spec['raw']!r}). Omit gpu_ids to use " + "the parent-visible devices." ) if len(set(requested_ids)) != len(requested_ids): @@ -2095,8 +2323,11 @@ def auto_select_gpu_ids( ) -> tuple[Optional[list[int]], Dict[str, Any]]: metadata: Dict[str, Any] = {"selection_mode": "auto"} - if get_device() != DeviceType.CUDA: - metadata["selection_mode"] = "non_cuda" + # Auto-selection needs per-device free-VRAM telemetry, available on CUDA + # (nvidia-smi) and XPU (torch.xpu) but not MLX/CPU, which fall + # through to inheriting parent visibility. + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + metadata["selection_mode"] = "non_accelerator" return None, metadata required_gb, estimate_metadata = estimate_required_model_memory_gb( @@ -2193,12 +2424,13 @@ def auto_select_gpu_ids( metadata["selection_mode"] = "auto" metadata["selected_gpu_ids"] = selected logger.debug( - "Selected GPUs automatically", - model_name = model_name, - selected_gpu_ids = selected, - usable_gb = metadata["usable_gb"], - required_gb = metadata.get("required_gb"), - multi_gpu_overhead = multi_gpu_overhead, + "Selected GPUs automatically: model=%s selected=%s usable_gb=%s " + "required_gb=%s multi_gpu_overhead=%s", + model_name, + selected, + metadata["usable_gb"], + metadata.get("required_gb"), + multi_gpu_overhead, ) return selected, metadata @@ -2214,12 +2446,13 @@ def auto_select_gpu_ids( metadata["usable_gb"] = round(fallback_usable, 3) metadata["selected_gpu_ids"] = fallback_all logger.warning( - "Falling back to all visible GPUs -- model may not fit", - model_name = model_name, - selected_gpu_ids = fallback_all, - usable_gb = metadata["usable_gb"], - required_gb = metadata.get("required_gb"), - multi_gpu_overhead = multi_gpu_overhead, + "Falling back to all visible GPUs; model may not fit: model=%s " + "selected=%s usable_gb=%s required_gb=%s multi_gpu_overhead=%s", + model_name, + fallback_all, + metadata["usable_gb"], + metadata.get("required_gb"), + multi_gpu_overhead, ) return fallback_all, metadata @@ -2253,10 +2486,10 @@ def prepare_gpu_selection( to a Hugging Face ``device_map`` string) and to ``apply_gpu_ids()`` in the worker subprocess (narrows ``CUDA_VISIBLE_DEVICES`` before torch/CUDA init). """ - if gpu_ids and get_device() != DeviceType.CUDA: + if gpu_ids and get_device() not in (DeviceType.CUDA, DeviceType.XPU): raise ValueError( - f"gpu_ids {list(gpu_ids)} is only supported on CUDA devices, " - f"but the current backend is '{get_device().value}'." + f"gpu_ids {list(gpu_ids)} is only supported on CUDA and Intel XPU " + f"devices, but the current backend is '{get_device().value}'." ) if gpu_ids: @@ -2329,11 +2562,14 @@ def get_physical_gpu_count() -> int: def _backend_visible_devices_env() -> Optional[str]: """Return the raw visibility env string that applies to this backend. - On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over - CUDA_VISIBLE_DEVICES; this mirrors ``_get_parent_visible_gpu_spec`` so + On XPU the control is ``ZE_AFFINITY_MASK`` (not ``CUDA_VISIBLE_DEVICES``); + on ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over + CUDA_VISIBLE_DEVICES. Mirrors ``_get_parent_visible_gpu_spec`` so ``backend_cuda_visible_devices`` reports the value actually narrowing the - visible device set. + visible device set on the current backend. """ + if get_device() == DeviceType.XPU: + return os.environ.get("ZE_AFFINITY_MASK") if IS_ROCM: return _get_parent_visible_gpu_spec().get("raw") return os.environ.get("CUDA_VISIBLE_DEVICES") @@ -2448,6 +2684,43 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count + # Prefer torch.xpu.device_count() on Intel XPU: the Level Zero runtime + # correctly interprets ZE_AFFINITY_MASK semantics (e.g. subdevice syntax + # "0.0,0.1" collapses onto one root GPU). Supersedes the torch fallback below. + if get_device() == DeviceType.XPU: + xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK") + xpu_mask_set = xpu_mask_raw is not None + xpu_visible = (xpu_mask_raw or "").strip() + if xpu_mask_set and xpu_visible == "": + _visible_gpu_count = 0 + return _visible_gpu_count + + try: + import torch + _visible_gpu_count = torch.xpu.device_count() + except Exception as e: + logger.debug( + "torch.xpu.device_count() failed, falling back to mask parsing: %s", + e, + ) + if xpu_visible: + # Fallback: count unique root device IDs from the mask. + # "device.subdevice" notation means "0.0,0.1" is 1 root, not 2. + # Without torch the hierarchy mode is unknown, so root-device + # counting is the conservative choice. + if xpu_visible == "*": + # Documented wildcard: all physical XPUs visible. + _visible_gpu_count = get_physical_gpu_count() + else: + roots = _parse_ze_mask_roots(xpu_visible) + # Non-parseable masks (",,,", "GPU-abc") yield an empty + # roots list, treated as 0 visible devices, not "all + # visible" -- no evidence the whole fleet was intended. + _visible_gpu_count = len(set(roots)) + else: + _visible_gpu_count = get_physical_gpu_count() + return _visible_gpu_count + # _get_parent_visible_gpu_spec() already handles HIP_VISIBLE_DEVICES / # ROCR_VISIBLE_DEVICES on ROCm. visible_spec = _get_parent_visible_gpu_spec() @@ -2461,20 +2734,18 @@ def get_visible_gpu_count() -> int: _visible_gpu_count = len([x for x in raw.split(",") if x.strip()]) return _visible_gpu_count - # No visibility env var set -- try torch, else physical count + # No visibility env var set -- try torch, else physical count. XPU is + # handled by the early return above, so only torch.cuda is needed here. try: import torch - if get_device() == DeviceType.XPU and hasattr(torch, "xpu"): - _visible_gpu_count = torch.xpu.device_count() - else: - _visible_gpu_count = torch.cuda.device_count() + _visible_gpu_count = torch.cuda.device_count() except Exception: _visible_gpu_count = get_physical_gpu_count() return _visible_gpu_count -def apply_gpu_ids(gpu_ids) -> None: +def apply_gpu_ids(gpu_ids, backend: Optional[str] = None) -> None: if gpu_ids is None: return @@ -2490,6 +2761,62 @@ def apply_gpu_ids(gpu_ids) -> None: else: value = str(gpu_ids) + # Intel XPU honors ZE_AFFINITY_MASK, not CUDA_VISIBLE_DEVICES; route XPU + # pinning through it so worker subprocesses are restricted to the intended GPU. + # Decide WITHOUT get_device(): workers call this before detect_hardware(), + # and a lazy detect would probe torch.cuda against the unmasked parent env, + # latching device enumeration before the mask below is written. Pre-detect, + # use env + torch BUILD attributes only (no runtime init, like the ROCm + # mirror below). + _is_xpu = DEVICE == DeviceType.XPU + if backend is not None: + # The spawning parent's detected backend (config["device_backend"]): + # exact and probe-free, so the mask target always matches what + # detect_hardware() decided in the parent, including its XPU + # availability check and CUDA fallback. + _is_xpu = backend == DeviceType.XPU.value + elif DEVICE is None: + # No parent backend passed (direct caller). version.xpu can be None + # on a working XPU build, so also accept torch.xpu._is_compiled() + # (a pure symbol-presence check, no runtime init). UNSLOTH_FORCE_XPU + # counts only on an XPU-capable build: detect_hardware() falls back + # to CUDA when XPU is missing, and the mask target must follow. + try: + import torch as _torch + + _ver = _torch.version + _is_comp = getattr(getattr(_torch, "xpu", None), "_is_compiled", None) + _xpu_build = (callable(_is_comp) and bool(_is_comp())) or ( + getattr(_ver, "xpu", None) is not None + ) + if os.environ.get("UNSLOTH_FORCE_XPU") == "1": + _is_xpu = _xpu_build + else: + # Mirror detect_hardware: hidden CUDA prefers XPU on an + # XPU-capable build (with or without a ZE mask -- detection + # falls through to XPU either way), where writing these ids + # to CUDA_VISIBLE_DEVICES would re-expose the deliberately + # hidden CUDA. + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + _cuda_hidden = _cvd is not None and _cvd.strip() in ("", "-1") + _is_xpu = _xpu_build and ( + _cuda_hidden + or (getattr(_ver, "cuda", None) is None and getattr(_ver, "hip", None) is None) + ) + except Exception as e: + logger.debug( + "apply_gpu_ids: torch XPU probe skipped (%s: %s)", + type(e).__name__, + e, + ) + if _is_xpu: + os.environ["ZE_AFFINITY_MASK"] = value + # Leave inherited CUDA_VISIBLE_DEVICES alone -- clearing it could let + # the worker flip back to CUDA on hybrid hosts. + _visible_gpu_count = None + logger.info("Applied gpu_ids: ZE_AFFINITY_MASK='%s'", value) + return + os.environ["CUDA_VISIBLE_DEVICES"] = value # Keep ROCm visibility env vars in sync. Workers may call apply_gpu_ids() # before detect_hardware() (IS_ROCM still False), so also mirror when the @@ -2534,26 +2861,41 @@ def get_device_map(gpu_ids: Optional[list[int]] = None) -> str: Returns ``"balanced"`` (shard evenly across GPUs) when: - ``gpu_ids`` explicitly lists >1 GPU, **or** - - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and - >1 GPU is visible (fallback: numeric IDs unresolvable, so assume - multi-GPU is intended). + - ``CUDA_VISIBLE_DEVICES``/``ZE_AFFINITY_MASK`` uses non-numeric + identifiers (UUID/MIG/wildcard) and >1 GPU is visible (fallback: + numeric IDs unresolvable, so assume multi-GPU is intended). - Returns ``"sequential"`` (single device) otherwise, including non-CUDA - backends (CPU, MLX). + Returns ``"sequential"`` (single device) otherwise, including CPU/MLX + backends. Use ``prepare_gpu_selection()`` upstream to determine ``gpu_ids`` -- it handles auto-selecting the minimum GPUs needed for a model. """ device = get_device() - if device == DeviceType.CUDA: + if device in (DeviceType.CUDA, DeviceType.XPU): multi_gpu = gpu_ids is not None and len(gpu_ids) > 1 if not multi_gpu: - # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU - # means multi-GPU sharding is intended. parent_visible_spec = _get_parent_visible_gpu_spec() - if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: - multi_gpu = True + if device == DeviceType.CUDA: + # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU + # means multi-GPU sharding is intended. + if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: + multi_gpu = True + elif device == DeviceType.XPU and gpu_ids is None: + # Shard across visible XPU ordinals via HF (no mask rewrite), + # only when no gpu_ids were passed -- an explicit gpu_ids=[0] + # means "use exactly device 0" and must stay sequential. + supports_physical = parent_visible_spec["supports_explicit_gpu_ids"] + has_multiple_numeric = ( + parent_visible_spec["numeric_ids"] is not None + and len(parent_visible_spec["numeric_ids"]) > 1 + ) + has_multiple_unresolved = ( + parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1 + ) + if has_multiple_unresolved or (not supports_physical and has_multiple_numeric): + multi_gpu = True if multi_gpu: return "balanced" @@ -2588,6 +2930,19 @@ def raise_if_offloaded( ) +def get_torch_device_str() -> str: + """ + Return the torch device string for the detected hardware. + E.g. "cuda", "xpu", or "cpu". + """ + device = get_device() + if device == DeviceType.CUDA: + return "cuda" + elif device == DeviceType.XPU: + return "xpu" + return "cpu" + + def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. @@ -2655,7 +3010,32 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: Returns ``None`` on spawn platforms (Windows, macOS) because ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``); only ``num_proc=None`` guarantees in-process execution. + + Also returns ``None`` on XPU once its runtime is initialized in this + process: ``os.fork()`` corrupts the Level-Zero context, making Triton + kernels fail with "Pointer argument doesn't reference XPU device memory". + Pre-init XPU hosts can still parallelize CPU-side preprocessing. """ if sys.platform in ("win32", "darwin"): return None + + if get_device() == DeviceType.XPU: + try: + import torch + except Exception: + # No torch means no active XPU runtime, so CPU-side dataset + # parallelism is still safe. + return safe_num_proc(desired) + + xpu = getattr(torch, "xpu", None) + is_initialized = getattr(xpu, "is_initialized", None) + if callable(is_initialized): + try: + if is_initialized(): + return None + except Exception as e: + # Treat a failing probe as "runtime not touched yet" so + # pre-init CPU preprocessing can still parallelize. + logger.debug("torch.xpu.is_initialized() probe failed: %s", e) + return safe_num_proc(desired) diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 05eb08067c..a264e06c85 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -5,7 +5,10 @@ from pathlib import Path from typing import Dict, Any, Optional +from functools import lru_cache import json +import math +import os import yaml import structlog from loggers import get_logger @@ -160,3 +163,137 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]: } return inference_config + + +# ── Effective sampling resolution for `unsloth run` / `unsloth start` ────────── +# +# Per-model recommended sampling is applied to a request only for the fields the +# client omitted; an operator can pin a field from the CLI via UNSLOTH_SAMPLING_* +# (a hard override that wins even over an explicit client value). Precedence per +# field: operator pin -> client explicit -> per-model recommendation -> the static +# schema default (mirroring ChatCompletionRequest, so behavior is unchanged when +# nothing is recommended or pinned). + +# field -> (env var, static default, min, max, is_int) +_SAMPLING_FIELDS = { + "temperature": ("UNSLOTH_SAMPLING_TEMPERATURE", 0.6, 0.0, 2.0, False), + "top_p": ("UNSLOTH_SAMPLING_TOP_P", 0.95, 0.0, 1.0, False), + "top_k": ("UNSLOTH_SAMPLING_TOP_K", 20, -1, 100, True), + "min_p": ("UNSLOTH_SAMPLING_MIN_P", 0.01, 0.0, 1.0, False), + "repetition_penalty": ("UNSLOTH_SAMPLING_REPETITION_PENALTY", 1.0, 1.0, 2.0, False), + "presence_penalty": ("UNSLOTH_SAMPLING_PRESENCE_PENALTY", 0.0, 0.0, 2.0, False), +} + +# Public, ordered tuple of the sampling fields callers resolve. +SAMPLING_FIELD_NAMES = tuple(_SAMPLING_FIELDS) + +# Fields the Studio Chat UI adopts as *per-model recommendations* from the backend +# `.inference` block. Its frontend `mergeBackendRecommendedInference` +# (presets/preset-policy.ts) seeds exactly these five and never reads repetition_penalty, +# so the server auto-recommends the same five for request parity. repetition_penalty stays a +# manual-only knob (client-sent or an UNSLOTH_SAMPLING_REPETITION_PENALTY operator pin), +# matching the UI where it is never auto-filled per model. +_UI_RECOMMENDED_FIELDS = ("temperature", "top_p", "top_k", "min_p", "presence_penalty") + + +def _clean_sampling_value(field: str, val: Any): + """Coerce ``val`` to the field's numeric type when it is a finite, in-range number, else None. + + Rejects bool, non-numeric, NaN/inf, and out-of-range values so neither a bad operator env + var nor a malformed model recommendation can reach llama-server. NaN matters because + ``nan < lo`` and ``nan > hi`` are both False, so a plain range check would let it through. + Coerce before the finiteness check: ``math.isfinite`` and ``float()`` raise ``OverflowError`` + on an int too big for a C double (an oversized UNSLOTH_SAMPLING_TOP_K would otherwise 500 the + request), while an in-range int is range-checked exactly and ``int()`` rejects a NaN/inf that + reached an int field. + """ + if isinstance(val, bool) or not isinstance(val, (int, float)): + return None + _env, _default, lo, hi, is_int = _SAMPLING_FIELDS[field] + try: + val = int(val) if is_int else float(val) + except (ValueError, OverflowError): + # int(nan)/int(inf) and float(oversized_int) raise; treat them as unusable. + return None + # After coercion an int is always finite; only a float can still be NaN/inf. + if isinstance(val, float) and not math.isfinite(val): + return None + if val < lo or val > hi: + return None + return val + + +def _operator_sampling_override(field: str): + """Operator-pinned value for a sampling field from UNSLOTH_SAMPLING_*, or None. + + An unparseable, non-finite, or out-of-range value is ignored so a bad env var can never + reach llama-server; the field then falls back to the client / recommended value. + """ + _env, _default, _lo, _hi, is_int = _SAMPLING_FIELDS[field] + raw = os.environ.get(_env) + if raw is None or raw.strip() == "": + return None + try: + val = int(raw) if is_int else float(raw) + except (TypeError, ValueError): + return None + return _clean_sampling_value(field, val) + + +@lru_cache(maxsize = 128) +def _recommended_sampling(model_id: str) -> Dict[str, Any]: + """Per-model recommended sampling, resolved through the SAME path the Studio Chat UI uses. + + The Chat UI seeds its sampling from the ``.inference`` block of the load/status responses, + which is exactly :func:`load_inference_config` (model-specific YAML -> family defaults + (inference_defaults.json) -> default.yaml). Sourcing recommendations here keeps the values + the server applies to a request identical to what the UI shows for the same model. Only the + fields the UI actually adopts (:data:`_UI_RECOMMENDED_FIELDS`) are recommended; each value + is validated (finite + in range) before use. Cached by model id. + """ + if not model_id: + return {} + try: + cfg = load_inference_config(model_id) or {} + except Exception as e: + logger.debug(f"Could not load recommended sampling for '{model_id}': {e}") + return {} + recommended: Dict[str, Any] = {} + for field in _UI_RECOMMENDED_FIELDS: + cleaned = _clean_sampling_value(field, cfg.get(field)) + if cleaned is not None: + recommended[field] = cleaned + return recommended + + +def resolve_effective_sampling( + model_id: Optional[str], + explicit: Dict[str, Any], + *, + fill_defaults: bool = True, +) -> Dict[str, Any]: + """Resolve the effective sampling params for a request. + + ``explicit`` maps each field in :data:`SAMPLING_FIELD_NAMES` to the client-sent + value, or ``None`` when the client omitted it. Precedence (highest first): an + operator ``UNSLOTH_SAMPLING_*`` pin, then the client's explicit value, then the + per-model recommendation, then the static schema default. + + When ``fill_defaults`` is False a field with no operator pin, client value, or + per-model recommendation is omitted from the result instead of set to the static + schema default, so a raw proxy body (``/v1/completions``) keeps llama-server's own + default for that field rather than being forced onto this schema's value. + """ + recommended = _recommended_sampling(model_id or "") + effective: Dict[str, Any] = {} + for field, (_env, default, _lo, _hi, _int) in _SAMPLING_FIELDS.items(): + override = _operator_sampling_override(field) + if override is not None: + effective[field] = override + elif explicit.get(field) is not None: + effective[field] = explicit[field] + elif field in recommended: + effective[field] = recommended[field] + elif fill_defaults: + effective[field] = default + return effective diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 91d7ad8f0e..3e12c15096 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -46,17 +46,6 @@ _PICKLE_WEIGHT_RE = re.compile( r"\.(bin|pt|pth|ckpt|pkl|pickle)$", re.IGNORECASE, ) -# Base-model safetensors set: HF names the base pickle pytorch_model.bin but the safetensors -# model.safetensors (stems differ), so a base pickle is replaced only by these, not an adapter's. -_BASE_SAFETENSORS_RE = re.compile( - r"^(model(-\d+-of-\d+)?\.safetensors|model\.safetensors\.index\.json)$", - re.IGNORECASE, -) -# Adapter (PEFT) safetensors set: adapter_model.safetensors, its shards, or index. -_ADAPTER_SAFETENSORS_RE = re.compile( - r"^(adapter_model(-\d+-of-\d+)?\.safetensors|adapter_model\.safetensors\.index\.json)$", - re.IGNORECASE, -) # Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/ # malicious or a future label) blocks, so Hub schema drift fails CLOSED. @@ -94,6 +83,13 @@ _INERT_SUFFIXES = frozenset( _SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"}) +# Torch-family weight indexes: from_pretrained feeds each shard they name to load_state_dict, which +# torch.load()s (pickle) any shard whose name does not end in .safetensors, whatever its stem. A +# pytorch index is superseded when a base safetensors is present (the loader prefers it); a +# safetensors index IS the chosen archive, so a non-safetensors target it names still loads. tf/flax +# indexes load via non-pickle loaders, so they are not a torch.load vector here. +_TORCH_INDEX_FILES = ("pytorch_model.bin.index.json", "model.safetensors.index.json") + # Root weight-index files. from_pretrained reads these to find sharded weights, so a # flagged subdir pickle is a load vector iff a root index references it. _TRANSFORMERS_INDEX_FILES = ( @@ -118,6 +114,28 @@ def _file_suffix(path: str) -> str: return "." + base.rsplit(".", 1)[1].lower() if "." in base else "" +def _hf_cache_snapshot_ref(local_path: str) -> Optional[tuple]: + """``(repo_id, revision)`` for an HF-cache snapshot path, else None. An inactive Studio + cache loads by its snapshot path but keeps the ``models--org--repo/snapshots/`` + layout, so the gate recovers its provenance and scans that exact commit instead of + exempting it (an older cached commit can hold a pickle since dropped from the branch).""" + try: + path = Path(local_path).resolve(strict = False) + except (OSError, ValueError): + return None + for parent in path.parents: + if parent.name != "snapshots": + continue + encoded = parent.parent.name + if not encoded.startswith("models--"): + return None + repo_id = encoded.removeprefix("models--").replace("--", "/") + if not repo_id: + return None + return repo_id, path.relative_to(parent).parts[0] # dir under snapshots/ + return None + + def _load_relative_path(norm: str, load_subdirs) -> str: """``norm`` relative to a ``from_pretrained`` load root. Some loads read from a snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``/LLM``), where a file @@ -145,13 +163,14 @@ def _indexed_shard_paths( model_name: str, hf_token: Optional[str], load_subdirs = (), + revision: Optional[str] = None, ): """Repo-relative weight paths a load could fetch via weight-index files. Returns a set (empty when the repo ships no index files -- a definitive "nothing sharded"), or None when the lookup was inconclusive (transient error) so the caller treats a flagged subdir pickle conservatively. Reads only small JSON indexes, never weights. Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map`` - entries re-prefixed to repo-relative paths. + entries re-prefixed to repo-relative paths. ``revision`` scopes to a cached commit. """ import json @@ -170,6 +189,7 @@ def _indexed_shard_paths( index_path = hf_hub_download( model_name, prefix + filename, + revision = revision, token = hf_token or None, cache_dir = active_hf_hub_cache(), ) @@ -264,9 +284,14 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: return model_name, load_subdirs -def _fetch_security_status(model_name: str, hf_token: Optional[str]): +def _fetch_security_status( + model_name: str, + hf_token: Optional[str], + revision: Optional[str] = None, +): """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; retries once on a transient error, then returns None so the caller fails open. + ``revision`` scopes the scan to a specific cached commit (else the default branch). """ from huggingface_hub import model_info as hf_model_info @@ -276,6 +301,7 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): try: info = hf_model_info( model_name, + revision = revision, token = token_arg, securityStatus = True, timeout = timeout, @@ -313,13 +339,72 @@ def _st_load_roots(snapshot: Path) -> list: return roots +def _indexed_pickle_shards(index_path: Path, root: Path, snapshot: Path) -> list: + """Shards a torch weight index points a ``from_pretrained`` load at that load_state_dict would + torch.load (pickle): every ``weight_map`` target NOT ending in ``.safetensors``, whatever its + stem (an arbitrary name like ``shards/payload`` still deserializes). Resolved relative to the + index dir (``root``) like the loader, so a shard in a nested dir is followed (iterdir misses it). + Lexical only, never ``Path.resolve()`` (HF snapshot files symlink into ``blobs/``, so resolving + escapes the snapshot and false-blocks every shard). Raises OSError -> caller fails CLOSED on an + unreadable/invalid index or a target escaping the snapshot.""" + import json + import os + + try: + # JSON is UTF-8 by spec; pin it so a non-ASCII index is not misdecoded (and needlessly + # blocked) under Windows' cp1252 default. + parsed = json.loads(index_path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + raise OSError(f"unreadable weight index: {index_path}") from exc + weight_map = parsed.get("weight_map") if isinstance(parsed, dict) else None + if not isinstance(weight_map, dict): + return [] # no dict weight_map -> the loader resolves no shards from this index + snapshot_norm = os.path.normpath(str(snapshot)) + shards = [] + for shard in weight_map.values(): + raw = str(shard) + if not raw: + continue + # Join the RAW weight_map value like from_pretrained's os.path.join: on POSIX a backslash is a + # literal filename char (not a separator), so normalizing it would probe a different path than + # the loader opens. normpath + containment stay platform-aware (os.sep) to block "..". + joined = os.path.normpath(os.path.join(str(root), raw)) + if joined != snapshot_norm and not joined.startswith(snapshot_norm + os.sep): + raise OSError(f"weight index escapes the snapshot: {index_path}") + shard_path = Path(joined) + # Case-SENSITIVE, mirroring load_state_dict's own endswith(".safetensors"): a shard named + # payload.SAFETENSORS is not treated as safetensors by the loader and falls to torch.load. + if not shard_path.name.endswith(".safetensors") and shard_path.is_file(): + shards.append(shard_path) + return shards + + +def _loader_resolves(root: Path, name: str) -> bool: + """True iff from_pretrained would open ``name`` under ``root``. ``is_file()`` honors the platform + (case-sensitive on Linux, case-insensitive on Windows/macOS), so it mirrors the loader's own + lookup: an oddly-cased decoy counts as an alternative only where the loader would truly open it. + A name-fold instead would let an uppercase MODEL.SAFETENSORS suppress the scan on Linux while the + loader, asking for the canonical lowercase name, silently falls through to a pickle index.""" + return (root / name).is_file() + + def _cached_pickle_weight_files(snapshot: Path) -> list: - """Pickle weight files in snapshot's ST load roots, EXCLUDING those whose weight family also - ships an inert safetensors in the same dir (the loader prefers it): a base pickle is suppressed - only by a base model.safetensors, an adapter pickle only by adapter_model.safetensors -- an - unrelated safetensors is no substitute. Load roots only. Raises OSError if the snapshot root is - unreadable (caller blocks).""" + """Pickle weight files a SentenceTransformer/Transformers load deserializes from snapshot's ST + load roots, EXCLUDING those whose weight family also ships an inert safetensors in the same dir + (the loader prefers it): a base pickle is suppressed only by a base model.safetensors, an adapter + pickle only by adapter_model.safetensors -- an unrelated safetensors is no substitute. Covers + both direct-child pickles AND pickle shards referenced by a local weight index (which the loader + follows into nested dirs, matching the online gate). Raises OSError -- caller fails CLOSED -- if + the snapshot root or a weight index is unreadable, or an index reference escapes the snapshot.""" blocked = [] + seen = set() + + def _add(path: Path): + key = str(path) + if key not in seen: + seen.add(key) + blocked.append(path) + for root in _st_load_roots(snapshot): try: entries = [p for p in root.iterdir() if p.is_file()] @@ -327,15 +412,35 @@ def _cached_pickle_weight_files(snapshot: Path) -> list: if root == snapshot: raise # top-level unreadable -> fail closed continue # unreadable module subdir: nothing loadable to attest here - has_base_safetensors = any(_BASE_SAFETENSORS_RE.match(p.name) for p in entries) - has_adapter_safetensors = any(_ADAPTER_SAFETENSORS_RE.match(p.name) for p in entries) + # Safetensors alternatives the loader would actually resolve (never a bare name-fold, which + # fails OPEN: see _loader_resolves). A base pickle is replaced only by a base safetensors, an + # adapter pickle only by an adapter one. A single model.safetensors also outranks BOTH indexes. + has_direct_base_safetensors = _loader_resolves(root, "model.safetensors") + has_base_safetensors = has_direct_base_safetensors or _loader_resolves( + root, "model.safetensors.index.json" + ) + has_adapter_safetensors = _loader_resolves(root, "adapter_model.safetensors") for path in entries: if not _PICKLE_WEIGHT_RE.match(path.name): continue is_adapter = path.name.lower().startswith("adapter_model") has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors if not has_alternative: - blocked.append(path) + _add(path) + # A torch weight index makes from_pretrained load nested shards iterdir never sees; the loader + # torch.loads any not ending in .safetensors. Probe the canonical index name with the loader's + # own lookup (_loader_resolves), so an oddly-cased artifact it would never open does not block. + # A direct model.safetensors wins over BOTH indexes; failing that a base safetensors still + # outranks the pytorch index, while a safetensors index is itself the chosen archive. + for index_name in _TORCH_INDEX_FILES: + if not _loader_resolves(root, index_name): + continue + if has_direct_base_safetensors: + continue + if index_name == "pytorch_model.bin.index.json" and has_base_safetensors: + continue + for shard_path in _indexed_pickle_shards(root / index_name, root, snapshot): + _add(shard_path) return blocked @@ -410,12 +515,17 @@ def evaluate_file_security( # fails open): the Spark-TTS "/LLM" alias is really unsloth/ from LLM/. model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs)) - # Local paths (including a local .gguf) have no Hub scan. A remote ref is scanned - # even if named "*.gguf", so a repo cannot dodge the scan via its name. + # Local paths have no Hub scan, EXCEPT an HF-cache snapshot whose canonical path + # encodes a repo id + commit: scan that exact commit so an inactive-cache load can't + # dodge the gate. A remote ref is scanned even if named "*.gguf" (name can't dodge it). + snapshot_revision = None try: from utils.paths import is_local_path if is_local_path(model_name): - return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan") + cache_ref = _hf_cache_snapshot_ref(model_name) + if cache_ref is None: + return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan") + model_name, snapshot_revision = cache_ref except Exception: # Cannot classify the path -> do not block on that account. return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked") @@ -424,7 +534,7 @@ def evaluate_file_security( if local_only_load: return _evaluate_local_only(model_name) - status = _fetch_security_status(model_name, hf_token) + status = _fetch_security_status(model_name, hf_token, revision = snapshot_revision) if not isinstance(status, dict): return FileSecurityDecision( model_name, False, reason = "scan unavailable; allowed (fail-open)" @@ -461,7 +571,9 @@ def evaluate_file_security( maybe_shard.append({"path": path, "level": level, "norm": norm}) if maybe_shard: - indexed = _indexed_shard_paths(model_name, hf_token, load_subdirs) + indexed = _indexed_shard_paths( + model_name, hf_token, load_subdirs, revision = snapshot_revision + ) for m in maybe_shard: # Block if a root index lists this shard, or if the lookup was inconclusive # (transient error -> stay conservative). A definitive "no index / not listed" diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index bf8348dd82..ce3d6704b7 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -294,17 +294,31 @@ def format_error_message(error: Exception, model_name: str) -> str: return "Invalid HF token. Please check your token and try again." if ( - "memory" in error_str - or "cuda" in error_str - or "mlx" in error_str - or "out of memory" in error_str + "out of memory" in error_str + or "out of device memory" in error_str + or "out_of_device_memory" in error_str # ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY + or "out_of_host_memory" in error_str # ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY + or "not enough memory" in error_str + or "cannot allocate memory" in error_str + or "memory allocation failed" in error_str + or "cublas_status_alloc_failed" in error_str # cuBLAS workspace OOM + or ("cuda error" in error_str and "alloc" in error_str) + or ("xpu" in error_str and ("alloc" in error_str or "memory" in error_str)) + or isinstance(error, MemoryError) + or ("mlx" in error_str and ("memory" in error_str or "allocate" in error_str)) ): + # Resolve get_device() at call time (not import time) so tests that + # monkey-patch utils.hardware.get_device after this module is loaded + # still see the patched backend. from utils.hardware import get_device device = get_device() - device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get( - device.value, "GPU" - ) + device_label = { + "cuda": "GPU", + "xpu": "Intel GPU", + "mlx": "Apple Silicon GPU", + "cpu": "system", + }.get(device.value, "GPU") return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory." return str(error) diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index e6c89b9cd7..9232defd70 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -524,8 +524,10 @@ export function AppProvider({ children }: AppProviderProps) { visibleToasts={2} expand={true} closeButton={true} - // Clear the chat header buttons on the right. - offset={{ top: 12, right: 64 }} + // Clear the chat header buttons on the right. On desktop, also drop + // below the ~34px custom window titlebar so toasts don't cover the + // minimize / maximize / close controls. + offset={{ top: isTauri ? 46 : 12, right: 64 }} /> diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index a6f64243ad..a31d9b6ced 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -105,7 +105,6 @@ import { archiveChatItem, ChatSearchDialog, clearNewChatDraft, - createChatProject, deleteChatProject, deleteChatItem, listStoredChatThreads, @@ -123,6 +122,7 @@ import { type ProjectRecord, type SidebarItem, } from "@/features/chat"; +import { NewProjectDialog } from "@/features/chat/components/new-project-dialog"; import { useAppearanceCustomStore, useSettingsDialogStore, @@ -696,7 +696,6 @@ export function AppSidebar() { }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); - const [projectNameDraft, setProjectNameDraft] = useState(""); const [projectCreateMoveTarget, setProjectCreateMoveTarget] = useState(null); const renameTrimmed = renameDraft.trim(); @@ -849,28 +848,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, }); } @@ -901,8 +898,8 @@ export function AppSidebar() { : "group/recent-item relative"; const actionClass = variant === "project" - ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" - : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; + ? "sidebar-row-action sidebar-touch-reveal group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + : "sidebar-row-action sidebar-touch-reveal group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; const buttonClass = cn( "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-ui-14p5 leading-ui-19 tracking-nav font-medium", // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the @@ -912,13 +909,14 @@ export function AppSidebar() { isPinned && variant !== "project" && "gap-[8.5px]", variant === "project" ? // Room for the hover pin quick-action plus the kebab. - "group-hover/project-chat-item:pr-14 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8" + "group-hover/project-chat-item:pr-14 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8 [@media(pointer:coarse)]:pr-14" : isPinned ? // Pinned rows show an extra unpin button on hover, so reserve more room // (pr-8 when the menu is open keeps the unpin button clear of the title). - "group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8" + "group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8 [@media(pointer:coarse)]:pr-16" : // Hover room for the kebab only; title keeps one more character. - "group-hover/recent-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-6", + // Touch rows clear the full always-visible kebab hit area (pr-10). + "group-hover/recent-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-6 [@media(pointer:coarse)]:pr-10", ); const isRenamingThis = @@ -987,7 +985,7 @@ export function AppSidebar() { togglePinnedChat(item.id); }} aria-label={isPinned ? "Unpin chat" : "Pin chat"} - className="sidebar-row-action is-unpin-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + className="sidebar-row-action sidebar-touch-reveal is-unpin-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" > @@ -1002,7 +1000,7 @@ export function AppSidebar() { togglePinnedChat(item.id); }} aria-label="Unpin chat" - className="sidebar-row-action is-unpin-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + className="sidebar-row-action sidebar-touch-reveal is-unpin-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" > @@ -1049,7 +1047,6 @@ export function AppSidebar() { { setProjectCreateMoveTarget(item); - setProjectNameDraft(""); setCreatingProject(true); }} > @@ -1392,7 +1389,6 @@ export function AppSidebar() { onClick={(e) => { e.stopPropagation(); 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" @@ -1854,6 +1850,18 @@ export function AppSidebar() { )} + {/* Collapsed rail has no room for the cog on the profile row, so it + sits above the avatar instead. */} + { + useSettingsDialogStore.getState().openDialog(); + closeMobileIfOpen(); + }} + /> @@ -2159,58 +2167,18 @@ export function AppSidebar() { - { setCreatingProject(open); - if (!open) { - setProjectNameDraft(""); - setProjectCreateMoveTarget(null); - } + if (!open) setProjectCreateMoveTarget(null); }} - > - - - - {projectCreateMoveTarget ? "Move to new project" : "New project"} - - - 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" - /> - - - - - - + title={ + projectCreateMoveTarget ? "Move to new project" : "Create project" + } + submitLabel={projectCreateMoveTarget ? "Create and move" : "Create project"} + onCreated={afterCreateProject} + /> ); } diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 97891e9358..09ca6d2530 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -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(0); const startTimeRef = useRef(null); @@ -368,6 +369,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ } }, [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 +390,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ if (isReasoningStreaming) { setDismissedWhileStreaming(!open); } else { + if (open) { + setRetainStreamingHeight(false); + } setManualOpen(open); } }, @@ -407,7 +422,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ aria-busy={isReasoningStreaming} streaming={isReasoningStreaming} > - + {children} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 43a0c43950..68f96c46c8 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -1556,6 +1556,7 @@ const Composer: FC<{ const draftThreadId = referenceThreadId; const draftKey = draftThreadId ? composerDraftKey(draftThreadId) : null; const lastDraftKeyRef = useRef(draftKey); + const draftSaveTimerRef = useRef | null>(null); useEffect(() => { const draft = draftKey ? (readComposerDraft(draftKey) ?? "") : ""; const composer = aui.composer(); @@ -1574,8 +1575,25 @@ const Composer: FC<{ return; } const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300); + draftSaveTimerRef.current = t; return () => clearTimeout(t); }, [composerText, draftKey]); + // Without this the restore effect above puts the sent text back when the + // runtime rebinds on the first message. + const draftKeyRef = useRef(draftKey); + useEffect(() => { + draftKeyRef.current = draftKey; + }, [draftKey]); + const clearStoredDraft = useCallback(() => { + if (draftSaveTimerRef.current !== null) { + clearTimeout(draftSaveTimerRef.current); + draftSaveTimerRef.current = null; + } + const key = draftKeyRef.current; + if (key) { + writeComposerDraft(key, ""); + } + }, []); // react-textarea-autosize re-measures only on value change or window resize, // not on the width swap from expanding, so it keeps the taller height and // leaves a stray blank row. Nudge a resize whenever input width changes. @@ -1726,9 +1744,10 @@ const Composer: FC<{ setPendingSend(false); dismissWaitToast(); if (text.trim().length > 0 || attachments.length > 0) { + clearStoredDraft(); aui.composer().send(); } - }, [pendingSend, indexingActive, aui, dismissWaitToast]); + }, [pendingSend, indexingActive, aui, clearStoredDraft, dismissWaitToast]); // Drop any queued send + toast on unmount (e.g. thread switch). useEffect( @@ -1771,6 +1790,7 @@ const Composer: FC<{ flushResourcesSync(() => { aui.composer().setText(""); }); + clearStoredDraft(); startPromptQueue( [queuedPrompt], createPromptQueueTarget(), @@ -1804,6 +1824,7 @@ const Composer: FC<{ closeOverlay(); return; } + clearStoredDraft(); setImageToolsEnabled(true); setPendingImageEditReference({ threadId: overlay.threadId ?? referenceThreadId, @@ -1821,11 +1842,15 @@ const Composer: FC<{ ); }); closeOverlay(); + return; } + + clearStoredDraft(); }, [ aui, canQueueCurrentPrompt, + clearStoredDraft, closeOverlay, composerText, createPromptQueueTarget, @@ -1947,6 +1972,7 @@ const Composer: FC<{ flushResourcesSync(() => { aui.composer().setText(""); }); + clearStoredDraft(); startPromptQueue([queuedPrompt], createPromptQueueTarget(), true); }} onSendClick={interceptSend} diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx index 1e816132b8..34e51b9d5a 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx @@ -83,7 +83,7 @@ function CopyBtn({ text }: { text: string }) { ); } -/** Save the executed script as a .py file via a client-side Blob (no server file serving). */ +/** Save the script as a .py file via a client-side Blob. */ function DownloadBtn({ code, name = "script.py" }: { code: string; name?: string }) { const download = useCallback(() => { if (typeof document === "undefined") { @@ -229,8 +229,8 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({ const authToken = getAuthToken(); return ( - // Run status and output collapse from history, but the script source is - // rendered outside ToolFallbackContent so it stays visible on reopen (#7165). + // Status/output collapse from history; the script source renders outside + // ToolFallbackContent so it stays visible on reopen (#7165). { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" }, ]; -// Optimizers the MLX trainer actually supports on Apple Silicon. Values must -// match SUPPORTED_MLX_OPTIMIZERS in unsloth-zoo's mlx/trainer.py; on MLX the -// bitsandbytes/torch names above have no meaning and are remapped to plain -// AdamW, so Studio offers this list instead when running on a Mac. +// MLX trainer optimizers (Apple Silicon); must match SUPPORTED_MLX_OPTIMIZERS in +// unsloth-zoo's mlx/trainer.py. The CUDA/torch names above are remapped to AdamW on MLX. export const MLX_OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ { value: "adamw", label: "AdamW" }, { value: "adam", label: "Adam" }, diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 2b02ffdd99..b089aa95da 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -105,6 +105,7 @@ import { import { resolveLoadMaxSeqLength } from "../presets/preset-policy"; import { generateAudio, + GenerationLengthError, listGgufVariants, loadModel, streamChatCompletions, @@ -4461,7 +4462,15 @@ export function createOpenAIStreamAdapter( ); if (!abortSignal.aborted) { const msg = err instanceof Error ? err.message : String(err); - if (err instanceof StreamInterruptedError) { + if (err instanceof GenerationLengthError) { + toast.error("Response ran out of tokens", { + description: + "The model used the full Max Tokens budget while thinking " + + "and did not produce a final answer. Increase Max Tokens in " + + "chat Settings or turn off thinking, then retry.", + duration: 8000, + }); + } else if (err instanceof StreamInterruptedError) { // Connection dropped mid-turn: surface it explicitly (the rethrow // below also marks the message with an inline error + Retry). toast.error("Response interrupted", { diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index ffaf099f29..4d123e98ab 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -50,6 +50,21 @@ export class StreamInterruptedError extends Error { } } +/** + * Thrown when a reasoning model consumes its output budget before emitting any + * standard content. Keeping this distinct from a dropped connection lets the + * chat UI explain why a completed stream contains only a thinking panel. + */ +export class GenerationLengthError extends Error { + constructor() { + super( + "The model reached the Max Tokens limit before producing a final answer. " + + "Increase Max Tokens or disable thinking, then retry.", + ); + this.name = "GenerationLengthError"; + } +} + export function notifyChatHistoryUpdated(): void { if (typeof window !== "undefined") { window.dispatchEvent(new Event(CHAT_HISTORY_UPDATED_EVENT)); @@ -982,6 +997,61 @@ function parseSseEvent(rawEvent: string): string[] { return dataLines; } +function hasNonWhitespaceText(value: unknown): boolean { + if (typeof value === "string") { + return value.trim().length > 0; + } + if (Array.isArray(value)) { + return value.some((item) => hasNonWhitespaceText(item)); + } + if (!value || typeof value !== "object") { + return false; + } + const record = value as Record; + return ["thinking", "text", "content", "reasoning", "summary"].some( + (key) => key in record && hasNonWhitespaceText(record[key]), + ); +} + +function classifyStructuredDeltaContent(content: unknown): { + hasAssistantContent: boolean; + hasReasoningContent: boolean; +} { + if (typeof content === "string") { + return { + hasAssistantContent: hasNonWhitespaceText(content), + hasReasoningContent: false, + }; + } + if (!Array.isArray(content)) { + return { + hasAssistantContent: false, + hasReasoningContent: false, + }; + } + + let hasAssistantContent = false; + let hasReasoningContent = false; + for (const part of content) { + if (typeof part === "string") { + hasAssistantContent ||= hasNonWhitespaceText(part); + continue; + } + if (!part || typeof part !== "object") { + continue; + } + const record = part as Record; + if (record.type === "thinking" || record.type === "reasoning") { + hasReasoningContent ||= hasNonWhitespaceText(record); + } else if (record.type === "text" || record.type === "output_text") { + const text = + typeof record.text === "string" ? record.text : record.content; + hasAssistantContent ||= hasNonWhitespaceText(text); + } + } + return { hasAssistantContent, hasReasoningContent }; +} + export async function* streamChatCompletions( payload: OpenAIChatCompletionsRequest, signal: AbortSignal, @@ -1009,6 +1079,19 @@ export async function* streamChatCompletions( // EOF without `[DONE]` or a finish_reason chunk means the stream was cut // mid-generation: surface as interrupted, not silent success. let sawTerminalSignal = false; + let terminalFinishReason: string | null = null; + let sawAssistantContent = false; + let sawReasoningContent = false; + + const throwIfReasoningOnlyLength = () => { + if ( + terminalFinishReason === "length" && + sawReasoningContent && + !sawAssistantContent + ) { + throw new GenerationLengthError(); + } + }; try { while (true) { @@ -1018,6 +1101,7 @@ export async function* streamChatCompletions( if (!sawTerminalSignal) { throw new StreamInterruptedError(); } + throwIfReasoningOnlyLength(); break; } @@ -1039,6 +1123,7 @@ export async function* streamChatCompletions( if (dataText === "[DONE]") { completed = true; sawTerminalSignal = true; + throwIfReasoningOnlyLength(); return; } @@ -1094,11 +1179,31 @@ export async function* streamChatCompletions( } // finish_reason is a valid terminal signal for providers that close // the stream without an explicit [DONE] sentinel. - const finishReason = ( + const parsedChoices = ( parsed as { - choices?: Array<{ finish_reason?: string | null }>; + choices?: Array<{ + delta?: Record; + finish_reason?: string | null; + }>; } - ).choices?.[0]?.finish_reason; + ).choices; + for (const choice of parsedChoices ?? []) { + const delta = choice.delta; + if (delta) { + const contentState = classifyStructuredDeltaContent(delta.content); + sawAssistantContent ||= contentState.hasAssistantContent; + sawReasoningContent ||= contentState.hasReasoningContent; + const reasoning = + delta.reasoning_content ?? + delta.reasoning ?? + delta.reasoning_details; + sawReasoningContent ||= hasNonWhitespaceText(reasoning); + } + if (choice.finish_reason) { + terminalFinishReason = choice.finish_reason; + } + } + const finishReason = parsedChoices?.[0]?.finish_reason; if (finishReason) { sawTerminalSignal = true; } diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts index 1e00357ea4..ab13130525 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -13,6 +13,7 @@ export type PersistedInferenceParams = Partial< export interface PersistedChatPreset { name: string; params: PersistedInferenceParams; + loadConfig?: Record; } export interface PersistedChatSettings { diff --git a/studio/frontend/src/features/chat/api/providers-api.ts b/studio/frontend/src/features/chat/api/providers-api.ts index 4ad996d54f..c6c5613272 100644 --- a/studio/frontend/src/features/chat/api/providers-api.ts +++ b/studio/frontend/src/features/chat/api/providers-api.ts @@ -23,6 +23,8 @@ export interface ProviderConfig { display_name: string; base_url: string; is_enabled: boolean; + models?: string[]; + available_models?: string[]; created_at: string; updated_at: string; } @@ -123,6 +125,8 @@ export async function createProviderConfig(payload: { providerType: string; displayName: string; baseUrl?: string | null; + models?: string[]; + availableModels?: string[]; }): Promise { const response = await authFetch("/api/providers/", { method: "POST", @@ -131,6 +135,8 @@ export async function createProviderConfig(payload: { provider_type: payload.providerType, display_name: payload.displayName, base_url: payload.baseUrl ?? null, + models: payload.models ?? [], + available_models: payload.availableModels ?? [], }), }); return parseJsonOrThrow(response); @@ -158,6 +164,8 @@ export async function updateProviderConfig( displayName?: string; baseUrl?: string | null; isEnabled?: boolean; + models?: string[]; + availableModels?: string[]; }, ): Promise { const response = await authFetch(`/api/providers/${providerId}`, { @@ -167,6 +175,10 @@ export async function updateProviderConfig( ...(payload.displayName === undefined ? {} : { display_name: payload.displayName }), ...(payload.baseUrl === undefined ? {} : { base_url: payload.baseUrl }), ...(payload.isEnabled === undefined ? {} : { is_enabled: payload.isEnabled }), + ...(payload.models === undefined ? {} : { models: payload.models }), + ...(payload.availableModels === undefined + ? {} + : { available_models: payload.availableModels }), }), }); return parseJsonOrThrow(response); diff --git a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx index 0f4c83e0db..1955c3aca1 100644 --- a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx +++ b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx @@ -12,6 +12,8 @@ import { import { MascotImg } from "@/components/mascot-img"; import { Button } from "@/components/ui/button"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { downloadFile, isDownloadCancelled } from "@/lib/native-files"; +import { toast } from "@/lib/toast"; import { cn } from "@/lib/utils"; import { CopyIcon, EyeIcon, Maximize2Icon, XIcon } from "lucide-react"; import { Download01Icon } from "@hugeicons/core-free-icons"; @@ -91,18 +93,6 @@ function ArtifactGeneratingPanel() { ); } -function downloadTextFile(filename: string, text: string): void { - const blob = new Blob([text], { type: "text/html;charset=utf-8" }); - const url = URL.createObjectURL(blob); - const anchor = document.createElement("a"); - anchor.href = url; - anchor.download = filename; - document.body.appendChild(anchor); - anchor.click(); - document.body.removeChild(anchor); - window.setTimeout(() => URL.revokeObjectURL(url), 0); -} - export function ArtifactSurface({ artifact, variant, @@ -205,7 +195,7 @@ export function ArtifactSurface({ className={cn( "relative flex min-h-0 flex-col bg-background", variant === "panel" - ? "artifact-panel-shell mx-2 mt-[72px] mb-8 h-[calc(100%_-_104px)] overflow-visible rounded-[28px] border-t border-border/70 bg-card/95" + ? "artifact-panel-shell mx-2 mt-[90px] mb-8 h-[calc(100%_-_122px)] overflow-visible rounded-[28px] border-t border-border/70 bg-card/95" : "h-[min(92vh,900px)] w-[min(96vw,1200px)] overflow-hidden rounded-2xl border border-border shadow-xl", )} aria-label={`${artifact.title} canvas`} @@ -265,7 +255,19 @@ export function ArtifactSurface({ size="icon" className="size-8" disabled={isLoadingArtifact || !hasArtifactCode} - onClick={() => downloadTextFile(filename, artifact.code)} + onClick={() => { + // Route through the native save dialog on desktop; the plain + // blob-anchor download is silently dropped by the Tauri WebView2. + void downloadFile( + artifact.code, + filename, + "text/html;charset=utf-8", + ).catch((err) => { + if (!isDownloadCancelled(err)) { + toast.error("Failed to save canvas HTML"); + } + }); + }} aria-label="Download canvas HTML" > diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index e46ea0ac46..7e0544e2d1 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -175,6 +175,7 @@ import { } from "./stores/chat-runtime-store"; import { useChatPreferencesStore } from "./stores/chat-preferences-store"; import { useExternalProvidersStore } from "./stores/external-providers-store"; +import { syncExternalProvidersFromBackend } from "./sync-external-providers"; import { buildChatTourSteps } from "./tour"; import type { ChatView, MessageRecord } from "./types"; import { @@ -184,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(() => @@ -997,7 +1002,14 @@ function ProjectLanding({ const active = useChatActive(); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const initialActiveThreadRef = useRef(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( null, ); @@ -1762,8 +1774,18 @@ export function ChatPage({ const externalProvidersForChat = connectionsEnabled ? externalProviders : []; useEffect(() => { - void hydratePersistedSettings(); - }, [hydratePersistedSettings]); + void (async () => { + await hydratePersistedSettings(); + try { + const synced = await syncExternalProvidersFromBackend( + useExternalProvidersStore.getState().providers, + ); + setExternalProviders(synced); + } catch { + // Silent on startup; Connections settings still surfaces load errors. + } + })(); + }, [hydratePersistedSettings, setExternalProviders]); useEffect(() => { // Skip while off-route: ChatPage stays mounted, and toast+navigate here would @@ -2278,7 +2300,7 @@ export function ChatPage({ } else if (outcome === "conflict") { toast.info("Resume this download from Models", { description: - "An earlier partial download used a different transport. Open the Models tab to resume or restart it.", + "An earlier partial download used a different transport. Open the Model hub tab to resume or restart it.", }); } else if (outcome === "busy") { toast.info("Download already in progress", { @@ -2399,7 +2421,7 @@ export function ChatPage({ // surface's onComplete auto-loads, mirroring the "started" branch. toast.info("Resume this download from Models", { description: - "An earlier partial download used a different transport. Open the Models tab to resume or restart it.", + "An earlier partial download used a different transport. Open the Model hub tab to resume or restart it.", }); return; } @@ -2665,7 +2687,7 @@ export function ChatPage({ config: meta?.config, nativePathToken: meta?.nativePathToken, nativePathExpiresAtMs: meta?.nativePathExpiresAtMs, - forceReload: isSameLoadedModel || undefined, + forceReload: meta?.forceReload ?? (isSameLoadedModel || undefined), }; await stageOrLoad(selection); })(); diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index bfe7c71918..4ab560e5ad 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -41,7 +41,6 @@ import { type ProviderRegistryEntry, createProviderConfig, deleteProviderConfig, - listProviderConfigs, listProviderModels, listProviderRegistry, testProviderConnection, @@ -49,7 +48,6 @@ import { } from "./api/providers-api"; import type { ExternalProviderConfig } from "./external-providers"; import { - CUSTOM_BACKEND_PROVIDER_TYPE, CUSTOM_PROVIDER_PRESETS, allowsManualModelIdsWithCatalog, customProviderBaseUrlPlaceholder, @@ -68,6 +66,10 @@ import { toExternalBackendProviderType, } from "./external-providers"; import { useExternalProvidersStore } from "./stores/external-providers-store"; +import { + pruneProviderModelIds, + syncExternalProvidersFromBackend, +} from "./sync-external-providers"; /** Matches navbar / thread layout easing (see index.css --ease-out-quart) */ const PROVIDER_FORM_EASE: [number, number, number, number] = [ @@ -76,58 +78,7 @@ const PROVIDER_FORM_EASE: [number, number, number, number] = [ const PROVIDER_FORM_DURATION = 0.2; const CUSTOM_PROVIDER_MISSING_KEY_MESSAGE = "No API key found. Add a valid API key for this connection."; -const ANTHROPIC_DATED_SNAPSHOT_SUFFIX = /-\d{8}$/; -const OPENAI_DEPRECATED_MODELS = new Set(["gpt-5.3"]); const HIDDEN_PROVIDER_TYPES = new Set(["qwen"]); -const OPENROUTER_EXCLUDED_MODELS = new Set([ - "google/chirp-3", - "kwaivgi/kling-v3.0-pro", - "openai/whisper-1", - "openai/gpt-4o-mini-transcribe", - "recraft/recraft-v4-pro", -]); - -function normalizeUrl(input: string): string { - return input.trim().replace(/\/+$/, ""); -} - -function resolveUiProviderTypeFromConfig( - configProviderType: string, - configDisplayName: string | null | undefined, - configBaseUrl: string | null | undefined, - registryRows: ProviderRegistryEntry[], - existingProviderType: string | undefined, -): string { - if (existingProviderType && isCustomProviderType(existingProviderType)) { - return existingProviderType; - } - if (configProviderType !== CUSTOM_BACKEND_PROVIDER_TYPE) { - return configProviderType; - } - const displayName = (configDisplayName ?? "").trim().toLowerCase(); - const matchingCustomPreset = CUSTOM_PROVIDER_PRESETS.find( - (preset) => preset.displayName.toLowerCase() === displayName, - ); - if (matchingCustomPreset) { - return matchingCustomPreset.providerType; - } - const openAiRegistry = registryRows.find( - (entry) => entry.provider_type === CUSTOM_BACKEND_PROVIDER_TYPE, - ); - if (!openAiRegistry) { - return configProviderType; - } - const openAiDisplayName = openAiRegistry.display_name.trim().toLowerCase(); - if (displayName.length > 0 && displayName !== openAiDisplayName) { - return LEGACY_CUSTOM_PROVIDER_TYPE; - } - const configUrl = normalizeUrl(configBaseUrl ?? ""); - const defaultUrl = normalizeUrl(openAiRegistry.base_url ?? ""); - if (configUrl.length > 0 && configUrl !== defaultUrl) { - return LEGACY_CUSTOM_PROVIDER_TYPE; - } - return configProviderType; -} function parseManualModelIds(text: string): string[] { const seen = new Set(); @@ -182,19 +133,6 @@ function shouldAppendOpenAiVersionPath(providerType: string): boolean { ); } -function pruneProviderModelIds(providerType: string, modelIds: string[]): string[] { - if (providerType === "anthropic") { - return modelIds.filter((id) => !ANTHROPIC_DATED_SNAPSHOT_SUFFIX.test(id)); - } - if (providerType === "openai") { - return modelIds.filter((id) => !OPENAI_DEPRECATED_MODELS.has(id)); - } - if (providerType === "openrouter") { - return modelIds.filter((id) => !OPENROUTER_EXCLUDED_MODELS.has(id)); - } - return modelIds; -} - function formatModelSummary(models: string[]): string { if (models.length === 0) { return "No models enabled"; @@ -360,9 +298,9 @@ export function ChatProvidersSettings({ } let syncSucceeded = false; try { - const [registryRows, configRows] = await Promise.all([ + const [registryRows, syncedProviders] = await Promise.all([ listProviderRegistry(), - listProviderConfigs(), + syncExternalProvidersFromBackend(providersRef.current), ]); if (!isMounted) return; syncSucceeded = true; @@ -377,61 +315,6 @@ export function ChatProvidersSettings({ } return registryRows[0]?.provider_type ?? ""; }); - const existingById = new Map(); - for (const provider of providersRef.current) { - existingById.set(provider.id, provider); - } - const syncedProviders: ExternalProviderConfig[] = configRows - .filter((config) => config.is_enabled) - .map((config) => { - const existing = existingById.get(config.id); - const uiProviderType = resolveUiProviderTypeFromConfig( - config.provider_type, - config.display_name, - config.base_url, - registryRows, - existing?.providerType, - ); - const createdAt = Number.isFinite(Date.parse(config.created_at)) - ? Date.parse(config.created_at) - : Date.now(); - const updatedAt = Number.isFinite(Date.parse(config.updated_at)) - ? Date.parse(config.updated_at) - : Date.now(); - const registryEntry = - registryRows.find((entry) => entry.provider_type === uiProviderType) ?? - registryRows.find((entry) => entry.provider_type === config.provider_type); - const defaultModels = pruneProviderModelIds( - uiProviderType, - registryEntry?.default_models ?? [], - ); - const savedModels = existing?.models ?? []; - const savedAvailableModels = existing?.availableModels ?? []; - const existingModels = pruneProviderModelIds( - uiProviderType, - savedModels.length > 0 ? savedModels : defaultModels, - ); - const existingAvailableModels = pruneProviderModelIds( - uiProviderType, - savedAvailableModels.length > 0 ? savedAvailableModels : defaultModels, - ); - return { - id: config.id, - providerType: uiProviderType, - name: config.display_name, - baseUrl: config.base_url ?? "", - models: existingModels, - availableModels: existingAvailableModels, - enablePromptCaching: supportsProviderPromptCaching(uiProviderType) - ? (existing?.enablePromptCaching ?? true) - : undefined, - isReasoningModel: supportsProviderReasoningToggle(uiProviderType) - ? existing?.isReasoningModel === true - : undefined, - createdAt: existing?.createdAt ?? createdAt, - updatedAt, - }; - }); // Trust the backend response. An empty array means every connection was // removed (often from another tab); mirror that locally, else stale // entries become un-removable here until localStorage is cleared. @@ -699,6 +582,10 @@ export function ChatProvidersSettings({ providerType: backendProviderType, displayName, baseUrl, + models: modelsToSave, + availableModels: manualOnly + ? [] + : pruneProviderModelIds(providerType, availableModels), }); const createdAt = Number.isFinite(Date.parse(created.created_at)) ? Date.parse(created.created_at) @@ -814,6 +701,10 @@ export function ChatProvidersSettings({ customProviderDisplayName(existing.providerType) : existing.name, baseUrl, + models: modelsToSave, + availableModels: manualOnly + ? [] + : pruneProviderModelIds(existing.providerType, availableModels), }); if (apiKey.trim()) { setExternalProviderApiKey(editingProviderId, apiKey.trim()); diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 672ab2b2d3..7b310c50d4 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -75,6 +75,12 @@ import { isSamePresetConfig, toPresetParams, } from "./presets/preset-policy"; +import { + applyPresetLoadConfig, + capturePresetLoadConfig, + formatPresetLoadConfigSummary, + isSamePresetLoadConfig, +} from "./presets/preset-load-config"; import { type ProviderCapabilities, getExternalMaxOutputTokens, @@ -385,6 +391,12 @@ export function ChatSettingsPanel({ (s) => s.ggufMaxContextLength, ); const customContextLength = useChatRuntimeStore((s) => s.customContextLength); + const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); + const gpuMemoryMode = useChatRuntimeStore((s) => s.gpuMemoryMode); + const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers); + const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe); + const tensorParallel = useChatRuntimeStore((s) => s.tensorParallel); + const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax); const speculativeType = useChatRuntimeStore((s) => s.speculativeType); const specFallbackReason = useChatRuntimeStore((s) => s.specFallbackReason); const mtpUpdatable = @@ -469,11 +481,50 @@ export function ChatSettingsPanel({ if (activePresetDefinition == null) { return false; } - if (activePresetDefinition.name === "Default") { - return activePresetSource === "modified"; - } - return !isSamePresetConfig(activePresetDefinition.params, params); - }, [activePresetDefinition, activePresetSource, params]); + const samplingChanged = + activePresetDefinition.name === "Default" + ? activePresetSource === "modified" + : !isSamePresetConfig(activePresetDefinition.params, params); + const currentLoadConfig = capturePresetLoadConfig(); + const loadChanged = !isSamePresetLoadConfig( + activePresetDefinition.loadConfig, + currentLoadConfig, + ); + return samplingChanged || loadChanged; + }, [ + activePresetDefinition, + activePresetSource, + params, + customContextLength, + ggufContextLength, + kvCacheDtype, + gpuMemoryMode, + gpuLayers, + nCpuMoe, + tensorParallel, + speculativeType, + specDraftNMax, + params.maxSeqLength, + ]); + const activePresetLoadSummary = useMemo( + () => formatPresetLoadConfigSummary(activePresetDefinition?.loadConfig), + [activePresetDefinition], + ); + const currentLoadSummary = useMemo( + () => formatPresetLoadConfigSummary(capturePresetLoadConfig()), + [ + customContextLength, + ggufContextLength, + kvCacheDtype, + gpuMemoryMode, + gpuLayers, + nCpuMoe, + tensorParallel, + speculativeType, + specDraftNMax, + params.maxSeqLength, + ], + ); const presetSaveState = useMemo( () => getPresetSaveState({ @@ -549,8 +600,14 @@ export function ChatSettingsPanel({ onParamsChange({ ...applyPresetParams(params, p.params), }); + if (p.loadConfig) { + applyPresetLoadConfig(p.loadConfig); + } setActivePreset(name); setActivePresetSource(getPresetSource(name)); + if (p.loadConfig && params.checkpoint) { + toast.info("Reload the model to apply load settings from this preset."); + } } } @@ -571,9 +628,14 @@ export function ChatSettingsPanel({ ? getBuiltinVariantName(trimmed, usedNames) : trimmed; const next = customPresets.filter((p) => p.name !== saveName); + const loadConfig = capturePresetLoadConfig(); const merged = [ ...next, - { name: saveName, params: toPresetParams(params) }, + { + name: saveName, + params: toPresetParams(params), + ...(loadConfig ? { loadConfig } : {}), + }, ]; setCustomPresets(merged); setActivePreset(saveName); @@ -598,8 +660,11 @@ export function ChatSettingsPanel({ if (activePreset === name) { if (fallbackPreset) { onParamsChange({ - ...applyPresetParams(params, fallbackPreset.params), + ... applyPresetParams(params, fallbackPreset.params), }); + if (fallbackPreset.loadConfig) { + applyPresetLoadConfig(fallbackPreset.loadConfig); + } setActivePreset(fallbackPreset.name); setActivePresetSource("builtin-default"); } @@ -901,6 +966,23 @@ export function ChatSettingsPanel({ Delete
+

+ Saving a preset also stores current load settings (context length, + KV cache dtype, speculative decoding, GPU layers). + {currentLoadSummary ? ( + <> + {" "} + Active now: {currentLoadSummary}. + + ) : null} + {activePresetLoadSummary && + activePresetLoadSummary !== currentLoadSummary ? ( + <> + {" "} + Saved in preset: {activePresetLoadSummary}. + + ) : null} +

diff --git a/studio/frontend/src/features/chat/components/new-project-dialog.tsx b/studio/frontend/src/features/chat/components/new-project-dialog.tsx index 880129ac6c..6aca3d36c3 100644 --- a/studio/frontend/src/features/chat/components/new-project-dialog.tsx +++ b/studio/frontend/src/features/chat/components/new-project-dialog.tsx @@ -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; }) { const navigate = useNavigate(); const [name, setName] = useState(""); + const [staged, setStaged] = useState([]); + 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({ { - if (!next) setName(""); - onOpenChange(next); + if (next) { + onOpenChange(true); + return; + } + close(); }} > - + - New project + {title} - 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. */} +
+ + + +
+ - diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 3e9ae4bbc7..4906a9cdc0 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -62,6 +62,7 @@ import { import { isExternalModelId } from "../external-providers"; import { applyPerModelConfigToRuntime, + normalizeMaxSeqLength, type PerModelConfig, } from "@/features/model-picker"; import type { @@ -604,12 +605,19 @@ export function useChatModelRuntime() { async function performLoad(): Promise { if (abortCtrl.signal.aborted) throw new Error("Cancelled"); let previousWasUnloaded = false; + const pendingLoadConfig = + typeof selection !== "string" ? selection.config : undefined; + if (pendingLoadConfig) { + applyPerModelConfigToRuntime(pendingLoadConfig); + } const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const stateBeforeUnload = useChatRuntimeStore.getState(); let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; let approvedRemoteCodeFingerprint: string | null = null; - const maxSeqLength = stateBeforeUnload.params.maxSeqLength; + const maxSeqLength = + normalizeMaxSeqLength(pendingLoadConfig?.maxSeqLength) ?? + stateBeforeUnload.params.maxSeqLength; const previousActiveNativePathToken = stateBeforeUnload.activeNativePathToken; const previousIsGguf = @@ -643,34 +651,54 @@ export function useChatModelRuntime() { const previousActiveNativePathExpiresAtMs = stateBeforeUnload.activeNativePathExpiresAtMs; // Snapshot the load settings at click time, before the awaits below - // (validation, the trust dialog, unload). - const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride; - const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype; + // (validation, the trust dialog, unload). When the picker staged a + // config payload, prefer it over the store: React may not have + // flushed NumericValueInput's blur commit into state yet. + const loadChatTemplateOverride = + pendingLoadConfig?.chatTemplateOverride?.trim() + ? pendingLoadConfig.chatTemplateOverride + : stateBeforeUnload.chatTemplateOverride; + const loadKvCacheDtype = + pendingLoadConfig?.kvCacheDtype ?? stateBeforeUnload.kvCacheDtype; // gpuMemoryMode is a standing preference (kept across a model switch); // the rest are per-model knobs the reset below clears, so they are // re-baselined there in lock-step with the store. - let loadCustomContextLength = stateBeforeUnload.customContextLength; + let loadCustomContextLength = + pendingLoadConfig?.customContextLength ?? + stateBeforeUnload.customContextLength; const loadGgufContextLength = stateBeforeUnload.ggufContextLength; - const loadTensorParallel = stateBeforeUnload.tensorParallel; + const loadTensorParallel = + pendingLoadConfig?.tensorParallel ?? stateBeforeUnload.tensorParallel; const loadActivePresetSource = stateBeforeUnload.activePresetSource; const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant; - const loadGpuMemoryMode = stateBeforeUnload.gpuMemoryMode; - let loadGpuLayers = stateBeforeUnload.gpuLayers; - let loadNCpuMoe = stateBeforeUnload.nCpuMoe; + const loadGpuMemoryMode = + pendingLoadConfig?.gpuMemoryMode ?? stateBeforeUnload.gpuMemoryMode; + let loadGpuLayers = + pendingLoadConfig?.gpuLayers ?? stateBeforeUnload.gpuLayers; + let loadNCpuMoe = + pendingLoadConfig?.nCpuMoe ?? stateBeforeUnload.nCpuMoe; let loadSplitRatio = stateBeforeUnload.splitRatio; // Reconcile the persisted pick against the GPUs present now, so a stale // cross-host / now-hidden pick is dropped before /load rather than // rejected there. Warm the device cache first: load-on-selection can // run before any GPU hook mounted, and a cold cache would pass the // pick through unvalidated. validateGpuIds derives from this too. - if (stateBeforeUnload.selectedGpuIds != null) { + if ( + pendingLoadConfig?.selectedGpuIds !== undefined || + stateBeforeUnload.selectedGpuIds != null + ) { await ensureGpuDeviceCache(); } - let loadSelectedGpuIds = reconcilePersistedGpuIds( - stateBeforeUnload.selectedGpuIds, - ); - let loadSpeculativeType = stateBeforeUnload.speculativeType; - let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax; + let loadSelectedGpuIds = + pendingLoadConfig?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds) + : reconcilePersistedGpuIds(stateBeforeUnload.selectedGpuIds); + let loadSpeculativeType = + pendingLoadConfig?.speculativeType != null + ? normalizeSpeculativeType(pendingLoadConfig.speculativeType) + : stateBeforeUnload.speculativeType; + let loadSpecDraftNMax = + pendingLoadConfig?.specDraftNMax ?? stateBeforeUnload.specDraftNMax; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). @@ -810,15 +838,23 @@ export function useChatModelRuntime() { // model loads at Auto/native, not the previous model's pin. customContextLength: null, }); - loadSpeculativeType = persistedSpeculativeType; - loadSpecDraftNMax = null; + loadSpeculativeType = + pendingLoadConfig?.speculativeType != null + ? normalizeSpeculativeType(pendingLoadConfig.speculativeType) + : persistedSpeculativeType; + loadSpecDraftNMax = pendingLoadConfig?.specDraftNMax ?? null; // Keep the click-time snapshot in lock-step with the store reset so // the load below sizes against the cleared per-model knobs, not the // previous model's (gpuMemoryMode is standing, so left as captured). - loadCustomContextLength = null; - loadSelectedGpuIds = null; - loadGpuLayers = GPU_LAYERS_AUTO; - loadNCpuMoe = 0; + // An explicit staged config from run-settings still wins. + loadCustomContextLength = + pendingLoadConfig?.customContextLength ?? null; + loadSelectedGpuIds = + pendingLoadConfig?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds) + : null; + loadGpuLayers = pendingLoadConfig?.gpuLayers ?? GPU_LAYERS_AUTO; + loadNCpuMoe = pendingLoadConfig?.nCpuMoe ?? 0; loadSplitRatio = null; } @@ -1289,12 +1325,19 @@ export function useChatModelRuntime() { prog.expected_bytes, dlSamples, ); - setLoadProgress({ - percent: pct, - label: progressLabel, - phase: "downloading", - }); - if (loadToastDismissedRef.current) return; + // loadProgress state is only read by the dismissed-toast inline + // status. Writing it while the toast is visible re-renders the + // whole chat page every poll — cheap in Chrome, janky in the + // desktop WebView2 (laggy typing). Feed the toast directly and + // only touch state when the inline view is actually live. + if (loadToastDismissedRef.current) { + setLoadProgress({ + percent: pct, + label: progressLabel, + phase: "downloading", + }); + return; + } toast(null, { id: toastId, ...modelLoadToastOptions( @@ -1316,19 +1359,23 @@ export function useChatModelRuntime() { const est = estimate(dlSamples, prog.downloaded_bytes, 0); const rateSuffix = est.stable ? ` • ${formatRate(est.rate)}` : ""; - setLoadProgress({ - percent: null, - label: `${dlGb.toFixed(1)} GB downloaded${rateSuffix}`, - phase: "downloading", - }); + // Inline-status-only state; skip the chat-page re-render unless it's shown. + if (loadToastDismissedRef.current) { + setLoadProgress({ + percent: null, + label: `${dlGb.toFixed(1)} GB downloaded${rateSuffix}`, + phase: "downloading", + }); + } } else if (prog.progress >= 1 && hasShownProgress) { downloadComplete = true; - setLoadProgress({ - percent: 100, - label: "Download complete", - phase: "starting", - }); - if (!loadToastDismissedRef.current) { + if (loadToastDismissedRef.current) { + setLoadProgress({ + percent: 100, + label: "Download complete", + phase: "starting", + }); + } else { toast(null, { id: toastId, ...modelLoadToastOptions( @@ -1382,12 +1429,17 @@ export function useChatModelRuntime() { formatEta(est.eta) !== "--" ? ` • ${formatEta(est.eta)} left` : "" }` : base; - setLoadProgress({ - percent: pct, - label, - phase: "starting", - }); - if (loadToastDismissedRef.current) return; + // Inline-status-only state (see pollDownload): while the toast is + // up, skip the state write so the chat page doesn't re-render every + // poll during "Starting model" — the desktop WebView2 typing-lag fix. + if (loadToastDismissedRef.current) { + setLoadProgress({ + percent: pct, + label, + phase: "starting", + }); + return; + } toast(null, { id: toastId, ...modelLoadToastOptions( diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 547c3f374e..f85ff3246b 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -222,7 +222,9 @@ export function applyActiveModelStatusToStore( incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null; const incomingSplit = incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null; - const incomingGpuIds = status.is_gguf ? (status.gpu_ids ?? null) : null; + const incomingGpuIds = status.is_gguf + ? (status.requested_gpu_ids ?? status.gpu_ids ?? null) + : null; const gpuStatusChanged = prevState.loadedGpuMemoryMode !== incomingGpuMode || prevState.loadedGpuLayers !== incomingGpuLayers || diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx index f3e1594795..2fafeab7d6 100644 --- a/studio/frontend/src/features/chat/permission-mode-select.tsx +++ b/studio/frontend/src/features/chat/permission-mode-select.tsx @@ -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 { - ChevronDown, - CircleAlert, - CircleOff, - Hand, - ShieldCheck, -} from "lucide-react"; +import { ChevronDown, CircleAlert, Hand, ShieldCheck } from "lucide-react"; +import type { ComponentType } from "react"; import { useState } from "react"; import { @@ -29,6 +24,7 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; +import { SparklesGlyph } from "@/lib/sparkles-icon"; import { Tick02Icon } from "@/lib/tick-icon"; import { cn } from "@/lib/utils"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -45,7 +41,7 @@ export const PERMISSION_MODE_OPTIONS: readonly { value: PermissionMode; label: string; description: string; - icon: typeof Hand; + icon: ComponentType<{ className?: string; strokeWidth?: number }>; }[] = [ { value: "ask", @@ -63,7 +59,7 @@ export const PERMISSION_MODE_OPTIONS: readonly { value: "off", label: "Run automatically", description: "Run tool calls without approval prompts inside the sandbox", - icon: CircleOff, + icon: SparklesGlyph, }, { value: "full", diff --git a/studio/frontend/src/features/chat/presets/preset-load-config.ts b/studio/frontend/src/features/chat/presets/preset-load-config.ts new file mode 100644 index 0000000000..1083655cf2 --- /dev/null +++ b/studio/frontend/src/features/chat/presets/preset-load-config.ts @@ -0,0 +1,244 @@ +// 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 { + applyPerModelConfigToRuntime, + currentRuntimePerModelConfig, + perModelConfigsEqual, +} from "@/features/model-picker"; +import { + CONTEXT_LENGTH_MIN, + DEFAULT_PER_MODEL_CONFIG, + DEFAULT_MAX_SEQ_LENGTH, + KV_CACHE_DTYPES, + MTP_SPECULATIVE_TYPES, + SPECULATIVE_TYPES, + normalizeMaxSeqLength, + type PerModelConfig, +} from "@/features/model-picker/model-config/per-model-config"; +import { + GPU_LAYERS_AUTO, + useChatRuntimeStore, + normalizeSpeculativeType, +} from "../stores/chat-runtime-store"; + +/** Load/runtime knobs saved in a chat preset (excludes per-model-only blobs). */ +export type PresetLoadConfig = Pick< + PerModelConfig, + | "customContextLength" + | "maxSeqLength" + | "kvCacheDtype" + | "speculativeType" + | "specDraftNMax" + | "tensorParallel" + | "gpuMemoryMode" + | "gpuLayers" + | "nCpuMoe" +>; + +const VALID_KV_CACHE_DTYPES = new Set(KV_CACHE_DTYPES); +const VALID_SPECULATIVE_TYPES = new Set(SPECULATIVE_TYPES); + +export const EMPTY_PRESET_LOAD_CONFIG: PresetLoadConfig = { + customContextLength: null, + maxSeqLength: null, + kvCacheDtype: null, + speculativeType: null, + specDraftNMax: null, + tensorParallel: false, +}; + +function toComparablePerModelConfig( + config: PresetLoadConfig, +): PerModelConfig { + return { + ...DEFAULT_PER_MODEL_CONFIG, + ...config, + chatTemplateOverride: null, + selectedGpuIds: null, + }; +} + +export function normalizePresetLoadConfig( + raw: unknown, +): PresetLoadConfig | undefined { + if (raw == null || typeof raw !== "object" || Array.isArray(raw)) { + return undefined; + } + const partial = raw as Record; + const rawSpecType = + typeof partial.speculativeType === "string" + ? normalizeSpeculativeType(partial.speculativeType) + : null; + const speculativeType = rawSpecType ?? null; + const specDraftNMax = + speculativeType != null && + MTP_SPECULATIVE_TYPES.has(speculativeType) && + typeof partial.specDraftNMax === "number" && + Number.isFinite(partial.specDraftNMax) + ? Math.max(1, Math.min(16, Math.round(partial.specDraftNMax))) + : null; + const gpuMemoryMode = + partial.gpuMemoryMode === "manual" ? ("manual" as const) : undefined; + let gpuLayers: number | undefined; + if (typeof partial.gpuLayers === "number" && Number.isFinite(partial.gpuLayers)) { + gpuLayers = partial.gpuLayers < 0 ? GPU_LAYERS_AUTO : Math.floor(partial.gpuLayers); + } + let nCpuMoe: number | undefined; + if (typeof partial.nCpuMoe === "number" && Number.isFinite(partial.nCpuMoe)) { + nCpuMoe = Math.max(0, Math.floor(partial.nCpuMoe)); + } + + const normalized: PresetLoadConfig = { + customContextLength: + typeof partial.customContextLength === "number" && + Number.isFinite(partial.customContextLength) && + partial.customContextLength > 0 + ? Math.max(CONTEXT_LENGTH_MIN, Math.floor(partial.customContextLength)) + : null, + maxSeqLength: normalizeMaxSeqLength(partial.maxSeqLength as number | null), + kvCacheDtype: + typeof partial.kvCacheDtype === "string" && + VALID_KV_CACHE_DTYPES.has(partial.kvCacheDtype) + ? partial.kvCacheDtype + : null, + speculativeType: + speculativeType && VALID_SPECULATIVE_TYPES.has(speculativeType) + ? speculativeType + : null, + specDraftNMax, + tensorParallel: + typeof partial.tensorParallel === "boolean" + ? partial.tensorParallel + : false, + ...(gpuMemoryMode ? { gpuMemoryMode } : {}), + ...(gpuLayers !== undefined ? { gpuLayers } : {}), + ...(nCpuMoe !== undefined ? { nCpuMoe } : {}), + }; + + return hasPresetLoadConfig(normalized) ? normalized : undefined; +} + +export function hasPresetLoadConfig( + config?: PresetLoadConfig | null, +): boolean { + return !isSamePresetLoadConfig(config, EMPTY_PRESET_LOAD_CONFIG); +} + +export function isSamePresetLoadConfig( + a?: PresetLoadConfig | null, + b?: PresetLoadConfig | null, +): boolean { + return perModelConfigsEqual( + toComparablePerModelConfig({ ...EMPTY_PRESET_LOAD_CONFIG, ...a }), + toComparablePerModelConfig({ ...EMPTY_PRESET_LOAD_CONFIG, ...b }), + ); +} + +export function capturePresetLoadConfig(): PresetLoadConfig | undefined { + const snapshot = currentRuntimePerModelConfig({ includeMaxSeqLength: true }); + const store = useChatRuntimeStore.getState(); + const isGguf = + store.activeGgufVariant != null || + store.ggufContextLength != null || + (store.params.checkpoint?.toLowerCase().endsWith(".gguf") ?? false); + const effectiveContextLength = + snapshot.customContextLength ?? + (isGguf ? store.ggufContextLength : null); + const captured: PresetLoadConfig = { + customContextLength: effectiveContextLength ?? null, + maxSeqLength: normalizeMaxSeqLength(snapshot.maxSeqLength), + kvCacheDtype: snapshot.kvCacheDtype ?? null, + speculativeType: normalizeSpeculativeType(snapshot.speculativeType), + specDraftNMax: snapshot.specDraftNMax ?? null, + tensorParallel: snapshot.tensorParallel ?? false, + ...(snapshot.gpuMemoryMode === "manual" + ? { gpuMemoryMode: "manual" as const } + : {}), + ...(snapshot.gpuLayers != null && snapshot.gpuLayers >= 0 + ? { gpuLayers: snapshot.gpuLayers } + : snapshot.gpuMemoryMode === "manual" + ? { gpuLayers: GPU_LAYERS_AUTO } + : {}), + ...(snapshot.nCpuMoe != null && snapshot.nCpuMoe > 0 + ? { nCpuMoe: snapshot.nCpuMoe } + : {}), + }; + return hasPresetLoadConfig(coalesceDefaultLoadKnobs(captured)) + ? coalesceDefaultLoadKnobs(captured) + : undefined; +} + +function coalesceDefaultLoadKnobs( + captured: PresetLoadConfig, +): PresetLoadConfig { + const result: PresetLoadConfig = { ...captured }; + if (normalizeMaxSeqLength(result.maxSeqLength) === DEFAULT_MAX_SEQ_LENGTH) { + result.maxSeqLength = null; + } + const speculativeType = normalizeSpeculativeType(result.speculativeType); + if (speculativeType == null || speculativeType === "auto") { + result.speculativeType = null; + } + if ( + (result.gpuLayers == null || result.gpuLayers < 0) && + result.gpuMemoryMode !== "manual" + ) { + delete result.gpuLayers; + } + if ((result.nCpuMoe ?? 0) === 0) { + delete result.nCpuMoe; + } + return result; +} + +export function applyPresetLoadConfig( + config?: PresetLoadConfig | null, +): void { + if (config == null) { + return; + } + const store = useChatRuntimeStore.getState(); + applyPerModelConfigToRuntime({ + ...DEFAULT_PER_MODEL_CONFIG, + maxSeqLength: normalizeMaxSeqLength(config.maxSeqLength) ?? DEFAULT_MAX_SEQ_LENGTH, + customContextLength: config.customContextLength ?? null, + kvCacheDtype: config.kvCacheDtype ?? null, + speculativeType: config.speculativeType ?? null, + specDraftNMax: config.specDraftNMax ?? null, + tensorParallel: config.tensorParallel ?? false, + chatTemplateOverride: null, + gpuMemoryMode: config.gpuMemoryMode, + gpuLayers: config.gpuLayers, + nCpuMoe: config.nCpuMoe, + selectedGpuIds: store.selectedGpuIds, + }); +} + +export function formatPresetLoadConfigSummary( + config?: PresetLoadConfig | null, +): string | null { + if (!config || !hasPresetLoadConfig(config)) { + return null; + } + const parts: string[] = []; + if (config.customContextLength != null) { + parts.push(`Ctx ${config.customContextLength.toLocaleString()}`); + } + if (config.kvCacheDtype) { + parts.push(`KV ${config.kvCacheDtype}`); + } + if (config.speculativeType && config.speculativeType !== "auto") { + parts.push(`Spec ${config.speculativeType}`); + } + if (config.gpuMemoryMode === "manual") { + parts.push("GPU manual"); + } + if (config.gpuLayers != null && config.gpuLayers >= 0) { + parts.push(`${config.gpuLayers} layers`); + } + if (config.tensorParallel) { + parts.push("TP"); + } + return parts.length > 0 ? parts.join(" · ") : null; +} diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts index 23d79a35e1..61c94707b4 100644 --- a/studio/frontend/src/features/chat/presets/preset-policy.ts +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -5,12 +5,15 @@ import { DEFAULT_INFERENCE_PARAMS, type InferenceParams, } from "../types/runtime"; +import type { PresetLoadConfig } from "./preset-load-config"; export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS; export interface Preset { name: string; params: InferenceParams; + /** Optional GGUF/load knobs captured with the preset. */ + loadConfig?: PresetLoadConfig; } export type PresetOwnedParams = Pick< @@ -85,6 +88,7 @@ export function normalizeCustomPresets(presets: Preset[]): Preset[] { return { name, params: preset.params, + ...(preset.loadConfig ? { loadConfig: preset.loadConfig } : {}), }; }) .filter((preset): preset is Preset => preset !== null); diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx index e20e517787..192c4e2331 100644 --- a/studio/frontend/src/features/chat/projects-page.tsx +++ b/studio/frontend/src/features/chat/projects-page.tsx @@ -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(null); const [renameDraft, setRenameDraft] = useState(""); const [deleting, setDeleting] = useState(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() { - + @@ -511,10 +488,7 @@ export function ProjectsPage() { - - - - + {/* Create project (name + drag-and-drop sources) */} + {/* Rename project */} ({ - name: preset.name, - params: { - ...DEFAULT_INFERENCE_PARAMS, - ...preset.params, - }, - })) ?? state.customPresets + settings.customPresets?.map((preset) => { + const loadConfig = normalizePresetLoadConfig(preset.loadConfig); + return { + name: preset.name, + params: { + ...DEFAULT_INFERENCE_PARAMS, + ...preset.params, + }, + ...(loadConfig ? { loadConfig } : {}), + }; + }) ?? state.customPresets ); } diff --git a/studio/frontend/src/features/chat/sync-external-providers.ts b/studio/frontend/src/features/chat/sync-external-providers.ts new file mode 100644 index 0000000000..cd966f7fd4 --- /dev/null +++ b/studio/frontend/src/features/chat/sync-external-providers.ts @@ -0,0 +1,221 @@ +// 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 { + type ProviderRegistryEntry, + listProviderConfigs, + listProviderRegistry, + updateProviderConfig, +} from "./api/providers-api"; +import { + CUSTOM_BACKEND_PROVIDER_TYPE, + CUSTOM_PROVIDER_PRESETS, + type ExternalProviderConfig, + isCustomProviderType, + isPromptCacheTtl, + LEGACY_CUSTOM_PROVIDER_TYPE, + supportsProviderPromptCaching, + supportsProviderPromptCacheTtl, + supportsProviderReasoningToggle, +} from "./external-providers"; + +const ANTHROPIC_DATED_SNAPSHOT_SUFFIX = /-\d{8}$/; +const OPENAI_DEPRECATED_MODELS = new Set(["gpt-5.3"]); +const OPENROUTER_EXCLUDED_MODELS = new Set([ + "google/chirp-3", + "kwaivgi/kling-v3.0-pro", + "openai/whisper-1", + "openai/gpt-4o-mini-transcribe", + "recraft/recraft-v4-pro", +]); + +function normalizeUrl(input: string): string { + return input.trim().replace(/\/+$/, ""); +} + +export function resolveUiProviderTypeFromConfig( + configProviderType: string, + configDisplayName: string | null | undefined, + configBaseUrl: string | null | undefined, + registryRows: ProviderRegistryEntry[], + existingProviderType: string | undefined, +): string { + if (existingProviderType && isCustomProviderType(existingProviderType)) { + return existingProviderType; + } + if (configProviderType !== CUSTOM_BACKEND_PROVIDER_TYPE) { + return configProviderType; + } + const displayName = (configDisplayName ?? "").trim().toLowerCase(); + const matchingCustomPreset = CUSTOM_PROVIDER_PRESETS.find( + (preset) => preset.displayName.toLowerCase() === displayName, + ); + if (matchingCustomPreset) { + return matchingCustomPreset.providerType; + } + const openAiRegistry = registryRows.find( + (entry) => entry.provider_type === CUSTOM_BACKEND_PROVIDER_TYPE, + ); + if (!openAiRegistry) { + return configProviderType; + } + const openAiDisplayName = openAiRegistry.display_name.trim().toLowerCase(); + if (displayName.length > 0 && displayName !== openAiDisplayName) { + return LEGACY_CUSTOM_PROVIDER_TYPE; + } + const configUrl = normalizeUrl(configBaseUrl ?? ""); + const defaultUrl = normalizeUrl(openAiRegistry.base_url ?? ""); + if (configUrl.length > 0 && configUrl !== defaultUrl) { + return LEGACY_CUSTOM_PROVIDER_TYPE; + } + return configProviderType; +} + +export function pruneProviderModelIds( + providerType: string, + modelIds: string[], +): string[] { + if (providerType === "anthropic") { + return modelIds.filter((id) => !ANTHROPIC_DATED_SNAPSHOT_SUFFIX.test(id)); + } + if (providerType === "openai") { + return modelIds.filter((id) => !OPENAI_DEPRECATED_MODELS.has(id)); + } + if (providerType === "openrouter") { + return modelIds.filter((id) => !OPENROUTER_EXCLUDED_MODELS.has(id)); + } + return modelIds; +} + +/** Carry browser-local provider knobs through a backend sync rebuild. */ +export function mergeLocalProviderOptions( + existing: ExternalProviderConfig | undefined, + synced: ExternalProviderConfig, +): ExternalProviderConfig { + if (!existing) { + return synced; + } + const providerType = synced.providerType; + return { + ...synced, + enablePromptCaching: supportsProviderPromptCaching(providerType) + ? (existing.enablePromptCaching ?? synced.enablePromptCaching ?? true) + : undefined, + promptCacheTtl: + supportsProviderPromptCacheTtl(providerType) && + isPromptCacheTtl(existing.promptCacheTtl) + ? existing.promptCacheTtl + : synced.promptCacheTtl, + isReasoningModel: supportsProviderReasoningToggle(providerType) + ? (existing.isReasoningModel ?? synced.isReasoningModel) + : undefined, + openaiContainerTtlMinutes: + providerType === "openai" && + typeof existing.openaiContainerTtlMinutes === "number" && + existing.openaiContainerTtlMinutes >= 1 + ? Math.min(existing.openaiContainerTtlMinutes, 20) + : synced.openaiContainerTtlMinutes, + }; +} + +/** Merge enabled backend provider configs with local store state. */ +export async function syncExternalProvidersFromBackend( + existingProviders: ExternalProviderConfig[], +): Promise { + const [registryRows, configRows] = await Promise.all([ + listProviderRegistry(), + listProviderConfigs(), + ]); + + const existingById = new Map(); + for (const provider of existingProviders) { + existingById.set(provider.id, provider); + } + + const backfillTasks: Promise[] = []; + const syncedProviders = configRows + .filter((config) => config.is_enabled) + .map((config) => { + const existing = existingById.get(config.id); + const uiProviderType = resolveUiProviderTypeFromConfig( + config.provider_type, + config.display_name, + config.base_url, + registryRows, + existing?.providerType, + ); + const createdAt = Number.isFinite(Date.parse(config.created_at)) + ? Date.parse(config.created_at) + : Date.now(); + const updatedAt = Number.isFinite(Date.parse(config.updated_at)) + ? Date.parse(config.updated_at) + : Date.now(); + const registryEntry = + registryRows.find((entry) => entry.provider_type === uiProviderType) ?? + registryRows.find((entry) => entry.provider_type === config.provider_type); + const defaultModels = pruneProviderModelIds( + uiProviderType, + registryEntry?.default_models ?? [], + ); + const serverModels = pruneProviderModelIds( + uiProviderType, + config.models ?? [], + ); + const serverAvailableModels = pruneProviderModelIds( + uiProviderType, + config.available_models ?? [], + ); + const savedModels = existing?.models ?? []; + const savedAvailableModels = existing?.availableModels ?? []; + const resolvedModels = pruneProviderModelIds( + uiProviderType, + serverModels.length > 0 + ? serverModels + : savedModels.length > 0 + ? savedModels + : defaultModels, + ); + const resolvedAvailableModels = pruneProviderModelIds( + uiProviderType, + serverAvailableModels.length > 0 + ? serverAvailableModels + : savedAvailableModels.length > 0 + ? savedAvailableModels + : defaultModels, + ); + const needsModelBackfill = + serverModels.length === 0 && savedModels.length > 0; + const needsAvailableBackfill = + serverAvailableModels.length === 0 && savedAvailableModels.length > 0; + if (needsModelBackfill || needsAvailableBackfill) { + backfillTasks.push( + updateProviderConfig(config.id, { + models: resolvedModels, + availableModels: resolvedAvailableModels, + }), + ); + } + const synced: ExternalProviderConfig = { + id: config.id, + providerType: uiProviderType, + name: config.display_name, + baseUrl: config.base_url ?? "", + models: resolvedModels, + availableModels: resolvedAvailableModels, + enablePromptCaching: supportsProviderPromptCaching(uiProviderType) + ? (existing?.enablePromptCaching ?? true) + : undefined, + isReasoningModel: supportsProviderReasoningToggle(uiProviderType) + ? existing?.isReasoningModel === true + : undefined, + createdAt: existing?.createdAt ?? createdAt, + updatedAt, + }; + return mergeLocalProviderOptions(existing, synced); + }); + + if (backfillTasks.length > 0) { + await Promise.allSettled(backfillTasks); + } + return syncedProviders; +} diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 0d0b92de67..b03046f6e7 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -190,7 +190,10 @@ export interface LoadModelResponse { n_layers?: number | null; /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ n_moe_layers?: number; + /** Effective GPU placement after fit-time narrowing. */ gpu_ids?: number[] | null; + /** User-requested GPU placement pool before fit-time narrowing. */ + requested_gpu_ids?: number[] | null; } export interface UnloadModelRequest { @@ -242,7 +245,10 @@ export interface InferenceStatusResponse { /** n_ctx the active GGUF load was invoked with (0 = Auto); re-seeds a * Manual + Auto-layers context pin on hydration. Null for non-GGUF. */ requested_context_length?: number | null; + /** Effective GPU placement after fit-time narrowing. */ gpu_ids?: number[] | null; + /** User-requested GPU placement pool before fit-time narrowing. */ + requested_gpu_ids?: number[] | null; n_layers?: number | null; /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ n_moe_layers?: number; diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index 4899cb9c83..48acfdee40 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -8,6 +8,7 @@ import { type PersistedChatSettings, type PersistedInferenceParams, } from "../api/chat-settings-api"; +import { normalizePresetLoadConfig } from "../presets/preset-load-config"; import { BUILTIN_PRESETS, defaultInferenceParams, @@ -152,6 +153,7 @@ function sanitizeInferenceParams( } function toFullPreset(preset: PersistedChatPreset): Preset { + const loadConfig = normalizePresetLoadConfig(preset.loadConfig); return { name: preset.name, params: { @@ -159,6 +161,7 @@ function toFullPreset(preset: PersistedChatPreset): Preset { ...preset.params, checkpoint: defaultInferenceParams.checkpoint, }, + ...(loadConfig ? { loadConfig } : {}), }; } @@ -174,7 +177,12 @@ function sanitizeCustomPresets( const name = item.name.trim(); if (!name) return null; const params = sanitizeInferenceParams(item.params); - return { name, params: params ?? {} }; + const loadConfig = normalizePresetLoadConfig(item.loadConfig); + return { + name, + params: params ?? {}, + ...(loadConfig ? { loadConfig } : {}), + }; }) .filter((preset): preset is PersistedChatPreset => preset !== null); @@ -183,6 +191,7 @@ function sanitizeCustomPresets( (preset, index) => ({ name: preset.name, params: presets[index]?.params ?? {}, + ...(preset.loadConfig ? { loadConfig: preset.loadConfig } : {}), }), ); } diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index f0e0950871..f844336d4a 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -64,7 +64,7 @@ export function ModelsHeader({ return (
void; + inputRef?: Ref; }) { return (
@@ -146,6 +158,7 @@ function MaxSeqLengthSetting({
void; displayValue?: string; info?: ReactNode; + inputRef?: Ref; }) { return (
@@ -199,6 +214,7 @@ function AdvancedGpuSlider({ {info && {info}}
) => void; layerCount: number | null; moeLayerCount: number | null; + gpuLayersInputRef?: Ref; + moeLayersInputRef?: Ref; }) { const gpuDevices = useGpuDevices(); const mode = config.gpuMemoryMode ?? "auto"; @@ -322,6 +342,7 @@ function GpuMemorySettings({ <> ) => void; @@ -407,6 +431,8 @@ function GgufAdvancedSettings({ onEditTemplate: () => void; layerCount: number | null; moeLayerCount: number | null; + gpuLayersInputRef?: Ref; + moeLayersInputRef?: Ref; }) { return ( <> @@ -415,7 +441,8 @@ function GgufAdvancedSettings({ KV Cache Dtype Lower KV cache precision to save VRAM at the cost of some quality. - f16/bf16 are full precision; q8_0/q5_1/q4_1 are quantized. + f16 is the default; bf16 and f32 are full precision; q8_0 through + iq4_nl are quantized. { + const files = Array.from(e.target.files ?? []); + e.target.value = ""; + addFiles(files); + }} + /> + {staged.length === 0 ? ( + + ) : ( +
+
    + {staged.map((entry) => ( +
  • + + + {entry.file.name} + + + {formatSize(entry.file.size)} + + +
    • + ))} +
+ +
+ )} + + + ); +} diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts index 8d6433d8c3..bdab7b0518 100644 --- a/studio/frontend/src/features/rag/components/use-rag-documents.ts +++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts @@ -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( diff --git a/studio/frontend/src/features/rag/components/vision-overrides.ts b/studio/frontend/src/features/rag/components/vision-overrides.ts new file mode 100644 index 0000000000..674484970b --- /dev/null +++ b/studio/frontend/src/features/rag/components/vision-overrides.ts @@ -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, + }; +} diff --git a/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx b/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx index 0a0eff849b..e777826315 100644 --- a/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx +++ b/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx @@ -1,10 +1,18 @@ // 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 { type ReactElement, useCallback } from "react"; -import { Lock, LockOpen, Maximize2, Minus, Plus } from "lucide-react"; -import { Panel, useReactFlow } from "@xyflow/react"; import { Button } from "@/components/ui/button"; +import { Panel, useReactFlow } from "@xyflow/react"; +import { + Focus, + Lock, + LockOpen, + Maximize2, + Minimize2, + Minus, + Plus, +} from "lucide-react"; +import { type ReactElement, useCallback } from "react"; import { buildFitViewOptions } from "../../utils/graph/fit-view"; import { RECIPE_FLOATING_ICON_BUTTON_CLASS } from "../recipe-floating-icon-button-class"; @@ -12,12 +20,16 @@ type ViewportControlsProps = { interactive: boolean; lockDisabled?: boolean; onToggleInteractive: () => void; + maximized: boolean; + onToggleMaximize: () => void; }; export function ViewportControls({ interactive, lockDisabled = false, onToggleInteractive, + maximized, + onToggleMaximize, }: ViewportControlsProps): ReactElement { const { zoomIn, zoomOut, fitView, getNodes } = useReactFlow(); @@ -61,9 +73,23 @@ export function ViewportControls({ size="icon" className={RECIPE_FLOATING_ICON_BUTTON_CLASS} onClick={handleFitView} - aria-label="Fit view" + aria-label="Center view" > - + + + ); diff --git a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx index 30302b86a7..f1f5e0958e 100644 --- a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx +++ b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx @@ -237,6 +237,7 @@ export function RecipeStudioPage({ }, [setActiveView]); const [processorsOpen, setProcessorsOpen] = useState(false); const [interactive, setInteractive] = useState(true); + const [maximized, setMaximized] = useState(false); const [runtimeIslandMinimized, setRuntimeIslandMinimized] = useState(false); const [recentCompletedExecution, setRecentCompletedExecution] = useState(null); @@ -569,6 +570,16 @@ export function RecipeStudioPage({ [reactFlowInstance], ); + const toggleMaximize = useCallback(() => { + // The maximized surface is a fixed z-50 overlay that already covers the + // app sidebar (z-10/z-20), so we don't touch the sidebar's own state — that + // state is persisted in pin mode and mutating it here would leak the + // temporary collapse into the next page/session. + setMaximized((prev) => !prev); + // Container size changes; refit once the layout settles. + scheduleFitView({ delayMs: TAB_SWITCH_FIT_DELAY_MS }); + }, [scheduleFitView]); + useEffect(() => { if ( previousActiveViewRef.current !== activeView && @@ -587,6 +598,15 @@ export function RecipeStudioPage({ } }, [activeView, reactFlowInstance]); + // The "Exit full view" control lives inside the editor canvas, which unmounts + // on other tabs. Drop full-view mode when leaving the editor so Easy/Runs + // aren't left under the fixed overlay. + useEffect(() => { + if (activeView !== "editor" && maximized) { + setMaximized(false); + } + }, [activeView, maximized]); + useEffect(() => { if ( !reactFlowInstance || @@ -732,6 +752,8 @@ export function RecipeStudioPage({ interactive={canvasInteractive} lockDisabled={executionLocked} onToggleInteractive={toggleInteractive} + maximized={maximized} + onToggleMaximize={toggleMaximize} /> {islandExecution && (isExecutionInProgress(islandExecution.status) || @@ -773,10 +795,25 @@ export function RecipeStudioPage({ } return ( -
-
+