Merge remote-tracking branch 'upstream/main' into fix/rocm-version-test-isolation

This commit is contained in:
LeoBorcherding 2026-07-27 18:42:11 -05:00
commit 5a94eab6eb
449 changed files with 46396 additions and 2835 deletions

View file

@ -7,7 +7,7 @@
#
# Why a separate workflow:
# - studio-backend-ci.yml's "Repo tests (CPU)" job already auto-discovers
# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 16
# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 17
# Bucket-A tests below live inside those --ignore dirs (CPU-runnable but
# historically excluded with their GPU siblings); pulling them out into
# a sibling job keeps the existing 760-passed baseline stable while we
@ -274,6 +274,7 @@ jobs:
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
@ -365,6 +366,7 @@ jobs:
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \
@ -2129,7 +2131,7 @@ jobs:
pip show unsloth_zoo
echo "::endgroup::"
echo "Consolidated job done. Coverage:"
echo " - 16 unsloth Bucket-A tests under tests/saving/ + tests/utils/"
echo " - 17 unsloth Bucket-A tests under tests/saving/ + tests/utils/"
echo " - unsloth_zoo @ ${UNSLOTH_ZOO_REF} pytest tests/ (5 GPU cases deselected)"
echo " - unsloth_zoo.compiler.test_apply_fused_lm_head"

View file

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

3
.gitignore vendored
View file

@ -238,4 +238,5 @@ package-lock.json
!studio/package-lock.json
llama.cpp/
# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo.
/~/
~/
/temp/

View file

