Merge remote-tracking branch 'upstream/main' into fix/rocm-version-test-isolation
This commit is contained in:
commit
5a94eab6eb
449 changed files with 46396 additions and 2835 deletions
6
.github/workflows/consolidated-tests-ci.yml
vendored
6
.github/workflows/consolidated-tests-ci.yml
vendored
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
46
.github/workflows/studio-backend-ci.yml
vendored
46
.github/workflows/studio-backend-ci.yml
vendored
|
|
@ -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
3
.gitignore
vendored
|
|
@ -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/
|
||||
|
|
|
|||
20
README.md
20
README.md
|
|
@ -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`):
|
||||
|
|
|
|||
18
install.ps1
18
install.ps1
|
|
@ -1917,12 +1917,14 @@ exit 0
|
|||
# (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU.
|
||||
elseif ($ROCmGpuLabel) {
|
||||
$nameArchTable = @(
|
||||
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
|
||||
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
|
||||
@{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: RX 9070 XT / 9070 GRE / 9070 / 9080)
|
||||
@{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: RX 9060 XT / 9060)
|
||||
@{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
|
||||
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
|
||||
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
|
||||
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
|
||||
@{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
|
||||
@{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
|
||||
@{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
|
||||
@{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32)
|
||||
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33)
|
||||
@{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point)
|
||||
@{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family
|
||||
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family
|
||||
|
|
@ -2203,6 +2205,7 @@ exit 0
|
|||
$archFamilyMap = @{
|
||||
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
|
||||
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
|
||||
"gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point)
|
||||
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
|
||||
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
|
||||
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
|
||||
|
|
@ -2224,6 +2227,7 @@ exit 0
|
|||
$torchFloorMap = @{
|
||||
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
|
||||
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
|
||||
"gfx1152" = "torch>=2.11.0,<2.12.0"
|
||||
}
|
||||
# Companion ranges track the torch ceiling so pip resolves a consistent
|
||||
# trio on AMD's per-arch index (each published independently). Mirrors
|
||||
|
|
@ -2231,10 +2235,12 @@ exit 0
|
|||
$torchvisionFloorMap = @{
|
||||
"gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0"
|
||||
"gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0"
|
||||
"gfx1152" = "torchvision>=0.26.0,<0.27.0"
|
||||
}
|
||||
$torchaudioFloorMap = @{
|
||||
"gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0"
|
||||
"gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0"
|
||||
"gfx1152" = "torchaudio>=2.11.0,<2.12.0"
|
||||
}
|
||||
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
|
||||
if ($archFamily) {
|
||||
|
|
@ -2264,7 +2270,7 @@ exit 0
|
|||
$_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
|
||||
}
|
||||
# Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare.
|
||||
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf
|
||||
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $_pinLeaf
|
||||
if ($_pinGfx211 -or $_pinRocm211) {
|
||||
$ROCmIndexUrl = $TorchIndexUrl
|
||||
$ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
|
||||
|
|
|
|||
247
install.sh
247
install.sh
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ lora:
|
|||
- "query"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ lora:
|
|||
- "value"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ lora:
|
|||
- "Wqkv"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ lora:
|
|||
- "shared_mlp.output_linear"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ lora:
|
|||
- "shared_mlp.output_linear"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ lora:
|
|||
- "v_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ lora:
|
|||
- "v_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "out_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ lora:
|
|||
- "all-linear"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ lora:
|
|||
- "out_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ lora:
|
|||
- "v_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ lora:
|
|||
- "gate_up_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ lora:
|
|||
- "gate_up_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ lora:
|
|||
- "down_proj"
|
||||
use_rslora: false
|
||||
use_loftq: false
|
||||
use_dora: false
|
||||
|
||||
logging:
|
||||
enable_wandb: false
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Add a link
Reference in a new issue