@ -103,7 +103,7 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
* **macOS:** Training, MLX and GGUF inference are ALL supported.
* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd).
* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819).
* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819). Vulkan accelerates GGUF inference only; training still requires a supported PyTorch or MLX backend.
* **Multi-GPU:** Available now, with a major upgrade on the way
#### macOS, Linux, WSL:
@ -112,12 +112,28 @@ curl -fsSL https://unsloth.ai/install.sh | sh
```
Use the same command to update.
To force the Vulkan llama.cpp backend, set `UNSLOTH_FORCE_VULKAN=1` **before installing or updating**. The setting selects the llama.cpp binary bundle, so setting it only when launching Studio cannot replace an existing CPU bundle:
```bash
export UNSLOTH_FORCE_VULKAN=1
curl -fsSL https://unsloth.ai/install.sh | sh
```
#### Windows:
```powershell
irm https://unsloth.ai/install.ps1 | iex
```
Use the same command to update.
To force the Vulkan llama.cpp backend, set the environment variable before running the installer or updater:
```powershell
$env:UNSLOTH_FORCE_VULKAN=1
irm https://unsloth.ai/install.ps1 | iex
```
Re-running the current installer replaces a previously selected CPU bundle when the backend differs. A separate Vulkan SDK is not required; the GPU driver must provide a working Vulkan runtime.
#### Launch
```bash
unsloth studio -p 8888
@ -263,6 +279,8 @@ unsloth studio -H 0.0.0.0 -p 8888
```
The Cloudflare tunnel is **off by default**: `-H 0.0.0.0` exposes the raw port only, not a public internet URL. Pair the wildcard bind with `--cloudflare` (`unsloth studio -H 0.0.0.0 --cloudflare`) to also publish a public `https://*.trycloudflare.com` link, or prefer `--secure` (above), which keeps the raw port private. `--cloudflare` has no effect on a loopback bind.
On a wildcard bind Unsloth works out the address to share by asking `ifconfig.me` for the public IP, then asks `check-host.net` whether that port is reachable so it can tell you if a firewall is in the way. Both contact a third party. Set `UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK=1` to skip them; the banner then shows the LAN address and no reachability line.
The first time Unsloth is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Unsloth shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`):

View file

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

View file

@ -257,6 +257,51 @@ run_install_cmd_retry() {
done
}
# True when the runtime target is gfx906 (MI50/Radeon VII): the prebuilt AMD
# bitsandbytes wheel carries no gfx906 kernels, and force-reinstalling it would
# clobber a user's source-built bnb (the only 4-bit path on this arch) on every
# `studio update`. So skip the auto-install and leave whatever bnb is present.
# _gfx906_target is set during torch-index resolution; also honor an explicit
# UNSLOTH_ROCM_GFX_ARCH so a pinned-index install still skips. The override is
# normalized (gfx906:sramecc-:xnack- -> gfx906) so a copied HIP gcnArchName counts.
_is_gfx906_bnb_skip() {
[ "${_gfx906_target:-false}" = true ] && return 0
_bnb_gfx_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]')
_bnb_gfx_env=${_bnb_gfx_env%%:*}
[ "$_bnb_gfx_env" = "gfx906" ] && return 0
# A pinned index (UNSLOTH_TORCH_INDEX_URL/_FAMILY) skips the reroute block that
# sets _gfx906_target, so a real gfx906 host with a pinned rocm6.3 index and no
# UNSLOTH_ROCM_GFX_ARCH would otherwise clobber a source-built bnb. Probe here
# in that gap; skip only when gfx906 is the SOLE distinct arch (mixed hosts
# opt in via the env var, mirroring the reroute block's de-dup rule).
if [ -z "$_bnb_gfx_env" ] && [ "${_torch_index_pinned:-false}" = true ]; then
_bnb_gfx_probe=$(_probe_amd_gfx_arch | awk 'NF && !seen[$0]++')
[ "$_bnb_gfx_probe" = "gfx906" ] && return 0
fi
return 1
}
# `pip install unsloth` resolves its unconditional bitsandbytes dep to a generic
# CUDA wheel (no gfx906 kernels) once we skip the prebuilt one. Snapshot bnb before
# the unsloth install, then drop a freshly pulled wheel afterwards while leaving a
# pre-existing source build in place.
_gfx906_bnb_installed() {
"$_VENV_PY" -c "import importlib.util as u, sys; sys.exit(0 if u.find_spec('bitsandbytes') else 1)" >/dev/null 2>&1
}
_gfx906_bnb_snapshot() {
_gfx906_bnb_absent_before=false
_is_gfx906_bnb_skip || return 0
_gfx906_bnb_installed || _gfx906_bnb_absent_before=true
}
_gfx906_bnb_prune() {
_is_gfx906_bnb_skip || return 0
[ "${_gfx906_bnb_absent_before:-false}" = true ] || return 0
_gfx906_bnb_installed || return 0
substep "gfx906: removing generic bitsandbytes pulled in as a dependency (no gfx906 kernels; build from source for 4-bit QLoRA)" "$C_WARN"
uv pip uninstall --python "$_VENV_PY" bitsandbytes >/dev/null 2>&1 \
|| "$_VENV_PY" -m pip uninstall -y bitsandbytes >/dev/null 2>&1 || true
}
# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
@ -655,6 +700,15 @@ _apt_distro_description() {
)
}
# ── Helper: can the controlling terminal actually be opened for reading? ──
# `test -r` only checks permission bits, which look fine in containers and
# systemd units where open() then fails with ENXIO. Probe with a real open.
# The subshell is required: in dash a failed redirection on the special
# builtin `:` exits the whole script.
_can_read_tty() {
( : </dev/tty ) >/dev/null 2>&1
}
# ── Helper: install packages via apt, escalating to sudo only if needed ──
# Usage: _smart_apt_install pkg1 pkg2 pkg3 ...
_smart_apt_install() {
@ -695,24 +749,63 @@ _smart_apt_install() {
echo " from your distro's official repositories (not a third-party tarball)."
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo ""
printf " Accept? [Y/n] "
if [ -r /dev/tty ]; then
read -r REPLY </dev/tty || REPLY="y"
else
REPLY="y"
fi
case "$REPLY" in
[nN]*)
if _can_read_tty; then
printf " Accept? [Y/n] "
# The device opened, so a failed read is EOF, not consent: decline,
# as the autostart prompt below does. Enter is still yes (a
# successful read of an empty line).
read -r REPLY </dev/tty || REPLY="n"
case "$REPLY" in
[nN]*)
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
;;
esac
# Mirror the headless branch: on a sudoers denial, a wrong password
# or an apt error, say what to run by hand instead of letting set -e
# abort on a bare sudo/apt message.
if sudo apt-get update -y </dev/null &&
sudo apt-get install -y $_STILL_MISSING </dev/null; then
:
else
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " Could not install these packages: $_STILL_MISSING"
echo " See the error above."
echo " Please install them first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
;;
*)
sudo apt-get update -y </dev/null
sudo apt-get install -y $_STILL_MISSING </dev/null
;;
esac
fi
else
# Nobody can answer a prompt or type a password here. -n makes sudo
# refuse rather than prompt into a closed stdin, which is how #7307
# died. Probe with the real commands: `sudo -l` answers whether they
# are *authorized*, not whether running them needs authentication.
# -k ignores any cached timestamp, so only a real NOPASSWD rule gets
# through, not someone's sudo in another shell minutes ago. Per
# sudo(8), -k alongside a command ignores the cached credentials and
# "will not update" them, so other sessions keep theirs.
echo " No terminal to confirm on; trying passwordless sudo."
if sudo -n -k apt-get update -y </dev/null &&
sudo -n -k apt-get install -y $_STILL_MISSING </dev/null; then
echo " Installed with passwordless sudo."
else
echo ""
echo " Could not install these packages: $_STILL_MISSING"
echo " Detected ${_ad_desc}."
# Either sudo refused, or apt failed on a bad repo, dpkg lock or
# network outage. sudo exits 1 on an auth/config problem and
# when the command cannot be executed, but otherwise passes the
# command's own status through, so state both causes.
echo " Either sudo needs a password here, or apt-get itself"
echo " failed; see the error above. With no terminal to"
echo " authenticate on, this cannot be done unattended."
echo " Please install them first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
fi
fi
else
echo ""
echo " sudo is not available on this system."
@ -2260,6 +2353,7 @@ _amd_arch_index_family_for_gfx() {
gfx1201|gfx1200) echo gfx120X-all ;;
gfx1151) echo gfx1151 ;;
gfx1150) echo gfx1150 ;;
gfx1152) echo gfx1152 ;;
gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;;
gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;;
gfx90a) echo gfx90a ;;
@ -2271,12 +2365,14 @@ _amd_arch_index_family_for_gfx() {
# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable).
_infer_amd_gfx_arch_from_gpu_name() {
case "$1" in
*"9070 XT"*|*9080*) echo gfx1201 ;;
*9070*|*9060*) echo gfx1200 ;;
*9070*|*9080*) echo gfx1201 ;;
*9060*) echo gfx1200 ;;
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;;
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;;
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;;
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;;
*"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) echo gfx1150 ;;
*"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1152 ;;
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) echo gfx1102 ;;
*"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) echo gfx1101 ;;
*"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) echo gfx1100 ;;
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;;
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;;
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;;
@ -2316,10 +2412,14 @@ _infer_linux_amd_gfx_arch() {
echo gfx1151
return 0
fi
if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]' /proc/cpuinfo 2>/dev/null; then
echo gfx1150
return 0
fi
if [ -n "$_gpu_evidence" ] && grep -qiE '860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
echo gfx1152
return 0
fi
if command -v lspci >/dev/null 2>&1; then
# A non-AMD controller can enumerate first (Intel/ASPEED before an AMD
# dGPU), so scan every display-class line and take the first AMD one
@ -3055,7 +3155,7 @@ if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
# whole handoff (a user-set override re-exports unchanged).
export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"
case "$_linux_inferred_gfx" in
gfx1201|gfx1200|gfx1151|gfx1150)
gfx1201|gfx1200|gfx1151|gfx1150|gfx1152)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
@ -3124,7 +3224,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"
@ -3241,10 +3341,20 @@ case "$_torch_index_leaf" in
if (n > 0) print vals[idx]
}')
fi
# An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the
# MI50 / Radeon VII path and must win over Strix probe-order detection on a
# mixed Strix + MI50 host, so the Strix reroute is suppressed when it is set.
# Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) and
# trim whitespace (mirrors the Python .strip()) so the feature-flag suffix or
# a stray newline does not defeat the exact gfx906 comparisons below.
_gfx906_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]')
_gfx906_env=${_gfx906_env%%:*}
_strix_gfx=""
case "$_runtime_gfx" in
gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
esac
if [ "$_gfx906_env" != "gfx906" ]; then
case "$_runtime_gfx" in
gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;;
esac
fi
# Skip rocm7.13+ generic indexes: they already ship the fixes, so the
# arch build (rocm7.13) would be a downgrade rather than a rescue.
if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then
@ -3272,6 +3382,57 @@ case "$_torch_index_leaf" in
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
# ── MI50 / Radeon VII (gfx906, Vega 20): legacy community-supported path ──
# Newer rocm wheel families bundle ROCm libraries whose Tensile kernels
# dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read for gfx906",
# ROCm/TheRock#1844), so a rocm6.4+/7.x index installs a torch that fails
# at the first BLAS call. The rocm6.3 index is the last one whose wheels
# run on gfx906 (torch 2.7.0 verified on MI50 32GB; up to 2.9 in community
# use). Reroute any newer picked index; leave rocm6.0-6.3 alone.
#
# Target resolution: an explicit UNSLOTH_ROCM_GFX_ARCH wins (lets a host
# whose rocminfo/amd-smi emit no gfx token still opt in; _gfx906_env was
# lowercased above, before the Strix block it suppresses). Otherwise only
# treat gfx906 as the target when it is the SOLE distinct arch present:
# _gfx_all is de-duplicated by visible index, which loses per-device
# ordinals on a mixed host, so a non-gfx906 selection must never be
# downgraded to rocm6.3 -- such hosts set UNSLOTH_ROCM_GFX_ARCH to opt in.
_gfx906_target=false
if [ -n "$_gfx906_env" ]; then
[ "$_gfx906_env" = "gfx906" ] && _gfx906_target=true
elif [ -n "$_gfx_all" ]; then
_gfx906_uniq=$(printf '%s\n' "$_gfx_all" | awk 'NF && !seen[$0]++')
[ "$_gfx906_uniq" = "gfx906" ] && _gfx906_target=true
fi
# gfx906 always trains from the PyTorch rocm6.3 wheels, never the Radeon repo
# (repo.radeon.com wheels carry no gfx906 BLAS kernels). Clear the Radeon
# marketing-name flag as soon as gfx906 is the target -- even when the host
# already picks rocm6.0-6.3 and the reroute below is a no-op -- so a Radeon VII
# does not divert to the radeon branch on those versions.
if [ "$_gfx906_target" = true ]; then
_amd_gpu_radeon=false
fi
if [ "$_gfx906_target" = true ] && ! _rocm_leaf_below "$_torch_index_leaf" 6 4; then
echo "" >&2
echo " [WARN] gfx906 (MI50 / Radeon VII / Vega 20) detected -- routing torch to the" >&2
echo " [WARN] rocm6.3 index: it is the last wheel family that runs on gfx906 (newer" >&2
echo " [WARN] rocm wheels ship without gfx906 BLAS kernels and fail at first use)." >&2
echo " [WARN] gfx906 is a community-maintained legacy path: 16-bit LoRA and full" >&2
echo " [WARN] finetuning work out of the box; bitsandbytes 4-bit QLoRA requires a" >&2
echo " [WARN] source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd)." >&2
echo "" >&2
_amd_gfx906_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
while [ "${_amd_gfx906_base%/}" != "$_amd_gfx906_base" ]; do
_amd_gfx906_base="${_amd_gfx906_base%/}"
done
TORCH_INDEX_URL="${_amd_gfx906_base}/rocm6.3"
# Reset to the default (<2.11) window: a rocm7.2 pick raised the floor
# to 2.11 above, which the rocm6.3 index (torch <= 2.9.x) cannot satisfy.
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
# (_amd_gpu_radeon already cleared above for every gfx906 target.)
fi
;;
esac
fi # _torch_index_pinned guard (Radeon + Strix reroute)
@ -3339,12 +3500,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)
@ -3496,6 +3659,7 @@ for _p in ('torch', 'torchvision', 'torchaudio'):
if [ "$_MIGRATED" = true ]; then
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the ROCm repair below fires.
_gfx906_bnb_snapshot
substep "upgrading unsloth in migrated environment..."
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
@ -3537,13 +3701,18 @@ if [ "$_MIGRATED" = true ]; then
# existing ROCm installs gain the AMD bitsandbytes build without a
# fresh reinstall.
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
if _is_gfx906_bnb_skip; then
substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
else
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
fi
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
_gfx906_bnb_prune
fi
elif [ -n "$TORCH_INDEX_URL" ]; then
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
@ -3734,8 +3903,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# host stays in GGUF-only mode rather than pulling in bitsandbytes,
# which is only useful once torch is present for training.
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
if _is_gfx906_bnb_skip; then
substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
else
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
fi
fi
_gfx906_bnb_snapshot
# Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
@ -3786,6 +3960,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
_gfx906_bnb_prune
fi
else
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
@ -4046,9 +4221,11 @@ echo ""
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then
echo ""
printf " Start Unsloth Studio now? [Y/n] "
# No readable answer (closed/EOF tty) defaults to no; Enter is still yes.
if [ -r /dev/tty ]; then
# Prompt only when something can answer: `test -r` passes on the unopenable
# /dev/tty found in containers, leaving a dangling question in the log.
if _can_read_tty; then
printf " Start Unsloth Studio now? [Y/n] "
read -r _reply </dev/tty || _reply="n"
else
_reply="n"

View file

@ -52,14 +52,14 @@ def _normalise_on(on_field):
def _load_workflow(path: Path):
try:
return yaml.safe_load(path.read_text())
return yaml.safe_load(path.read_text(encoding = "utf-8"))
except Exception as exc:
print(f"ERROR: failed to parse {path}: {exc}", file = sys.stderr)
sys.exit(2)
def _extract_cache_keys(path: Path) -> list[str]:
text = path.read_text()
text = path.read_text(encoding = "utf-8")
keys: list[str] = []
for m in re.finditer(r"(?:^|\n)\s*key:\s*([^\n]+)", text):
keys.append(m.group(1).strip())
@ -104,7 +104,7 @@ def main() -> int:
for t in RESTRICTED_TRIGGERS:
if t in triggers:
text = path.read_text()
text = path.read_text(encoding = "utf-8")
if "lint:workflow_triggers-allow-workflow_run" not in text:
findings.append(
f"{path.name}: RESTRICTED trigger '{t}' requires an "

View file

@ -98,6 +98,14 @@
"evidence": "L587: while True: sha256:06c2c7f15d73bf192e5e3272c5ff5fcaeff7f6774fef5f4eca6ef473ae50e2b3",
"evidence_hash": "57acd497f404c203e4450d0580ad85aa8a33406e8d64ad06fbac6cf47d97b24d"
},
{
"package": "fastapi",
"file": "fastapi/routing.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L592: while True: sha256:84283c09277ded3296998b2a6a838744457b606829cf5ab5d0da6f222ff020a0",
"evidence_hash": "a7295004315e26a8f3c64fb837521e9fdd7268219bb43e000fb0236ab0259223"
},
{
"package": "fastmcp-slim",
"file": "fastmcp/cli/apps_dev.py",

View file

@ -1,134 +1,145 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/github/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/github/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"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",
"<div class=\"align-center\">\n",
"<a href=\"https://unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
"<a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord button.png\" width=\"145\"></a>\n",
"<a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a> Join Discord if you need help + ⭐ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐\n",
"</div>\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": [
"<p align=\"left\"><img src=\"https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png\" width=\"600\"></p>"
],
"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",
"<div class=\"align-center\">\n",
" <a href=\"https://unsloth.ai\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
" <a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord.png\" width=\"145\"></a>\n",
" <a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a>\n",
"\n",
" Join Discord if you need help + ⭐️ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐️\n",
"\n",
" <b>This notebook is licensed <a href=\"https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0\">AGPL-3.0</a></b>\n",
"</div>"
],
"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",
"<div class=\"align-center\">\n",
"<a href=\"https://unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
"<a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord button.png\" width=\"145\"></a>\n",
"<a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a> Join Discord if you need help + ⭐ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐\n",
"</div>\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": [
"<p align=\"left\"><img src=\"https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png\" width=\"600\"></p>"
]
},
{
"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",
"<div class=\"align-center\">\n",
" <a href=\"https://unsloth.ai\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
" <a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord.png\" width=\"145\"></a>\n",
" <a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a>\n",
"\n",
" Join Discord if you need help + ⭐️ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐️\n",
"\n",
" <b>This notebook is licensed <a href=\"https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0\">AGPL-3.0</a></b>\n",
"</div>"
]
}
],
"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
}

View file

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

View file

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

View file

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

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -30,6 +30,7 @@ lora:
- "query"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -30,6 +30,7 @@ lora:
- "value"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -29,6 +29,7 @@ lora:
- "Wqkv"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -37,6 +37,7 @@ lora:
- "shared_mlp.output_linear"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -37,6 +37,7 @@ lora:
- "shared_mlp.output_linear"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -30,6 +30,7 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

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

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -33,6 +33,7 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -38,6 +38,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -37,6 +37,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "out_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -37,6 +37,7 @@ lora:
- "out_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -38,6 +38,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -33,6 +33,7 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -36,6 +36,7 @@ lora:
- "gate_up_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -36,6 +36,7 @@ lora:
- "gate_up_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
use_dora: false
logging:
enable_wandb: false

View file

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

View file

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

View file

@ -164,6 +164,21 @@ async def get_current_subject_allow_password_change(
)
# The literal the examples ship with; pasted unedited more often than a revoked key.
API_KEY_PLACEHOLDER = f"{API_KEY_PREFIX}YOUR_KEY"
def _invalid_api_key_detail(token: str) -> str:
"""Why the key failed. Only the example placeholder is called out; every real
key gets one indistinguishable message, so this leaks no key existence."""
if token == API_KEY_PLACEHOLDER:
return (
"This is the placeholder key from the example. Create an API key in "
f"Unsloth Studio under Settings > API and use it in place of {API_KEY_PLACEHOLDER}."
)
return "Invalid or expired API key"
async def _get_current_subject(
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
) -> str:
@ -176,7 +191,7 @@ async def _get_current_subject(
if username is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid or expired API key",
detail = _invalid_api_key_detail(token),
)
return username

View file

@ -44,7 +44,7 @@ def generate_bootstrap_password() -> str:
# Persisted from a previous run?
if _BOOTSTRAP_PW_PATH.is_file():
_bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
_bootstrap_password = _BOOTSTRAP_PW_PATH.read_text(encoding = "utf-8").strip()
if _bootstrap_password:
return _bootstrap_password
@ -57,7 +57,7 @@ def generate_bootstrap_password() -> str:
# Persist so the same passphrase survives restarts until password change.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password, encoding = "utf-8")
try:
os.chmod(_BOOTSTRAP_PW_PATH, 0o600)
except OSError:
@ -76,7 +76,7 @@ def _load_bootstrap_password() -> Optional[str]:
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
bootstrap_password = _BOOTSTRAP_PW_PATH.read_text(encoding = "utf-8").strip()
if bootstrap_password:
_bootstrap_password = bootstrap_password
return _bootstrap_password
@ -99,7 +99,7 @@ def clear_bootstrap_password() -> None:
# stale plaintext can't be re-seeded by generate_bootstrap_password()
# if a later reset-password deletes auth.db and re-validates it.
try:
_BOOTSTRAP_PW_PATH.write_text("")
_BOOTSTRAP_PW_PATH.write_text("", encoding = "utf-8")
cleared = True
except OSError:
cleared = False

View file

@ -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", encoding = "utf-8")
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(encoding = "utf-8").splitlines()
if len(lines) >= 2 and lines[0] and lines[1]:
return lines[0], lines[1]
except (OSError, UnicodeDecodeError) 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"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 22px; font-weight: 800;">
Unsloth Studio Login (Colab)
</h2>
<p style="color: #333333; margin: 0 0 12px 0; font-size: 14px; font-weight: bold;">
Log in as <code>{username}</code> with this password. This cell is visible only in
your notebook session.
</p>
<p style="color: #333333; margin: 0; font-size: 14px; font-family: monospace; font-weight: bold;">
Password: <code>{password}</code>
</p>
</div>
"""
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 "
'<code style="background:#f3f3f3;padding:2px 6px;border-radius:4px;">'
"cloudflare=True</code> 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 <code style="background:#f3f3f3;padding:2px 6px;'
'border-radius:4px;">start(cloudflare=True)</code> for a working link.'
)
return f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 26px; font-weight: 800;
display: flex; align-items: center; gap: 12px;">
<img src="https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/unsloth-gem.png"
height="48" style="display:block;">
Unsloth Studio is Ready!
</h2>
<p style="color: #333333; margin: 0 0 8px 0; font-size: 15px; font-weight: bold;">
{embed_note}
</p>
<p style="color: #666666; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
{short_url}
</p>
</div>
"""
return f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 26px; font-weight: 800;
@ -100,15 +311,52 @@ def show_link(port: int = 8888, *, _url: "str | None" = None):
</p>
</div>
"""
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"""
<p style="color: #000000; margin: 16px 0 0 0; font-size: 20px; font-weight: 800;">
Password
</p>
<p style="margin: 6px 0 0 0;"><code style="display: inline-block; font-size: 24px;
font-weight: 800; text-decoration: underline; background: #f3f3f3;
padding: 4px 10px; border-radius: 6px;">{password}</code></p>
<p style="color: #666666; margin: 6px 0 0 0; font-size: 12px;">
Log in as <code>{username}</code> with this password. Shown only in your
notebook session, and never included in the shared link.
</p>"""
return f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
@ -213,40 +480,55 @@ def _shareable_link_html(cloudflare_url: str) -> str:
Open Unsloth Studio
</a>
<p style="color: #333333; margin: 12px 0 0 0; font-size: 14px; font-weight: bold;">
This Cloudflare HTTPS link works from any device share it with anyone. The Colab view below only works in this tab.
This Cloudflare HTTPS link works from any device, so you can share it with anyone.
</p>
<p style="color: #333333; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
🔗 {cloudflare_url}
</p>
🔗 <a href="{cloudflare_url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
style="color: #000000; text-decoration: underline; cursor: pointer;">{cloudflare_url}</a>
</p>{login_block}
</div>
"""
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"""
<div style="font-family:system-ui,-apple-system,sans-serif;margin:8px 0;
@ -266,41 +548,110 @@ def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None):
</div>
""")
)
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):

View file

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

View file

@ -6,12 +6,14 @@
Run in a short-lived subprocess (``python _vulkan_probe.py <bindir>``) so the
Vulkan instance never lives in the long-running backend process. Loads the
bundled ggml Vulkan backend from ``<bindir>`` and prints one
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` line per device to stdout.
Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>\\t<name>`` line per device to
stdout. Indices are ggml's own Vulkan device ordinals, which need not match
nvidia-smi order. ``is_igpu`` (from ggml's device type) is ``1`` for an
integrated GPU sharing system RAM. ``total_bytes`` is the device-local heap;
the reader uses it to reserve absolute headroom on a discrete card (parity
with the CUDA/ROCm fit) and ignores it for an iGPU, whose "VRAM" is shared
system RAM. ``name`` is ggml's device description (the marketing name, e.g.
"AMD Radeon RX 9070 XT"); empty when the registry lookup fails.
Uses only the standard library so it stays runnable as a bare script.
"""
@ -24,15 +26,30 @@ import sys
_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
def _igpu_flags(base, lib, count: int) -> list[bool]:
"""Per-device integrated-GPU flags via ggml's backend registry.
def _igpu_flags_and_names(base, lib, count: int) -> tuple[list[bool], list[str]]:
"""Per-device integrated-GPU flags and descriptions via ggml's backend registry.
The Vulkan reg enumerates devices in the same order as
``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
i``), so reg index == device ordinal. Returns all-False on any failure so
the reader never over-caps a discrete card.
i``), so reg index == device ordinal. Returns all-False / empty-name on any
failure so the reader never over-caps a discrete card and the memory
readings still get through.
"""
flags = [False] * count
names = [""] * count
# The name lookup is bound OUTSIDE the type-detection try: a ggml-base
# without ggml_backend_dev_description (older/custom build) must degrade to
# unnamed devices, not abort before the iGPU flags are read (which would
# count an iGPU's shared RAM as VRAM).
describe = None
try:
base.ggml_backend_dev_description.restype = ctypes.c_char_p
base.ggml_backend_dev_description.argtypes = [ctypes.c_void_p]
describe = base.ggml_backend_dev_description
except Exception:
pass
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
@ -45,17 +62,31 @@ def _igpu_flags(base, lib, count: int) -> list[bool]:
reg = lib.ggml_backend_vk_reg()
if not reg:
return flags
return flags, names
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
if describe is not None:
try:
desc = describe(dev)
if desc:
# Tabs/newlines would corrupt the line protocol;
# spaces are safe.
names[i] = (
desc.decode("utf-8", errors = "replace")
.replace("\t", " ")
.replace("\n", " ")
.strip()
)
except Exception:
pass
except Exception:
# Best-effort: any failure degrades to "discrete" so the memory
# readings still get through instead of crashing the probe.
# Best-effort: any failure degrades to "discrete"/"unnamed" so the
# memory readings still get through instead of crashing the probe.
pass
return flags
return flags, names
def main() -> int:
@ -63,6 +94,14 @@ def main() -> int:
return 0
bindir = sys.argv[1]
# Device names can be non-ASCII (localized drivers); the platform-default
# stdout encoding (e.g. cp1252) would raise on them and lose the whole
# inventory. The reader decodes UTF-8 with the same error mode.
try:
sys.stdout.reconfigure(encoding = "utf-8", errors = "replace")
except Exception:
pass
# Hold add_dll_directory's handle for the rest of main() (the documented
# idiom) so bindir stays on the search path while the sibling ggml DLLs
# resolve below.
@ -96,12 +135,12 @@ def main() -> int:
]
count = lib.ggml_backend_vk_get_device_count()
igpu = _igpu_flags(base, lib, count)
igpu, names = _igpu_flags_and_names(base, lib, count)
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
rows.append("%d\t%d\t%d\t%d\t%s" % (i, free.value, int(igpu[i]), total.value, names[i]))
sys.stdout.write("\n".join(rows))
return 0

View file

@ -52,6 +52,13 @@ class ApiMonitorEntry:
total_tokens: Optional[int] = None
total_tokens_authoritative: bool = False
error: Optional[str] = None
# "request" (HTTP call) or "lifecycle" (model load/unload: event/reason, not a prompt; shared).
kind: str = "request"
event: Optional[str] = None
reason: Optional[str] = None
shared: bool = False
# 0-100 for a running download row; None when not applicable.
progress: Optional[float] = None
def snapshot(self, *, include_details: bool = True) -> dict[str, Any]:
duration_ms = None
@ -85,6 +92,10 @@ class ApiMonitorEntry:
"completion_tokens": self.completion_tokens,
"total_tokens": self.total_tokens,
"error": self.error,
"kind": self.kind,
"event": self.event,
"reason": self.reason,
"progress": self.progress,
}
if include_details:
payload["prompt"] = self.prompt
@ -127,6 +138,73 @@ class ApiMonitor:
self._trim_terminal_locked()
return entry.id
def record_lifecycle(
self,
*,
event: str,
model: str,
reason: Optional[str] = None,
running: bool = False,
) -> str:
"""Record a model load/unload alongside the request traffic that caused it.
``running=True`` opens the row for the caller to close with :meth:`finish` /
:meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to
every subject) and share the request retention budget.
"""
now = time.time()
entry = ApiMonitorEntry(
id = f"apievt_{uuid.uuid4().hex[:12]}",
endpoint = f"model.{event}",
method = "",
model = model or "default",
prompt = "",
status = "running" if running else "completed",
started_at = now,
updated_at = now,
started_monotonic = time.monotonic(),
finished_at = None if running else now,
finished_monotonic = None if running else time.monotonic(),
kind = "lifecycle",
event = event,
reason = reason,
shared = True,
)
with self._lock:
self._entries.appendleft(entry)
self._trim_terminal_locked()
return entry.id
def relabel(self, entry_id: Optional[str], model: str) -> None:
"""Rename an open lifecycle row once the load resolves its real id: up front
the caller only has the load path, which may be an HF snapshot dir."""
if not entry_id or not model:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is not None:
entry.model = model
entry.updated_at = time.time()
def set_progress(self, entry_id: Optional[str], progress: Optional[float]) -> None:
"""Update an open download row's percentage (clamped to 0-100)."""
if not entry_id or progress is None:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is not None and entry.status == "running":
entry.progress = min(100.0, max(0.0, float(progress)))
entry.updated_at = time.time()
def discard(self, entry_id: Optional[str]) -> None:
"""Drop a row that turned out not to be an event (an already-satisfied load)."""
if not entry_id:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is not None:
self._entries.remove(entry)
def append_reply(self, entry_id: Optional[str], text: str) -> None:
if not entry_id or not text:
return
@ -212,6 +290,18 @@ class ApiMonitor:
self._entries.appendleft(entry)
self._trim_terminal_locked()
def fail_open(self, entry_id: Optional[str], error: str) -> None:
"""Fail only a still-open row: unlike :meth:`fail`, a catch-all in a
``finally`` cannot stamp an error onto a request that already succeeded."""
if not entry_id:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is None or entry.finished_at is not None:
return
# Same lock as the check, so a finish() cannot land in between.
self._fail_locked(entry, error)
def fail(self, entry_id: Optional[str], error: str) -> None:
if not entry_id:
return
@ -224,15 +314,18 @@ class ApiMonitor:
if error:
entry.error = _trim(error, 1000)
return
now = time.time()
entry.status = "error"
entry.error = _trim(error, 1000)
entry.updated_at = now
entry.finished_at = now
entry.finished_monotonic = time.monotonic()
self._entries.remove(entry)
self._entries.appendleft(entry)
self._trim_terminal_locked()
self._fail_locked(entry, error)
def _fail_locked(self, entry: ApiMonitorEntry, error: str) -> None:
now = time.time()
entry.status = "error"
entry.error = _trim(error, 1000)
entry.updated_at = now
entry.finished_at = now
entry.finished_monotonic = time.monotonic()
self._entries.remove(entry)
self._entries.appendleft(entry)
self._trim_terminal_locked()
def snapshot(
self,
@ -244,7 +337,7 @@ class ApiMonitor:
return [
entry.snapshot(include_details = include_details)
for entry in self._entries
if subject is None or entry.subject == subject
if self._visible(entry, subject)
]
def get(
@ -257,22 +350,29 @@ class ApiMonitor:
entry = self._find_locked(entry_id)
if entry is None:
return None
if subject is not None and entry.subject != subject:
if not self._visible(entry, subject):
return None
return entry.snapshot(include_details = True)
def active_count(self, *, subject: Optional[str] = None) -> int:
# Lifecycle rows show as "running" while loading but are not in-flight API requests.
with self._lock:
return sum(
1
for entry in self._entries
if entry.status == "running" and (subject is None or entry.subject == subject)
if entry.status == "running"
and entry.kind != "lifecycle"
and (subject is None or entry.subject == subject)
)
def clear(self) -> None:
with self._lock:
self._entries.clear()
@staticmethod
def _visible(entry: ApiMonitorEntry, subject: Optional[str]) -> bool:
return subject is None or entry.subject == subject or entry.shared
def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]:
for entry in self._entries:
if entry.id == entry_id:

View file

@ -326,6 +326,58 @@ def _normalize_tool_call_arguments(messages: list) -> list:
return out if mutated else messages
def _take_tool_result(pending: list, call_id) -> Optional[dict]:
if call_id:
for i, result in enumerate(pending):
if result.get("tool_call_id") == call_id:
return pending.pop(i)
for i, result in enumerate(pending):
if not result.get("tool_call_id"):
return pending.pop(i)
return None
def _split_parallel_tool_calls(messages: list) -> list:
"""Llama 3.x templates render one call per message, so split parallel calls
into consecutive single-call messages, each followed by its own result."""
if not any(isinstance(m, dict) and len(m.get("tool_calls") or ()) > 1 for m in messages):
return messages
out: list = []
i = 0
total = len(messages)
while i < total:
msg = messages[i]
calls = msg.get("tool_calls") if isinstance(msg, dict) else None
if not calls or len(calls) <= 1:
out.append(msg)
i += 1
continue
# Tool results right after this message answer its calls.
j = i + 1
pending: list = []
while (
j < total
and isinstance(messages[j], dict)
and messages[j].get("role") in ("tool", "ipython")
):
pending.append(messages[j])
j += 1
for idx, call in enumerate(calls):
piece = {**msg, "tool_calls": [call]}
if idx:
piece["content"] = ""
out.append(piece)
result = _take_tool_result(pending, call.get("id") if isinstance(call, dict) else None)
if result is not None:
out.append(result)
out.extend(pending)
i = j
return out
def apply_chat_template_for_generation(
tokenizer,
messages: list,
@ -378,13 +430,21 @@ def apply_chat_template_for_generation(
try:
return _render(messages)
except Exception:
# Strict tool templates reject the JSON-string ``arguments`` form via
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
# Original messages render first, so working templates stay byte-identical.
# Retry with repairs applied cumulatively. Originals render first, so
# working templates stay byte-identical.
candidates: list = []
normalized = _normalize_tool_call_arguments(messages)
if normalized is messages:
raise
return _render(normalized)
if normalized is not messages:
candidates.append(normalized)
split = _split_parallel_tool_calls(normalized)
if split is not normalized:
candidates.append(split)
for candidate in candidates:
try:
return _render(candidate)
except Exception:
continue
raise
def render_native_template(

View file

@ -567,7 +567,7 @@ class InferenceBackend:
_meta_path = Path(config.path) / "export_metadata.json"
try:
if _meta_path.exists():
_meta = json.loads(_meta_path.read_text())
_meta = json.loads(_meta_path.read_text(encoding = "utf-8"))
if _meta.get("base_model"):
processor_source = _meta["base_model"]
except Exception:

View file

@ -13,37 +13,85 @@ from __future__ import annotations
import asyncio
import os
import sys
import threading
from collections import deque
from dataclasses import dataclass
from typing import Deque, Optional
ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL"
ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT"
ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL"
ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE"
# dataclass(slots = True) halves per-instance overhead. Measured as perf-neutral
# here, not a speed win: it costs a little on construction and gains it back on
# access. It is 3.10+ and this package declares >=3.9, so gate it rather than
# dropping it outright. Empty on 3.9 means a plain dataclass.
_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {}
ADMISSION_CONTROL_ENV = "UNSLOTH_LLAMA_ADMISSION_CONTROL"
ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_TIMEOUT"
ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_LLAMA_ADMISSION_KEEPALIVE_INTERVAL"
ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_LLAMA_ADMISSION_MAX_QUEUE"
ADMISSION_QUEUE_PER_SLOT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_PER_SLOT"
# The UNSLOTH_OPENAI_COMPAT_* spellings predate this queue being shared with the
# Anthropic /v1/messages route (same llama-server slots). Still honored; the
# neutral name above wins when both are set.
_LEGACY_ENV = {
ADMISSION_CONTROL_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL",
ADMISSION_QUEUE_TIMEOUT_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT",
ADMISSION_KEEPALIVE_INTERVAL_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL",
ADMISSION_MAX_QUEUE_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE",
}
DEFAULT_ADMISSION_ENABLED = True
# None: a queued request waits for its slot indefinitely rather than timing out.
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0
DEFAULT_ADMISSION_MAX_QUEUE = 64
# None: no absolute cap, the wait line is sized from the pool instead.
DEFAULT_ADMISSION_MAX_QUEUE = None
# Wait line = 16 x the serving slots, so it tracks --parallel (4 slots -> 64
# waiters, 8 -> 128). Purely a memory guard; waiting itself is never timed out.
DEFAULT_ADMISSION_QUEUE_PER_SLOT = 16
# Floor for the scaled line, so a 1-slot backend (plain `unsloth studio`, or any
# load downshifted to fit VRAM) keeps the depth it had before scaling existed
# rather than dropping to 16 and rejecting callers that used to queue.
DEFAULT_ADMISSION_MIN_QUEUE = 64
@dataclass(frozen = True)
@dataclass(frozen = True, **_SLOTS)
class LlamaAdmissionConfig:
enabled: bool = DEFAULT_ADMISSION_ENABLED
queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE
queue_per_slot: Optional[int] = DEFAULT_ADMISSION_QUEUE_PER_SLOT
# Unconditional floor on the scaled line. The env path clears it when the
# operator sets QUEUE_PER_SLOT, so only the default multiplier is floored.
min_queue: Optional[int] = DEFAULT_ADMISSION_MIN_QUEUE
def queue_limit(self, capacity: int) -> Optional[int]:
"""How many callers may line up for a pool of ``capacity`` slots.
An explicit ``max_queue`` wins; otherwise the line scales with the slots
so it follows ``--parallel``. The default multiplier is floored, so a
1-slot backend does not end up shallower than it was before scaling. None
(or any non-positive setting) means an unbounded line.
"""
if self.max_queue is not None:
return self.max_queue if self.max_queue > 0 else None
if not self.queue_per_slot or self.queue_per_slot <= 0:
return None
scaled = self.queue_per_slot * max(1, capacity)
return max(self.min_queue, scaled) if self.min_queue else scaled
@dataclass(frozen = True)
@dataclass(frozen = True, **_SLOTS)
class LlamaAdmissionSnapshot:
key: str
capacity: int
active: int
queued: int
free: int = 0
class LlamaAdmissionError(Exception):
@ -69,8 +117,17 @@ class LlamaAdmissionCancelled(LlamaAdmissionError):
pass
def _bool_env(name: str, default: bool) -> bool:
def _raw_env(name: str) -> Optional[str]:
"""Value for a canonical name, falling back to its legacy spelling."""
value = os.environ.get(name)
if value is None or not value.strip():
legacy = _LEGACY_ENV.get(name)
value = os.environ.get(legacy) if legacy else None
return value
def _bool_env(name: str, default: bool) -> bool:
value = _raw_env(name)
if value is None or not value.strip():
return default
value = value.strip().lower()
@ -82,7 +139,7 @@ def _bool_env(name: str, default: bool) -> bool:
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
value = os.environ.get(name)
value = _raw_env(name)
if value is None or not value.strip():
return default
try:
@ -93,7 +150,7 @@ def _optional_positive_float_env(name: str, default: Optional[float]) -> Optiona
def _positive_float_env(name: str, default: float) -> float:
value = os.environ.get(name)
value = _raw_env(name)
if value is None or not value.strip():
return default
try:
@ -103,19 +160,38 @@ def _positive_float_env(name: str, default: float) -> float:
return parsed if parsed > 0 else default
def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
def _queue_limits_from_env() -> tuple[Optional[int], Optional[int], Optional[int]]:
"""(max_queue, queue_per_slot, min_queue) from the environment.
An absolute MAX_QUEUE wins outright; MAX_QUEUE=0 asks for an unbounded line.
Unset leaves the per-slot multiplier in charge (itself 0 for unbounded). The
floor applies only to the default multiplier: setting QUEUE_PER_SLOT means
the operator wants that exact depth, however shallow.
"""
# Explicit means it parsed, not just that something was set: a typo falls back
# to the default multiplier, so it has to keep the default's floor too.
raw_per_slot = _raw_env(ADMISSION_QUEUE_PER_SLOT_ENV)
try:
parsed = int(value.strip())
per_slot = int((raw_per_slot or "").strip())
except ValueError:
return default
return parsed if parsed > 0 else None
per_slot, min_queue = DEFAULT_ADMISSION_QUEUE_PER_SLOT, DEFAULT_ADMISSION_MIN_QUEUE
else:
per_slot, min_queue = (per_slot if per_slot > 0 else None), None
raw = _raw_env(ADMISSION_MAX_QUEUE_ENV)
if raw is None or not raw.strip():
return None, per_slot, min_queue
try:
parsed = int(raw.strip())
except ValueError:
return None, per_slot, min_queue
return (parsed, None, None) if parsed > 0 else (None, None, None)
def llama_admission_config_from_env() -> LlamaAdmissionConfig:
max_queue, queue_per_slot, min_queue = _queue_limits_from_env()
return LlamaAdmissionConfig(
queue_per_slot = queue_per_slot,
min_queue = min_queue,
enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED),
queue_timeout_s = _optional_positive_float_env(
ADMISSION_QUEUE_TIMEOUT_ENV,
@ -125,14 +201,11 @@ def llama_admission_config_from_env() -> LlamaAdmissionConfig:
ADMISSION_KEEPALIVE_INTERVAL_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
),
max_queue = _optional_positive_int_env(
ADMISSION_MAX_QUEUE_ENV,
DEFAULT_ADMISSION_MAX_QUEUE,
),
max_queue = max_queue,
)
@dataclass
@dataclass(**_SLOTS)
class _Waiter:
loop: asyncio.AbstractEventLoop
future: asyncio.Future
@ -141,11 +214,23 @@ class _Waiter:
class LlamaAdmissionLease:
def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
__slots__ = ("_queue", "_slot", "_released", "_release_lock")
def __init__(
self,
queue: Optional["LlamaAdmissionQueue"],
slot: Optional[int] = None,
):
self._queue = queue
self._slot = slot
self._released = False
self._release_lock = threading.Lock()
@property
def slot(self) -> Optional[int]:
"""Pool slot this lease holds, or None when admission is disabled."""
return self._slot
def release(self) -> None:
queue = None
with self._release_lock:
@ -154,7 +239,7 @@ class LlamaAdmissionLease:
self._released = True
queue = self._queue
if queue is not None:
queue.release()
queue.release(self._slot)
async def __aenter__(self) -> "LlamaAdmissionLease":
return self
@ -164,6 +249,8 @@ class LlamaAdmissionLease:
class LlamaAdmissionReservation:
__slots__ = ("_queue", "_lease", "_waiter", "snapshot")
def __init__(
self,
*,
@ -195,6 +282,13 @@ class LlamaAdmissionReservation:
return self._lease
async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]:
"""Wait up to ``timeout_s`` for a slot.
A timeout leaves this reservation queued so the caller can poll again.
Any exit that abandons the wait for good must call ``cancel()``, or the
slot granted later is delivered to a future nobody reads and is never
released.
"""
lease = self.lease_nowait()
if lease is not None:
return lease
@ -229,35 +323,80 @@ class LlamaAdmissionReservation:
class LlamaAdmissionQueue:
"""A fixed pool of generation slots for one llama-server, plus a FIFO wait line.
The pool mirrors llama-server's own ``--parallel`` slots: ``capacity`` slot ids
are each either free or held by exactly one caller. A caller that finds every
slot busy waits in arrival order and is handed the next slot to free, so no
caller is starved. This bounds only the callers that reserve: chat completions
and messages do, while /v1/completions, Studio's own chat endpoint and RAG
captioning all reach llama-server directly, so it is not a global cap.
Waiting is unbounded in time by default (``queue_timeout_s``
None); the wait line itself is bounded, and only how many may line up before
new arrivals are rejected. By default that is ``16 x slots`` floored at 64,
not unlimited: an unbounded line takes ``max_queue`` or ``queue_per_slot``
set to 0. See ``LlamaAdmissionConfig.queue_limit``.
"""
__slots__ = ("key", "_lock", "_capacity", "_free", "_in_use", "_held", "_waiters")
def __init__(self, key: str):
self.key = key
self._lock = threading.Lock()
self._active = 0
self._capacity = 1
self._free: list[int] = [0]
# Held slots as a bitmask: one int instead of a set, so the pool costs the
# same whether it is idle or saturated. _held is its popcount, kept as a
# counter because int.bit_count() is 3.10+ and this package targets 3.9.
self._in_use = 0
self._held = 0
self._waiters: Deque[_Waiter] = deque()
def _resize_pool_locked(self, capacity: int) -> None:
# Slots past a shrunk capacity retire when their holder releases them.
if capacity == self._capacity:
return
self._capacity = capacity
self._free = [slot for slot in range(capacity) if not self._in_use >> slot & 1]
def _can_admit_locked(self) -> bool:
# Slots still held above a shrunk capacity keep occupying the backend, so
# count every held slot against the ceiling, not just the ids below it.
return bool(self._free) and self._held < self._capacity
def _take_slot_locked(self) -> Optional[int]:
if not self._can_admit_locked():
return None
slot = self._free.pop()
self._in_use |= 1 << slot
self._held += 1
return slot
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
capacity = max(1, int(capacity or 1))
if not config.enabled:
return LlamaAdmissionReservation(
queue = None,
lease = LlamaAdmissionLease(None),
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0),
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0, capacity),
)
loop = asyncio.get_running_loop()
with self._lock:
self._capacity = capacity
self._prune_waiters_locked()
self._resize_pool_locked(capacity)
self._grant_waiters_locked()
if self._active < self._capacity and not self._waiters:
self._active += 1
return LlamaAdmissionReservation(
queue = self,
lease = LlamaAdmissionLease(self),
snapshot = self._snapshot_locked(),
)
if config.max_queue is not None and len(self._waiters) >= config.max_queue:
if not self._waiters:
slot = self._take_slot_locked()
if slot is not None:
# No snapshot here: callers read it through snapshot_now(),
# which re-reads the queue, so building one per admitted
# request would be pure allocation on the hot path.
return LlamaAdmissionReservation(
queue = self,
lease = LlamaAdmissionLease(self, slot),
)
limit = config.queue_limit(self._capacity)
if limit is not None and self._live_waiters_locked() >= limit:
raise LlamaAdmissionQueueFull(
"llama-server generation queue is full",
snapshot = self._snapshot_locked(),
@ -270,13 +409,20 @@ class LlamaAdmissionQueue:
return LlamaAdmissionReservation(
queue = self,
waiter = waiter,
snapshot = self._snapshot_locked(),
)
def release(self) -> None:
def _release_slot_locked(self, slot: Optional[int]) -> None:
# A slot id at or past a shrunk capacity retires instead of returning.
if slot is None or not self._in_use >> slot & 1:
return
self._in_use &= ~(1 << slot)
self._held -= 1
if slot < self._capacity:
self._free.append(slot)
def release(self, slot: Optional[int]) -> None:
with self._lock:
if self._active > 0:
self._active -= 1
self._release_slot_locked(slot)
self._grant_waiters_locked()
def cancel(self, waiter: _Waiter) -> None:
@ -291,7 +437,13 @@ class LlamaAdmissionQueue:
lease_to_release = waiter.granted_lease
waiter.granted_lease = None
if not waiter.future.done():
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
try:
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
except RuntimeError:
# Loop gone. Routes call cancel() from finally blocks, so
# raising here would both mask their exception and skip the
# release below, stranding the slot for the process lifetime.
pass
if lease_to_release is not None:
lease_to_release.release()
@ -303,20 +455,30 @@ class LlamaAdmissionQueue:
def is_idle(self) -> bool:
with self._lock:
self._prune_waiters_locked()
return self._active == 0 and not self._waiters
return self._in_use == 0 and not self._waiters
def _grant_waiters_locked(self) -> None:
self._prune_waiters_locked()
while self._waiters and self._active < self._capacity:
# Dead waiters are skipped as they are popped, so no prune is needed here.
while self._waiters and self._can_admit_locked():
waiter = self._waiters.popleft()
if waiter.cancelled or waiter.future.done():
continue
self._active += 1
lease = LlamaAdmissionLease(self)
slot = self._take_slot_locked()
lease = LlamaAdmissionLease(self, slot)
waiter.granted_lease = lease
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
try:
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
except RuntimeError:
# Waiter's loop is gone. Reclaim the slot; leaving the bit set
# would strand it, since _free is rebuilt from the bitmask.
waiter.granted_lease = None
self._release_slot_locked(slot)
def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None:
# Runs on the waiter's own loop thread, which is also the only thread that
# cancels that reservation, so waiter state is safe to touch unlocked here.
# release() may be called from any thread, but only reaches this via
# call_soon_threadsafe. Cancelling off-loop would need this under _lock.
if waiter.cancelled or waiter.future.done():
waiter.granted_lease = None
if not waiter.future.done():
@ -331,16 +493,32 @@ class LlamaAdmissionQueue:
lease.release()
def _prune_waiters_locked(self) -> None:
# Rebuilding the deque on every reserve/release dominated the hot path, so
# only pay it when a waiter actually died out of band (an externally
# cancelled future); cancel() already drops its own waiter eagerly.
for waiter in self._waiters:
if waiter.cancelled or waiter.future.done():
break
else:
return
self._waiters = deque(
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
)
def _live_waiters_locked(self) -> int:
self._prune_waiters_locked()
return len(self._waiters)
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:
return LlamaAdmissionSnapshot(
key = self.key,
capacity = self._capacity,
active = self._active,
active = self._held,
queued = len(self._waiters),
# What another caller could actually take, so the admission log never
# shows free slots next to queued requests: after a shrink, ids below
# the new capacity can be free while holdovers still fill the ceiling.
free = min(len(self._free), max(0, self._capacity - self._held)),
)

File diff suppressed because it is too large Load diff

Some files were not shown because too many files have changed in this diff Show more