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

Merge remote-tracking branch 'origin/main' into woa-nvidia-wsl-fallback

Browse source 
# Conflicts:
#	scripts/uninstall.sh
#	studio/setup.sh
This commit is contained in:
Daniel Han 2026-06-11 20:20:53 -07:00
commit 9bdd5436b2
234 changed files with 23233 additions and 3046 deletions
Download patch file Download diff file Expand all files Collapse all files

20
.github/CODEOWNERS vendored
View file

@ -6,10 +6,10 @@
/unsloth/models/rl_replacements.py @Datta0 @pluesclues @danielhanchen
/unsloth/trainer.py @danielhanchen
/unsloth/models/sentence_transformer.py @Etherll @danielhanchen
/unsloth/save.py @rolandtannous @danielhanchen
/unsloth/save.py @danielhanchen
/unsloth/tokenizer_utils.py @mmathew23 @danielhanchen
/unsloth/chat_templates.py @rolandtannous @danielhanchen
/unsloth/ollama_template_mappers.py @rolandtannous @danielhanchen
/unsloth/chat_templates.py @danielhanchen
/unsloth/ollama_template_mappers.py @danielhanchen
/unsloth/kernels/moe/*.py @Datta0
/unsloth/import_fixes.py @danielhanchen
/unsloth/device_type.py @danielhanchen
@ -45,14 +45,14 @@
/unsloth/utils/hf_hub.py @mmathew23
/unsloth/utils/packing.py @mmathew23
/cli/ @rolandtannous @Manan17
/studio/frontend/ @Shine1i @rolandtannous @Manan17
/cli/ @Manan17
/studio/frontend/ @Shine1i @Manan17
/studio/frontend/public/ @Shine1i
/studio/backend/ @rolandtannous
/studio/backend/core/data_recipe/ @rolandtannous
/studio/backend/tests/ @rolandtannous @danielhanchen
/tests/ @rolandtannous @danielhanchen
/scripts/ @rolandtannous @danielhanchen
/studio/backend/
/studio/backend/core/data_recipe/
/studio/backend/tests/ @danielhanchen
/tests/ @danielhanchen
/scripts/ @danielhanchen
# Snapshot data for the notebook linter / Colab oracle. Drift in these
# files changes the pin floor for every Unsloth notebook, so refreshes

11
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -333,6 +333,17 @@ jobs:
run: |
python -m pytest -v --tb=short tests/test_callback_signature_drift.py
- name: generation correctness guards (HARD GATE)
# Deterministic CPU guards, each validated to fail on its pre-fix code:
# leftpad = batched left-padded generation (#1066/#3699, fixed by
# #2216 + #4100; staging proof: unsloth-staging-2 PRs 170/172);
# rope_scaling_drift = config.rope_scaling dropped by replaced rotary
# classes (#2405). AST checks run first so import breakage cannot mask them.
run: |
python -m pytest -v --tb=short \
tests/utils/test_prepare_inputs_leftpad.py \
tests/utils/test_rope_scaling_drift.py
- name: unsloth Bucket-A — CPU tests not in Repo tests (CPU)
# CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/
# that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model

28
.github/workflows/lint-ci.yml vendored
View file

@ -105,29 +105,37 @@ jobs:
# exactly where a hoist refactor lives, and it skips brand-new
# files whose re-export imports would otherwise look "unused".
#
# actions/checkout uses fetch-depth: 1, so the base branch is not
# present locally. Fetch the single base commit with an explicit
# refspec so origin/<base> is reliably created (a bare
# `git fetch origin <ref>` only updates FETCH_HEAD in some
# configs). Two-dot diff avoids needing a merge-base on a shallow
# clone.
# Diff against the true merge-base, not the base tip. A two-dot
# diff against the tip re-lints every file the base branch
# changed after the PR branched, comparing newer base code
# (BEFORE) against the PR's older snapshot (AFTER) - a
# time-reversed comparison that flags the base branch's own
# refactors as blockers on PRs that never touched those files.
# The compare API returns the merge-base without needing local
# history, and fetching that single commit by SHA keeps the
# shallow (fetch-depth: 1) clone.
if: github.event_name == 'pull_request'
env:
GH_TOKEN: ${{ github.token }}
run: |
git fetch --no-tags --depth=1 origin \
"${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}"
MERGE_BASE=$(gh api \
"repos/${{ github.repository }}/compare/${{ github.event.pull_request.base.sha }}...${{ github.event.pull_request.head.sha }}" \
--jq .merge_base_commit.sha)
git fetch --no-tags --depth=1 origin "$MERGE_BASE"
mapfile -t CHANGED < <(
git diff --name-only --diff-filter=M \
"origin/${{ github.base_ref }}" HEAD -- '*.py' \
"$MERGE_BASE" HEAD -- '*.py' \
| grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true
)
if [ "${#CHANGED[@]}" -eq 0 ]; then
echo "no in-place-modified Python files to check"
exit 0
fi
printf 'merge base: %s\n' "$MERGE_BASE"
printf 'checking %d file(s):\n' "${#CHANGED[@]}"
printf ' %s\n' "${CHANGED[@]}"
python scripts/verify_import_hoist.py \
--before "origin/${{ github.base_ref }}" --after HEAD "${CHANGED[@]}"
--before "$MERGE_BASE" --after HEAD "${CHANGED[@]}"
- name: No leftover debugger / pdb / breakpoint calls
# Catches the "I'll just stick a breakpoint() here" mistake

51
.github/workflows/mlx-ci.yml vendored
View file

@ -231,33 +231,14 @@ jobs:
tests/studio/test_is_mlx_dispatch_gate.py \
tests/studio/test_mlx_training_worker_behaviors.py
# Studio prebuilt llama.cpp install + GGUF inference. Drives the
# exact path Studio's setup.sh takes on macOS: invokes
# studio/install_llama_prebuilt.py with --published-repo
# ggml-org/llama.cpp and --published-release-tag b9049 (the
# latest llama.cpp release at the time this step was added; bump
# via UNSLOTH_LLAMA_TAG / DEFAULT_LLAMA_TAG when refreshing).
# The installer downloads llama-b9049-bin-macos-arm64.tar.gz,
# which is the universal Apple Silicon (arm64) build -- the
# same artifact works on M1/M2/M3/M4 because llama.cpp compiles
# against the ARMv8.2 baseline.
#
# The b9049 release also publishes:
# - llama-b9049-bin-macos-arm64-kleidiai.tar.gz
# KleidiAI dispatches at runtime; on M1 it falls back where
# ISA features (e.g. I8MM) are missing, so this asset also
# runs on M1 -- Studio just doesn't choose it by default.
# - llama-b9049-bin-macos-x64.tar.gz
# Intel-only; would only run on M1 via Rosetta 2 emulation,
# which we explicitly avoid.
# - iOS XCFramework
# iOS-app build artifact, unrelated to a macOS desktop CI.
#
# After install, downloads a small published GGUF
# (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) from HuggingFace and
# runs the prebuilt llama-cli on it. Asserts the prompt echo
# appears in stdout. If the install fails OR the binary exits
# non-zero, that's an Unsloth/Studio bug.
# Studio prebuilt llama.cpp install + GGUF inference. Mirrors the
# path Studio's setup.sh takes on macOS since #5963: plan against
# the unslothai/llama.cpp fork's latest release, which ships the
# bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the
# default policy reads. After install, downloads a small published
# GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates
# llama-server /completion end to end. An install failure or a
# non-zero binary exit is an Unsloth/Studio bug.
- name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
@ -272,20 +253,12 @@ jobs:
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# --simple-policy is required when --published-repo points
# at upstream ggml-org/llama.cpp; that repo doesn't ship the
# llama-prebuilt-manifest.json asset Studio's default policy
# expects, so the simple platform-specific policy maps
# Darwin+arm64 -> bin-macos-arm64 directly. studio/setup.sh
# passes both --published-repo ggml-org/llama.cpp AND
# --simple-policy automatically on macOS, so this CI step
# exercises the same code path users hit when they run
# `curl -fsSL https://unsloth.ai/install.sh | sh`.
# Mirror studio/setup.sh on macOS (the install.sh user path):
# it plans against the unslothai/llama.cpp fork's latest
# release with no policy or tag flags.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo ggml-org/llama.cpp \
--published-release-tag b9049 \
--simple-policy
--published-repo unslothai/llama.cpp
# Studio bundles only llama-server + llama-quantize from the
# prebuilt (not llama-cli) -- inference goes through

6
.github/workflows/studio-api-smoke.yml vendored
View file

@ -77,7 +77,7 @@ jobs:
path: hf-cache
# Same key as studio-ui-smoke.yml so the two jobs share a
# single GGUF download across CI.
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -88,17 +88,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

38
.github/workflows/studio-frontend-ci.yml vendored
View file

@ -17,6 +17,8 @@ on:
- 'studio/frontend/**'
- 'scripts/check_frontend_dep_removal.py'
- 'tests/studio/test_frontend_dep_removal.py'
- 'scripts/sync_allow_scripts_pins.py'
- 'tests/studio/test_sync_allow_scripts_pins.py'
- '.github/workflows/studio-frontend-ci.yml'
push:
branches: [main, pip]
@ -60,6 +62,19 @@ jobs:
with:
node-version: '22'
# node 22 bundles npm 10.x, which predates allowScripts. Move to the
# 11.x line and fail loudly if the gate is still missing, so the
# strict flag below can never silently degrade into a warning.
- name: Upgrade npm to 11.x (allowScripts enforcement)
working-directory: ${{ github.workspace }}
run: |
npm install -g npm@^11 --no-fund --no-audit
V=$(npm -v)
case "$V" in
11.1[6-9].*|11.[2-9][0-9].*|1[2-9].*) echo "npm $V has allowScripts" ;;
*) echo "::error::npm $V lacks allowScripts (need >=11.16)"; exit 1 ;;
esac
# Run the structural lockfile scan BEFORE npm ci. A compromised
# tarball runs its `prepare` / `postinstall` during `npm ci`,
# so any catch has to fire upstream of that. The scanner is
@ -68,14 +83,23 @@ jobs:
working-directory: ${{ github.workspace }}
run: python3 scripts/lockfile_supply_chain_audit.py
# Dependency bumps strand the version-pinned allowScripts entries.
# The paired pre-commit hook auto-fixes PRs; this is the backstop.
- name: allowScripts pins must match the lockfile
working-directory: ${{ github.workspace }}
run: |
python3 tests/studio/test_sync_allow_scripts_pins.py
python3 scripts/sync_allow_scripts_pins.py --check
- name: Lockfile must agree with package.json (npm ci is strict)
# Lifecycle scripts (esbuild native-binary postinstall, etc.) are
# required for `vite build`. The pre-install lockfile structural
# audit (lockfile_supply_chain_audit.py) is the practical defence
# against the npm postinstall-dropper class -- it fires BEFORE any
# tarball runs, on the injection pattern itself rather than an
# advisory-DB lookup.
run: npm ci --no-fund --no-audit
# The vite 8 chain (rolldown, lightningcss, tailwind oxide) ships napi
# binaries with no install scripts. The only script-bearing deps are
# covered by `allowScripts` in package.json (npm >=11.16, default in
# npm 12). The pre-install lockfile audit above stays the first line
# of defence -- it fires before any tarball can run code.
# --strict-allow-scripts: any unreviewed install script hard-fails
# the job; the sync hook keeps the pins fresh after bumps.
run: npm ci --strict-allow-scripts --no-fund --no-audit
- name: npm ci must not have modified the working tree
working-directory: ${{ github.workspace }}

25
.github/workflows/studio-inference-smoke.yml vendored
View file

@ -20,7 +20,7 @@
# enable_tools / enabled_tools, and enable_thinking on/off.
#
# 3. JSON, images
# Qwen3-VL-2B-Instruct UD-IQ2_XXS (~570 MiB) + mmproj-F16 (~780 MiB).
# Qwen3-VL-2B-Instruct UD-Q4_K_XL (~1.1 GiB) + mmproj-F16 (~780 MiB).
# response_format JSON-schema decoding and OpenAI image_url
# (data URI) plus Anthropic source/base64 image inputs.
#
@ -91,7 +91,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -102,17 +102,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -375,6 +377,7 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -792,8 +795,14 @@ jobs:
timeout-minutes: 30
env:
GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF
GGUF_VARIANT: UD-IQ2_XXS
GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf
# UD-Q4_K_XL, not UD-IQ2_XXS: at 2-bit the temp-0 answer to the JSON
# step's capital-of-France probe flips with the host's SIMD kernels
# (GitHub runners deterministically answered France while other CPUs
# answer Paris; seeds do not rescue it, 1/5 Paris at temp 0.7). The
# Q4 quant answered Paris 13/13 across temps and seeds on the same
# runners, so the hard Paris assertion below stays reliable.
GGUF_VARIANT: UD-Q4_K_XL
GGUF_FILE: Qwen3-VL-2B-Instruct-UD-Q4_K_XL.gguf
MMPROJ_FILE: mmproj-F16.gguf
STUDIO_PORT: '18890'
HF_HOME: ${{ github.workspace }}/hf-cache
@ -823,7 +832,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Prime HF_HOME with the GGUF + mmproj
id: prime-hf
@ -835,17 +844,19 @@ jobs:
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj)
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

6
.github/workflows/studio-mac-api-smoke.yml vendored
View file

@ -62,7 +62,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -73,17 +73,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

8
.github/workflows/studio-mac-inference-smoke.yml vendored
View file

@ -85,7 +85,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -96,6 +96,7 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
# Save partial caches on cancel/timeout -- hf download resumes by
# content hash. `outcome != skipped` keeps cache-hit a no-op.
@ -104,11 +105,12 @@ jobs:
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -361,6 +363,7 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -749,6 +752,7 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

1
.github/workflows/studio-mac-install-matrix.yml vendored
View file

@ -63,6 +63,7 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

6
.github/workflows/studio-mac-ui-smoke.yml vendored
View file

@ -62,7 +62,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -73,17 +73,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

3
.github/workflows/studio-mac-update-smoke.yml vendored
View file

@ -62,6 +62,7 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -73,6 +74,7 @@ jobs:
- name: First update should be a no-op (prebuilt already validated)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -91,6 +93,7 @@ jobs:
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

6
.github/workflows/studio-ui-smoke.yml vendored
View file

@ -76,7 +76,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -87,17 +87,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

3
.github/workflows/studio-update-smoke.yml vendored
View file

@ -71,6 +71,7 @@ jobs:
# prebuilt path falls back to source build.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -85,6 +86,7 @@ jobs:
# idempotency regressed.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -107,6 +109,7 @@ jobs:
# the first one.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

6
.github/workflows/studio-windows-api-smoke.yml vendored
View file

@ -69,7 +69,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -80,13 +80,14 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -123,6 +124,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;

13
.github/workflows/studio-windows-inference-smoke.yml vendored
View file

@ -101,7 +101,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -114,6 +114,7 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME cache for ${{ env.GGUF_REPO }}
# Only write a fresh cache entry when we actually rebuilt the
@ -123,7 +124,7 @@ jobs:
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -160,6 +161,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -504,6 +506,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -879,7 +882,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Prime HF_HOME with the GGUF + mmproj
id: prime-hf
@ -891,13 +894,14 @@ jobs:
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj)
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -934,6 +938,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;

6
.github/workflows/studio-windows-ui-smoke.yml vendored
View file

@ -85,7 +85,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -96,13 +96,14 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -143,6 +144,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 redirects ALL PowerShell streams (stdout, stderr,

3
.github/workflows/studio-windows-update-smoke.yml vendored
View file

@ -133,6 +133,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -179,6 +180,7 @@ jobs:
- name: First update should be a no-op (prebuilt already validated)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -197,6 +199,7 @@ jobs:
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

12
.pre-commit-config.yaml
View file

@ -19,3 +19,15 @@ repos:
exclude: '(chat_templates|ollama_template_mappers|_auto_install|mapper)\.py$'
additional_dependencies:
- ruff==0.6.9
# Re-pins allowScripts entries after dependency bumps. pre-commit.ci
# pushes the fix to PR branches, Dependabot's included, so stale pins
# heal without a human in the loop.
- id: sync-allow-scripts-pins
name: Sync allowScripts pins with the frontend lockfile
# `python <script>` not a direct exec: autofix commits can drop the
# executable bit, which kills shebang-style entries.
entry: python scripts/sync_allow_scripts_pins.py
args: [--fix]
language: python
files: ^studio/frontend/(package\.json|package-lock\.json)$
pass_filenames: false

441
install.ps1
View file

@ -158,6 +158,10 @@ function Install-UnslothStudio {
# UNSLOTH_PYTHON pins the version (mirrors install.sh --python); default 3.13.
$PythonVersion = if ($env:UNSLOTH_PYTHON) { $env:UNSLOTH_PYTHON } else { "3.13" }
# python.org fallback patch, used only when winget is unavailable/broken AND
# the live python.org listing can't be fetched. The installer URL scheme is
# stable so an older patch still installs. Bump alongside $PythonVersion.
$PythonFallbackFullVersion = "3.13.13"
# Resolve install destinations. Priority: UNSLOTH_STUDIO_HOME, then
# STUDIO_HOME alias, then USERPROFILE-redirect, then default.
@ -868,6 +872,7 @@ shell.Run cmd, 0, False
try {
$wshell = New-Object -ComObject WScript.Shell
$createdShortcutCount = 0
$createdShortcutPaths = @()
foreach ($linkPath in @($desktopLink, $startMenuLink)) {
if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue }
try {
@ -881,12 +886,46 @@ shell.Run cmd, 0, False
}
$shortcut.Save()
$createdShortcutCount++
$createdShortcutPaths += $linkPath
} catch {
substep "could not create shortcut at ${linkPath}: $($_.Exception.Message)" "Yellow"
}
}
if ($createdShortcutCount -gt 0) {
substep "Created Unsloth Studio shortcut"
# Force Explorer to re-read each new shortcut's icon so it renders
# immediately instead of a stale/generic entry (a same-name .lnk
# recreated across reinstalls keeps Explorer's cached per-item icon).
# The reliable, non-disruptive fix (no explorer restart) is a per-item
# SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW per .lnk; the global
# SHCNE_ASSOCCHANGED broadcast alone does NOT recover a stale item.
# Also clear the on-disk icon cache (covers heavier staleness).
try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {}
try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {}
try {
Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);' -ErrorAction SilentlyContinue
# SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut
foreach ($scPath in $createdShortcutPaths) {
try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, $scPath, [System.IntPtr]::Zero) } catch {}
}
# SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders)
[UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero)
} catch {}
# Win11's Start Menu (StartMenuExperienceHost) keeps its OWN
# pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT
# invalidate, so a rewritten same-name shortcut shows the old tile
# until the host restarts. Drop only the render caches (NEVER
# start2.bin -- the pinned layout) and let the host rebuild.
# Best-effort; Win10 has no such host (Test-Path skips it).
try {
$smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
if (Test-Path -LiteralPath $smehTemp) {
Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch {}
} else {
substep "no Unsloth Studio shortcuts were created" "Yellow"
}
@ -958,10 +997,13 @@ shell.Run cmd, 0, False
function Find-CompatiblePython {
# Try the Python Launcher first (most reliable on Windows)
# py.exe resolves to the standard CPython install, not conda.
$pyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue
if ($pyLauncher -and $pyLauncher.Source -notmatch $script:CondaSkipPattern) {
# Prefer the requested $PythonVersion, then newest-first fallback.
$minors = @($PythonVersion) + (@("3.13", "3.12", "3.11") | Where-Object { $_ -ne $PythonVersion })
# Prefer the requested $PythonVersion, then newest-first fallback.
$minors = @($PythonVersion) + (@("3.13", "3.12", "3.11") | Where-Object { $_ -ne $PythonVersion })
# Enumerate every py.exe on PATH with -All (Windows PowerShell 5.1
# returns only the first launcher without it) and search each for a
# supported, non-conda interpreter.
foreach ($pyLauncher in @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue)) {
if ($pyLauncher.Source -match $script:CondaSkipPattern) { continue }
foreach ($minor in $minors) {
try {
$out = & $pyLauncher.Source "-$minor" --version 2>&1 | Out-String
@ -999,6 +1041,79 @@ shell.Run cmd, 0, False
return $null
}
# ── Fallback: install CPython directly from python.org ──
# Used when winget is unavailable or fails (notably msstore cert-pinning error
# 0x8a15005e, which aborts `winget install` unless --source winget is given).
# Downloads the official installer and runs it silently as a per-user install
# (no UAC), putting python.exe + the py launcher on PATH. Mirrors the uv ->
# astral.sh fallback below. Returns @{ Version; Path } or $null.
function Install-PythonFromPythonOrg {
# python.org ships one installer per architecture.
$archSuffix = switch (Get-TauriDiagArch) {
"x86_64" { "-amd64" }
"arm64" { "-arm64" }
"x86" { "" }
default { $null }
}
if ($null -eq $archSuffix) {
substep "No python.org installer is available for this architecture." "Yellow"
return $null
}
# Resolve the latest $PythonVersion.x patch from the python.org listing,
# falling back to a same-minor version if the listing cannot be fetched.
# Use the pinned full version only when it matches the requested minor so a
# non-default UNSLOTH_PYTHON (e.g. 3.12) doesn't silently install 3.13.
$full = if ($PythonFallbackFullVersion -like "$PythonVersion.*") { $PythonFallbackFullVersion } else { "$PythonVersion.0" }
try {
$listing = [string](Invoke-RestMethod -Uri "https://www.python.org/ftp/python/" -UseBasicParsing -TimeoutSec 20)
$patches = [regex]::Matches($listing, ([regex]::Escape($PythonVersion) + '\.(\d+)/')) |
ForEach-Object { [int]$_.Groups[1].Value } | Sort-Object -Descending
if ($patches.Count -gt 0) { $full = "$PythonVersion.$($patches[0])" }
} catch {}
$file = "python-$full$archSuffix.exe"
$url = "https://www.python.org/ftp/python/$full/$file"
$dest = Join-Path ([System.IO.Path]::GetTempPath()) $file
substep "downloading Python $full from python.org..." "Yellow"
try {
Invoke-WebRequest -Uri $url -OutFile $dest -UseBasicParsing
} catch {
substep "python.org download failed: $($_.Exception.Message)" "Yellow"
return $null
}
# Per-user install => no UAC. PrependPath puts python + py on PATH;
# Include_launcher installs py.exe (preferred by Find-CompatiblePython).
substep "installing Python $full (silent, per-user)..."
$installArgs = @(
"/quiet",
"InstallAllUsers=0",
"PrependPath=1",
"Include_launcher=1",
# Launcher per-user too: Include_launcher defaults InstallLauncherAllUsers=1,
# which needs admin and would break this non-admin per-user fallback.
"InstallLauncherAllUsers=0",
"Include_pip=1",
"AssociateFiles=0",
"Shortcuts=0"
)
$rc = 1
try {
$proc = Start-Process -FilePath $dest -ArgumentList $installArgs -Wait -PassThru
$rc = $proc.ExitCode
} catch {
substep "python.org installer failed to start: $($_.Exception.Message)" "Yellow"
} finally {
Remove-Item -LiteralPath $dest -Force -ErrorAction SilentlyContinue
}
if ($rc -ne 0) {
substep "python.org installer exited with code $rc." "Yellow"
}
Refresh-SessionPath
return (Find-CompatiblePython)
}
# ── Install Python if no compatible version (3.11-3.13) found ──
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
Write-TauriLog "STEP" "Installing Python"
@ -1008,47 +1123,62 @@ shell.Run cmd, 0, False
step "python" "Python $($DetectedPython.Version) already installed"
}
if (-not $DetectedPython) {
if (-not $script:WingetAvailable) {
Write-Host "[ERROR] No compatible Python (3.11-3.13) found and winget is unavailable on this host." -ForegroundColor Red
Write-Host " Install Python $PythonVersion from https://www.python.org/downloads/" -ForegroundColor Yellow
Write-Host " and re-run this installer (make sure 'Add Python to PATH' is checked)." -ForegroundColor Yellow
return (Exit-InstallFailure "winget required to install Python on this host")
}
substep "installing Python ${PythonVersion}..."
$pythonPackageId = "Python.Python.$PythonVersion"
# Temporarily lower ErrorActionPreference so that winget stderr
# (progress bars, warnings) does not become a terminating error
# on PowerShell 5.1 where native-command stderr is ErrorRecord.
$prevEAP = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements
$wingetExit = $LASTEXITCODE
} catch { $wingetExit = 1 }
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
$wingetExit = $null
# Re-detect after install (PATH may have changed)
$DetectedPython = Find-CompatiblePython
if (-not $DetectedPython) {
# Python still not functional after winget -- force reinstall.
# This handles both real failures AND "already installed" codes where
# winget thinks Python is present but it's not actually on PATH
# (e.g. user partially uninstalled, or installed via a different method).
substep "Python not found on PATH after winget. Retrying with --force..." "Yellow"
if ($script:WingetAvailable) {
# --source winget avoids the msstore source, which can fail with
# cert-pinning error 0x8a15005e and abort the whole `winget install`
# (winget then demands --source). Python and uv both live in the
# winget source, so pinning it is correct and faster.
#
# Lower ErrorActionPreference so winget stderr (progress/warnings) is
# not a terminating error on PS 5.1 (native stderr is ErrorRecord).
$prevEAP = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements --force
winget install -e --id $pythonPackageId --source winget --accept-package-agreements --accept-source-agreements
$wingetExit = $LASTEXITCODE
} catch { $wingetExit = 1 }
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
# Re-detect after install (PATH may have changed)
$DetectedPython = Find-CompatiblePython
if (-not $DetectedPython) {
# Python still not functional after winget -- force reinstall.
# This handles both real failures AND "already installed" codes where
# winget thinks Python is present but it's not actually on PATH
# (e.g. user partially uninstalled, or installed via a different method).
substep "Python not found on PATH after winget. Retrying with --force..." "Yellow"
$ErrorActionPreference = "Continue"
try {
winget install -e --id $pythonPackageId --source winget --accept-package-agreements --accept-source-agreements --force
$wingetExit = $LASTEXITCODE
} catch { $wingetExit = 1 }
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
$DetectedPython = Find-CompatiblePython
}
}
# Fall back to python.org if winget is unavailable OR couldn't install a
# working Python (missing/broken winget, msstore cert errors --source
# winget can't fix). Keeps the install automatic instead of failing out.
if (-not $DetectedPython) {
if ($script:WingetAvailable) {
substep "winget could not install Python -- falling back to python.org..." "Yellow"
} else {
substep "winget is unavailable -- installing Python from python.org..." "Yellow"
}
$DetectedPython = Install-PythonFromPythonOrg
}
if (-not $DetectedPython) {
Write-Host "[ERROR] Python installation failed (exit code $wingetExit)" -ForegroundColor Red
$exitNote = if ($null -ne $wingetExit) { " (winget exit code $wingetExit)" } else { "" }
Write-Host "[ERROR] Python installation failed$exitNote" -ForegroundColor Red
Write-Host " Please install Python $PythonVersion manually from https://www.python.org/downloads/" -ForegroundColor Yellow
Write-Host " Make sure to check 'Add Python to PATH' during installation." -ForegroundColor Yellow
Write-Host " Then re-run this installer." -ForegroundColor Yellow
@ -1068,7 +1198,7 @@ shell.Run cmd, 0, False
if ($script:WingetAvailable) {
$prevEAP = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {}
try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {}
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
}
@ -1244,14 +1374,107 @@ shell.Run cmd, 0, False
try { [System.IO.File]::WriteAllText((Join-Path $VenvDir ".unsloth-studio-owned"), "") } catch {}
}
# ── Helper: run amd-smi without triggering a UAC elevation prompt ──
# amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing
# DiskPart UAC prompt mid-install (Studio backend amd.py hits the same).
# __COMPAT_LAYER=RunAsInvoker forces it (and helpers it spawns) to run
# un-elevated; on failure the WMI name -> gfx fallback still resolves the arch.
function Invoke-AmdSmiNoElevate {
param(
[Parameter(Mandatory = $true, Position = 0)][string]$Exe,
[Parameter(Position = 1)][string[]]$SmiArgs = @(),
[int]$TimeoutSec = 30
)
# RunAsInvoker blocks the auto-elevation/UAC prompt; the timeout bounds a
# flaky amd-smi that can otherwise spin for minutes (30s mirrors amd.py).
$prevCompat = [Environment]::GetEnvironmentVariable('__COMPAT_LAYER', 'Process')
$env:__COMPAT_LAYER = 'RunAsInvoker'
try {
# [Process]::Start, NOT Start-Process -PassThru: the latter leaves
# .ExitCode $null after WaitForExit on PS 5.1, so $LASTEXITCODE (checked
# by callers) reads non-zero and kills detection. Async reads drain the
# pipes (no deadlock); amd-smi args have no spaces so a plain join is safe.
$psi = New-Object System.Diagnostics.ProcessStartInfo
$psi.FileName = $Exe
$psi.Arguments = ($SmiArgs -join ' ')
$psi.UseShellExecute = $false
$psi.RedirectStandardOutput = $true
$psi.RedirectStandardError = $true
$psi.CreateNoWindow = $true
$proc = [System.Diagnostics.Process]::Start($psi)
$outTask = $proc.StandardOutput.ReadToEndAsync()
$errTask = $proc.StandardError.ReadToEndAsync()
if (-not $proc.WaitForExit($TimeoutSec * 1000)) {
try { $proc.Kill() } catch {}
$global:LASTEXITCODE = 124
return ""
}
$global:LASTEXITCODE = $proc.ExitCode
return ($outTask.Result + "`n" + $errTask.Result)
} catch {
$global:LASTEXITCODE = 1
return ""
} finally {
if ($null -eq $prevCompat) {
Remove-Item Env:__COMPAT_LAYER -ErrorAction SilentlyContinue
} else {
$env:__COMPAT_LAYER = $prevCompat
}
}
}
# ── Helper: run nvidia-smi under a timeout ──
# A wedged NVIDIA driver can make nvidia-smi block during init or after a
# reset; WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate) so detection
# cannot hang the installer. No RunAsInvoker compat layer: nvidia-smi does
# not auto-elevate. Returns combined stdout+stderr; "" on timeout/failure.
function Invoke-NvidiaSmiBounded {
param(
[Parameter(Mandatory = $true, Position = 0)][string]$Exe,
[Parameter(Position = 1)][string[]]$SmiArgs = @(),
[int]$TimeoutSec = 10
)
try {
$psi = New-Object System.Diagnostics.ProcessStartInfo
$psi.FileName = $Exe
$psi.Arguments = ($SmiArgs -join ' ')
$psi.UseShellExecute = $false
$psi.RedirectStandardOutput = $true
$psi.RedirectStandardError = $true
$psi.CreateNoWindow = $true
$proc = [System.Diagnostics.Process]::Start($psi)
$outTask = $proc.StandardOutput.ReadToEndAsync()
$errTask = $proc.StandardError.ReadToEndAsync()
if (-not $proc.WaitForExit($TimeoutSec * 1000)) {
try { $proc.Kill() } catch {}
$global:LASTEXITCODE = 124
return ""
}
$global:LASTEXITCODE = $proc.ExitCode
return ($outTask.Result + "`n" + $errTask.Result)
} catch {
$global:LASTEXITCODE = 1
return ""
}
}
# ── Helper: nvidia-smi -L lists at least one real GPU ──
# Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0
# while listing no GPU, which would mark an AMD host NVIDIA and suppress
# ROCm detection. Require a "GPU <n>:" data row.
function Test-NvidiaSmiHasGpu {
param([Parameter(Mandatory = $true)][string]$Exe)
$out = Invoke-NvidiaSmiBounded $Exe @('-L')
return ($LASTEXITCODE -eq 0 -and $out -match '(?m)^GPU\s+\d+:')
}
# ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ──
$HasNvidiaSmi = $false
$NvidiaSmiExe = $null
try {
$nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
if ($nvSmiCmd) {
& $nvSmiCmd.Source *> $null
if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source }
if ($nvSmiCmd -and (Test-NvidiaSmiHasGpu $nvSmiCmd.Source)) {
$HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source
}
} catch {}
if (-not $HasNvidiaSmi) {
@ -1261,8 +1484,7 @@ shell.Run cmd, 0, False
)) {
if (Test-Path $p) {
try {
& $p *> $null
if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break }
if (Test-NvidiaSmiHasGpu $p) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break }
} catch {}
}
}
@ -1317,11 +1539,21 @@ shell.Run cmd, 0, False
}
} catch {}
}
if (-not $HasROCm) {
# On hosts without a working HIP runtime amd-smi elevates a child at runtime,
# popping a UAC/DiskPart prompt RunAsInvoker can't suppress (manifest is
# asInvoker). So only probe when a HIP SDK is present (hipinfo found ->
# un-elevated) or the user opts in; else fall through to WMI name inference
# (enough to pick ROCm wheels + lemonade llama.cpp).
# An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the
# HIP-SDK heuristic: a HIP SDK binary with a broken runtime can still pop the
# prompt, so $HipSdkInstalled must NOT silently re-enable it.
$amdSmiOptOut = $env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(0|false|no|off)$'
$amdSmiAllowed = (-not $amdSmiOptOut) -and ($HipSdkInstalled -or ($env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(1|true|yes|on)$'))
if (-not $HasROCm -and $amdSmiAllowed) {
$amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiExe) {
try {
$smiOut = & $amdSmiExe.Source list 2>&1 | Out-String
$smiOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('list')
if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") {
$HasROCm = $true
# Mirror the hipinfo path: collect all gfx tokens in enumeration
@ -1336,7 +1568,7 @@ shell.Run cmd, 0, False
# Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+,
# including the GFX target needed for wheel index selection.
$smiAsicOut = ""
try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {}
try { $smiAsicOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('static','--asic') } catch {}
$_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() })
if ($_asicGfxTokens.Count -gt 0) {
$ROCmGfxArch = if ($_smiVisIdx -lt $_asicGfxTokens.Count) { $_asicGfxTokens[$_smiVisIdx] } else { $_asicGfxTokens[0] }
@ -1360,9 +1592,12 @@ shell.Run cmd, 0, False
} catch {}
}
# ── Arch resolution: env-var override → name inference ──────────────
# Covers users whose amd-smi is too old to report the GFX target and
# who don't have hipinfo (HIP-runtime-only, common on Strix Halo / iGPU).
if ($HasROCm -and -not $ROCmGfxArch) {
# Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime
# ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the
# studio setup forward --rocm-gfx and pull a GPU-accelerated (lemonade)
# llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels
# still require a confirmed HIP SDK -- they stay gated on $HasROCm below.
if (-not $ROCmGfxArch) {
# 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running.
if ($env:UNSLOTH_ROCM_GFX_ARCH) {
$ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower()
@ -1370,15 +1605,20 @@ shell.Run cmd, 0, False
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan"
}
# 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI).
# Targets only arches the lemonade-sdk ROCm prebuilts cover
# (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU.
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4
@{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail)
@{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point)
@{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop
@{ P = "RX 7600"; A = "gfx1102" } # RDNA 3
@{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix)
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
@{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI 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 = "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) -- lemonade gfx103X
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X
@{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X
)
foreach ($row in $nameArchTable) {
if ($ROCmGpuLabel -match $row.P) {
@ -1419,11 +1659,11 @@ shell.Run cmd, 0, False
}
} catch {}
}
if (-not $ROCmVersion) {
if (-not $ROCmVersion -and $amdSmiAllowed) {
$amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiVer) {
try {
$smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String
$smiVerOut = Invoke-AmdSmiNoElevate $amdSmiVer.Source @('version')
if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') {
$ROCmVersion = $Matches[1]
}
@ -1433,6 +1673,46 @@ shell.Run cmd, 0, False
}
}
# ── Optional WSL-ROCm driver hint ────────────────────────────────────────
# An AMD GPU can also be used inside WSL2, but only with Adrenalin >= 26.2.2
# (first production ROCDXG/WSL release); native Windows GPU works with any
# recent driver. We can't auto-install it (AMD referrer-gates downloads, no
# winget package), so just point at AMD's page. Shown only when the installed
# driver predates 26.2.2 (Feb 2026); suppress with UNSLOTH_SKIP_AMD_DRIVER_HINT=1.
function Show-AmdWslDriverHint {
if ($env:UNSLOTH_SKIP_AMD_DRIVER_HINT) { return }
try {
$amd = Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue |
Where-Object { $_.Name -match 'AMD|Radeon' } | Select-Object -First 1
if (-not $amd) { return }
$drvDate = $null
try {
if ($amd.DriverDate -is [datetime]) {
# Get-CimInstance returns DriverDate already parsed.
$drvDate = $amd.DriverDate
} elseif ($amd.DriverDate) {
# Get-WmiObject style WMI datetime string.
$drvDate = [Management.ManagementDateTimeConverter]::ToDateTime([string]$amd.DriverDate)
}
} catch {}
# Older than 26.2.2 (Feb 2026) => can't expose the GPU to WSL ROCm.
# Unreadable date => still show the hint (informational, suppressible).
if ($drvDate -and $drvDate -ge (Get-Date '2026-02-01')) { return }
substep "Tip: to use this GPU inside WSL too, install AMD Adrenalin 26.2.2+ (for WSL2)." "Cyan"
substep " Your current driver predates it; native Windows GPU is unaffected. Get it from AMD:" "Cyan"
substep " https://www.amd.com/en/resources/support-articles/release-notes/RN-RAD-WIN-26-2-2.html" "Cyan"
substep " Then reboot and run this installer inside an Ubuntu-24.04 WSL distro." "Cyan"
# If WSL isn't installed yet, point at the command that provisions it
# (best-effort; wsl.exe absent => no WSL).
$hasWsl = $false
try { $hasWsl = [bool](Get-Command wsl.exe -ErrorAction SilentlyContinue) } catch {}
if (-not $hasWsl) {
substep " No WSL yet? Install it first: wsl --install -d Ubuntu-24.04" "Cyan"
}
substep " (suppress: set UNSLOTH_SKIP_AMD_DRIVER_HINT=1)" "Cyan"
} catch {}
}
if ($HasNvidiaSmi) {
step "gpu" "NVIDIA GPU detected"
} elseif ($HasROCm) {
@ -1449,15 +1729,24 @@ shell.Run cmd, 0, False
substep " This is a driver issue, not an SDK issue." "Yellow"
substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} elseif ($ROCmGfxArch) {
# Known arch: Studio setup installs AMD's bundled-runtime ROCm PyTorch wheels
# (repo.amd.com), which ship their own runtime -- HIP SDK optional.
step "gpu" "AMD ROCm ($ROCmGfxArch)" "Cyan"
substep "Detected: $ROCmGpuLabel" "Cyan"
substep "GPU PyTorch uses AMD's bundled-runtime ROCm wheels -- HIP SDK not required (optional)." "Cyan"
} elseif ($ROCmGpuLabel) {
step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow"
step "gpu" "AMD GPU detected -- arch unknown" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "Install the HIP SDK for ROCm GPU inference:" "Yellow"
substep "Could not determine the GPU arch -- install the HIP SDK or set" "Yellow"
substep "UNSLOTH_ROCM_GFX_ARCH to enable GPU ROCm PyTorch:" "Yellow"
substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} else {
step "gpu" "none (chat-only / GGUF)" "Yellow"
substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow"
}
# On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint.
if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint }
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
# Mirrors Get-PytorchCudaTag in setup.ps1.
@ -1465,7 +1754,7 @@ shell.Run cmd, 0, False
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = & $NvidiaSmiExe 2>&1 | Out-String
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
# Newer NVIDIA drivers (e.g. 610.x on Windows) print
# "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y".
# Accept both spellings so we don't fall through to the cu126 default.
@ -1873,16 +2162,24 @@ shell.Run cmd, 0, False
# ── Print CPU-only hint when no GPU detected ──
if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") {
Write-Host ""
if ($HipSdkInstalled -and -not $HasROCm) {
substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow"
} elseif ($ROCmGpuLabel) {
substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow"
if ($ROCmGfxArch) {
# Known AMD arch: install.ps1 lays down CPU PyTorch as a base, then
# setup.ps1 swaps in AMD's bundled-runtime GPU ROCm wheels (no HIP SDK).
substep "Installing CPU PyTorch as a base -- Studio setup installs GPU ROCm" "Cyan"
substep "wheels for $ROCmGfxArch next (bundled runtime; HIP SDK not required)." "Cyan"
} else {
substep "No NVIDIA GPU detected." "Yellow"
if ($HipSdkInstalled -and -not $HasROCm) {
substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow"
} elseif ($ROCmGpuLabel) {
substep "Installing CPU-only PyTorch (AMD GPU arch unknown -- install the HIP SDK" "Yellow"
substep "or set UNSLOTH_ROCM_GFX_ARCH to enable GPU ROCm)." "Yellow"
} else {
substep "No NVIDIA GPU detected." "Yellow"
}
substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow"
substep "re-run with --no-torch for a faster, lighter install:" "Yellow"
substep ".\install.ps1 --no-torch" "Yellow"
}
substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow"
substep "re-run with --no-torch for a faster, lighter install:" "Yellow"
substep ".\install.ps1 --no-torch" "Yellow"
Write-Host ""
}
@ -1923,7 +2220,7 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.1" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -1937,7 +2234,7 @@ shell.Run cmd, 0, False
}
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.1" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -1984,7 +2281,7 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.1" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
@ -1996,7 +2293,7 @@ shell.Run cmd, 0, False
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.1" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2024,7 +2321,7 @@ shell.Run cmd, 0, False
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.1" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.3" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@ -2139,6 +2436,11 @@ shell.Run cmd, 0, False
$studioArgs = @('studio', 'setup')
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
$env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1"
# Hand the venv interpreter to setup.ps1 so it reuses the Python we already
# resolved and built the venv with, instead of re-probing the system (which
# can trip over an unsupported `python` 3.14 or a Store stub on PATH even
# though the venv is fine). setup.ps1 Test-Path-guards this before use.
$env:UNSLOTH_SETUP_PYTHON = Join-Path $VenvDir "Scripts\python.exe"
try {
& $UnslothExe @studioArgs
$setupExit = $LASTEXITCODE
@ -2149,6 +2451,7 @@ shell.Run cmd, 0, False
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
}
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue
}
if ($setupExit -ne 0) {
Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red

336
install.sh
View file

@ -1207,6 +1207,15 @@ STUB_EOF
# Escape single quotes for PowerShell single-quoted string embedding
_css_sc_args_ps=$(printf '%s' "$_css_sc_args" | sed "s/'/''/g")
# DISTINCT shortcut name so the WSL launcher never clobbers a native
# install's "Unsloth Studio.lnk" in the same folder. Per-distro suffix.
if [ -n "$_css_distro" ]; then
_css_lnk_name="Unsloth Studio (WSL - ${_css_distro}).lnk"
else
_css_lnk_name="Unsloth Studio (WSL).lnk"
fi
_css_lnk_name_ps=$(printf '%s' "$_css_lnk_name" | sed "s/'/''/g")
# Create shortcuts via a temp PowerShell script to avoid escaping issues
_css_ps1_tmp=$(mktemp /tmp/unsloth-shortcut-XXXXXX.ps1 2>/dev/null) || true
if [ -n "$_css_ps1_tmp" ]; then
@ -1214,19 +1223,57 @@ STUB_EOF
\$WshShell = New-Object -ComObject WScript.Shell
\$targetExe = (Get-Command '$_css_sc_target' -ErrorAction SilentlyContinue).Source
if (-not \$targetExe) { exit 1 }
# Best-effort: fetch the Unsloth icon to a stable Windows path (shared with a
# native install if one exists) so the WSL shortcut shows the proper icon.
\$iconDir = Join-Path \$env:LOCALAPPDATA 'Unsloth Studio'
\$iconPath = Join-Path \$iconDir 'unsloth.ico'
if (-not (Test-Path -LiteralPath \$iconPath)) {
try {
New-Item -ItemType Directory -Force -Path \$iconDir | Out-Null
Invoke-WebRequest -Uri 'https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth.ico' -OutFile \$iconPath -UseBasicParsing -ErrorAction Stop
} catch {}
}
\$hasIcon = \$false
if (Test-Path -LiteralPath \$iconPath) {
try { \$b = [System.IO.File]::ReadAllBytes(\$iconPath); if (\$b.Length -ge 4 -and \$b[0] -eq 0 -and \$b[1] -eq 0 -and \$b[2] -eq 1 -and \$b[3] -eq 0) { \$hasIcon = \$true } } catch {}
}
\$locations = @(
[Environment]::GetFolderPath('Desktop'),
(Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs')
)
\$created = @()
foreach (\$dir in \$locations) {
if (-not \$dir -or -not (Test-Path \$dir)) { continue }
\$linkPath = Join-Path \$dir 'Unsloth Studio.lnk'
\$linkPath = Join-Path \$dir '$_css_lnk_name_ps'
\$shortcut = \$WshShell.CreateShortcut(\$linkPath)
\$shortcut.TargetPath = \$targetExe
\$shortcut.Arguments = '$_css_sc_args_ps'
\$shortcut.Description = 'Launch Unsloth Studio'
\$shortcut.Description = 'Launch Unsloth Studio (WSL)'
if (\$hasIcon) { \$shortcut.IconLocation = "\$iconPath,0" }
\$shortcut.Save()
\$created += \$linkPath
}
# Force Explorer to re-read EACH new shortcut's icon so it renders immediately
# instead of a stale/blank (generic) icon. The reliable, NON-disruptive fix
# (no explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM,
# SHCNF_PATHW, <lnk>) -- the global SHCNE_ASSOCCHANGED alone does not recover a
# stale item. Also clear the on-disk icon cache for heavier staleness.
try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {}
try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {}
try {
Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int e, uint f, string a, System.IntPtr b);' -ErrorAction SilentlyContinue
foreach (\$p in \$created) { try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, \$p, [System.IntPtr]::Zero) } catch {} }
[UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, \$null, [System.IntPtr]::Zero)
} catch {}
# Win11 Start Menu keeps its own tile-icon cache (preserve start2.bin).
try {
\$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState'
if (Test-Path -LiteralPath \$smeh) {
Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch {}
WSLPS1_EOF
# Convert WSL path to Windows path for powershell.exe
@ -1236,6 +1283,13 @@ WSLPS1_EOF
fi
rm -f "$_css_ps1_tmp"
fi
# If WSL interop is disabled (powershell.exe "Exec format error"), the
# shortcut wasn't created; tell the user how to launch / re-enable it.
if [ "$_css_created" -ne 1 ]; then
substep "Couldn't create the Windows shortcut (WSL interop may be disabled)." "$C_WARN"
substep " Launch Studio from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN"
substep " (re-enable shortcuts: turn WSL interop back on, e.g. run 'wsl --shutdown' then reopen WSL.)" "$C_WARN"
fi
fi
if [ "$_css_created" -eq 1 ]; then
@ -1650,9 +1704,28 @@ _find_no_torch_runtime() {
}
# ── AMD ROCm GPU detection helper ──
# WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg only when
# HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and /opt/rocm/bin can be
# off PATH outside login shells (the profile.d drop-in). Seed both before any
# rocminfo probe or a ROCDXG WSL host is misdetected as CPU-only.
_ensure_rocm_probe_env() {
export HSA_ENABLE_DXG_DETECTION="${HSA_ENABLE_DXG_DETECTION:-1}"
if ! command -v rocminfo >/dev/null 2>&1 && [ -x /opt/rocm/bin/rocminfo ]; then
PATH="$PATH:/opt/rocm/bin"
fi
}
# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs
# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices).
# Always returns 1 (false) when an NVIDIA GPU is present: blocks every
# detection path (rocminfo, amd-smi, KFD sysfs) from producing a false
# positive on NVIDIA-only or NVIDIA-primary hosts, even when ROCm tools
# are co-installed.
_has_amd_rocm_gpu() {
_ensure_rocm_probe_env
if _has_usable_nvidia_gpu; then
return 1
fi
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
return 0
@ -1660,27 +1733,71 @@ _has_amd_rocm_gpu() {
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
return 0
elif [ -e /dev/kfd ] && \
awk '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \
awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \
gpu && amd { found=1 } END{ exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
# vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver
# 560+) can register KFD topology nodes with non-zero gpu_id but
# vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting
# NVIDIA-only hosts to the ROCm install path.
return 0
fi
return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) only if nvidia-smi is present AND actually lists a GPU.
# Prevents AMD-only hosts with a stale nvidia-smi on PATH from being routed
# into the CUDA branch.
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
else
return 1
fi
"$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# ── Detect GPU and choose PyTorch index URL ──
@ -1697,14 +1814,16 @@ get_torch_index_url() {
# packages) is not sufficient: otherwise an AMD-only host would
# silently install CUDA wheels.
_smi=""
_nvidia_detected=0
if _has_usable_nvidia_gpu; then
_nvidia_detected=1
if command -v nvidia-smi >/dev/null 2>&1; then
_smi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_smi="/usr/bin/nvidia-smi"
fi
fi
if [ -z "$_smi" ]; then
if [ "$_nvidia_detected" -eq 0 ]; then
# No NVIDIA GPU -- check for AMD ROCm GPU.
# PyTorch only publishes ROCm wheels for linux-x86_64; skip the
# ROCm branch entirely on aarch64 / arm64 / other architectures
@ -1781,7 +1900,11 @@ get_torch_index_url() {
# of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions
# (POSIX sed does not support "?" without -E). The two patterns are
# mutually exclusive per line, so head -1 picks the first emitted match.
_cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \
# Bound the call (a wedged nvidia-smi would otherwise hang here) and force
# the C locale for stable parsing. LC_ALL is exported inside this command
# substitution subshell so it reaches nvidia-smi through _run_bounded
# without depending on `env`; the export is scoped to the subshell.
_cuda_ver=$(export LC_ALL=C; _run_bounded "$_smi" 2>/dev/null \
| sed -n \
-e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
-e 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
@ -1914,8 +2037,139 @@ _pick_radeon_wheel() {
esac
}
# ── ROCm-on-WSL bootstrap for AMD Strix Halo (gfx1151) ───────────────────────
# No-op everywhere except: WSL + GPU wanted + no usable GPU yet + /dev/dxg +
# Strix Halo APU. Every other config (NVIDIA, native-Linux ROCm, macOS, Windows,
# CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts
# the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 +
# librocdxg), then sources the env it persisted so detection finds the GPU.
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
# Leave any already-usable GPU completely alone (NVIDIA, or working ROCm).
if _has_usable_nvidia_gpu; then return 0; fi
# "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the
# generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and
# would skip this bootstrap while the real GPU is still unusable. awk consumes
# all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail.
_ensure_rocm_probe_env
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
return 0
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
[ -e /dev/dxg ] || return 0
# Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match
# the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S").
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
command -v bash >/dev/null 2>&1 || return 0
# Fast path: already configured (librocdxg present) but launched from a
# non-login shell so the persisted env wasn't loaded -- just load it.
if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then
if [ -r /etc/profile.d/unsloth-rocm-wsl.sh ]; then
# shellcheck disable=SC1091
. /etc/profile.d/unsloth-rocm-wsl.sh || true
else
# librocdxg present but the env drop-in is gone (e.g. a Studio
# uninstall removed it while keeping shared ROCm). Restore the FULL
# env inline (so rocminfo is on PATH) and recreate the drop-in.
_rw_rocm=/opt/rocm
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${_rw_rocm}/bin:${PATH}"
export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}"
# Persist the drop-in so later non-login Studio launches get the env
# too. /etc/profile.d is root-owned: a plain redirect fails for a
# non-root reinstall (ROCm would silently disappear after this shell),
# so tee through sudo when not root. Best-effort -- the current shell
# already has the env, so the install proceeds either way.
_rw_dropin="$(
printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n'
printf 'export HSA_ENABLE_DXG_DETECTION=1\n'
printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n'
printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}"
printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}"
printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n'
)"
if [ "$(id -u)" = "0" ]; then
printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
elif command -v sudo >/dev/null 2>&1; then
printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true
fi
fi
return 0
fi
echo ""
substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN"
substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU."
substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)"
# Locate the helper: prefer the copy shipped beside install.sh, else fetch it.
_rw_helper="${_REPO_ROOT:-.}/scripts/install_rocm_wsl_strixhalo.sh"
_rw_tmp=""
if [ ! -r "$_rw_helper" ]; then
_rw_tmp="$(mktemp 2>/dev/null || echo /tmp/_unsloth_rocm_wsl.sh)"
if download "https://raw.githubusercontent.com/unslothai/unsloth/main/scripts/install_rocm_wsl_strixhalo.sh" "$_rw_tmp" 2>/dev/null; then
_rw_helper="$_rw_tmp"
else
substep "Could not fetch the ROCm-on-WSL helper; using CPU fallback." "$C_WARN"
[ -n "$_rw_tmp" ] && rm -f "$_rw_tmp"
return 0
fi
fi
# Consent: the narrow guarded case is exactly the GPU setup the user ran the
# installer for, so it proceeds AUTOMATICALLY by default (works with no TTY,
# e.g. `curl ... | sh`). Opt out via UNSLOTH_SKIP_ROCM_WSL_SETUP=1 (top of
# function). The Tauri app drives its own consent UI, so under TAURI_MODE it
# only runs when the app passes UNSLOTH_ROCM_WSL_AUTO=1; else surface and wait.
_rw_go=1
if [ "${TAURI_MODE:-false}" = "true" ] && [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" != "1" ]; then
tauri_log "ROCM_WSL_AVAILABLE" "strixhalo"
substep "Enable the GPU from the desktop app (or set UNSLOTH_ROCM_WSL_AUTO=1)." "$C_WARN"
_rw_go=0
fi
if [ "$_rw_go" = "1" ]; then
# Helper does its own sudo + is idempotent. SMOKE_TEST=0: install.sh
# installs torch itself right after, into the real venv.
if UNSLOTH_WSL_SMOKE_TEST=0 bash "$_rw_helper"; then
# Pull the helper's persisted env into THIS shell so detection
# (rocminfo) now enumerates the GPU and routes to gfx1151.
if [ -r /etc/profile.d/unsloth-rocm-wsl.sh ]; then
# shellcheck disable=SC1091
. /etc/profile.d/unsloth-rocm-wsl.sh || true
fi
substep "ROCm-on-WSL ready; continuing with GPU install." "$C_OK"
else
substep "ROCm-on-WSL setup did not complete; falling back to CPU-only." "$C_WARN"
fi
fi
[ -n "$_rw_tmp" ] && rm -f "$_rw_tmp"
return 0
}
_maybe_bootstrap_rocm_wsl || true
TORCH_INDEX_URL=$(get_torch_index_url)
# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that
# downstream scripts (setup.sh -> install_python_stack.py) know what was
# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts.
# Classify on the FINAL path segment only: a custom UNSLOTH_PYTORCH_MIRROR
# whose base path happens to contain "rocm" or "gfx" must not mislabel a
# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix
# overrides in gfxNNNN/, so the trailing slash is stripped first).
_torch_index_leaf="${TORCH_INDEX_URL%/}"
_torch_index_leaf="${_torch_index_leaf##*/}"
case "$_torch_index_leaf" in
rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;;
cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;;
*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
esac
# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
# All other ROCm tags and CUDA stay within <2.11.0.
case "$TORCH_INDEX_URL" in
@ -2018,6 +2272,7 @@ if _has_usable_nvidia_gpu; then
step "gpu" "NVIDIA GPU detected"
elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
# Probe gfx arch for the display label, honouring HIP_VISIBLE_DEVICES
_ensure_rocm_probe_env
_gpu_disp_gfx_all=""
_gpu_disp_mkt=""
if command -v rocminfo >/dev/null 2>&1; then
@ -2048,14 +2303,20 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx"
# Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable)
elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then
# Kept in sync with the nameArchTable in install.ps1 / setup.ps1.
# 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
*"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 iGPU
*"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 iGPU
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop
*"RX 7600"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU
*"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI 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)
*"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)
*"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) _gpu_disp_gfx="gfx1034" ;; # RDNA 2 (Navi 24)
esac
if [ -n "$_gpu_disp_gfx" ]; then
substep "gfx arch inferred from GPU name: $_gpu_disp_gfx"
@ -2089,7 +2350,34 @@ case "$TORCH_INDEX_URL" in
*/cpu)
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
if [ "$OS" = "wsl" ]; then
# WSL + no GPU detected (detection above found nothing). Common
# cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet --
# /dev/dxg present (graphics) but no ROCm runtime.
_wsl_ubu_ver=""
[ -r /etc/os-release ] && _wsl_ubu_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}")
if [ -e /dev/dxg ]; then
substep "A GPU is plumbed into WSL (/dev/dxg) but no ROCm runtime is exposed to it." "$C_WARN"
fi
substep "For an AMD GPU, ROCm-on-WSL currently needs ALL of:"
substep " 1. AMD Adrenalin Edition 26.1.1+ on Windows (26.2.2+ for Strix Halo / Ryzen AI Max+)."
substep " Older drivers lack production ROCDXG/WSL support, so ROCm can't see the GPU."
substep " Get it from AMD (open in a browser -- direct downloads are referrer-gated):"
substep " https://www.amd.com/en/resources/support-articles/release-notes/RN-RAD-WIN-26-2-2.html"
substep " 2. ROCm 7.2.1 + librocdxg inside WSL (with HSA_ENABLE_DXG_DETECTION=1)."
substep " 3. A WSL distro AMD supports for ROCm -- Ubuntu 24.04 is the known-good one."
if [ -n "$_wsl_ubu_ver" ] && [ "$_wsl_ubu_ver" != "24.04" ]; then
substep " This distro is Ubuntu $_wsl_ubu_ver, which AMD may not support for ROCm-on-WSL yet." "$C_WARN"
fi
substep "Set up the GPU in WSL with a dedicated Ubuntu 24.04 distro:"
substep " wsl --install Ubuntu-24.04 # run in Windows PowerShell, then reopen WSL"
substep " # then re-run this installer inside Ubuntu-24.04 -- it will detect the GPU."
substep "AMD ROCm-on-WSL docs: https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/"
substep "Strix Halo (gfx1151): this installer auto-offers ROCm-on-WSL setup once the"
substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself."
else
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
fi
substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):"
substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
fi
@ -2117,7 +2405,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.1" unsloth-zoo
"unsloth>=2026.6.3" unsloth-zoo
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2130,7 +2418,7 @@ if [ "$_MIGRATED" = true ]; then
else
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.1" unsloth-zoo
"unsloth>=2026.6.3" unsloth-zoo
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2334,7 +2622,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.6.1" unsloth-zoo
"unsloth>=2026.6.3" unsloth-zoo
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2352,7 +2640,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.1" unsloth-zoo
--upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2408,7 +2696,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.1" --torch-backend=auto
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.3" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."

7
pyproject.toml
View file

@ -26,6 +26,7 @@ classifiers = [
]
dependencies = [
"typer",
"rich",
"pydantic",
"pyyaml",
"nest-asyncio",
@ -70,7 +71,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.6.1",
"unsloth_zoo>=2026.6.3",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -91,7 +92,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.6.1",
"unsloth_zoo>=2026.6.3",
"torchvision",
"unsloth[triton]",
]
@ -581,7 +582,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.6.1",
"unsloth_zoo>=2026.6.3",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

289
scripts/install_rocm_wsl_strixhalo.sh Normal file
View file

@ -0,0 +1,289 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# ──────────────────────────────────────────────────────────────────────────────
# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151)
# ──────────────────────────────────────────────────────────────────────────────
# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime
# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG
# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04
# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via
# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies).
#
# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with
# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once
# installed + rebooted, /dev/dxg is exposed to WSL and this script builds the rest.
#
# HOW ROCDXG WORKS (and why older /usr/lib/wsl/lib notes are wrong): librocdxg.so
# is AMD's user-mode bridge between the Linux HSA runtime and the Windows driver
# over /dev/dxg. The STANDARD hsa-rocr runtime (NOT the gone "roc4wsl" package)
# loads it when HSA_ENABLE_DXG_DETECTION=1. No hsa/rocm libs need injecting into
# /usr/lib/wsl/lib (it holds only d3d12/dxcore), yet rocminfo enumerates gfx1151
# fine -- so we gate on /dev/dxg, not on WSL lib injection.
#
# KNOWN CAVEAT (ROCm/ROCm#6022): librocdxg can cap usable ROCm VRAM at the WSL
# VM's RAM (.wslconfig [wsl2] memory=) on some BIOS UMA layouts, and amd-smi
# doesn't work in WSL. On OOM below capacity, raise memory= (then wsl --shutdown)
# and watch GPU use from Windows. Large-UMA BIOS exposes the full pool regardless.
#
# Verified on Ryzen AI Max+ PRO 395 / Radeon 8060S (gfx1151) with ROCm 7.2.1 +
# Ubuntu 24.04 + WSL2 + Adrenalin. These pins MOVE; bump + re-verify on newer ROCm.
# ──────────────────────────────────────────────────────────────────────────────
set -euo pipefail
# ── Tunables (override via env) ──────────────────────────────────────────────
ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install
GFX="gfx1151"
LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build
# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test.
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/"
# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs
# torch itself into the real venv right after, so a duplicate download is wasteful.
SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}"
# REQUIRED constraint -- without it pip prefers PyPI's newer CUDA torch over the
# gfx1151 ROCm wheel. 2.11 carries AMD's real gfx1151 fix (matches install.sh).
TORCH_CONSTRAINT="${UNSLOTH_WSL_TORCH_CONSTRAINT:-torch>=2.11.0,<2.12.0}"
ROCM_DIR="" # resolved after install
say() { printf '\n\033[1;36m== %s\033[0m\n' "$*"; }
note() { printf ' %s\n' "$*"; }
die() { printf '\n\033[1;31m[BLOCKED] %s\033[0m\n' "$*" >&2; exit 1; }
# sudo only if not already root (WSL distros often run as root)
SUDO=""
if [ "$(id -u)" -ne 0 ]; then
command -v sudo >/dev/null 2>&1 || die "Need root or sudo to install ROCm."
SUDO="sudo"
fi
# ── Windows 11 SDK (headers for the librocdxg build) ─────────────────────────
# librocdxg's cmake build needs the Windows SDK 'shared' headers, which live on
# the Windows HOST under C:\Program Files (x86)\Windows Kits\10\Include\<ver>\.
_WIN_SDK_INC_BASE="/mnt/c/Program Files (x86)/Windows Kits/10/Include"
# Print the newest installed SDK include dir with 'shared' headers, or nothing.
# find + read loop (not `for ... in $(ls)`) since the base path has a space.
_find_win_sdk() {
[ -d "$_WIN_SDK_INC_BASE" ] || return 0
while IFS= read -r _inc; do
[ -n "$_inc" ] || continue
if [ -d "$_inc/shared" ]; then printf '%s' "$_inc"; return 0; fi
done < <(find "$_WIN_SDK_INC_BASE" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort -Vr)
return 0
}
# Best-effort: install the Windows 11 SDK on the Windows HOST via winget so the
# build has its headers with no manual step. Elevates -> ONE UAC prompt; headers
# appear under /mnt/c immediately (no reboot). Never fatal -- failure falls
# through to a manual-install message. Opt out: UNSLOTH_SKIP_WIN_SDK_INSTALL=1.
_install_windows_sdk_via_winget() {
[ "${UNSLOTH_SKIP_WIN_SDK_INSTALL:-0}" = "1" ] && { note "Skipping Windows SDK auto-install (UNSLOTH_SKIP_WIN_SDK_INSTALL=1)."; return 0; }
command -v powershell.exe >/dev/null 2>&1 || return 0
# `command -v` succeeds even with WSL interop OFF (.exe on PATH but fails
# with "Exec format error"); verify it actually executes.
powershell.exe -NoProfile -Command "exit 0" >/dev/null 2>&1 || return 0
if ! powershell.exe -NoProfile -Command "if (Get-Command winget -ErrorAction SilentlyContinue) { exit 0 } else { exit 1 }" >/dev/null 2>&1; then
note "winget not available on the Windows host -- cannot auto-install the Windows SDK."
return 0
fi
say "Installing the Windows 11 SDK on the Windows host via winget"
note "librocdxg needs its headers. Approve the UAC prompt on the Windows desktop."
note "One-time (~1-3 GB download); opt out with UNSLOTH_SKIP_WIN_SDK_INSTALL=1."
# Newest SDK first, then a fallback. Header presence is the source of truth
# (re-check each attempt), not winget's exit code. </dev/null so winget never
# consumes a piped `curl | sh` stdin.
for _sdk_id in Microsoft.WindowsSDK.10.0.26100 Microsoft.WindowsSDK.10.0.22621; do
note "winget install ${_sdk_id} ..."
# --source winget: pin the community source so a broken default msstore
# source (the cert failure this PR fixes) can't abort SDK resolution.
powershell.exe -NoProfile -Command "winget install --id ${_sdk_id} -e --source winget --accept-source-agreements --accept-package-agreements --disable-interactivity" </dev/null || true
if [ -n "$(_find_win_sdk)" ]; then
note "Windows SDK headers present after install."
return 0
fi
done
note "Automatic Windows SDK install did not complete."
return 0
}
# ── PREFLIGHT ────────────────────────────────────────────────────────────────
say "Preflight checks"
# shellcheck disable=SC1091
. /etc/os-release 2>/dev/null || true
if [ "${VERSION_ID:-}" != "24.04" ]; then
die "This targets Ubuntu 24.04 (found '${VERSION_ID:-unknown}'). AMD's ROCm-on-WSL supports 24.04; create a dedicated distro: wsl --install Ubuntu-24.04 (do not run on 26.04 -- ROCm 7.2 does not target it yet)."
fi
if [ ! -e /dev/dxg ]; then
die "/dev/dxg missing -- WSL GPU paravirtualization not present. Ensure this is WSL2 (not WSL1) on a recent Windows build, and that an AMD GPU + ROCDXG-capable Adrenalin driver is installed on the Windows host (then reboot)."
fi
note "Ubuntu 24.04 + /dev/dxg present."
# Don't block on hsa/rocm libs in /usr/lib/wsl/lib: a working ROCDXG setup
# doesn't need them (only d3d12/dxcore). Real readiness is checked via rocminfo.
# ── Step 1: build/runtime prerequisites ──────────────────────────────────────
say "Installing build prerequisites"
export DEBIAN_FRONTEND=noninteractive
$SUDO apt-get update -y
# `make` is explicit: cmake shells out to it but Ubuntu only *recommends* it, so
# minimal images lack it and the librocdxg `make -j` build would fail.
$SUDO apt-get install -y cmake make gcc g++ git wget gpg ca-certificates python3-venv python3-pip
# ── Step 2: ROCm ${ROCM_VER} userspace (no DKMS -- WSL uses the Windows driver) ─
say "Installing ROCm ${ROCM_VER} userspace"
if ! command -v rocminfo >/dev/null 2>&1 && [ ! -x /opt/rocm/bin/rocminfo ]; then
# Direct apt-repo install (leaner than amdgpu-install; repo is indexed by
# ROCm version, e.g. .../apt/7.2.1).
$SUDO mkdir -p /etc/apt/keyrings
wget -qO- https://repo.radeon.com/rocm/rocm.gpg.key \
| gpg --dearmor | $SUDO tee /etc/apt/keyrings/rocm.gpg >/dev/null
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VER} noble main" \
| $SUDO tee /etc/apt/sources.list.d/rocm.list >/dev/null
printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' \
| $SUDO tee /etc/apt/preferences.d/rocm-pin-600 >/dev/null
$SUDO apt-get update -y
# rocm-libs pulls everything torch links at runtime (rocblas, hipblas,
# miopen-hip, rccl, ...); hsa-rocr + rocminfo come as deps. Large (~5 GB
# download / ~23 GB installed).
$SUDO apt-get install -y rocm-libs rocminfo hip-runtime-amd
else
note "ROCm already present -- skipping apt install."
fi
# Resolve the real ROCm dir and ensure the canonical /opt/rocm symlink. apt lays
# ROCm under /opt/rocm-<ver> and rocm-core symlinks /opt/rocm -> that; repair if
# an earlier partial run left /opt/rocm as a real dir blocking the symlink.
_real="$(ls -d /opt/rocm-* 2>/dev/null | sort -V | tail -1 || true)"
if [ -n "$_real" ] && [ ! -L /opt/rocm ] && [ -d /opt/rocm ]; then
# /opt/rocm is a real dir blocking the symlink. Only treat it as a removable
# stray stub if it's NOT a real ROCm install (a real one has bin/rocminfo /
# bin/hipcc / .info/version) -- this protects a user's pre-existing ROCm. Even
# then we MOVE IT ASIDE, never rm -rf, so a wrong guess can't lose data.
if [ -e /opt/rocm/bin/rocminfo ] || [ -e /opt/rocm/bin/hipcc ] || [ -e /opt/rocm/.info/version ]; then
note "/opt/rocm is a real ROCm install -- leaving it untouched (will install librocdxg into it)."
else
note "Moving stray /opt/rocm stub aside -> $_real (not deleting it)"
$SUDO cp -an /opt/rocm/. "$_real"/ 2>/dev/null || true
$SUDO mv /opt/rocm "/opt/rocm.unsloth-stub-bak.$(date +%s)" 2>/dev/null || true
[ -e /opt/rocm ] || $SUDO ln -s "$_real" /opt/rocm
fi
elif [ -n "$_real" ] && [ ! -e /opt/rocm ]; then
$SUDO ln -s "$_real" /opt/rocm
fi
if [ -L /opt/rocm ] || [ -d /opt/rocm ]; then ROCM_DIR="/opt/rocm"; else ROCM_DIR="$_real"; fi
{ [ -n "$ROCM_DIR" ] && [ -d "$ROCM_DIR" ]; } || die "ROCm not found under /opt after install."
note "ROCm at ${ROCM_DIR}"
# ── Step 3: build librocdxg (DXG <-> HSA bridge; not yet an apt package) ──────
say "Building librocdxg (${LIBROCDXG_REF})"
if [ -e "${ROCM_DIR}/lib/librocdxg.so" ]; then
note "librocdxg already installed -- skipping build."
else
# Discover the newest installed Win11 SDK (version differs per machine). If
# absent, auto-install via winget (one UAC prompt) and re-discover; only if
# that ALSO fails do we stop with manual instructions.
_win_sdk="$(_find_win_sdk)"
if [ -z "$_win_sdk" ]; then
note "Windows 11 SDK headers not found -- attempting automatic install..."
_install_windows_sdk_via_winget
_win_sdk="$(_find_win_sdk)"
fi
[ -n "$_win_sdk" ] || die "Windows 11 SDK headers not found under 'C:\\Program Files (x86)\\Windows Kits\\10\\Include\\*\\shared', and the automatic winget install did not complete. Install it on the Windows host (e.g. 'winget install Microsoft.WindowsSDK.10.0.26100') and re-run."
note "Windows SDK: ${_win_sdk}"
_src="${HOME}/.unsloth/librocdxg"
rm -rf "$_src"
git clone --depth 1 --branch "$LIBROCDXG_REF" https://github.com/ROCm/librocdxg.git "$_src" \
|| git clone "https://github.com/ROCm/librocdxg.git" "$_src"
(
cd "$_src"
git checkout "$LIBROCDXG_REF" 2>/dev/null || true
mkdir -p build && cd build
cmake .. -DWIN_SDK="${_win_sdk}/shared"
make -j"$(nproc)"
$SUDO make install
)
fi
# Ensure soname symlinks resolve to whatever version was built (e.g. 1.2.0).
_dxg_real="$(ls -1 "${ROCM_DIR}"/lib/librocdxg.so.*.* 2>/dev/null | sort -V | tail -1 || true)"
if [ -n "$_dxg_real" ]; then
_dxg_base="$(basename "$_dxg_real")" # librocdxg.so.1.2.0
_dxg_major="$(printf '%s' "$_dxg_base" | sed -E 's/librocdxg\.so\.([0-9]+).*/\1/')"
$SUDO ln -sf "$_dxg_base" "${ROCM_DIR}/lib/librocdxg.so.${_dxg_major}"
$SUDO ln -sf "librocdxg.so.${_dxg_major}" "${ROCM_DIR}/lib/librocdxg.so"
fi
echo "${ROCM_DIR}/lib" | $SUDO tee /etc/ld.so.conf.d/rocm.conf >/dev/null
$SUDO ldconfig
# ── Step 4: persist environment (system-wide so Studio's worker inherits it) ──
say "Persisting ROCm-on-WSL environment"
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
$SUDO tee "$_envfile" >/dev/null <<EOF
# >>> Unsloth ROCm-on-WSL (gfx1151) >>>
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${ROCM_DIR}/bin:\${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}"
# <<< Unsloth ROCm-on-WSL (gfx1151) <<<
EOF
# also drop into ~/.bashrc for interactive shells
if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then
cat "$_envfile" >> "${HOME}/.bashrc"
fi
# export into the current process so verification below works immediately
export HSA_ENABLE_DXG_DETECTION=1
export PATH="${ROCM_DIR}/bin:${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}"
# ── Step 5: verify the runtime enumerates the GPU ────────────────────────────
say "Verifying rocminfo sees ${GFX}"
# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs
# rocminfo on first match, which under `set -o pipefail` turns a successful match
# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a
# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass.
_rocminfo_out="$(rocminfo 2>/dev/null || true)"
if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then
printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true
die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
fi
# Display-only summary: best-effort (|| true) so head's early pipe-close under
# `set -o pipefail` can't fail the bootstrap after verification already passed.
printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true
note "ROCm-on-WSL runtime is live for ${GFX}."
# ── Step 6 (optional): torch smoke test from the gfx1151 index ───────────────
if [ "$SMOKE_TEST" = "1" ]; then
say "Smoke-testing PyTorch on ${GFX} (throwaway venv)"
_venv="${HOME}/.unsloth/rocm-smoketest"
rm -rf "$_venv"; python3 -m venv "$_venv"
"$_venv/bin/pip" install --quiet --upgrade pip
# gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py
# deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch.
"$_venv/bin/pip" install --index-url "$TORCH_INDEX" \
--extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \
die "torch install from ${TORCH_INDEX} failed."
"$_venv/bin/python" - <<'PY'
import torch
ok = torch.cuda.is_available()
print("torch:", torch.__version__, "| cuda(rocm) available:", ok)
if ok:
print("device:", torch.cuda.get_device_name(0))
free, total = torch.cuda.mem_get_info(0)
print(f"mem: free={free/1e9:.1f} GB total={total/1e9:.1f} GB")
import time
a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
torch.cuda.synchronize(); t0 = time.time()
for _ in range(10): c = a @ b
torch.cuda.synchronize()
print(f"matmul ok ({(time.time()-t0)/10*1e3:.1f} ms/iter)")
raise SystemExit(0 if ok else 1)
PY
rm -rf "$_venv"
fi
say "Done."
note "ROCm-on-WSL is ready for ${GFX}. If you ran this standalone, install Unsloth"
note "in THIS distro and it will detect the GPU automatically:"
note " curl -fsSL https://unsloth.ai/install.sh | sh"

143
scripts/sync_allow_scripts_pins.py Normal file
View file

@ -0,0 +1,143 @@
#!/usr/bin/env python3
"""Keep `allowScripts` pins in studio/frontend/package.json in sync with
package-lock.json.
`npm approve-scripts` writes version-pinned entries ("pkg@1.2.3": true).
A dependency bump strands the pin, so the approval (or denial) silently
stops matching and the package's install scripts fall back to
"unreviewed". This tool re-pins existing entries to the versions the
lockfile actually resolves; it never adds or removes entries, so
approving a brand-new script-bearing package stays a human decision.
Usage:
python scripts/sync_allow_scripts_pins.py --check # CI: exit 1 on drift
python scripts/sync_allow_scripts_pins.py --fix # rewrite package.json
Pinned keys follow npm's allowScripts grammar: "name@1.2.3" or
"name@1.2.3 || 1.2.4". Bare names (no version) match every version and
are left alone. Entries whose range is not an exact-version disjunction
(wildcards, tags) are left alone too.
"""
from __future__ import annotations
import argparse
import json
import re
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_DIR = REPO_ROOT / "studio" / "frontend"
EXACT_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?$")
def split_spec(key: str) -> tuple[str, str | None]:
"""'@scope/name@1.2.3' -> ('@scope/name', '1.2.3'); bare names -> (key, None)."""
if key.startswith("@"):
rest = key[1:]
if "@" not in rest:
return key, None
name, rng = rest.split("@", 1)
return "@" + name, rng
if "@" not in key:
return key, None
name, rng = key.split("@", 1)
return name, rng
def is_exact_disjunction(rng: str) -> bool:
parts = [p.strip() for p in rng.split("||")]
return all(EXACT_VERSION_RE.match(p) for p in parts) and bool(parts)
def version_sort_key(version: str) -> tuple:
release = version.split("-", 1)[0].split("+", 1)[0]
return tuple(int(x) for x in release.split(".")), version
def script_versions_from_lock(lock: dict) -> dict[str, list[str]]:
"""Map package name -> sorted versions that carry install scripts."""
out: dict[str, set[str]] = {}
for path, meta in (lock.get("packages") or {}).items():
if not path or not meta.get("hasInstallScript"):
continue
name = path.rsplit("node_modules/", 1)[-1]
version = meta.get("version")
if name and version:
out.setdefault(name, set()).add(version)
return {n: sorted(vs, key = version_sort_key) for n, vs in out.items()}
def desired_key(name: str, versions: list[str]) -> str:
return f"{name}@{' || '.join(versions)}"
def compute_renames(policy: dict, lock_versions: dict[str, list[str]]) -> dict[str, str]:
renames: dict[str, str] = {}
for key in policy:
name, rng = split_spec(key)
if rng is None or not is_exact_disjunction(rng):
continue # bare name or non-exact spec: matches by name, never stale
versions = lock_versions.get(name)
if not versions:
continue # package gone or script-free now: stale pin is inert
want = desired_key(name, versions)
if key != want:
renames[key] = want
return renames
def main(argv: list[str] | None = None) -> int:
ap = argparse.ArgumentParser(description = __doc__)
mode = ap.add_mutually_exclusive_group(required = True)
mode.add_argument("--check", action = "store_true", help = "exit 1 if pins are stale")
mode.add_argument("--fix", action = "store_true", help = "rewrite package.json in place")
ap.add_argument(
"--dir",
type = Path,
default = DEFAULT_DIR,
help = "directory holding package.json + package-lock.json",
)
args = ap.parse_args(argv)
pkg_path = args.dir / "package.json"
lock_path = args.dir / "package-lock.json"
if not pkg_path.exists() or not lock_path.exists():
print(f"sync-allow-scripts: nothing to do ({args.dir} has no package.json + lockfile)")
return 0
pkg = json.loads(pkg_path.read_text(encoding = "utf-8"))
policy = pkg.get("allowScripts")
if not isinstance(policy, dict) or not policy:
print("sync-allow-scripts: no allowScripts policy in package.json, nothing to do")
return 0
lock = json.loads(lock_path.read_text(encoding = "utf-8"))
renames = compute_renames(policy, script_versions_from_lock(lock))
if not renames:
print(f"sync-allow-scripts: {len(policy)} allowScripts entries in sync with the lockfile")
return 0
for old, new in renames.items():
print(f' stale pin: "{old}" -> "{new}"')
if args.check:
print(
"sync-allow-scripts: pins are stale; run "
"`python scripts/sync_allow_scripts_pins.py --fix` and commit the result"
)
return 1
pkg["allowScripts"] = {renames.get(k, k): v for k, v in policy.items()}
pkg_path.write_text(json.dumps(pkg, indent = 2, ensure_ascii = False) + "\n", encoding = "utf-8")
print(
f"sync-allow-scripts: re-pinned {len(renames)} entr{'y' if len(renames) == 1 else 'ies'} in {pkg_path}"
)
return 0
if __name__ == "__main__":
sys.exit(main())

82
scripts/uninstall.ps1
View file

@ -16,15 +16,19 @@ function Uninstall-UnslothStudio {
function _Step { param([string]$Msg) Write-Host $Msg }
function _Substep { param([string]$Msg, [string]$Color = "Gray") Write-Host " $Msg" -ForegroundColor $Color }
# Remove a file/dir/symlink only if it exists. Idempotent.
# Remove a file/dir/symlink if present. Idempotent; retries since a just-killed
# process can briefly hold a handle (Windows refuses the delete until released).
function _RemovePath {
param([string]$Path)
if ([string]::IsNullOrWhiteSpace($Path)) { return }
if (Test-Path -LiteralPath $Path) {
if (-not (Test-Path -LiteralPath $Path)) { return }
for ($attempt = 1; $attempt -le 3; $attempt++) {
try {
Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
_Substep "removed: $Path" "Green"
return
} catch {
if ($attempt -lt 3) { Start-Sleep -Milliseconds 700; continue }
_Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow"
}
}
@ -236,9 +240,58 @@ function Uninstall-UnslothStudio {
} catch { }
}
# Stop processes that would block deleting the paths we remove. Unlike
# _StopStudioProcesses (venv exe only), this also catches llama-server/llama-cli,
# the unsloth.exe shim, and orphaned mp workers under SYSTEM python holding a
# venv DLL (an open DLL handle blocks the dir delete) -- found by scanning each
# candidate's loaded modules, not just its image path.
function _StopProcessesLockingRoots {
param([string[]]$Roots)
$clean = @($Roots | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\','/') })
if ($clean.Count -eq 0) { return }
$underRoot = {
param($p)
if (-not $p) { return $false }
foreach ($r in $clean) { if ($p -ieq $r -or $p -ilike "$r\*") { return $true } }
return $false
}
# 1. Image path under a target root (venv python, shim, llama-server).
try {
foreach ($proc in (Get-CimInstance Win32_Process -ErrorAction SilentlyContinue)) {
if ((& $underRoot $proc.ExecutablePath)) {
try { Stop-Process -Id $proc.ProcessId -Force -ErrorAction SilentlyContinue } catch { }
}
}
} catch { }
# 2. A loaded module under a target root (orphaned mp-fork python holding a
# venv DLL). Scoped to names that load our DLLs to keep the scan fast.
try {
$cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli -ErrorAction SilentlyContinue
foreach ($proc in $cands) {
$hit = $false
try {
foreach ($m in $proc.Modules) { if ((& $underRoot $m.FileName)) { $hit = $true; break } }
} catch { } # access denied enumerating modules -> skip
if ($hit) { try { Stop-Process -Id $proc.Id -Force -ErrorAction SilentlyContinue } catch { } }
}
} catch { }
}
# Default install root + default data dir.
$defaultStudioHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth\studio" } else { $null }
$defaultDataDir = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth Studio" } else { $null }
# Default-mode ~/.unsloth holds a SHARED llama.cpp build + .cache that are
# siblings of studio (not under it), so deleting <studio> misses them -- handle
# explicitly. No-op in env/custom mode (nested under the custom root, removed
# with it). A user-set UNSLOTH_LLAMA_CPP_PATH is left alone.
$defaultUnslothHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth" } else { $null }
$defaultLlamaCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "llama.cpp" } else { $null }
$defaultCache = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".cache" } else { $null }
# llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging,
# sibling of the install dir). Usually pruned after activate, but an interrupted
# build can leave a "<name>.staging-XXXX" tree; removing it lets the empty-dir
# cleanup of ~/.unsloth below succeed. No-op in env/custom mode and when absent.
$defaultStaging = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".staging" } else { $null }
# Build known-root list FIRST so the port-file kill can verify ownership.
$customRoots = @(_CustomStudioRoots)
@ -255,6 +308,9 @@ function Uninstall-UnslothStudio {
_StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots
}
_StopStudioProcesses -KnownRoots $knownRoots
# Also stop anything holding a handle on the exact paths we delete (llama-server,
# the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused.
_StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache))
# ── Remove custom-root install trees ──
_Step "Removing data and install directories..."
@ -273,6 +329,16 @@ function Uninstall-UnslothStudio {
if ($defaultStudioHome) { _RemovePath $defaultStudioHome }
# Default data dir.
if ($defaultDataDir) { _RemovePath $defaultDataDir }
# Default-mode shared llama.cpp build + cache (siblings of studio under
# ~/.unsloth). No-op in env/custom mode and when absent.
if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp }
if ($defaultCache) { _RemovePath $defaultCache }
if ($defaultStaging) { _RemovePath $defaultStaging }
# Drop ~/.unsloth itself, but ONLY if now empty -- never nuke unrelated content.
if ($defaultUnslothHome -and (Test-Path -LiteralPath $defaultUnslothHome) -and
-not (Get-ChildItem -LiteralPath $defaultUnslothHome -Force -ErrorAction SilentlyContinue)) {
_RemovePath $defaultUnslothHome
}
# ── Remove desktop and Start Menu shortcuts ──
_Step "Removing desktop and Start Menu shortcuts..."
@ -283,6 +349,18 @@ function Uninstall-UnslothStudio {
if ($env:APPDATA) {
_RemovePath (Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk")
}
# Invalidate the Win11 Start Menu tile cache so the removed shortcut's tile
# disappears promptly instead of lingering stale (mirrors install.ps1's
# New-StudioShortcuts). Preserves start2.bin (the pin layout).
try {
$smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
if (Test-Path -LiteralPath $smehTemp) {
Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch { }
# ── Clean user PATH and registry backup ──
_Step "Cleaning user PATH and registry..."

119
scripts/uninstall.sh
View file

@ -212,10 +212,24 @@ _custom_studio_roots | while IFS= read -r _custom_root; do
_remove_path "$_custom_root"
done
_remove_path "$HOME/.unsloth/studio"
# CUDA llama.cpp from provision_llama_cuda.sh (+ the fetched script). Clears the
# native-Linux build dir, or on WSL the symlink to the build install.ps1 removes.
# Default-mode shared llama.cpp build + cache are siblings of studio (not removed
# by deleting it). No-op in env/custom mode (they nest under the custom root) and
# when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept.
_remove_path "$HOME/.unsloth/llama.cpp"
# provision_llama_cuda.sh fetched by the WoA/Spark CUDA-build path (install.ps1
# background build + direct-WSL setup.sh). No-op when absent.
_remove_path "$HOME/.unsloth/provision_llama_cuda.sh"
_remove_path "$HOME/.unsloth/.cache"
# llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging).
# Normally pruned after activate, but an interrupted build can leave it behind;
# removing it lets the rmdir below succeed. No-op in env/custom mode and absent.
_remove_path "$HOME/.unsloth/.staging"
# ROCm-on-WSL helper artifacts (librocdxg build clone + smoke-test venv). No-op
# where they don't exist; removing them lets the rmdir below succeed.
_remove_path "$HOME/.unsloth/librocdxg"
_remove_path "$HOME/.unsloth/rocm-smoketest"
# Drop ~/.unsloth only if now empty (rmdir refuses non-empty, so user content is kept).
rmdir "$HOME/.unsloth" 2>/dev/null || true
_remove_path "$HOME/.local/share/unsloth"
# CLI shim: only the symlink Studio created, never a pip-installed file.
_remove_cli_shim
@ -248,22 +262,50 @@ case "$_os" in
Linux)
if [ "$_is_wsl" = "1" ]; then
echo "Removing WSL Windows-side shortcuts..."
# install.sh creates 'Unsloth Studio.lnk' on the Windows Desktop and
# Start Menu Programs folder via powershell.exe; mirror that path.
if command -v powershell.exe >/dev/null 2>&1; then
# install.sh creates per-distro 'Unsloth Studio (WSL - <distro>).lnk'
# on the Windows Desktop + Start Menu via powershell.exe. Scope removal
# to THIS distro (passed as $args[0]) so a multi-distro install keeps the
# other distros' launchers; the TARGET=wsl.exe check still spares a
# native install's "Unsloth Studio.lnk". Prefer powershell.exe; test it
# can EXECUTE (`command -v` succeeds even with interop OFF -- .exe then
# fails "Exec format error", common on systemd-enabled distros).
_wsl_distro="${WSL_DISTRO_NAME:-}"
_ps_ran=0
if command -v powershell.exe >/dev/null 2>&1 && \
powershell.exe -NoProfile -Command "exit 0" >/dev/null 2>&1; then
_ps_ran=1
# Inject the distro into the command: a -Command string does not
# receive trailing tokens as $args. WSL distro names are safe to
# embed (no quotes/$/backtick).
# shellcheck disable=SC2016
# $env:APPDATA is a PowerShell expansion; intentionally literal at shell level.
powershell.exe -NoProfile -Command '
# $env:APPDATA/$distro are PowerShell-side; $_wsl_distro is injected from shell.
powershell.exe -NoProfile -Command '$distro = "'"$_wsl_distro"'";
$dirs = @(
[Environment]::GetFolderPath("Desktop"),
(Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs")
);
$ws = New-Object -ComObject WScript.Shell;
foreach ($d in $dirs) {
if (-not $d) { continue }
$p = Join-Path $d "Unsloth Studio.lnk";
if (Test-Path -LiteralPath $p) { Remove-Item -LiteralPath $p -Force }
if (-not $d -or -not (Test-Path -LiteralPath $d)) { continue }
Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio*.lnk" -ErrorAction SilentlyContinue | ForEach-Object {
try {
$sc = $ws.CreateShortcut($_.FullName);
if ("$($sc.TargetPath) $($sc.Arguments)" -notmatch "wsl\.exe") { return }
# When the distro is known, require the per-distro
# name for this distro or its -d "<distro>" argument
# so launchers for other distros are not removed.
if ($distro) {
$nameMatch = ($_.Name -eq "Unsloth Studio (WSL - $distro).lnk");
$argMatch = ($sc.Arguments -match ("-d\s+`"?" + [regex]::Escape($distro) + "`"?"));
if (-not ($nameMatch -or $argMatch)) { return }
}
Remove-Item -LiteralPath $_.FullName -Force -ErrorAction SilentlyContinue
} catch { }
}
}
# WSL-fallback native shim/launcher dir (%LOCALAPPDATA%\Unsloth) + its PATH entry.
# WoA WSL-fallback (install.ps1) native shim/launcher dir
# (%LOCALAPPDATA%\Unsloth) + its PATH entry. install.ps1 created the
# shim; clean it here too so a WSL-side bash uninstall is complete.
$ud = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth" } else { $null };
if ($ud) {
$shim = (Join-Path $ud "bin").TrimEnd("\","/");
@ -272,6 +314,61 @@ case "$_os" in
if (Test-Path -LiteralPath $ud) { Remove-Item -LiteralPath $ud -Recurse -Force -ErrorAction SilentlyContinue }
}' >/dev/null 2>&1 || true
fi
# Fallback when powershell.exe can't run (interop disabled): remove the
# WSL .lnk files via drvfs. The "Unsloth Studio (WSL..." name is
# WSL-specific, so a native install's "Unsloth Studio.lnk" never matches.
if [ "$_ps_ran" = "0" ]; then
for _drive in /mnt/c /mnt/d /mnt/e; do
[ -d "$_drive/Users" ] || continue
for _udir in "$_drive"/Users/*; do
[ -d "$_udir" ] || continue
for _scdir in \
"$_udir/Desktop" \
"$_udir/OneDrive/Desktop" \
"$_udir"/OneDrive*/Desktop \
"$_udir/AppData/Roaming/Microsoft/Windows/Start Menu/Programs"; do
[ -d "$_scdir" ] || continue
if [ -n "$_wsl_distro" ]; then
# Exact per-distro name (no glob) so other distros survive.
_lnk="$_scdir/Unsloth Studio (WSL - ${_wsl_distro}).lnk"
[ -e "$_lnk" ] && rm -f "$_lnk" 2>/dev/null && echo " removed: $_lnk" || true
else
# Distro unknown: fall back to the broad WSL prefix.
for _lnk in "$_scdir"/"Unsloth Studio (WSL"*.lnk; do
[ -e "$_lnk" ] && rm -f "$_lnk" 2>/dev/null && echo " removed: $_lnk" || true
done
fi
done
done
done
fi
# ── ROCm-on-WSL config (install_rocm_wsl_strixhalo.sh) ──
# Remove Unsloth's own ROCDXG config (the env it persisted). The system
# ROCm userspace is a shared prereq (like CUDA) and is LEFT IN PLACE by
# default; set UNSLOTH_UNINSTALL_ROCM=1 to remove it too.
echo "Removing ROCm-on-WSL config..."
_sudo=""
if [ "$_uid" != "0" ] && command -v sudo >/dev/null 2>&1; then _sudo="sudo"; fi
$_sudo rm -f /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
if [ -f "$HOME/.bashrc" ] && grep -q "Unsloth ROCm-on-WSL" "$HOME/.bashrc" 2>/dev/null; then
_bk=$(mktemp 2>/dev/null || echo "$HOME/.bashrc.unsloth.tmp")
if sed '/# >>> Unsloth ROCm-on-WSL/,/# <<< Unsloth ROCm-on-WSL/d' "$HOME/.bashrc" > "$_bk" 2>/dev/null; then
cat "$_bk" > "$HOME/.bashrc" 2>/dev/null || true
echo " cleaned ROCm-on-WSL block from ~/.bashrc"
fi
rm -f "$_bk" 2>/dev/null || true
fi
if [ "${UNSLOTH_UNINSTALL_ROCM:-0}" = "1" ]; then
echo " removing system ROCm (UNSLOTH_UNINSTALL_ROCM=1)..."
$_sudo rm -f /etc/apt/sources.list.d/rocm.list /etc/apt/preferences.d/rocm-pin-600 \
/etc/apt/keyrings/rocm.gpg /etc/ld.so.conf.d/rocm.conf 2>/dev/null || true
$_sudo sh -c 'rm -rf /opt/rocm /opt/rocm-*' 2>/dev/null || true
if command -v ldconfig >/dev/null 2>&1; then $_sudo ldconfig 2>/dev/null || true; fi
elif [ -d /opt/rocm ]; then
echo " Note: ROCm userspace (/opt/rocm*) left in place (shared prereq)."
echo " Remove it by re-running with UNSLOTH_UNINSTALL_ROCM=1, or manually:"
echo " sudo rm -rf /opt/rocm /opt/rocm-* && sudo ldconfig"
fi
fi
echo "Removing Linux .desktop entry..."
_remove_path "$HOME/.local/share/applications/unsloth-studio.desktop"

302
studio/backend/cloudflare_tunnel.py Normal file
View file

@ -0,0 +1,302 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches.
The raw http://<ip>:<port> is often unreachable (https-vs-http, blocked ports,
closed security groups); a cloudflared quick tunnel gives a free
https://*.trycloudflare.com URL that works anywhere, with no account or domain.
Best-effort throughout: any failure collapses to "no URL" and Studio keeps
running. Stdlib only (back-end imports are lazy) so it is safe to import early.
"""
from __future__ import annotations
import os
import platform
import re
import shutil
import subprocess
import sys
import threading
from pathlib import Path
from typing import Optional, Tuple
# cloudflared logs the quick-tunnel URL; match only the URL so we do not depend
# on the surrounding wording, which Cloudflare may change.
_URL_RE = re.compile(r"https://[A-Za-z0-9-]+\.trycloudflare\.com")
_RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download"
_URL_TIMEOUT = 15.0 # seconds to wait for the public URL before giving up
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
if sys.platform != "win32":
return {}
flags = getattr(subprocess, "CREATE_NO_WINDOW", 0)
return {"creationflags": flags} if flags else {}
def _asset_name() -> Optional[Tuple[str, bool]]:
"""(release asset filename, is_tgz) for this OS/arch, or None if unsupported."""
system = platform.system().lower()
machine = platform.machine().lower()
is_x64 = machine in ("x86_64", "amd64", "x64")
is_arm64 = machine in ("aarch64", "arm64")
is_x86 = machine in ("i386", "i686", "x86")
if system == "linux":
if is_x64:
return ("cloudflared-linux-amd64", False)
if is_arm64:
return ("cloudflared-linux-arm64", False)
elif system == "darwin":
if is_arm64:
return ("cloudflared-darwin-arm64.tgz", True)
if is_x64:
return ("cloudflared-darwin-amd64.tgz", True)
elif system == "windows":
if is_x64:
return ("cloudflared-windows-amd64.exe", False)
if is_x86:
return ("cloudflared-windows-386.exe", False)
return None
def _cache_path() -> Optional[Path]:
"""studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable."""
try:
from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import
except Exception:
return None
name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared"
return studio_bin_root() / name
def find_cloudflared() -> Optional[str]:
"""Locate an existing cloudflared: PATH first, then the Studio bin cache."""
on_path = shutil.which("cloudflared")
if on_path:
return on_path
cached = _cache_path()
if cached is not None and cached.is_file() and os.access(cached, os.X_OK):
return str(cached)
return None
def _download(url: str, dest: Path) -> bool:
"""Download url to dest via urllib (temp file + atomic rename). Best-effort -> bool."""
import tempfile
import urllib.request
tmp_path: Optional[Path] = None
try:
dest.parent.mkdir(parents = True, exist_ok = True)
with tempfile.NamedTemporaryFile(
prefix = dest.name + ".tmp-", dir = dest.parent, delete = False
) as handle:
tmp_path = Path(handle.name)
# GitHub's CDN 403s the default Python-urllib User-Agent.
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _DOWNLOAD_TIMEOUT) as response:
shutil.copyfileobj(response, handle)
if tmp_path.stat().st_size == 0:
raise RuntimeError("empty download")
os.replace(tmp_path, dest)
return True
except Exception:
if tmp_path is not None:
try:
tmp_path.unlink(missing_ok = True)
except Exception:
pass
return False
def _extract_tgz_member(tgz_path: Path, dest: Path) -> bool:
"""Extract just the `cloudflared` member from a darwin .tgz to dest.
Rejects absolute paths and `..` traversal so a hostile archive cannot write
outside dest. Best-effort -> bool.
"""
import tarfile
try:
with tarfile.open(tgz_path, "r:gz") as tar:
member = None
for m in tar.getmembers():
if not m.isfile() or os.path.basename(m.name) != "cloudflared":
continue
if m.name.startswith("/") or ".." in Path(m.name).parts:
continue
member = m
break
if member is None:
return False
src = tar.extractfile(member)
if src is None:
return False
with src, open(dest, "wb") as out:
shutil.copyfileobj(src, out)
return True
except Exception:
return False
def ensure_cloudflared() -> Optional[str]:
"""Return a cloudflared path, downloading + caching the binary once if missing."""
existing = find_cloudflared()
if existing:
return existing
asset = _asset_name()
cached = _cache_path()
if asset is None or cached is None:
return None
name, is_tgz = asset
url = f"{_RELEASE_BASE}/{name}"
try:
cached.parent.mkdir(parents = True, exist_ok = True)
if is_tgz:
tgz = cached.with_suffix(".tgz")
if not _download(url, tgz) or not _extract_tgz_member(tgz, cached):
tgz.unlink(missing_ok = True)
return None
tgz.unlink(missing_ok = True)
elif not _download(url, cached):
return None
if sys.platform != "win32":
os.chmod(cached, 0o755)
return str(cached)
except Exception:
return None
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:<port>. Best-effort throughout.
Use localhost (not the wildcard bind) as the tunnel origin so cloudflared's
upstream stays local-only.
"""
def __init__(self, port: int, binary: str):
self.port = port
self.binary = binary
self._proc: Optional[subprocess.Popen] = None
self._lock = threading.Lock()
self._url_event = threading.Event()
self.url: Optional[str] = None
self.error: Optional[str] = None
def start(self) -> None:
cmd = [
self.binary,
"tunnel",
"--url",
f"http://localhost:{self.port}",
"--no-autoupdate",
]
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
stdin = subprocess.DEVNULL,
text = True,
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
)
with self._lock:
self._proc = proc
threading.Thread(
target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True
).start()
def _reader(self, proc: subprocess.Popen) -> None:
# Drain cloudflared's output, capture the first trycloudflare URL, and
# keep draining so it never blocks on a full pipe.
try:
if proc.stdout is not None:
for line in proc.stdout:
if self.url is None:
match = _URL_RE.search(line)
if match:
self.url = match.group(0)
self._url_event.set()
except Exception:
pass
finally:
if self.url is None:
self.error = "cloudflared exited before emitting a tunnel URL"
self._url_event.set()
def wait_for_url(self, timeout: float = _URL_TIMEOUT) -> Optional[str]:
self._url_event.wait(timeout)
return self.url
def stop(self) -> None:
"""Terminate the tunnel. Idempotent and safe to call from a signal handler."""
with self._lock:
proc, self._proc = self._proc, None
if proc is None:
return
try:
if proc.poll() is None:
proc.terminate()
try:
proc.wait(timeout = 5)
except subprocess.TimeoutExpired:
proc.kill()
try:
proc.wait(timeout = 5)
except Exception:
pass
except Exception:
pass
# Single serving process per Studio launch, so one module-level tunnel handle is
# enough; the lock guards the start/stop/shutdown races.
_active_tunnel: Optional[CloudflareTunnel] = None
_active_lock = threading.Lock()
def start_studio_tunnel(port: int, timeout: float = _URL_TIMEOUT) -> Optional[str]:
"""Start a quick tunnel and return its public URL, or None (best-effort).
On any failure (no binary, no URL within timeout, early crash) the tunnel is
stopped and None is returned, so the caller prints a hint and continues.
"""
global _active_tunnel
binary = ensure_cloudflared()
if not binary:
return None
tunnel = CloudflareTunnel(port, binary)
# Register before start/wait so a shutdown during the URL wait can stop it.
with _active_lock:
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
try:
tunnel.start()
url = tunnel.wait_for_url(timeout)
except Exception:
url = None
if url:
return url
# No URL (or crash): drop it unless a concurrent shutdown already replaced it.
with _active_lock:
if _active_tunnel is tunnel:
_active_tunnel = None
tunnel.stop()
return None
def stop_studio_tunnel() -> None:
"""Terminate the active tunnel, if any. Idempotent."""
global _active_tunnel
with _active_lock:
tunnel, _active_tunnel = _active_tunnel, None
if tunnel is not None:
tunnel.stop()

21
studio/backend/core/inference/external_provider.py
View file

@ -469,9 +469,24 @@ def _apply_mistral_reasoning_controls(
# Shared client reused across all requests for HTTP connection pooling.
# Auth headers and timeouts are per-request, so one client handles every
# provider without storing credentials.
_http_client = httpx.AsyncClient()
# Auth headers and timeouts are passed per-request, so a single client
# handles every provider without storing credentials.
def _create_shared_http_client() -> httpx.AsyncClient:
# Unsupported env proxy schemes (socks:// etc) raise at construction and
# would crash Studio startup (#6090); retry ignoring env proxies instead.
try:
return httpx.AsyncClient()
except (ImportError, ValueError) as exc:
exc_str = str(exc)
if "Unknown scheme for proxy URL" not in exc_str and "socksio" not in exc_str:
raise
logger.warning(
"Ignoring unsupported environment proxy for the shared HTTP client: %s", exc_str
)
return httpx.AsyncClient(trust_env = False)
_http_client = _create_shared_http_client()
# Cap per-image fetch well below Gemini's ~20 MB total request budget.

9
studio/backend/core/inference/inference.py
View file

@ -25,6 +25,7 @@ from utils.hardware import (
get_visible_gpu_count,
)
from core.inference.audio_codecs import AudioCodecManager
from core.inference.runtime_context import runtime_context_length
from io import StringIO
import structlog
from loggers import get_logger
@ -405,6 +406,10 @@ class InferenceBackend:
# Reject CPU/disk offload for audio models too
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
self.models[model_name]["context_length"] = runtime_context_length(
self.models[model_name].get("model"),
max_seq_length,
)
self.active_model_name = model_name
self.loading_models.discard(model_name)
@ -485,6 +490,10 @@ class InferenceBackend:
self.models[model_name]["tokenizer"] = tokenizer
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
self.models[model_name]["context_length"] = runtime_context_length(
self.models[model_name].get("model"),
max_seq_length,
)
self._load_chat_template_info(model_name)

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

File diff suppressed because it is too large Load diff

8
studio/backend/core/inference/llama_server_args.py
View file

@ -132,7 +132,13 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
"--spec-ngram-size",
"--draft-min",
"--draft-max",
# MTP path (llama.cpp #22673).
# MTP path (llama.cpp #22673). --model-draft and aliases are
# Studio-managed since the separate-drafter support (Gemma 4): an
# inherited copy must not last-wins-override the auto-detected
# drafter. Explicit extras for the current load are never stripped.
"--model-draft",
"-md",
"--spec-draft-model",
"--spec-draft-n-max",
"--spec-draft-n-min",
"--spec-draft-p-min",

81
studio/backend/core/inference/mcp_client.py
View file

@ -25,16 +25,85 @@ def is_stdio(address: str) -> bool:
return not address.strip().lower().startswith(("http://", "https://"))
def _split_windows_command_line(address: str) -> list[str]:
"""Parse a Windows command line using the same backslash/quote rules that
subprocess.list2cmdline() writes. This keeps trailing backslashes before a
closing quote from being doubled in the resulting argv."""
parts: list[str] = []
current: list[str] = []
in_quotes = False
backslashes = 0
arg_started = False
i = 0
while i < len(address):
ch = address[i]
if ch == "\\":
backslashes += 1
i += 1
continue
if ch == '"':
current.extend("\\" * (backslashes // 2))
if backslashes % 2:
current.append('"')
else:
in_quotes = not in_quotes
arg_started = True
backslashes = 0
i += 1
continue
if ch.isspace() and not in_quotes:
if backslashes:
current.extend("\\" * backslashes)
arg_started = True
backslashes = 0
if arg_started or current:
parts.append("".join(current))
current = []
arg_started = False
i += 1
while i < len(address) and address[i].isspace():
i += 1
continue
if backslashes:
current.extend("\\" * backslashes)
arg_started = True
backslashes = 0
current.append(ch)
arg_started = True
i += 1
if backslashes:
current.extend("\\" * backslashes)
arg_started = True
if in_quotes:
raise ValueError("No closing quotation")
if arg_started or current:
parts.append("".join(current))
return parts
def parse_stdio_command(address: str) -> list[str]:
"""Split a stdio command line into argv. Shared by route validation and the
transport so both agree on quoting (notably Windows backslash paths)."""
posix = sys.platform != "win32"
parts = shlex.split(address, posix = posix)
if not posix:
# posix=False keeps backslash paths but also keeps surrounding quotes;
# strip a matched pair so argv reaches the subprocess clean.
parts = [p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p for p in parts]
return parts
if posix:
return shlex.split(address, posix = posix)
if address.lstrip().startswith("'"):
raise ValueError("Single-quoted executables are not supported on Windows")
return _split_windows_command_line(address)
def join_stdio_command(parts: list[str]) -> str:
"""Inverse of parse_stdio_command: join argv into a single command string
that parse_stdio_command() splits back into ``parts`` on this platform.
Config files (issue #5936) carry structured command + args; storage holds
one string in the url field. Windows uses list2cmdline so spaced/backslash
paths round-trip through the posix=False quote-strip; posix uses shlex."""
if sys.platform == "win32":
import subprocess
return subprocess.list2cmdline(parts)
return shlex.join(parts)
def stdio_mcp_enabled() -> bool:

169
studio/backend/core/inference/mcp_config_import.py Normal file
View file

@ -0,0 +1,169 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Parse a standard ``mcpServers`` JSON config (Claude Desktop / Cursor / Cline
/ VS Code) into entries the existing MCP storage understands. See issue #5936.
A stdio entry (``command`` + ``args`` + ``env``) is joined into the single
command string the ``url`` field already stores; a remote entry (``url`` +
``headers``) maps straight through. Parsing never raises on a single bad entry:
it returns ``(entries, errors)`` so one malformed server can't sink the import.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from core.inference.mcp_client import join_stdio_command
_SCALAR = (str, int, float, bool)
_UNSUPPORTED_STDIO_FIELDS = ("cwd", "envFile")
_UNSUPPORTED_TIMEOUT_FIELDS = ("timeout", "timeoutMs", "timeoutSeconds")
_HTTP_REMOTE_TYPES = ("http", "streamableHttp")
@dataclass
class ParsedMcpEntry:
display_name: str
url: str # joined command (stdio) or http(s) url (remote)
headers: Optional[dict[str, str]] # env vars (stdio) or http headers (remote)
is_stdio: bool
is_enabled: bool = True
use_oauth: bool = False
def _coerce_str_dict(value: dict) -> dict[str, str]:
return {str(k): str(v) for k, v in value.items()}
def _has_variable_reference(value: object) -> bool:
if isinstance(value, str):
return "${" in value
if isinstance(value, list):
return any(_has_variable_reference(item) for item in value)
if isinstance(value, dict):
return any(_has_variable_reference(item) for item in value.values())
return False
def _has_null_value(value: object) -> bool:
return isinstance(value, dict) and any(item is None for item in value.values())
def _enabled_from_spec(label: str, spec: dict) -> tuple[Optional[bool], Optional[str]]:
disabled = spec.get("disabled")
if disabled is None:
return True, None
if not isinstance(disabled, bool):
return None, f"{label}: 'disabled' must be true or false."
return not disabled, None
def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Optional[str]]:
label = str(name).strip()
if not label:
return None, "Server entry has an empty name."
if not isinstance(spec, dict):
return None, f"{label}: entry must be an object."
if _has_variable_reference(spec):
return None, f"{label}: VS Code variable references are not supported by import."
is_enabled, error = _enabled_from_spec(label, spec)
if error:
return None, error
has_command = bool(spec.get("command"))
has_url = bool(spec.get("url"))
if has_command and has_url:
return None, f"{label}: entry has both 'command' and 'url'; use one."
if not has_command and not has_url:
return None, f"{label}: entry needs a 'command' (stdio) or 'url' (remote)."
if has_command:
command = spec["command"]
if not isinstance(command, str):
return None, f"{label}: 'command' must be a string."
entry_type = spec.get("type")
if entry_type is not None and entry_type != "stdio":
return None, f"{label}: stdio entry has unsupported type {entry_type!r}."
sandbox_enabled = spec.get("sandboxEnabled")
if sandbox_enabled is not None and not isinstance(sandbox_enabled, bool):
return None, f"{label}: 'sandboxEnabled' must be true or false."
if sandbox_enabled:
return None, f"{label}: sandboxed stdio servers cannot be preserved by import."
unsupported = [field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None]
if unsupported:
return None, f"{label}: import cannot preserve {', '.join(unsupported)}."
if spec.get("oauth") is not None:
return None, f"{label}: 'oauth' is only supported for remote servers."
args = spec.get("args") or []
if not isinstance(args, list) or not all(isinstance(a, _SCALAR) for a in args):
return None, f"{label}: 'args' must be a list of strings."
env = spec.get("env")
if env is not None and not isinstance(env, dict):
return None, f"{label}: 'env' must be an object."
if _has_null_value(env):
return None, f"{label}: null environment values are not supported by import."
url = join_stdio_command([command, *(str(a) for a in args)])
headers = _coerce_str_dict(env) if env else None
return ParsedMcpEntry(label, url, headers, True, is_enabled = is_enabled), None
url = spec["url"]
if not isinstance(url, str):
return None, f"{label}: 'url' must be a string."
url = url.strip()
entry_type = spec.get("type")
if entry_type is not None and entry_type not in (*_HTTP_REMOTE_TYPES, "sse"):
return None, f"{label}: remote entry has unsupported type {entry_type!r}."
unsupported_timeout = [
field for field in _UNSUPPORTED_TIMEOUT_FIELDS if spec.get(field) is not None
]
if unsupported_timeout:
return None, f"{label}: import cannot preserve {', '.join(unsupported_timeout)}."
url_infers_sse = url.rstrip("/").endswith("/sse")
if entry_type == "sse" and not url_infers_sse:
return None, f"{label}: explicit SSE transport cannot be preserved for this URL."
if entry_type in _HTTP_REMOTE_TYPES and url_infers_sse:
return None, f"{label}: explicit HTTP transport cannot be preserved for this URL."
oauth_raw = spec.get("oauth")
if oauth_raw is not None and not isinstance(oauth_raw, dict):
return None, f"{label}: 'oauth' must be an object."
headers_raw = spec.get("headers")
if headers_raw is not None and not isinstance(headers_raw, dict):
return None, f"{label}: 'headers' must be an object."
if _has_null_value(headers_raw):
return None, f"{label}: null header values are not supported by import."
headers = _coerce_str_dict(headers_raw) if headers_raw else None
return ParsedMcpEntry(
label,
url,
headers,
False,
is_enabled = is_enabled,
use_oauth = oauth_raw is not None,
), None
def parse_mcp_config(config: object) -> tuple[list[ParsedMcpEntry], list[str]]:
"""Parse a Claude-Desktop/Cursor/Cline/VS Code config. Accepts the
``mcpServers`` key (primary) or ``servers`` (VS Code alias). Returns
``(entries, errors)``; a bad entry adds an error rather than raising."""
if not isinstance(config, dict):
return [], ["Config must be a JSON object."]
servers_key = "mcpServers" if "mcpServers" in config else "servers"
servers = config.get(servers_key)
if servers is None:
return [], ["Config has no 'mcpServers' (or 'servers') object."]
if not isinstance(servers, dict):
return [], [f"'{servers_key}' must be an object mapping name -> server."]
entries: list[ParsedMcpEntry] = []
errors: list[str] = []
for name, spec in servers.items():
entry, error = _parse_entry(name, spec)
if error:
errors.append(error)
elif entry:
entries.append(entry)
return entries, errors

2
studio/backend/core/inference/mlx_inference.py
View file

@ -7,6 +7,7 @@ instead of torch/transformers for model loading and generation.
import threading
from typing import Optional, Generator
from core.inference.runtime_context import runtime_context_length
from loggers import get_logger
logger = get_logger(__name__)
@ -175,6 +176,7 @@ class MLXInferenceBackend:
"is_audio": False,
"audio_type": None,
"has_audio_input": False,
"context_length": runtime_context_length(self._model, max_seq_length),
}
# Capture chat_template_info so the worker IPC reply ships it back and
# the route layer classifies capabilities like the other paths.

49
studio/backend/core/inference/orchestrator.py
View file

@ -17,6 +17,7 @@ Pattern follows core/training/training.py.
import atexit
import base64
import os
import signal
import structlog
from loggers import get_logger
import multiprocessing as mp
@ -239,6 +240,43 @@ class InferenceOrchestrator:
"""True if the subprocess is alive."""
return self._proc is not None and self._proc.is_alive()
def _subprocess_crash_message(self, context: str) -> str:
"""Return a user-facing crash message with the worker exit status."""
context_label = {
"wait": "loading the model",
"generation": "generating a response",
"audio generation": "generating audio",
"audio input generation": "processing audio input",
}.get(context, context)
message = f"The inference worker stopped unexpectedly while {context_label}."
if self._proc is None:
return f"{message} Details: process missing."
exitcode = self._proc.exitcode
pid = self._proc.pid
if exitcode is None:
return f"{message} Details: pid={pid}."
if exitcode < 0:
signum = -exitcode
try:
sig_name = signal.Signals(signum).name
except ValueError:
sig_name = f"SIG{signum}"
suffix = ""
if sig_name == "SIGKILL":
suffix = (
" This usually means the system killed it under memory pressure. "
"Try a smaller model, lower context length, or close other GPU-heavy apps."
)
return (
f"{message}{suffix} " f"Details: pid={pid}, signal={sig_name}, exitcode={exitcode}."
)
return f"{message} Details: pid={pid}, exitcode={exitcode}."
# ------------------------------------------------------------------
# Queue helpers
# ------------------------------------------------------------------
@ -286,7 +324,7 @@ class InferenceOrchestrator:
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess crashed during wait")
raise RuntimeError(self._subprocess_crash_message("wait"))
continue
rtype = resp.get("type", "")
@ -510,7 +548,7 @@ class InferenceOrchestrator:
except queue.Empty:
# Timeout — check subprocess health
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess crashed during generation"
yield f"Error: {self._subprocess_crash_message('generation')}"
return
continue
@ -689,6 +727,7 @@ class InferenceOrchestrator:
"is_audio": model_info.get("is_audio", False),
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
"context_length": model_info.get("context_length"),
}
# Mirror chat_template_info so routes can classify caps
# without re-entering the subprocess.
@ -1028,7 +1067,7 @@ class InferenceOrchestrator:
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess crashed during generation"
yield f"Error: {self._subprocess_crash_message('generation')}"
return
continue
@ -1125,7 +1164,7 @@ class InferenceOrchestrator:
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess crashed during audio generation")
raise RuntimeError(self._subprocess_crash_message("audio generation"))
continue
rtype = resp.get("type", "")
@ -1247,7 +1286,7 @@ class InferenceOrchestrator:
if resp is None:
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess crashed during audio input generation"
yield ("Error: " + self._subprocess_crash_message("audio input generation"))
return
continue

22
studio/backend/core/inference/runtime_context.py Normal file
View file

@ -0,0 +1,22 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Runtime context length helpers shared by inference backends."""
from __future__ import annotations
from typing import Any, Optional
def runtime_context_length(model: Any, fallback: Optional[int] = None) -> Optional[int]:
"""Return the effective context length Unsloth attached to a loaded model."""
for value in (getattr(model, "max_seq_length", None), fallback):
if isinstance(value, bool):
continue
try:
value_int = int(value)
except (TypeError, ValueError):
continue
if value_int > 0:
return value_int
return None

13
studio/backend/core/inference/worker.py
View file

@ -315,6 +315,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
_context_length = _entry.get("context_length")
if _context_length is not None:
model_info["context_length"] = int(_context_length)
except Exception as _ctx_exc:
logger.warning("context_length forward failed: %s", _ctx_exc)
# Forward chat_template_info so the parent can classify capabilities.
try:
_bm = getattr(backend, "models", {}) or {}
@ -881,6 +893,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
name: {
"is_vision": info.get("is_vision", False),
"is_lora": info.get("is_lora", False),
"context_length": info.get("context_length"),
}
for name, info in backend.models.items()
},

3
studio/backend/core/training/trainer.py
View file

@ -58,7 +58,8 @@ from pathlib import Path
from typing import Optional, Callable
from dataclasses import dataclass
import pandas as pd
from datasets import Dataset, load_dataset
from datasets import Dataset
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
from core.inference.llama_cpp import _hf_offline_if_dns_dead
from utils.models import is_vision_model, detect_audio_type

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

@ -534,8 +534,19 @@ class TrainingBackend:
except (TypeError, ValueError):
logger.debug("Could not convert loss to float: %s", _raw_loss)
_safe_loss = None
if _safe_loss is not None and not math.isfinite(_safe_loss):
_loss_is_nonfinite = _safe_loss is not None and not math.isfinite(_safe_loss)
if _loss_is_nonfinite:
# Drop the value rather than laundering it back to the last
# finite loss; clients see loss=None at this step so the NaN
# is not hidden behind a stale value. Training continues.
_safe_loss = None
if not getattr(self._progress, "_nonfinite_loss_warned", False):
self._progress._nonfinite_loss_warned = True
logger.warning(
"Training produced non-finite loss at step %s; "
"loss field will report null until it recovers.",
event.get("step", "?"),
)
try:
_safe_lr = float(_raw_lr) if _raw_lr is not None else None
except (TypeError, ValueError):
@ -545,6 +556,10 @@ class TrainingBackend:
_safe_lr = None
if _safe_loss is not None:
self._progress.loss = _safe_loss
elif _loss_is_nonfinite:
# Clear stale finite loss so the API doesn't keep
# reporting the last good value while NaN is happening.
self._progress.loss = None
if _safe_lr is not None:
self._progress.learning_rate = _safe_lr
self._progress.total_steps = event.get("total_steps", self._progress.total_steps)

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

@ -26,6 +26,22 @@ import subprocess as _sp
from pathlib import Path
from typing import Any, Callable
# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ──────
# Mirrors main.py. In WSL the AMD GPU is reached via the ROCDXG bridge
# (librocdxg.so over /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_
# DETECTION=1 is set before torch touches the GPU. A worker spawned outside a
# login shell misses the installer's persisted env and falls back to CPU.
# Gated to no-op unless BOTH /dev/dxg and librocdxg.so exist, so native Linux
# ROCm, NVIDIA, macOS and Windows are unaffected.
if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ:
try:
if os.path.exists("/dev/dxg") and any(
os.path.exists(_p + "/librocdxg.so") for _p in ("/opt/rocm/lib", "/opt/rocm/lib64")
):
os.environ["HSA_ENABLE_DXG_DETECTION"] = "1"
except Exception:
pass
logger = get_logger(__name__)
from utils.hardware import apply_gpu_ids
from utils.wheel_utils import (
@ -678,8 +694,11 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
``set_per_process_memory_fraction`` cap to leave OS headroom.
Classification priority:
1. ``gcnArchName`` / variant spellings (stable, naming-independent).
2. Device-name substring match (last resort when all arch attrs absent;
1. ``props.is_integrated`` truthy (hipDeviceProp_t.integrated -- the
driver's own unified-memory answer; covers APUs beyond the hardcoded
arch set, e.g. gfx1103 Phoenix iGPUs). Only ever upgrades to unified.
2. ``gcnArchName`` / variant spellings (stable, naming-independent).
3. Device-name substring match (last resort when all arch attrs absent;
AMD SDK / Radeon wheels may not populate them):
- gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
- gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
@ -692,6 +711,16 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
gcn_arch = _v
break
# Driver's own answer first: hipDeviceProp_t.integrated (exposed as
# props.is_integrated; same gate PR #5988's UMA safetensors fast-load
# uses). Strictly additive -- only a truthy value upgrades to unified;
# 0/absent falls through to the arch/name logic below, so a wheel that
# omits or zeroes the field can never downgrade the known APU set. This
# covers unified APUs outside the hardcoded arches (gfx1103 Phoenix
# iGPUs, future parts) with one universal signal.
if getattr(props, "is_integrated", 0):
return gcn_arch, True
if gcn_arch:
return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"}
@ -1069,7 +1098,74 @@ def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int
return new_w, new_h
def _resize_mlx_vlm_image(image, resize):
_MLX_VLM_RESIZED_IMAGE_LAYOUT_CACHE = {}
def _mlx_vlm_resized_image_layout(processor = None) -> str | None:
"""Return the numpy image layout expected after Studio-side VLM resizing."""
image_processor = getattr(processor, "image_processor", None)
if image_processor is None:
return None
cls = image_processor.__class__
key = (getattr(cls, "__module__", ""), getattr(cls, "__qualname__", cls.__name__))
if key in _MLX_VLM_RESIZED_IMAGE_LAYOUT_CACHE:
return _MLX_VLM_RESIZED_IMAGE_LAYOUT_CACHE[key]
copied_image_processor = _copy_mlx_vlm_image_processor(image_processor)
layout = (
_probe_mlx_vlm_numpy_image_layout(copied_image_processor)
if copied_image_processor is not None
else None
)
_MLX_VLM_RESIZED_IMAGE_LAYOUT_CACHE[key] = layout
return layout
def _copy_mlx_vlm_image_processor(image_processor):
import copy
try:
return copy.deepcopy(image_processor)
except Exception:
try:
return copy.copy(image_processor)
except Exception:
return None
def _probe_mlx_vlm_numpy_image_layout(image_processor) -> str | None:
try:
import numpy as np
except ImportError:
return None
def _accepts(candidate) -> bool:
try:
image_processor(images = [candidate])
return True
except TypeError:
try:
image_processor([candidate])
return True
except Exception:
return False
except Exception:
return False
# Use an asymmetric image so CHW-vs-HWC mistakes are visible to processors
# that skip conversion for 3D numpy arrays.
hwc = np.zeros((64, 96, 3), dtype = np.uint8)
chw = np.ascontiguousarray(hwc.transpose(2, 0, 1))
if _accepts(hwc):
return None
if _accepts(chw):
return "chw"
return None
def _resize_mlx_vlm_image(
image,
resize,
image_layout = None,
):
if resize is None:
return image
try:
@ -1087,16 +1183,27 @@ def _resize_mlx_vlm_image(image, resize):
# On resize, hand mlx-vlm a writable RGB ndarray so its PIL-path
# square-resize is skipped and HF processors don't warn on non-writable
# views. resize=None above keeps the original PIL.
return np.array(image, copy = True)
array = np.array(image, copy = True)
if image_layout == "chw":
return np.ascontiguousarray(array.transpose(2, 0, 1))
return array
def _resize_mlx_vlm_images(value, resize):
def _resize_mlx_vlm_images(
value,
resize,
image_layout = None,
):
if isinstance(value, list):
return [_resize_mlx_vlm_image(image, resize) for image in value]
return _resize_mlx_vlm_image(value, resize)
return [_resize_mlx_vlm_image(image, resize, image_layout = image_layout) for image in value]
return _resize_mlx_vlm_image(value, resize, image_layout = image_layout)
def _adapt_for_mlx_vlm(items, resize = None):
def _adapt_for_mlx_vlm(
items,
resize = None,
image_layout = None,
):
"""Adapt GPU-path VLM dataset output for mlx-vlm.
The GPU path embeds PIL images in message content as
@ -1116,7 +1223,13 @@ def _adapt_for_mlx_vlm(items, resize = None):
if isinstance(part, dict) and part.get("type") == "image":
img = part.get("image")
if img is not None:
images.append(_resize_mlx_vlm_image(img, resize))
images.append(
_resize_mlx_vlm_image(
img,
resize,
image_layout = image_layout,
)
)
new_content.append({"type": "image"})
else:
new_content.append(part)
@ -1127,9 +1240,17 @@ def _adapt_for_mlx_vlm(items, resize = None):
if images:
out["image"] = images[0] if len(images) == 1 else images
elif "image" in item:
out["image"] = _resize_mlx_vlm_images(item["image"], resize)
out["image"] = _resize_mlx_vlm_images(
item["image"],
resize,
image_layout = image_layout,
)
elif "images" in item:
out["images"] = _resize_mlx_vlm_images(item["images"], resize)
out["images"] = _resize_mlx_vlm_images(
item["images"],
resize,
image_layout = image_layout,
)
adapted.append(out)
return adapted
@ -1256,7 +1377,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
"(unsloth_zoo.mlx.loader / unsloth_zoo.mlx.trainer). Reinstall via "
"install.sh on Apple Silicon."
) from e
from datasets import load_dataset
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
if mx.metal.is_available():
info = mx.device_info()
@ -1276,6 +1397,14 @@ def _run_mlx_training(event_queue, stop_queue, config):
message = "LoftQ is not supported for MLX training yet."
_send("error", error = message)
raise NotImplementedError(message)
if config.get("is_embedding"):
message = "Embedding model training is not supported for MLX training yet."
_send("error", error = message)
raise NotImplementedError(message)
if config.get("training_type") == "Continued Pretraining":
message = "Continued Pretraining is not supported for MLX training yet."
_send("error", error = message)
raise NotImplementedError(message)
optim_name = _normalize_mlx_studio_optimizer(config.get("optim", "adamw_8bit"))
lr_scheduler_type = _normalize_mlx_studio_scheduler(config.get("lr_scheduler_type", "linear"))
@ -1284,6 +1413,12 @@ def _run_mlx_training(event_queue, stop_queue, config):
# Force text-only for non-image datasets even on vision-capable models
# (e.g. Qwen3.5-VL trained on plain alpaca text).
_send("status", status_message = f"Loading {model_name}...")
# Pull through resume_from_checkpoint so MLXTrainer.train() can restore
# optimizer + step state and continue cleanly. Was previously dropped on
# the floor for the MLX path, so the Resume UI button silently restarted
# from step 0 (the CUDA path at lines 2729 / 3108 has been forwarding
# this all along).
resume_from_checkpoint = config.get("resume_from_checkpoint") or None
is_dataset_image = bool(config.get("is_dataset_image", False))
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
@ -1317,7 +1452,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
"status",
status_message = f"MLX vision image resize: {vision_image_size} (max dimension)",
)
# ── 2. Apply LoRA / full FT ──
# gradient_checkpointing stays a string ("mlx"/"unsloth"/"none"/etc.);
# get_peft_model and MLXTrainer both accept and handle strings.
@ -1426,6 +1560,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
# Reuse the GPU format pipeline for VLM (auto-detects OCR/caption/llava/
# sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column).
format_type = config.get("format_type", "")
custom_format_mapping = config.get("custom_format_mapping")
try:
from utils.datasets import format_and_template_dataset
def _fmt_progress(status_message = "", **_kw):
@ -1439,12 +1574,19 @@ def _run_mlx_training(event_queue, stop_queue, config):
tokenizer = tokenizer,
is_vlm = True,
dataset_name = hf_dataset or "local",
custom_format_mapping = custom_format_mapping,
progress_callback = _fmt_progress,
)
if vlm_info.get("success"):
vision_image_layout = (
_mlx_vlm_resized_image_layout(tokenizer)
if vision_image_size is not None
else None
)
dataset = _adapt_for_mlx_vlm(
vlm_info["dataset"],
resize = vision_image_size,
image_layout = vision_image_layout,
)
else:
errors = vlm_info.get("errors", [])
@ -1456,11 +1598,18 @@ def _run_mlx_training(event_queue, stop_queue, config):
tokenizer = tokenizer,
is_vlm = True,
dataset_name = hf_dataset or "local",
custom_format_mapping = custom_format_mapping,
)
if ev_info.get("success"):
vision_image_layout = (
_mlx_vlm_resized_image_layout(tokenizer)
if vision_image_size is not None
else None
)
eval_dataset = _adapt_for_mlx_vlm(
ev_info["dataset"],
resize = vision_image_size,
image_layout = vision_image_layout,
)
elif format_type:
@ -1472,6 +1621,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
is_vlm = False,
format_type = format_type,
dataset_name = hf_dataset or "local",
custom_format_mapping = custom_format_mapping,
progress_callback = _fmt_progress,
)
if info.get("success", True):
dataset = info.get("dataset", dataset)
@ -1483,6 +1634,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
is_vlm = False,
format_type = format_type,
dataset_name = hf_dataset or "local",
custom_format_mapping = custom_format_mapping,
)
if ev.get("success", True):
eval_dataset = ev.get("dataset", eval_dataset)
@ -1730,7 +1882,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
# ── 11. Run training ──
gc.collect()
mx.synchronize()
trainer.train()
trainer.train(resume_from_checkpoint = resume_from_checkpoint)
# ── 12. Save and finalize ──
if trainer.stop_requested and not _stop_save[0]:
@ -1829,7 +1981,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
}
)
return
# Activate correct transformers version (Gemma-4 needs 5.5.0, etc.)
# Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.)
# before any transformers/mlx-lm imports in _run_mlx_training.
try:
_activate_transformers_version(model_name)
@ -2006,12 +2158,29 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.info("Windows ROCm: torch.compile (dynamo) disabled")
# bitsandbytes' import-time get_rocm_gpu_arch() probe runs
# `hipinfo.exe` from PATH; the AMD torch wheel ships it in the venv
# Scripts dir, which is on PATH only for activated venvs. Prepend
# it so the probe succeeds instead of logging a scary (harmless)
# "Could not detect ROCm GPU architecture" ERROR on every import.
# Normally inherited from main.py's env, but workers can also be
# spawned standalone (tests, CLI) -- keep the guard here too.
_scripts_dir = os.path.dirname(sys.executable)
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")):
import shutil as _shutil
if not _shutil.which("hipinfo.exe"):
os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "")
# BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB
# wheel may ship a DLL whose suffix doesn't match. Detect the actual
# DLL name and override; "72" is a safe fallback. Callers may
# pre-set the var to override.
if "BNB_ROCM_VERSION" not in os.environ:
# DLL name and override. Values seeded by the installer are
# redetectable defaults, while caller overrides remain authoritative.
if (
"BNB_ROCM_VERSION" not in os.environ
or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"
):
_bnb_rocm_ver = None
_found_rocm_bnb = False
try:
import glob as _glob
import importlib.util as _ilu
@ -2024,6 +2193,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
_found_rocm_bnb = True
_m = _re.search(
r"libbitsandbytes_rocm(\d+)\.dll",
os.path.basename(_dll),
@ -2035,14 +2205,20 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
_bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
except Exception:
pass
_bnb_rocm_ver = _bnb_rocm_ver or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
logger.info(
"Windows ROCm: set BNB_ROCM_VERSION=%s "
"(detected from installed BNB wheel; "
"overrides torch.version.hip auto-detection)",
_bnb_rocm_ver,
)
# Only when a ROCm bnb DLL actually exists (mirrors main.py):
# without one the seeded value and its marker stay untouched,
# so later import fixes can still redetect or opt out. DLL
# with unparsable name -> seeded value or "72".
if _found_rocm_bnb:
_bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
logger.info(
"Windows ROCm: set BNB_ROCM_VERSION=%s "
"(detected from installed BNB wheel; "
"overrides torch.version.hip auto-detection)",
_bnb_rocm_ver,
)
# Parse HIP version for the kernel-fix gate below, falling back to
# the rocm version embedded in torch.__version__ when version.hip is
@ -2202,10 +2378,25 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
if _is_unified and not _gcn_arch:
logger.debug(
"ROCm OOM guard: gcnArchName absent -- inferred "
"unified memory from device name %r; applying 0.80 cap",
"unified memory from device name %r; applying unified cap",
_dev_name,
)
_mem_fraction = 0.80 if _is_unified else 0.90
# Unified hosts on native Windows: mem_get_info's total is the
# WDDM budget the driver grants HIP (BIOS carve + ~half of the
# remaining RAM) -- the OS share is already outside it, so the
# Linux 0.80 starve-protection double-taxes (48.49 GiB budget →
# 38.79 allowed) and blocks loads that fit in free memory.
# 1.0 removes the double-tax. Current AMD Windows wheels only
# enforce sub-1.0 fractions (measured on gfx1151: 0.5 caps,
# 1.0 still allocates past the budget via WDDM overcommit), so
# 1.0 behaves like torch's uncapped default, with WDDM
# arbitrating residency; on wheels that do enforce it, it caps
# at exactly the driver-granted budget. On Linux the total
# spans nearly all RAM, so keep the 0.80 OS headroom there.
if _is_unified:
_mem_fraction = 1.0 if sys.platform == "win32" else 0.80
else:
_mem_fraction = 0.90
_torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction)
logger.info(
"ROCm OOM guard: set_per_process_memory_fraction(%.2f) — "
@ -2215,6 +2406,28 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
_dev_name,
_gcn_arch or "unknown arch",
)
# Unified Windows APUs: the WDDM budget is user-raisable, but
# nothing on the box says so -- users see "48 GB VRAM" on a
# 96 GB machine and assume a Studio bug. Say where the limit
# comes from and how to raise it.
if _is_unified and sys.platform == "win32":
try:
import psutil as _psutil
_phys = _psutil.virtual_memory().total
_granted = _torch_mem.cuda.mem_get_info(0)[1]
if _granted < 0.75 * _phys:
logger.info(
"Windows grants the GPU %.1f GiB of %.1f GiB "
"system RAM (driver/WDDM budget). To raise it: "
"increase the BIOS UMA frame buffer size, or "
"AMD Software > Performance > Tuning > "
"Variable Graphics Memory.",
_granted / 1024**3,
_phys / 1024**3,
)
except Exception:
pass
except Exception as _oom_guard_err:
logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)
@ -2768,7 +2981,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
)
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
from datasets import load_dataset, Dataset
from datasets import Dataset
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
from transformers import TrainerCallback
from utils.paths import datasets_root, resolve_output_dir
except ImportError as e:

5
studio/backend/hub/services/datasets/formatting.py
View file

@ -223,7 +223,10 @@ def _load_processed_hf_preview_slice(
if not _is_valid_repo_id(request.dataset_name):
return None
try:
from datasets import DownloadConfig, load_dataset
from datasets import DownloadConfig
# Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
except Exception:
return None

3
studio/backend/hub/services/datasets/local.py
View file

@ -224,7 +224,8 @@ def _stream_file_preview_slice(path: Path, preview_size: int):
def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
from datasets import load_dataset
# Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
if dataset_path.is_dir():
parquet_dir = (

7
studio/backend/hub/services/models/common.py
View file

@ -21,6 +21,7 @@ from hub.utils.gguf import (
extract_quant_label,
is_gguf_filename as _is_gguf_filename,
is_mmproj_filename as _is_mmproj_filename,
is_mtp_drafter_path as _is_mtp_drafter_path,
)
from hub.utils.paths import is_valid_repo_id as _is_valid_repo_id
@ -71,7 +72,7 @@ def _is_model_directory(d: Path) -> bool:
if suffix == ".safetensors":
return True
if suffix == ".gguf":
return "mmproj" not in f.name.lower()
return "mmproj" not in f.name.lower() and not _is_mtp_drafter_path(f.name)
if suffix == ".bin":
name = f.name.lower()
return (
@ -276,7 +277,9 @@ def _classify_non_gguf_model_format(
def _is_main_gguf_filename(name: str) -> bool:
return _is_gguf_filename(name) and not _is_mmproj_filename(name)
return (
_is_gguf_filename(name) and not _is_mmproj_filename(name) and not _is_mtp_drafter_path(name)
)
def _iter_gguf_paths(root: Path):

6
studio/backend/hub/services/models/deletion.py
View file

@ -32,6 +32,7 @@ from hub.services.models.common import (
_is_gguf_filename,
_is_main_gguf_filename,
_is_mmproj_filename,
_is_mtp_drafter_path,
)
logger = get_logger(__name__)
@ -139,7 +140,10 @@ def _delete_gguf_variant_from_repos(
if matched and not sibling_active and not _has_remaining_main_gguf(target_repo):
companion_matches = _repo_file_matches(
target_repo,
lambda name: _is_gguf_filename(name) and _is_mmproj_filename(name),
# Companions: mmproj and the MTP drafter -- downloaded with
# every variant, so the last variant's delete reclaims them.
lambda name: _is_gguf_filename(name)
and (_is_mmproj_filename(name) or _is_mtp_drafter_path(name)),
)
for snap, _blob, name in companion_matches:
try:

9
studio/backend/hub/services/models/gguf_variants.py
View file

@ -38,6 +38,7 @@ from hub.utils.paths import (
)
from hub.services.models.common import (
_is_mmproj_filename,
_is_mtp_drafter_path,
_iter_gguf_paths,
)
from hub.utils.gguf_plan import (
@ -479,7 +480,7 @@ async def get_gguf_variants_response(
continue
key = rel.lower()
by_filename[key] = max(by_filename.get(key, 0), size)
if _is_mmproj_filename(f.name):
if _is_mmproj_filename(f.name) or _is_mtp_drafter_path(rel):
continue
q = extract_quant_label(rel).lower()
by_quant[q] = by_quant.get(q, 0) + size
@ -593,7 +594,11 @@ async def get_gguf_variants_response(
requirement = requirements_by_quant.get(variant.quant.lower())
if requirement is None:
continue
if requirement.mmproj_hashes & incomplete_hashes and _filenames_cached(
# companion_hashes adds the MTP drafter (mmproj_hashes covers
# every mmproj precision in the repo, not just the planned one).
if (
(requirement.mmproj_hashes | requirement.companion_hashes) & incomplete_hashes
) and _filenames_cached(
requirement.main_filenames,
requirement.main_size_bytes,
):

36
studio/backend/hub/utils/gguf.py
View file

@ -83,6 +83,28 @@ def is_mmproj_filename(filename: str) -> bool:
return "mmproj" in filename.lower()
def is_mtp_drafter_path(path: str) -> bool:
"""True for a separate-file MTP drafter (speculative head), a companion to
the main model rather than a selectable quant.
Covers the repo-root ``mtp-*.gguf`` (the Q8_0 copy unsloth ships for
llama.cpp ``-hf`` auto-discovery) and the ``MTP/`` subdir copies (Gemma 4).
Repos that bake the head into the main GGUF (Qwen) have no such file, so
this is False for them. Must be excluded from main-model selection
everywhere mmproj is.
CANONICAL COPY. Layering keeps two mirrors that must change in lockstep:
utils/models/model_config.py ``_is_mtp_drafter`` (utils cannot import
hub) and core/inference/llama_cpp.py ``_is_companion_gguf_path`` (core
avoids hub imports; bundles the mmproj check).
"""
p = path.lower()
if not p.endswith(".gguf"):
return False
name = p.rsplit("/", 1)[-1]
return name.startswith("mtp-") or "/mtp/" in f"/{p}"
def is_gguf_filename(filename: str) -> bool:
return filename.lower().endswith(".gguf")
@ -119,7 +141,9 @@ def iter_gguf_files(directory: Path, recursive: bool = False):
def pick_best_gguf(filenames: list[str]) -> Optional[str]:
gguf_files = [
name for name in filenames if is_gguf_filename(name) and not is_mmproj_filename(name)
name
for name in filenames
if is_gguf_filename(name) and not is_mmproj_filename(name) and not is_mtp_drafter_path(name)
]
if not gguf_files:
return None
@ -270,6 +294,12 @@ def list_partial_gguf_variants_from_state(
for expected in manifest.expected_files:
if not is_gguf_filename(expected.path):
continue
if is_mtp_drafter_path(expected.path):
# Downloaded with every variant (like mmproj) but not a
# selectable quant; count it so the shown download size
# matches what is fetched.
companion_bytes += max(0, int(expected.size or 0))
continue
if is_mmproj_filename(expected.path):
has_vision = True
companion_bytes += max(0, int(expected.size or 0))
@ -336,6 +366,8 @@ def list_gguf_variants(
filename = getattr(sibling, "rfilename", None)
if not isinstance(filename, str) or not is_gguf_filename(filename):
continue
if is_mtp_drafter_path(filename):
continue
if is_mmproj_filename(filename):
has_vision = True
continue
@ -389,6 +421,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
except OSError:
size = 0
rel = file.relative_to(root).as_posix()
if is_mtp_drafter_path(rel):
continue
quant = extract_quant_label(rel)
quant_totals[quant] = quant_totals.get(quant, 0) + size
quant_first_file.setdefault(quant, rel)

77
studio/backend/hub/utils/gguf_plan.py
View file

@ -7,7 +7,12 @@ from dataclasses import dataclass
from typing import Optional, Sequence
from hub.utils.download_manifest import ExpectedFile
from hub.utils.gguf import extract_quant_label, is_gguf_filename, is_mmproj_filename
from hub.utils.gguf import (
extract_quant_label,
is_gguf_filename,
is_mmproj_filename,
is_mtp_drafter_path,
)
@dataclass(frozen = True)
@ -53,24 +58,30 @@ def expected_file_from_sibling(sibling) -> Optional[ExpectedFile]:
def is_companion_gguf_path(path: str) -> bool:
return is_gguf_filename(path) and is_mmproj_filename(path)
"""Companion (non-main) GGUF downloaded alongside a variant: the vision
mmproj or the separate MTP drafter (Gemma 4)."""
return is_gguf_filename(path) and (is_mmproj_filename(path) or is_mtp_drafter_path(path))
def is_main_gguf_variant_path(path: str, variant: str) -> bool:
return (
is_gguf_filename(path)
and not is_mmproj_filename(path)
and not is_mtp_drafter_path(path)
and extract_quant_label(path).lower() == variant.lower()
)
def _gguf_rfilename(sibling) -> Optional[str]:
"""The sibling's rfilename when it is a GGUF, else None."""
name = getattr(sibling, "rfilename", None)
if isinstance(name, str) and is_gguf_filename(name):
return name
return None
def mmproj_siblings(siblings: Sequence) -> list:
return [
s
for s in siblings
if isinstance(getattr(s, "rfilename", None), str)
and is_companion_gguf_path(getattr(s, "rfilename"))
]
return [s for s in siblings if (name := _gguf_rfilename(s)) and is_mmproj_filename(name)]
def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
@ -83,6 +94,25 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
)
def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
"""The separate MTP drafter to fetch with every variant: the repo-root
``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
(Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
prefix, first in sort order) so download and load resolve the same file;
the higher-precision ``MTP/`` subdir copies are for explicit selection and
are not auto-fetched. None for repos with the head baked into the main
GGUF (Qwen)."""
candidates = sorted(
(
s
for s in siblings
if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
),
key = lambda s: getattr(s, "rfilename"),
)
return candidates[0] if candidates else None
def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]:
main: dict[str, list] = {}
all_mmproj = mmproj_siblings(siblings)
@ -94,12 +124,20 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]:
all_mmproj_hashes = frozenset(h for h in (sibling_sha256(s) for s in all_mmproj) if h)
companion = preferred_mmproj_sibling(siblings)
companion_expected = expected_file_from_sibling(companion) if companion is not None else None
mtp_sibling = preferred_mtp_sibling(siblings)
mtp_expected = expected_file_from_sibling(mtp_sibling) if mtp_sibling is not None else None
companions_expected = tuple(
file for file in (companion_expected, mtp_expected) if file is not None
)
for sibling in siblings:
name = getattr(sibling, "rfilename", None)
if not isinstance(name, str) or not is_gguf_filename(name):
name = _gguf_rfilename(sibling)
if name is None:
continue
if is_mmproj_filename(name):
# Companions are folded into every plan below; keep them out of the
# quant grouping so a drafter never lands in a variant's main files
# (the root mtp-*.gguf carries a quant label, e.g. Q8_0).
if is_mmproj_filename(name) or is_mtp_drafter_path(name):
continue
quant = extract_quant_label(name).lower()
main.setdefault(quant, []).append(sibling)
@ -111,11 +149,7 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]:
for sibling in target_main_siblings
if (file := expected_file_from_sibling(sibling)) is not None
)
expected_files = (
(*main_expected, companion_expected)
if companion_expected is not None
else main_expected
)
expected_files = (*main_expected, *companions_expected)
plans[quant] = plan_from_expected_files(
quant,
expected_files,
@ -135,6 +169,9 @@ def plan_from_expected_files(
expected = tuple(expected_files)
main_files = tuple(file for file in expected if is_main_gguf_variant_path(file.path, variant))
companion_files = tuple(file for file in expected if is_companion_gguf_path(file.path))
# Manifest-resume fallback for the mmproj fields below: companion_files
# also holds the MTP drafter, so keep an mmproj-only view.
mmproj_files = tuple(file for file in companion_files if is_mmproj_filename(file.path))
main_hashes = frozenset(file.sha256 for file in main_files if file.sha256)
companion_hashes = frozenset(file.sha256 for file in companion_files if file.sha256)
required_hashes = frozenset(file.sha256 for file in expected if file.sha256)
@ -149,9 +186,13 @@ def plan_from_expected_files(
mmproj_filenames = (
all_mmproj_filenames
if all_mmproj_filenames is not None
else frozenset(file.path for file in companion_files)
else frozenset(file.path for file in mmproj_files)
),
mmproj_hashes = (
all_mmproj_hashes
if all_mmproj_hashes is not None
else frozenset(file.sha256 for file in mmproj_files if file.sha256)
),
mmproj_hashes = (all_mmproj_hashes if all_mmproj_hashes is not None else companion_hashes),
expected_files = expected,
main_size_bytes = main_size,
download_size_bytes = download_size,

9
studio/backend/hub/utils/inventory_scan.py
View file

@ -25,7 +25,12 @@ from loggers import get_logger
logger = get_logger(__name__)
from hub.utils.gguf import extract_quant_label, is_gguf_filename, is_mmproj_filename
from hub.utils.gguf import (
extract_quant_label,
is_gguf_filename,
is_mmproj_filename,
is_mtp_drafter_path,
)
from hub.utils.state_dir import RepoType
from hub.utils.hf_cache_state import (
@ -336,7 +341,7 @@ def _completed_gguf_variants(snapshot_dir: Optional[Path]) -> set[str]:
except OSError:
continue
rel = path.relative_to(snapshot_dir).as_posix()
if not is_gguf_filename(rel) or is_mmproj_filename(rel):
if not is_gguf_filename(rel) or is_mmproj_filename(rel) or is_mtp_drafter_path(rel):
continue
quant = extract_quant_label(rel)
split = _GGUF_SPLIT_RE.search(path.name)

66
studio/backend/main.py
View file

@ -63,17 +63,38 @@ if sys.platform == "win32":
_add_rocm_dll_dirs()
del _add_rocm_dll_dirs
# ── Windows AMD ROCm: make hipInfo.exe resolvable for subprocess probes ──
# bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via PATH at import
# time; the AMD torch wheel ships it in the venv Scripts dir, which is on
# PATH only when the venv is activated -- Studio launches python directly.
# Without this, every bitsandbytes import logs a scary (but harmless)
# "Could not detect ROCm GPU architecture: [WinError 2]" ERROR + WARNING.
# Gated on the file existing: only AMD ROCm wheels ship hipInfo.exe, so
# NVIDIA/CPU hosts are untouched. os.add_dll_directory above does not help
# here -- subprocess PATH resolution ignores DLL search directories.
_scripts_dir = os.path.dirname(sys.executable)
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")):
import shutil as _shutil
if not _shutil.which("hipinfo.exe"):
os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "")
del _shutil
del _scripts_dir
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
# bitsandbytes derives the rocm<ver>.dll name from torch.version.hip, but the
# wheel ships rocm72.dll, so the server crashes ("Configured ROCm binary not
# found") without this. Detect the shipped DLL and fall back to "72" (mirrors
# worker.py). Gate on the rocm bnb DLL / HIP_PATH rather than torch.version.hip
# to avoid importing torch on every Windows host.
if "BNB_ROCM_VERSION" not in os.environ:
# found") without this. Detect the shipped DLL (mirrors worker.py); gate on
# the rocm bnb DLL rather than torch.version.hip to avoid importing torch on
# every Windows host.
# Values seeded by the installer's sitecustomize.py are redetectable
# defaults; explicit caller values remain authoritative.
if (
"BNB_ROCM_VERSION" not in os.environ
or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"
):
import glob as _glob
import logging as _logging
_hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"))
_bnb_rocm_ver = None
_found_rocm_bnb = False
try:
@ -96,18 +117,43 @@ if sys.platform == "win32":
_bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
except Exception as _e:
_logging.getLogger(__name__).warning(
"Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'",
"Windows ROCm: BNB DLL detection failed (%s); leaving BNB_ROCM_VERSION as is",
_e,
)
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72")
if _found_rocm_bnb or _hip_env:
_bnb_rocm_ver_final = _bnb_rocm_ver or "72"
# Only when a ROCm bnb DLL actually exists: HIP_PATH/ROCM_PATH alone
# (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a
# non-ROCm bitsandbytes, which raises at import. DLL unparsable -> "72".
if _found_rocm_bnb:
_bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
_logging.getLogger(__name__).info(
"Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)",
_bnb_rocm_ver_final,
)
# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ──────
# In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over
# /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE
# torch touches the GPU. A worker launched outside a login shell (e.g.
# `wsl.exe -d Ubuntu-24.04 python ...`) misses the installer's persisted env
# and silently falls back to CPU. Set it here, gated to no-op unless BOTH
# /dev/dxg AND librocdxg.so exist -- native Linux ROCm, NVIDIA, macOS and
# Windows are unaffected.
elif sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ:
try:
if os.path.exists("/dev/dxg") and any(
os.path.exists(os.path.join(_p, "librocdxg.so"))
for _p in ("/opt/rocm/lib", "/opt/rocm/lib64")
):
os.environ["HSA_ENABLE_DXG_DETECTION"] = "1"
import logging as _logging
_logging.getLogger(__name__).info(
"WSL ROCm: set HSA_ENABLE_DXG_DETECTION=1 (librocdxg bridge present)"
)
except Exception:
pass
# Put backend dir on sys.path so _platform_compat is importable when main.py
# is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)
@ -221,6 +267,7 @@ from routes import (
training_history_router,
training_router,
)
from routes.llama import router as llama_router
from hub.routes import (
inventory_router as hub_inventory_router,
datasets_router as hub_datasets_router,
@ -759,6 +806,7 @@ app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp
app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])

34
studio/backend/models/inference.py
View file

@ -72,7 +72,7 @@ class LoadRequest(BaseModel):
"ngram-mod+draft-mtp chain on both platforms), 'off' (disabled). "
"Legacy values 'default' (-> auto), 'draft-mtp' (-> mtp), "
"'ngram-mod' (-> ngram), and 'ngram-simple' (kept as-is) are "
"still accepted. Ignored for non-GGUF and vision models."
"still accepted. Ignored for non-GGUF models."
),
)
spec_draft_n_max: Optional[int] = Field(
@ -171,7 +171,7 @@ class LoadResponse(BaseModel):
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
context_length: Optional[int] = Field(
None, description = "Model's native context length (from GGUF metadata)"
None, description = "Runtime context length in tokens for the loaded model"
)
max_context_length: Optional[int] = Field(
None, description = "Maximum context length currently available on this hardware"
@ -346,6 +346,17 @@ class InferenceStatusResponse(BaseModel):
"False -> recommend `unsloth studio update`."
),
)
spec_fallback_reason: Optional[str] = Field(
None,
description = (
"Why MTP was disabled on the loaded model despite being requested "
"(auto on an MTP model, or forced mtp / mtp+ngram). "
"'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would "
"re-enable it (show the update affordance); 'runtime_error' -> the "
"current build could not run it. None when MTP engaged or was not "
"requested."
),
)
llama_cpp_prebuilt_stale: bool = Field(
False,
description = (
@ -533,8 +544,10 @@ class ChatMessage(BaseModel):
if self.role == "tool":
# tool_call_id resolution happens at ChatCompletionRequest scope.
if not self.content:
raise ValueError('role="tool" messages require non-empty "content".')
# OpenAI accepts empty tool results (commands with no output);
# normalize to "" instead of a 400 agentic clients treat as fatal.
if self.content is None or self.content == []:
self.content = ""
elif self.role == "assistant":
# Post-Stop sentinel: collapse content="" / [] to None.
if (self.content == "" or self.content == []) and not self.tool_calls:
@ -631,7 +644,8 @@ class ChatCompletionRequest(BaseModel):
None, description = "[x-unsloth] Base64-encoded image for vision models"
)
audio_base64: Optional[str] = Field(
None, description = "[x-unsloth] Base64-encoded WAV for audio-input models (ASR)"
None,
description = "[x-unsloth] Base64-encoded audio (wav/mp3/ogg/flac/m4a) for audio-input models",
)
use_adapter: Optional[Union[bool, str]] = Field(
None,
@ -680,6 +694,16 @@ class ChatCompletionRequest(BaseModel):
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
)
context_overflow: Optional[Literal["error", "truncate_middle"]] = Field(
None,
description = (
"[x-unsloth] Passthrough behavior when the prompt exceeds the real "
"context window. 'error' (default) returns a 400 with "
"code=context_length_exceeded. 'truncate_middle' drops middle "
"turn-groups (system prompt, first turn, and recent turns kept; "
"tool calls stay paired with their results) and retries."
),
)
max_tool_calls_per_message: Optional[int] = Field(
25,
ge = 0,

11
studio/backend/models/mcp_servers.py
View file

@ -44,3 +44,14 @@ class McpServerProbeResult(BaseModel):
ok: bool
tool_count: int = 0
error: Optional[str] = None
class McpServerImportRequest(BaseModel):
# A standard mcpServers JSON config (Claude Desktop / Cursor / Cline / VS Code).
config: dict
class McpServerImportResult(BaseModel):
created: list[McpServerResponse] = Field(default_factory = list)
skipped: list[str] = Field(default_factory = list) # display names skipped as duplicates
errors: list[str] = Field(default_factory = list)

3
studio/backend/routes/datasets.py
View file

@ -412,7 +412,8 @@ def _build_local_dataset_items() -> list[LocalDatasetItem]:
def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
from datasets import load_dataset
# Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
if dataset_path.is_dir():
parquet_dir = (

3
studio/backend/routes/export.py
View file

@ -78,8 +78,7 @@ async def load_checkpoint(
for _ in range(60): # up to 30s
if not trn.is_training_active():
break
import time
time.sleep(0.5)
await asyncio.sleep(0.5)
else:
logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
except Exception as e:

669
studio/backend/routes/inference.py
View file

@ -29,6 +29,16 @@ from utils.models import extract_model_size_b as _extract_model_size_b
from utils.api_errors import openai_error_body, anthropic_error_body
def _positive_int_or_none(value: Any) -> Optional[int]:
if isinstance(value, bool):
return None
try:
value_int = int(value)
except (TypeError, ValueError):
return None
return value_int if value_int > 0 else None
def _install_httpcore_asyncgen_silencer() -> None:
"""Silence benign httpx/httpcore asyncgen GC noise on Python 3.13.
@ -242,6 +252,169 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException":
)
_OVERFLOW_TRUNCATE_MAX_RETRIES = 3
# Truncated-prompt share of the real window; the rest is generation headroom
# so a near-full prompt cannot cut a tool call mid-JSON at the wall.
_OVERFLOW_PROMPT_TARGET_FRACTION = 0.75
def _overflow_truncation_requested(payload) -> bool:
"""True when the request (or the UNSLOTH_CONTEXT_OVERFLOW server default,
for clients that cannot send custom fields) opted into truncation."""
requested = getattr(payload, "context_overflow", None)
if requested is not None:
return requested == "truncate_middle"
return os.environ.get("UNSLOTH_CONTEXT_OVERFLOW", "").strip().lower() == "truncate_middle"
def _parse_overflow_counts(err_text: str):
"""(n_prompt_tokens, n_ctx) from an exceed_context_size_error body, or
None. Tolerates \\" around keys (body may be a re-wrapped JSON string)."""
m_prompt = _re.search(r'n_prompt_tokens\\?"?\s*:\s*(\d+)', err_text)
m_ctx = _re.search(r'n_ctx\\?"?\s*:\s*(\d+)', err_text)
if m_prompt and m_ctx:
return int(m_prompt.group(1)), int(m_ctx.group(1))
return None
def _estimate_message_tokens(msg: dict) -> int:
try:
return max(1, len(json.dumps(msg, ensure_ascii = False)) // 4)
except Exception:
return 1
def _truncate_middle_messages(messages: list, keep_ratio: float):
"""Drop whole turn-groups from the middle of an OpenAI message list.
Always kept: leading system message(s), the first group (task anchor),
and the trailing groups. A group is a user message, or an assistant
message plus its following tool results, so surviving tool_calls stay
paired with their results as chat templates require.
Returns (new_messages, dropped_message_count).
"""
if not messages or keep_ratio >= 1.0:
return messages, 0
head: list = []
idx = 0
while idx < len(messages) and messages[idx].get("role") in ("system", "developer"):
head.append(messages[idx])
idx += 1
groups: list[list] = []
for msg in messages[idx:]:
role = msg.get("role")
if role == "tool" and groups:
groups[-1].append(msg)
elif role == "tool":
groups.append([msg]) # orphan tool result; treat as its own group
else:
groups.append([msg])
# Anchor group plus the last 3 groups stay.
protected_tail = min(3, max(1, len(groups) - 1))
if len(groups) <= 1 + protected_tail:
return messages, 0
total_est = sum(_estimate_message_tokens(m) for m in messages)
target_est = int(total_est * keep_ratio)
anchor = groups[0]
middle = groups[1:-protected_tail]
tail = groups[-protected_tail:]
current_est = total_est
kept_middle: list[list] = list(middle)
dropped = 0
# Drop oldest-first until the estimate fits the target.
while kept_middle and current_est > target_est:
victim = kept_middle.pop(0)
dropped += len(victim)
current_est -= sum(_estimate_message_tokens(m) for m in victim)
if dropped == 0:
return messages, 0
new_messages = head + anchor
for grp in kept_middle:
new_messages.extend(grp)
for grp in tail:
new_messages.extend(grp)
return new_messages, dropped
_CLIP_MARKER = "\n[... truncated by context_overflow=truncate_middle ...]\n"
# Generous head+tail first; cut harder if the estimate still misses the target.
_CLIP_KEEP_CHARS = (1500, 400)
def _clip_long_contents(messages: list, target_est: int) -> int:
"""Clip oversized string contents middle-out until ``target_est`` is met.
Tool results first, then earlier user turns, the final message last.
Message count and roles never change, so tool pairing holds even when
group-dropping could not free enough. Returns messages clipped.
"""
def _candidates():
tools = [m for m in messages if m.get("role") == "tool"]
users = [m for m in messages[:-1] if m.get("role") == "user"]
last = [messages[-1]] if messages else []
return tools + users + last
clipped = 0
for keep in _CLIP_KEEP_CHARS:
for msg in _candidates():
if sum(_estimate_message_tokens(m) for m in messages) <= target_est:
return clipped
content = msg.get("content")
if not isinstance(content, str) or len(content) <= 2 * keep + len(_CLIP_MARKER):
continue
msg["content"] = content[:keep] + _CLIP_MARKER + content[-keep:]
clipped += 1
return clipped
def _apply_overflow_truncation(body: dict, err_text: str) -> bool:
"""Shrink a passthrough body after an upstream context overflow: drop
middle turn-groups, clip still-oversized contents, clamp ``max_tokens``
to the generation headroom. Returns False when nothing could shrink."""
counts = _parse_overflow_counts(err_text)
messages = body.get("messages") or []
total_est = sum(_estimate_message_tokens(m) for m in messages)
if counts:
n_prompt, n_ctx = counts
keep_ratio = min(0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt))
# Scale the server-token target into char-estimate units.
target_est = int(total_est * keep_ratio)
else:
n_ctx = None
keep_ratio = 0.6 # no counts in the error; cut conservatively
target_est = int(total_est * keep_ratio)
new_messages, dropped = _truncate_middle_messages(messages, keep_ratio)
if dropped:
body["messages"] = new_messages
clipped = 0
if sum(_estimate_message_tokens(m) for m in body.get("messages") or []) > target_est:
clipped = _clip_long_contents(body.get("messages") or [], target_est)
if not dropped and not clipped:
return False
if n_ctx:
headroom = max(1024, int(n_ctx * (1.0 - _OVERFLOW_PROMPT_TARGET_FRACTION)))
cur_max = body.get("max_tokens")
body["max_tokens"] = min(cur_max, headroom) if cur_max else headroom
logger.warning(
"context_overflow=truncate_middle: dropped %d middle messages, clipped "
"%d contents (keep_ratio %.2f); retrying within the real window",
dropped,
clipped,
keep_ratio,
)
return True
def _anthropic_stream_error_event(exc):
"""Anthropic in-band SSE ``error`` event for a mid-stream failure, or ``None``
to fall through to a normal message_delta finish. Returns an event only for a
@ -418,6 +591,7 @@ try:
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
@ -427,7 +601,10 @@ try:
)
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import load_model_defaults
from utils.models.model_config import (
detect_mtp_file,
load_model_defaults,
)
from utils.native_path_leases import (
NativePathLeaseError,
display_label_for_native_path,
@ -445,6 +622,7 @@ except ImportError:
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
@ -454,7 +632,10 @@ except ImportError:
)
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import load_model_defaults
from utils.models.model_config import (
detect_mtp_file,
load_model_defaults,
)
from utils.native_path_leases import (
NativePathLeaseError,
display_label_for_native_path,
@ -927,35 +1108,42 @@ def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str
logger = get_logger(__name__)
def _validate_native_mmproj_companion(mmproj_path: str | None, gguf_path: str | None) -> None:
if not mmproj_path or not gguf_path:
def _validate_native_gguf_companion(
companion_path: str | None, gguf_path: str | None, label: str
) -> None:
"""Reject a companion GGUF (mmproj / MTP drafter) that a native-lease load
would otherwise hand to llama-server: must be a regular file (no symlink
escaping the leased directory) living next to the selected GGUF."""
if not companion_path or not gguf_path:
return
import stat as _stat_module
mm = Path(mmproj_path)
companion = Path(companion_path)
gguf = Path(gguf_path)
try:
mm_lstat = os.lstat(mm)
companion_lstat = os.lstat(companion)
except OSError as exc:
raise HTTPException(
status_code = 400,
detail = "Native vision companion is no longer accessible.",
detail = f"Native {label} is no longer accessible.",
) from exc
if _stat_module.S_ISLNK(mm_lstat.st_mode) or not _stat_module.S_ISREG(mm_lstat.st_mode):
if _stat_module.S_ISLNK(companion_lstat.st_mode) or not _stat_module.S_ISREG(
companion_lstat.st_mode
):
raise HTTPException(
status_code = 400,
detail = "Native vision companion must be a regular file.",
detail = f"Native {label} must be a regular file.",
)
try:
if mm.resolve(strict = True).parent != gguf.resolve(strict = True).parent:
if companion.resolve(strict = True).parent != gguf.resolve(strict = True).parent:
raise HTTPException(
status_code = 400,
detail = "Native vision companion must live next to the selected GGUF.",
detail = f"Native {label} must live next to the selected GGUF.",
)
except OSError as exc:
raise HTTPException(
status_code = 400,
detail = "Native vision companion is no longer accessible.",
detail = f"Native {label} is no longer accessible.",
) from exc
@ -980,13 +1168,11 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC
llama_backend.cache_type_kv
):
return False
# Vision loads silently drop speculative decoding (llama_cpp.py gates spec
# on ``not is_vision``), so treat the request as ``off`` against the
# backend's ``None`` to avoid a redundant reload.
if llama_backend.is_vision:
req_mode = "off"
else:
req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto"
# Spec decoding works on vision models too (MTP is mmproj-compatible,
# llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare
# the real requested mode -- coercing vision to ``off`` here used to
# swallow every spec-mode change on a vision model as already_loaded.
req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto"
backend_mode = llama_backend.requested_spec_mode or "auto"
if req_mode != backend_mode:
return False
@ -1008,6 +1194,33 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC
else:
if list(request.llama_extra_args) != backend_extra:
return False
# A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing
# next to the loaded weights changes the launch command (--model-draft),
# so a duplicate /load must reload rather than dedupe. Always compare the
# detected vs stored drafter when the mode can use one and the user does
# not own --spec-type: the resolved-path compare is cheap and handles all
# four cases (both None -> match; one None -> reload; equal -> match;
# different -> reload), including a drafter deleted out from under a
# running server. Runs last: it stats the filesystem, so every pure-memory
# comparison above short-circuits first. Resolve both sides since the
# stored launch path may be a snapshot symlink while detect_mtp_file
# returns the resolved blob.
if req_mode in ("auto", "mtp", "mtp+ngram") and llama_backend.gguf_path:
effective_extras = (
request.llama_extra_args
if request.llama_extra_args is not None
else llama_backend.extra_args
)
if not _extra_args_set_spec_type(effective_extras):
detected = detect_mtp_file(llama_backend.gguf_path)
stored = llama_backend.mtp_draft_path
try:
detected_resolved = Path(detected).resolve() if detected else None
stored_resolved = Path(stored).resolve() if stored else None
except OSError:
return False
if detected_resolved != stored_resolved:
return False
return True
@ -1187,6 +1400,7 @@ async def load_model(
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
context_length = _positive_int_or_none(_model_info.get("context_length")),
chat_template = _chat_template,
)
@ -1313,12 +1527,26 @@ async def load_model(
)
else:
# Local mode: llama-server loads via -m <path>
if native_grant_backed and config.gguf_mmproj_file:
_validate_native_mmproj_companion(config.gguf_mmproj_file, config.gguf_file)
if native_grant_backed:
if config.gguf_mmproj_file:
_validate_native_gguf_companion(
config.gguf_mmproj_file, config.gguf_file, "vision companion"
)
if config.gguf_mtp_file:
# The drafter is optional (unlike mmproj for a vision
# model): drop it rather than fail the load.
try:
_validate_native_gguf_companion(
config.gguf_mtp_file, config.gguf_file, "MTP drafter"
)
except HTTPException as exc:
logger.warning("Dropping MTP drafter for native load: %s", exc.detail)
config.gguf_mtp_file = None
success = await asyncio.to_thread(
llama_backend.load_model,
gguf_path = config.gguf_file,
mmproj_path = config.gguf_mmproj_file,
mtp_draft_path = config.gguf_mtp_file,
# Pass the resolved variant so _extra_args_source keys off
# the same string the inheritance check at the top of /load
# uses (#5401 followup).
@ -1514,6 +1742,7 @@ async def load_model(
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
context_length = _positive_int_or_none(_model_info.get("context_length")),
chat_template = _chat_template,
)
@ -1852,6 +2081,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
llama_cpp_supports_mtp = _supports_mtp,
spec_fallback_reason = llama_backend.spec_fallback_reason,
llama_cpp_prebuilt_stale = _stale,
llama_cpp_installed_tag = _installed_tag,
llama_cpp_latest_tag = _latest_tag,
@ -1901,6 +2131,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
context_length = _positive_int_or_none(model_info.get("context_length")),
chat_template = chat_template,
llama_cpp_supports_mtp = _supports_mtp,
llama_cpp_prebuilt_stale = _stale,
@ -2064,6 +2295,172 @@ def _decode_audio_base64(b64: str) -> np.ndarray:
return waveform.squeeze(0).numpy()
# Reject oversized audio before decoding. base64 inflates raw bytes by ~4/3, so
# cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally
# bounds the *decoded* length, since a small compressed file (opus/flac/etc.)
# can expand to a far larger PCM array than the encoded-size cap implies.
_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024
_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3
_MAX_AUDIO_SECONDS = 30 * 60
_WAV_HEADER_BYTES = 44
_MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000
def _sniff_audio_container(raw: bytes) -> Optional[str]:
"""Return 'wav' or 'mp3' if the bytes are a container llama-server accepts
directly (so we can forward them untouched), else None (needs transcoding)."""
if len(raw) >= 12 and raw[:4] == b"RIFF" and raw[8:12] == b"WAVE":
return "wav"
# mp3: ID3 tag, or an MPEG audio frame sync (no other accepted format leads
# with 0xFF, so the simple sync check doesn't collide).
if raw[:3] == b"ID3" or (len(raw) >= 2 and raw[0] == 0xFF and (raw[1] & 0xE0) == 0xE0):
return "mp3"
return None
def _mono_f32_to_wav_bytes(arr: np.ndarray, sample_rate: int) -> bytes:
"""Encode a mono float32 array as 16-bit PCM WAV bytes.
Torch-free (numpy + stdlib only) so it works on no-torch GGUF-only installs;
the shared audio_codecs helper pulls in torch at import time.
"""
import io
import wave
arr = np.nan_to_num(np.asarray(arr, dtype = np.float32).flatten(), posinf = 0.0, neginf = 0.0)
if arr.size == 0:
raise ValueError("decoded audio is empty")
peak = float(np.abs(arr).max())
if peak > 1.0:
arr = arr / peak
pcm = (arr * 32767.0).astype(np.int16)
buf = io.BytesIO()
with wave.open(buf, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(int(sample_rate))
wf.writeframes(pcm.tobytes())
return buf.getvalue()
def _resample_mono_linear(arr: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray:
"""Small numpy-only resampler for upload size limiting."""
if source_rate <= 0 or target_rate <= 0 or source_rate == target_rate:
return arr
duration = len(arr) / float(source_rate)
target_len = max(1, int(round(duration * target_rate)))
if target_len == len(arr):
return arr
source_x = np.linspace(0.0, duration, num = len(arr), endpoint = False)
target_x = np.linspace(0.0, duration, num = target_len, endpoint = False)
return np.interp(target_x, source_x, arr).astype(np.float32)
def _fit_transcoded_audio_to_wav_cap(arr: np.ndarray, sample_rate: int) -> tuple[np.ndarray, int]:
"""Downsample only when needed so transcoded WAV stays within the upload cap."""
if sample_rate <= 0:
raise ValueError("decoded audio has an invalid sample rate")
wav_bytes = _WAV_HEADER_BYTES + len(arr) * 2
if wav_bytes <= _MAX_AUDIO_RAW_BYTES:
return arr, sample_rate
duration = len(arr) / float(sample_rate)
max_samples = max(1, (_MAX_AUDIO_RAW_BYTES - _WAV_HEADER_BYTES) // 2)
target_rate = int(max_samples // duration)
if target_rate < _MIN_TRANSCODE_AUDIO_SAMPLE_RATE:
raise ValueError("decoded audio exceeds the transcoded WAV size limit")
target_rate = min(sample_rate, target_rate)
fitted = _resample_mono_linear(arr, sample_rate, target_rate)
if _WAV_HEADER_BYTES + len(fitted) * 2 > _MAX_AUDIO_RAW_BYTES:
raise ValueError("decoded audio exceeds the transcoded WAV size limit")
return fitted, target_rate
def _decode_audio_mono(raw: bytes) -> tuple[np.ndarray, int]:
"""Decode audio bytes to (mono float32 array, native sample_rate).
soundfile (libsndfile) reads wav/mp3/ogg/flac straight from memory. librosa
(ffmpeg-backed) additionally covers m4a/webm but needs a real path and is
absent on no-torch GGUF-only installs. Both imports are inside the fallback
so a missing decoder degrades to the next one (and finally a clear error)
rather than crashing.
"""
import io
try:
import soundfile as sf
arr, sr = sf.read(io.BytesIO(raw), dtype = "float32")
except Exception:
try:
import librosa
except ModuleNotFoundError as e:
raise RuntimeError(
"this audio format needs librosa, which is not installed in "
"GGUF-only environments; use wav, mp3, ogg or flac"
) from e
import os
import tempfile
from utils.paths import ensure_dir, tmp_root
with tempfile.NamedTemporaryFile(
suffix = ".audio",
delete = False,
dir = str(ensure_dir(tmp_root())),
) as tmp:
tmp.write(raw)
tmp_path = tmp.name
try:
arr, sr = librosa.load(tmp_path, sr = None, mono = True)
finally:
os.unlink(tmp_path)
if arr.ndim > 1:
arr = arr.mean(axis = 1)
if sr > 0 and len(arr) > sr * _MAX_AUDIO_SECONDS:
raise ValueError(f"decoded audio exceeds the {_MAX_AUDIO_SECONDS // 60}-minute limit")
return arr, sr
def _prepare_audio_for_llama(b64: str) -> tuple[str, str]:
"""Return (base64, format) ready for llama-server's input_audio part.
llama-server's API only accepts wav/mp3, and decodes/resamples/down-mixes
them itself, so wav and mp3 uploads are forwarded untouched (no decode, no
PCM payload inflation). Other containers (m4a/ogg/webm/flac) are decoded to
a mono WAV. Blocking; call via a thread from async paths.
"""
if b64.startswith("data:"):
b64 = b64.split(",", 1)[1] if "," in b64 else ""
raw = base64.b64decode(b64)
passthrough = _sniff_audio_container(raw)
if passthrough is not None:
return b64, passthrough
arr, sr = _decode_audio_mono(raw)
arr, sr = _fit_transcoded_audio_to_wav_cap(arr, sr)
return base64.b64encode(_mono_f32_to_wav_bytes(arr, sr)).decode("ascii"), "wav"
def _inject_audio_part(messages: list[dict], audio_b64: str, audio_format: str) -> None:
"""Append an input_audio part to the last user message, in place.
Audio rides in the message list like image_url parts do, so it flows through
both the plain and tool-calling generation paths.
"""
part = {
"type": "input_audio",
"input_audio": {"data": audio_b64, "format": audio_format},
}
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content")
if isinstance(content, list):
content.append(part)
else:
msg["content"] = [{"type": "text", "text": content or ""}, part]
return
def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[str]"]:
"""
Parse OpenAI-format messages into components the inference backend expects.
@ -3021,9 +3418,12 @@ async def openai_chat_completions(
if _wants_multiple_choices(payload):
_raise_unsupported_n("GGUF tool or response_format passthrough")
if payload.audio_base64:
# This path forwards the request verbatim, so the transcoded audio
# never gets injected. (The agentic tool loop below does support
# audio.)
raise HTTPException(
status_code = 400,
detail = "Audio input is not supported for GGUF chat models yet.",
detail = "Audio input is not supported together with guided decoding or client-supplied tools yet.",
)
# Preserve the vision guard from the non-passthrough path below:
@ -3074,11 +3474,34 @@ async def openai_chat_completions(
# ── GGUF path: proxy to llama-server /v1/chat/completions ──
if using_gguf:
# Forward uploaded audio as an input_audio part. wav/mp3 pass through
# untouched (llama-server decodes and resamples them via the mmproj
# audio encoder); other containers are transcoded to WAV here. The part
# is injected into the message list below so it rides through both the
# plain and tool-calling paths, exactly like image_url parts.
audio_b64 = None
audio_format = "wav"
if payload.audio_base64:
raise HTTPException(
status_code = 400,
detail = "Audio input is not supported for GGUF chat models yet.",
)
if not getattr(llama_backend, "_has_audio_input", False):
raise HTTPException(
status_code = 400,
detail = "Audio provided but current GGUF model does not support audio input.",
)
if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS:
raise HTTPException(
status_code = 413,
detail = "Audio file is too large (max ~25 MB).",
)
try:
audio_b64, audio_format = await asyncio.to_thread(
_prepare_audio_for_llama, payload.audio_base64
)
except Exception as e:
logger.warning("Audio decode failed: %s", e, exc_info = True)
raise HTTPException(
status_code = 400,
detail = "Could not decode the provided audio file.",
)
gguf_messages, _ = _openai_messages_for_gguf_chat(
payload,
@ -3086,6 +3509,8 @@ async def openai_chat_completions(
)
gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt)
image_b64 = None
if audio_b64:
_inject_audio_part(gguf_messages, audio_b64, audio_format)
cancel_event = threading.Event()
@ -4207,26 +4632,45 @@ def _openai_model_objects() -> list[dict]:
# Check GGUF backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
models.append(
{
"id": llama_backend.model_identifier,
"object": "model",
"created": _created,
"owned_by": "local",
}
)
entry = {
"id": llama_backend.model_identifier,
"object": "model",
"created": _created,
"owned_by": "local",
}
_ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None))
if _ctx is not None:
entry["context_length"] = _ctx
_max_ctx = _positive_int_or_none(getattr(llama_backend, "max_context_length", None))
if _max_ctx is not None:
entry["max_context_length"] = _max_ctx
_native_ctx = _positive_int_or_none(getattr(llama_backend, "native_context_length", None))
if _native_ctx is not None:
entry["native_context_length"] = _native_ctx
models.append(entry)
# Check Unsloth backend
backend = get_inference_backend()
if backend.active_model_name:
models.append(
{
"id": backend.active_model_name,
"object": "model",
"created": _created,
"owned_by": "local",
}
)
model_info = backend.models.get(backend.active_model_name, {})
entry = {
"id": backend.active_model_name,
"object": "model",
"created": _created,
"owned_by": "local",
}
_ctx = _positive_int_or_none(model_info.get("context_length"))
if _ctx is None:
for _candidate in (
getattr(backend, "context_length", None),
getattr(backend, "max_seq_length", None),
):
_ctx = _positive_int_or_none(_candidate)
if _ctx is not None:
break
if _ctx is not None:
entry["context_length"] = _ctx
models.append(entry)
return models
@ -4476,6 +4920,17 @@ def _responses_message_text(content: Union[str, list]) -> str:
return "\n".join(parts)
def _responses_tool_output_text(output: Union[str, list]) -> str:
"""Return Chat Completions-safe content for a Responses tool result."""
if isinstance(output, str):
return output if output.strip() else "(no output)"
if output:
return json.dumps(output)
return "(no output)"
def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]:
"""Convert a ResponsesRequest's ``input`` into a Chat-format ``ChatMessage`` list.
@ -4535,10 +4990,9 @@ def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]:
if isinstance(item, ResponsesFunctionCallOutputInputItem):
# Chat Completions `role="tool"` requires string content; serialize
# a Responses content-array output.
output = item.output
if not isinstance(output, str):
output = json.dumps(output)
# a Responses content-array output and keep empty outputs from
# tripping the stricter ChatMessage role validator.
output = _responses_tool_output_text(item.output)
messages.append(
ChatMessage(
role = "tool",
@ -4783,7 +5237,9 @@ async def _responses_stream(
detail = "Image provided but current GGUF model does not support vision.",
)
body = _build_openai_passthrough_body(chat_req, backend_ctx = llama_backend.context_length)
body = _build_openai_passthrough_body(
chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
)
body["stream_options"] = {"include_usage": True}
target_url = f"{llama_backend.base_url}/v1/chat/completions"
@ -6489,7 +6945,11 @@ def _extract_response_format(payload):
return rf if isinstance(rf, dict) else None
def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict:
def _build_openai_passthrough_body(
payload,
backend_ctx = None,
llama_backend = None,
) -> dict:
"""Assemble the llama-server request body from a ChatCompletionRequest.
Only known OpenAI / llama-server fields are forwarded, so Studio-specific
@ -6500,12 +6960,19 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict:
system_prompt, _, _ = _extract_content_parts(payload.messages)
messages = _set_or_prepend_system_message(messages, system_prompt)
tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto"
# When the caller asked for a specific reasoning mode, forward it via
# chat_template_kwargs so the Jinja template renders with (or without) the
# reasoning preamble.
tpl_kwargs = None
if payload.enable_thinking is not None:
tpl_kwargs = {"enable_thinking": bool(payload.enable_thinking)}
# Forward per-request reasoning fields (enable_thinking / reasoning_effort /
# preserve_thinking) via chat_template_kwargs so the Jinja template renders
# in the caller's mode, gated on the active template's capabilities exactly
# like the non-passthrough paths.
tpl_kwargs = (
llama_backend._request_reasoning_kwargs(
payload.enable_thinking,
payload.reasoning_effort,
payload.preserve_thinking,
)
if llama_backend is not None
else None
)
return _build_passthrough_payload(
messages,
payload.tools,
@ -6540,7 +7007,9 @@ async def _openai_passthrough_stream(
the client sees a standard OpenAI response.
"""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length)
body = _build_openai_passthrough_body(
payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
)
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
@ -6558,27 +7027,32 @@ async def _openai_passthrough_stream(
limits = httpx.Limits(max_keepalive_connections = 0),
)
resp = None
try:
req = client.build_request("POST", target_url, json = body)
resp = await client.send(req, stream = True)
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable.
logger.error("openai passthrough stream: upstream unreachable: %s", e)
if resp is not None:
_truncate_budget = (
_OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0
)
while True:
try:
req = client.build_request("POST", target_url, json = body)
resp = await client.send(req, stream = True)
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable.
logger.error("openai passthrough stream: upstream unreachable: %s", e)
if resp is not None:
try:
await resp.aclose()
except Exception:
pass
try:
await resp.aclose()
await client.aclose()
except Exception:
pass
try:
await client.aclose()
except Exception:
pass
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
if resp.status_code != 200:
if resp.status_code == 200:
break
err_bytes = await resp.aread()
err_text = err_bytes.decode("utf-8", errors = "replace")
logger.error(
@ -6591,6 +7065,14 @@ async def _openai_passthrough_stream(
await resp.aclose()
except Exception:
pass
# Opt-in overflow policy: shrink and retry instead of a fatal 400.
if (
_truncate_budget > 0
and _classify_llama_generation_error(Exception(err_text))
and _apply_overflow_truncation(body, err_text)
):
_truncate_budget -= 1
continue
try:
await client.aclose()
except Exception:
@ -6685,22 +7167,37 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
``tool_calls``, and accurate ``usage`` token counts.
"""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length)
body = _build_openai_passthrough_body(
payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
)
try:
async with httpx.AsyncClient() as client:
resp = await client.post(target_url, json = body, timeout = 600)
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable. Surface the
# same friendly message the sync chat path emits so operators don't see
# a bare 500 with no diagnostic.
logger.error("openai passthrough non-streaming: upstream unreachable: %s", e)
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
_truncate_budget = (
_OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0
)
while True:
try:
async with httpx.AsyncClient() as client:
resp = await client.post(target_url, json = body, timeout = 600)
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable. Surface the
# same friendly message the sync chat path emits so operators don't see
# a bare 500 with no diagnostic.
logger.error("openai passthrough non-streaming: upstream unreachable: %s", e)
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
if resp.status_code != 200:
if resp.status_code == 200:
break
# Opt-in overflow policy: shrink and retry instead of a fatal 400.
if (
_truncate_budget > 0
and _classify_llama_generation_error(Exception(resp.text))
and _apply_overflow_truncation(body, resp.text)
):
_truncate_budget -= 1
continue
raise _openai_passthrough_error(resp.status_code, resp.text)
# The guided-decoding fence wraps each choice's JSON content in a

85
studio/backend/routes/llama.py Normal file
View file

@ -0,0 +1,85 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""llama.cpp prebuilt update endpoints.
GET /api/llama/update-status -> is a newer prebuilt available + job state
POST /api/llama/update -> download + atomically swap to the latest
Detection reuses utils.llama_cpp_freshness; the swap reuses
install_llama_prebuilt.py via utils.llama_cpp_update. Both fail open so the UI
never blocks on a missing marker / offline GitHub.
"""
from __future__ import annotations
import asyncio
from typing import Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject
from utils.llama_cpp_update import get_update_status, start_update
router = APIRouter()
class LlamaUpdateJob(BaseModel):
state: str = Field("idle", description = "idle | running | success | error")
message: str = ""
from_tag: Optional[str] = None
to_tag: Optional[str] = None
error: Optional[str] = None
progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.")
started_at: Optional[str] = None
finished_at: Optional[str] = None
class LlamaUpdateStatusResponse(BaseModel):
supported: bool = Field(
False,
description = "True when the install came from an Unsloth prebuilt (has a marker).",
)
update_available: bool = Field(
False, description = "True when the latest release is genuinely newer than the install."
)
stale: bool = Field(
False, description = "Update available AND install older than the staleness threshold."
)
installed_tag: Optional[str] = None
latest_tag: Optional[str] = None
published_repo: Optional[str] = None
installed_at_utc: Optional[str] = None
age_days: Optional[int] = None
source_build: bool = Field(
False, description = "True when there is no marker (source build) but a prebuilt is offered."
)
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
class LlamaUpdateActionResponse(BaseModel):
started: bool
reason: Optional[str] = None
message: Optional[str] = None
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
@router.get("/update-status", response_model = LlamaUpdateStatusResponse)
async def llama_update_status(
force_refresh: bool = Query(
False, description = "Bypass the 24h release cache for an explicit check."
),
current_subject: str = Depends(get_current_subject),
) -> LlamaUpdateStatusResponse:
# Off the event loop: detection may probe the host and read GitHub.
status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh)
return LlamaUpdateStatusResponse(**status)
@router.post("/update", response_model = LlamaUpdateActionResponse)
async def llama_update(
current_subject: str = Depends(get_current_subject),
) -> LlamaUpdateActionResponse:
action = await asyncio.to_thread(start_update)
return LlamaUpdateActionResponse(**action)

43
studio/backend/routes/mcp_servers.py
View file

@ -18,8 +18,11 @@ from core.inference.mcp_client import (
probe_timeout,
stdio_mcp_enabled,
)
from core.inference.mcp_config_import import parse_mcp_config
from models.mcp_servers import (
McpServerCreate,
McpServerImportRequest,
McpServerImportResult,
McpServerProbeResult,
McpServerResponse,
McpServerTestRequest,
@ -240,6 +243,46 @@ async def refresh_mcp_server_tools(
return McpServerProbeResult(ok = True, tool_count = len(tools))
@router.post("/import", response_model = McpServerImportResult)
async def import_mcp_servers(
payload: McpServerImportRequest, current_subject: str = Depends(get_current_subject)
):
"""Bulk-register servers from a standard mcpServers JSON config (issue
#5936). Each entry rides the existing create path: _validate_url applies
the same stdio gate (a stdio entry becomes a per-entry error when stdio is
off; http still imports), and entries whose url already exists are skipped
so re-importing the same file is idempotent. One bad entry never 400s the
whole batch -- failures are reported per entry."""
entries, errors = parse_mcp_config(payload.config)
created: list[McpServerResponse] = []
skipped: list[str] = []
seen_urls = {row["url"] for row in mcp_servers_db.list_servers()}
for entry in entries:
try:
url = _validate_url(entry.url)
except HTTPException as exc:
errors.append(f"{entry.display_name}: {exc.detail}")
continue
if url in seen_urls:
skipped.append(entry.display_name)
continue
headers = _normalize_headers(entry.headers)
server_id = uuid.uuid4().hex[:16]
mcp_servers_db.create_server(
id = server_id,
display_name = entry.display_name,
url = url,
headers_json = json.dumps(headers) if headers else None,
is_enabled = entry.is_enabled,
use_oauth = entry.use_oauth and not is_stdio(url),
)
seen_urls.add(url)
created.append(_row_to_response(mcp_servers_db.get_server(server_id)))
return McpServerImportResult(created = created, skipped = skipped, errors = errors)
@router.post("/test", response_model = McpServerProbeResult)
async def test_mcp_server(
payload: McpServerTestRequest, current_subject: str = Depends(get_current_subject)

31
studio/backend/routes/training.py
View file

@ -665,10 +665,17 @@ async def stream_training_progress(
# If not active, send final state and exit
if not is_active:
if backend.step_history:
final_step = backend.step_history[-1]
_live = (getattr(tp, "step", 0) or 0) if tp else 0
if backend.step_history or _live > 0:
final_step = backend.step_history[-1] if backend.step_history else 0
final_loss = backend.loss_history[-1] if backend.loss_history else None
final_lr = backend.lr_history[-1] if backend.lr_history else None
# Histories skip non-finite steps; report the live step with
# loss=None instead of the last finite pair.
if _live > final_step:
final_step = _live
final_loss = getattr(tp, "loss", None)
final_lr = getattr(tp, "learning_rate", final_lr)
final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step
final_epoch = getattr(tp, "epoch", None) if tp else None
payload = build_progress(
@ -697,11 +704,18 @@ async def stream_training_progress(
while backend.is_training_active():
try:
if backend.step_history:
current_step = backend.step_history[-1]
tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0
if backend.step_history or live_step > 0:
current_step = backend.step_history[-1] if backend.step_history else 0
current_loss = backend.loss_history[-1] if backend.loss_history else None
current_lr = backend.lr_history[-1] if backend.lr_history else None
tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
# Histories skip non-finite steps; follow the live progress
# step and report its loss (None until it recovers).
if live_step > current_step:
current_step = live_step
current_loss = getattr(tp_inner, "loss", None)
current_lr = getattr(tp_inner, "learning_rate", current_lr)
current_total_steps = (
getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step
)
@ -798,6 +812,13 @@ async def stream_training_progress(
final_loss = backend.loss_history[-1] if backend.loss_history else None
final_lr = backend.lr_history[-1] if backend.lr_history else None
final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
# If the run ended on a non-finite stretch, report the live step with
# loss=None instead of rolling back to the last finite pair.
_final_live_step = (getattr(final_tp, "step", 0) or 0) if final_tp else 0
if _final_live_step > (final_step if final_step is not None else -1):
final_step = _final_live_step
final_loss = getattr(final_tp, "loss", None)
final_lr = getattr(final_tp, "learning_rate", final_lr)
final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step
final_epoch = getattr(final_tp, "epoch", None) if final_tp else None
final_payload = build_progress(

167
studio/backend/run.py
View file

@ -8,6 +8,7 @@ Self-contained; can be moved to any directory.
import os
import sys
import time
from pathlib import Path
from typing import Optional
@ -231,6 +232,9 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
public internet. Synchronous so output lands between the banner URLs and the
stop hint. Bounded at ~15s; failures swallowed (verifier failing != Studio
failing). Only meaningful for a wildcard bind."""
global _public_reachable
# Reset to "unknown" each run; set True/False only when the probe decides.
_public_reachable = None
import ipaddress
import json
import time
@ -323,12 +327,14 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
print("", flush = True)
if ok_nodes:
_public_reachable = True
print(
f"{ok_c} Reachability check: {url}/ is reachable from the "
f"public internet ({ok_nodes}/{total} probe nodes connected).{reset}",
flush = True,
)
elif err_nodes:
_public_reachable = False
print(
f"{err_c} Reachability check: {url}/ is NOT reachable from "
f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}",
@ -413,9 +419,31 @@ def _emit_startup_output(host: str, port: int, display_host: str) -> None:
print_studio_stop_hint()
elif wildcard_bind:
_verify_global_reachability(display_host, port)
_print_cloudflare_line()
print_studio_stop_hint()
def _print_cloudflare_line() -> None:
"""Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up.
Reads the module-level URL set by ``run_server``. Prints nothing when the
tunnel is disabled or failed -- failures are silently ignored. When the public
reachability probe just failed (``_public_reachable is False``) but the tunnel
is up, reword to point the user at the Cloudflare link as the way in.
"""
if not _cloudflare_url:
return
from startup_banner import stdout_supports_color
accent = "\033[38;5;150;1m"
reset = "\033[0m"
if _public_reachable is False:
line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}"
else:
line = f" Secure link access via Cloudflare: {_cloudflare_url}"
print(f"{accent}{line}{reset}" if stdout_supports_color() else line)
def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
"""Return (pid, process_name) listening on *port*, or None.
@ -583,6 +611,13 @@ def _graceful_shutdown(server = None):
except Exception as e:
logger.warning("Error shutting down llama-server: %s", e)
# 6. Stop the Cloudflare tunnel (if started).
try:
from cloudflare_tunnel import stop_studio_tunnel
stop_studio_tunnel()
except Exception as e:
logger.warning("Error stopping Cloudflare tunnel: %s", e)
logger.info("All subprocesses cleaned up")
@ -593,6 +628,16 @@ _server = None
# Shutdown event -- wakes the main loop on signal.
_shutdown_event = None
# trycloudflare.com URL for 0.0.0.0 binds (set by run_server, read by the banner);
# None when there is no tunnel (loopback, disabled, or a silently-ignored failure).
_cloudflare_url = None
# Public reachability from the last _verify_global_reachability run, read by the
# Cloudflare banner line. True when the public ip:port probe confirmed reachable,
# False when it confirmed NOT reachable, None when the probe did not run or could
# not decide (timeout, blocked, private address).
_public_reachable = None
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
@ -673,6 +718,94 @@ def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Pa
return None, attempted
class _TeeStream:
"""Mirror writes to the original stream and a session log file.
Console behavior is unchanged (writes/returns delegate to the original
stream; Tauri's structured-stdout protocol and isatty probes see exactly
what they saw before). The file copy is best-effort: a full disk or a
closed handle must never break the console."""
def __init__(self, stream, log_fh):
self._stream = stream
self._log_fh = log_fh
def write(self, data):
try:
self._log_fh.write(data)
except Exception:
pass
return self._stream.write(data)
def flush(self):
try:
self._log_fh.flush()
except Exception:
pass
try:
self._stream.flush()
except Exception:
pass
def __getattr__(self, name):
return getattr(self._stream, name)
def _setup_server_disk_logging():
"""Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim
faulthandler at the same file so hard crashes (access violations /
SIGSEGV in the GPU runtime) leave a stack trace on disk.
Also exports PYTHONFAULTHANDLER=1 so child Python processes (training
workers) dump native-crash stacks to their captured stderr. Keeps the
newest 20 session logs. Opt out with UNSLOTH_STUDIO_NO_FILE_LOG=1.
Returns the log path, or None when disabled/unavailable.
"""
if os.environ.get("UNSLOTH_STUDIO_NO_FILE_LOG") == "1":
return None
try:
from utils.paths import studio_root
log_dir = Path(studio_root()) / "logs" / "server"
except Exception:
home = (
os.environ.get("UNSLOTH_STUDIO_HOME")
or os.environ.get("STUDIO_HOME")
or os.path.join(os.path.expanduser("~"), ".unsloth", "studio")
)
log_dir = Path(home) / "logs" / "server"
try:
log_dir.mkdir(parents = True, exist_ok = True)
stamp = time.strftime("%Y%m%d-%H%M%S")
log_path = log_dir / f"server-{stamp}-pid{os.getpid()}.log"
# Line-buffered so the tail survives a hard kill; errors="replace"
# so a console encoding quirk can never take the server down.
log_fh = open(log_path, "w", encoding = "utf-8", errors = "replace", buffering = 1)
except Exception:
return None
import faulthandler
try:
faulthandler.enable(file = log_fh, all_threads = True)
except Exception:
pass
# Children (training workers) inherit: their native-crash stacks land on
# the stderr the server already captures.
os.environ.setdefault("PYTHONFAULTHANDLER", "1")
sys.stdout = _TeeStream(sys.stdout, log_fh)
sys.stderr = _TeeStream(sys.stderr, log_fh)
# Best-effort retention: keep the newest 20 session logs.
try:
logs = sorted(log_dir.glob("server-*.log"), key = lambda p: p.stat().st_mtime)
for old in logs[:-20]:
old.unlink(missing_ok = True)
except Exception:
pass
return log_path
def run_server(
host: str = "127.0.0.1",
port: int = 8888,
@ -680,6 +813,7 @@ def run_server(
silent: bool = False,
api_only: bool = False,
llama_parallel_slots: int = 1,
cloudflare: bool = True,
):
"""
Start the FastAPI server.
@ -705,6 +839,16 @@ def run_server(
except Exception:
pass
# Persist a session log + native-crash stacks BEFORE importing main, so
# even import-time failures leave evidence on disk. Field report: Studio
# "terminates without a warning" -- a native crash in the GPU runtime
# kills the process with no Python traceback, and a desktop-shortcut
# console closes before anything can be read. Console-only logging made
# that undiagnosable.
_session_log = _setup_server_disk_logging()
if _session_log is not None and not silent:
print(f"Session log: {_session_log}")
# Set env var BEFORE importing main so CORS middleware picks it up.
if api_only:
os.environ["UNSLOTH_API_ONLY"] = "1"
@ -874,6 +1018,21 @@ def run_server(
if api_only:
print(f"TAURI_PORT={port}", flush = True)
# Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often
# unreachable). Started pre-banner and even when silent so the CLI banner can
# read app.state.cloudflare_url; torn down by _graceful_shutdown.
global _cloudflare_url
_cloudflare_url = None
app.state.cloudflare_url = None
_cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel
_cloudflare_url = start_studio_tunnel(port)
app.state.cloudflare_url = _cloudflare_url
except Exception as e:
logger.debug("Cloudflare tunnel skipped: %s", e)
if not silent:
_emit_startup_output(host, port, display_host)
@ -912,6 +1071,13 @@ if __name__ == "__main__":
action = "store_true",
help = "API server only, no frontend (for Tauri)",
)
parser.add_argument(
"--cloudflare",
action = argparse.BooleanOptionalAction,
default = True,
help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 "
"(default on; --no-cloudflare to disable)",
)
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct
# backend launches; `unsloth studio run` always passes its own value (4).
_PARALLEL_MIN = 1
@ -938,6 +1104,7 @@ if __name__ == "__main__":
silent = args.silent,
api_only = args.api_only,
llama_parallel_slots = args.parallel,
cloudflare = args.cloudflare,
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)

479
studio/backend/tests/test_cloudflare_tunnel.py Normal file
View file

@ -0,0 +1,479 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the Cloudflare quick-tunnel helper and run.py wiring.
cloudflare_tunnel.py is stdlib-only (storage_roots is imported lazily), so it
loads via spec_from_file_location without the studio venv. run.py defaults are
checked by AST so we never import its heavy deps (uvicorn/structlog).
"""
import ast
import importlib.util
import io
import sys
import tarfile
import types
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parent.parent
_CT_PY = _BACKEND / "cloudflare_tunnel.py"
_RUN_PY = _BACKEND / "run.py"
def _load_ct():
spec = importlib.util.spec_from_file_location("cloudflare_tunnel", _CT_PY)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
ct = _load_ct()
# ── URL parsing ──────────────────────────────────────────────────────
def test_url_regex_extracts_and_ignores_noise():
blob = (
"2026-06-11T10:00:00Z INF Thank you for trying Cloudflare Tunnel.\n"
"2026-06-11T10:00:01Z INF Requesting new quick Tunnel on trycloudflare.com...\n"
"2026-06-11T10:00:01Z INF | https://setting-democracy-gathering.trycloudflare.com |\n"
"2026-06-11T10:00:02Z INF Registered tunnel connection https://not-the-url.example.com\n"
)
m = ct._URL_RE.search(blob)
assert m is not None
assert m.group(0) == "https://setting-democracy-gathering.trycloudflare.com"
def test_url_regex_no_match_on_unrelated():
assert ct._URL_RE.search("INF connecting to https://api.cloudflare.com/v4") is None
# ── asset mapping ────────────────────────────────────────────────────
@pytest.mark.parametrize(
"system,machine,expected",
[
("Linux", "x86_64", ("cloudflared-linux-amd64", False)),
("Linux", "aarch64", ("cloudflared-linux-arm64", False)),
("Darwin", "arm64", ("cloudflared-darwin-arm64.tgz", True)),
("Darwin", "x86_64", ("cloudflared-darwin-amd64.tgz", True)),
("Windows", "AMD64", ("cloudflared-windows-amd64.exe", False)),
("Windows", "x86", ("cloudflared-windows-386.exe", False)),
("Linux", "mips", None),
("Plan9", "x86_64", None),
],
)
def test_asset_name(monkeypatch, system, machine, expected):
monkeypatch.setattr(ct.platform, "system", lambda: system)
monkeypatch.setattr(ct.platform, "machine", lambda: machine)
assert ct._asset_name() == expected
# ── binary discovery ─────────────────────────────────────────────────
def test_find_cloudflared_prefers_path(monkeypatch):
monkeypatch.setattr(ct.shutil, "which", lambda name: "/usr/local/bin/cloudflared")
assert ct.find_cloudflared() == "/usr/local/bin/cloudflared"
def test_find_cloudflared_falls_back_to_cache(monkeypatch, tmp_path):
cached = tmp_path / "cloudflared"
cached.write_text("#!/bin/sh\n")
cached.chmod(0o755)
monkeypatch.setattr(ct.shutil, "which", lambda name: None)
monkeypatch.setattr(ct, "_cache_path", lambda: cached)
assert ct.find_cloudflared() == str(cached)
def test_find_cloudflared_none_when_missing(monkeypatch, tmp_path):
monkeypatch.setattr(ct.shutil, "which", lambda name: None)
monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "absent")
assert ct.find_cloudflared() is None
# ── ensure / download ────────────────────────────────────────────────
def test_ensure_downloads_and_chmods_when_missing(monkeypatch, tmp_path):
cached = tmp_path / "cloudflared"
monkeypatch.setattr(ct, "find_cloudflared", lambda: None)
monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-linux-amd64", False))
monkeypatch.setattr(ct, "_cache_path", lambda: cached)
def fake_download(url, dest):
assert url.endswith("/cloudflared-linux-amd64")
dest.write_bytes(b"ELF-ish")
return True
monkeypatch.setattr(ct, "_download", fake_download)
monkeypatch.setattr(ct.sys, "platform", "linux")
path = ct.ensure_cloudflared()
assert path == str(cached)
assert cached.exists()
assert cached.stat().st_mode & 0o111 # executable bit set
def test_ensure_returns_none_on_download_failure(monkeypatch, tmp_path):
monkeypatch.setattr(ct, "find_cloudflared", lambda: None)
monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-linux-amd64", False))
monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "cloudflared")
monkeypatch.setattr(ct, "_download", lambda url, dest: False)
assert ct.ensure_cloudflared() is None
def test_ensure_returns_none_for_unsupported_arch(monkeypatch, tmp_path):
monkeypatch.setattr(ct, "find_cloudflared", lambda: None)
monkeypatch.setattr(ct, "_asset_name", lambda: None)
monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "cloudflared")
assert ct.ensure_cloudflared() is None
def test_download_sets_user_agent(monkeypatch, tmp_path):
import urllib.request
captured = {}
class _Resp:
_sent = False
def __enter__(self):
return self
def __exit__(self, *a):
return False
def read(self, n = -1):
if self._sent:
return b""
self._sent = True
return b"payload"
def fake_urlopen(req, timeout = None):
captured["ua"] = req.get_header("User-agent")
return _Resp()
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
dest = tmp_path / "cloudflared"
assert ct._download("https://github.com/cloudflare/cloudflared/x", dest) is True
assert captured["ua"] == "unsloth-studio" # GitHub CDN 403s the default UA
assert dest.read_bytes() == b"payload"
# ── cross-platform: Windows (.exe), macOS (.tgz) ─────────────────────
def test_cache_path_uses_exe_on_windows(monkeypatch, tmp_path):
import types
fake_sr = types.ModuleType("utils.paths.storage_roots")
fake_sr.studio_bin_root = lambda: tmp_path
monkeypatch.setitem(sys.modules, "utils.paths.storage_roots", fake_sr)
monkeypatch.setattr(ct.sys, "platform", "win32")
assert ct._cache_path() == tmp_path / "cloudflared.exe"
def test_ensure_windows_downloads_exe(monkeypatch, tmp_path):
cached = tmp_path / "cloudflared.exe"
monkeypatch.setattr(ct, "find_cloudflared", lambda: None)
monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-windows-amd64.exe", False))
monkeypatch.setattr(ct, "_cache_path", lambda: cached)
monkeypatch.setattr(ct.sys, "platform", "win32")
def fake_download(url, dest):
assert url.endswith("/cloudflared-windows-amd64.exe")
dest.write_bytes(b"MZ") # PE header magic
return True
monkeypatch.setattr(ct, "_download", fake_download)
# chmod is skipped on Windows; would raise on a path that does not exist yet.
monkeypatch.setattr(ct.os, "chmod", lambda *a, **k: pytest.fail("chmod called on win32"))
assert ct.ensure_cloudflared() == str(cached)
assert cached.read_bytes() == b"MZ"
def test_ensure_macos_extracts_tgz_and_chmods(monkeypatch, tmp_path):
cached = tmp_path / "cloudflared"
monkeypatch.setattr(ct, "find_cloudflared", lambda: None)
monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-darwin-arm64.tgz", True))
monkeypatch.setattr(ct, "_cache_path", lambda: cached)
monkeypatch.setattr(ct.sys, "platform", "darwin")
def fake_download(url, dest):
# dest is cached.with_suffix(".tgz"); write a real archive there.
assert url.endswith("/cloudflared-darwin-arm64.tgz")
with tarfile.open(dest, "w:gz") as tar:
data = b"mach-o"
info = tarfile.TarInfo(name = "cloudflared")
info.size = len(data)
tar.addfile(info, io.BytesIO(data))
return True
monkeypatch.setattr(ct, "_download", fake_download)
path = ct.ensure_cloudflared()
assert path == str(cached)
assert cached.read_bytes() == b"mach-o"
assert cached.stat().st_mode & 0o111 # chmod applied on posix
assert not cached.with_suffix(".tgz").exists() # temp archive cleaned up
# ── .tgz extraction (darwin) ─────────────────────────────────────────
def _make_tgz(
tmp_path,
member_name,
data = b"bin",
):
tgz = tmp_path / "cf.tgz"
with tarfile.open(tgz, "w:gz") as tar:
info = tarfile.TarInfo(name = member_name)
info.size = len(data)
tar.addfile(info, io.BytesIO(data))
return tgz
def test_tgz_extraction_extracts_clean_member(tmp_path):
tgz = _make_tgz(tmp_path, "cloudflared")
dest = tmp_path / "out"
assert ct._extract_tgz_member(tgz, dest) is True
assert dest.read_bytes() == b"bin"
def test_tgz_extraction_rejects_traversal(tmp_path):
tgz = _make_tgz(tmp_path, "../cloudflared")
dest = tmp_path / "out"
assert ct._extract_tgz_member(tgz, dest) is False
assert not dest.exists()
def test_tgz_extraction_missing_member(tmp_path):
tgz = _make_tgz(tmp_path, "README")
dest = tmp_path / "out"
assert ct._extract_tgz_member(tgz, dest) is False
# ── tunnel lifecycle ─────────────────────────────────────────────────
class _FakePopen:
def __init__(self):
self.terminated = False
self.killed = False
self._alive = True
def poll(self):
return None if self._alive else 0
def terminate(self):
self.terminated = True
self._alive = False
def wait(self, timeout = None):
if self._alive:
raise ct.subprocess.TimeoutExpired(cmd = "cloudflared", timeout = timeout)
return 0
def kill(self):
self.killed = True
self._alive = False
def test_stop_terminates_process():
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
fake = _FakePopen()
t._proc = fake
t.stop()
assert fake.terminated is True
assert t._proc is None
# second stop is a no-op (idempotent)
t.stop()
def test_wait_for_url_times_out_without_blocking():
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
assert t.wait_for_url(timeout = 0.05) is None
def test_start_studio_tunnel_no_binary(monkeypatch):
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None)
assert ct.start_studio_tunnel(8080) is None
def test_start_studio_tunnel_registers_before_wait(monkeypatch):
# The tunnel must be visible to stop_studio_tunnel() during the URL wait,
# else a shutdown in that window orphans cloudflared.
seen = {}
class _Stub:
def __init__(self, port, binary):
self.url = None
def start(self):
pass
def wait_for_url(self, timeout):
seen["active_during_wait"] = ct._active_tunnel is self
self.url = "https://x.trycloudflare.com"
return self.url
def stop(self):
seen["stopped"] = True
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
try:
assert ct.start_studio_tunnel(8080) == "https://x.trycloudflare.com"
assert seen["active_during_wait"] is True
finally:
ct.stop_studio_tunnel()
def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch):
seen = {}
class _Stub:
def __init__(self, port, binary):
self.url = None
def start(self):
pass
def wait_for_url(self, timeout):
return None
def stop(self):
seen["stopped"] = True
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
assert ct.start_studio_tunnel(8080) is None
assert seen.get("stopped") is True
assert ct._active_tunnel is None
def test_start_studio_tunnel_returns_url(monkeypatch):
class _StubTunnel:
def __init__(self, port, binary):
self.url = None
def start(self):
self.url = "https://stub-xyz.trycloudflare.com"
def wait_for_url(self, timeout):
return self.url
def stop(self):
pass
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _StubTunnel)
try:
assert ct.start_studio_tunnel(8080) == "https://stub-xyz.trycloudflare.com"
finally:
ct.stop_studio_tunnel()
# ── run.py source-level pins (AST / source, no heavy import) ─────────
def _func_param_defaults(source, func_name):
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
args = node.args.args
defaults = node.args.defaults
offset = len(args) - len(defaults)
out = {}
for i, d in enumerate(defaults):
if isinstance(d, ast.Constant):
out[args[offset + i].arg] = d.value
return out
return {}
def _argparse_default(source, option):
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
if node.func.attr == "add_argument" and node.args:
a0 = node.args[0]
if isinstance(a0, ast.Constant) and a0.value == option:
for kw in node.keywords:
if kw.arg == "default" and isinstance(kw.value, ast.Constant):
return kw.value.value
return None
def test_run_server_cloudflare_default_true():
defaults = _func_param_defaults(_RUN_PY.read_text(), "run_server")
assert defaults.get("cloudflare") is True
def test_argparse_cloudflare_default_true():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True
def test_run_server_gates_tunnel_on_wildcard():
# Guard against accidentally widening the trigger beyond 0.0.0.0.
source = _RUN_PY.read_text()
assert "_cloudflare_enabled" in source
assert 'host == "0.0.0.0"' in source
def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable):
"""Exec the real _print_cloudflare_line source in isolation (run.py has heavy
deps), with the two module globals injected and startup_banner stubbed."""
src = _RUN_PY.read_text()
tree = ast.parse(src)
func_src = next(
ast.get_source_segment(src, n)
for n in ast.walk(tree)
if isinstance(n, ast.FunctionDef) and n.name == "_print_cloudflare_line"
)
stub = types.ModuleType("startup_banner")
stub.stdout_supports_color = lambda: False
monkeypatch.setitem(sys.modules, "startup_banner", stub)
captured: list[str] = []
ns = {
"_cloudflare_url": cloudflare_url,
"_public_reachable": public_reachable,
"print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)),
}
exec(compile(func_src, "<print_cloudflare_line>", "exec"), ns)
ns["_print_cloudflare_line"]()
return "\n".join(captured)
def test_cloudflare_line_reworded_when_public_unreachable(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = False
)
assert "Use the secure link access via Cloudflare instead: https://x.trycloudflare.com" in out
def test_cloudflare_line_default_wording_when_reachable(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = True
)
assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out
assert "Use the secure link" not in out
def test_cloudflare_line_default_wording_when_unknown(monkeypatch):
# Probe did not run / could not decide -> keep the existing wording.
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = None
)
assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out
assert "Use the secure link" not in out
def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch):
out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False)
assert out == ""

277
studio/backend/tests/test_context_overflow_truncation.py Normal file
View file

@ -0,0 +1,277 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the opt-in ``context_overflow="truncate_middle"`` passthrough policy.
On ``exceed_context_size_error`` the passthrough drops middle turn-groups and
retries inside the real window instead of surfacing a fatal 400. Truncation
keeps the system prompt, the first turn, and recent turns, and never orphans
a tool result from its tool_calls turn. Also covers ``/v1/models`` exposing
the real post-readback context window.
"""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from routes.inference import (
_apply_overflow_truncation,
_clip_long_contents,
_CLIP_MARKER,
_estimate_message_tokens,
_openai_model_objects,
_overflow_truncation_requested,
_parse_overflow_counts,
_truncate_middle_messages,
)
import routes.inference as routes_mod
# Nick's actual error body from the Discord report logs.
_NICK_ERROR = (
'{"detail":"llama-server error: {\\"error\\":{\\"code\\":400,'
'\\"message\\":\\"request (70494 tokens) exceeds the available context size '
'(67584 tokens), try increasing it\\",\\"type\\":\\"exceed_context_size_error\\",'
'\\"n_prompt_tokens\\":70494,\\"n_ctx\\":67584}}"}'
)
def _tool_turn(i: int, result_chars: int = 400) -> list[dict]:
"""An assistant tool_calls turn paired with its tool result."""
return [
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": f"call_{i}",
"type": "function",
"function": {"name": "read", "arguments": f'{{"filePath":"/f{i}"}}'},
}
],
},
{"role": "tool", "tool_call_id": f"call_{i}", "content": "x" * result_chars},
]
def _conversation(n_tool_turns: int = 12) -> list[dict]:
msgs = [
{"role": "system", "content": "You are an agent." * 20},
{"role": "user", "content": "Do the big task." * 20},
]
for i in range(n_tool_turns):
msgs.extend(_tool_turn(i))
msgs.append({"role": "assistant", "content": "halfway summary"})
msgs.append({"role": "user", "content": "keep going"})
return msgs
# ---------------------------------------------------------------------------
# _parse_overflow_counts
# ---------------------------------------------------------------------------
def test_parse_overflow_counts_nick_error():
assert _parse_overflow_counts(_NICK_ERROR) == (70494, 67584)
def test_parse_overflow_counts_missing_fields():
assert _parse_overflow_counts('{"error":"something else"}') is None
# ---------------------------------------------------------------------------
# _truncate_middle_messages
# ---------------------------------------------------------------------------
def test_truncation_drops_middle_keeps_anchors():
msgs = _conversation()
new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.5)
assert dropped > 0
assert len(new) == len(msgs) - dropped
# System prompt and task anchor survive.
assert new[0]["role"] == "system"
assert new[1] == msgs[1]
# The most recent turns survive verbatim.
assert new[-1] == msgs[-1]
assert new[-2] == msgs[-2]
def test_truncation_never_orphans_tool_results():
msgs = _conversation()
new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.4)
assert dropped > 0
surviving_call_ids = {
tc["id"] for m in new if m.get("role") == "assistant" for tc in (m.get("tool_calls") or [])
}
for m in new:
if m.get("role") == "tool":
assert m["tool_call_id"] in surviving_call_ids
def test_truncation_reduces_estimated_size_toward_target():
msgs = _conversation()
total = sum(_estimate_message_tokens(m) for m in msgs)
new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.5)
new_total = sum(_estimate_message_tokens(m) for m in new)
assert dropped > 0
assert new_total < total
# Should land at or below the requested share, modulo one whole group.
biggest_group = max(
_estimate_message_tokens(a) + _estimate_message_tokens(b)
for a, b in zip(msgs[2:-2:2], msgs[3:-2:2])
)
assert new_total <= int(total * 0.5) + biggest_group
def test_truncation_noop_when_keep_ratio_full():
msgs = _conversation()
new, dropped = _truncate_middle_messages(msgs, keep_ratio = 1.0)
assert dropped == 0
assert new == msgs
def test_truncation_noop_when_only_protected_turns_remain():
msgs = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "task"},
*_tool_turn(0),
{"role": "user", "content": "latest"},
]
new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.1)
assert dropped == 0
assert new == msgs
# ---------------------------------------------------------------------------
# _apply_overflow_truncation
# ---------------------------------------------------------------------------
def test_apply_overflow_truncation_mutates_body_and_clamps_max_tokens():
body = {"messages": _conversation(), "max_tokens": 32000}
assert _apply_overflow_truncation(body, _NICK_ERROR) is True
assert len(body["messages"]) < len(_conversation())
# Generation headroom: max_tokens clamped to the non-prompt share of n_ctx.
assert body["max_tokens"] <= max(1024, int(67584 * 0.25))
def test_apply_overflow_truncation_returns_false_when_nothing_droppable():
body = {
"messages": [
{"role": "system", "content": "sys"},
{"role": "user", "content": "task"},
{"role": "user", "content": "latest"},
],
"max_tokens": 32000,
}
assert _apply_overflow_truncation(body, _NICK_ERROR) is False
def test_apply_overflow_truncation_clips_giant_protected_tool_results():
"""One giant burst (few turn-groups, all protected) must still shrink:
stage 2 clips oversized tool contents instead of giving up."""
msgs = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "task"},
*_tool_turn(0, result_chars = 60000),
*_tool_turn(1, result_chars = 60000),
]
body = {"messages": msgs, "max_tokens": 32000}
n_before = len(msgs)
assert _apply_overflow_truncation(body, _NICK_ERROR) is True
# No message disappeared (pairing intact), but contents were clipped.
assert len(body["messages"]) == n_before
clipped = [m for m in body["messages"] if _CLIP_MARKER in str(m.get("content"))]
assert clipped, "expected at least one clipped tool result"
surviving_call_ids = {
tc["id"]
for m in body["messages"]
if m.get("role") == "assistant"
for tc in (m.get("tool_calls") or [])
}
for m in body["messages"]:
if m.get("role") == "tool":
assert m["tool_call_id"] in surviving_call_ids
def test_clip_long_contents_reaches_target_and_keeps_structure():
msgs = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "task"},
*_tool_turn(0, result_chars = 40000),
{"role": "user", "content": "latest question"},
]
total = sum(_estimate_message_tokens(m) for m in msgs)
clipped = _clip_long_contents(msgs, target_est = total // 4)
assert clipped >= 1
assert sum(_estimate_message_tokens(m) for m in msgs) <= total // 4
# Roles and count unchanged; the short final user message untouched.
assert [m["role"] for m in msgs] == ["system", "user", "assistant", "tool", "user"]
assert msgs[-1]["content"] == "latest question"
def test_overflow_truncation_requested_reads_field(monkeypatch):
monkeypatch.delenv("UNSLOTH_CONTEXT_OVERFLOW", raising = False)
class _P:
context_overflow = "truncate_middle"
class _Q:
context_overflow = None
assert _overflow_truncation_requested(_P()) is True
assert _overflow_truncation_requested(_Q()) is False
assert _overflow_truncation_requested(object()) is False
def test_overflow_truncation_server_default_env(monkeypatch):
"""UNSLOTH_CONTEXT_OVERFLOW enables the policy for clients that cannot
send custom body fields; an explicit per-request 'error' still wins."""
class _Unset:
context_overflow = None
class _ExplicitError:
context_overflow = "error"
monkeypatch.setenv("UNSLOTH_CONTEXT_OVERFLOW", "truncate_middle")
assert _overflow_truncation_requested(_Unset()) is True
assert _overflow_truncation_requested(_ExplicitError()) is False
monkeypatch.setenv("UNSLOTH_CONTEXT_OVERFLOW", "error")
assert _overflow_truncation_requested(_Unset()) is False
# ---------------------------------------------------------------------------
# /v1/models context metadata
# ---------------------------------------------------------------------------
class _FakeLlamaBackend:
is_loaded = True
model_identifier = "unsloth/Qwen3.6-27B-GGUF"
context_length = 67584
max_context_length = 262144
class _FakeEmptyBackend:
active_model_name = None
def test_v1_models_exposes_real_context_window(monkeypatch):
monkeypatch.setattr(routes_mod, "get_llama_cpp_backend", lambda: _FakeLlamaBackend())
monkeypatch.setattr(routes_mod, "get_inference_backend", lambda: _FakeEmptyBackend())
models = _openai_model_objects()
assert len(models) == 1
entry = models[0]
assert entry["id"] == "unsloth/Qwen3.6-27B-GGUF"
# The REAL (post /props readback) window, not the requested one.
assert entry["context_length"] == 67584
assert entry["max_context_length"] == 262144

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

@ -419,6 +419,8 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
routes_module.__path__ = []
settings_module = ModuleType("routes.settings")
settings_module.router = APIRouter()
llama_module = ModuleType("routes.llama")
llama_module.router = APIRouter()
prompts_module = ModuleType("routes.prompts")
prompts_module.router = APIRouter()
@ -440,9 +442,11 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
}.items():
setattr(routes_module, name, router)
routes_module.settings = settings_module
routes_module.llama = llama_module
monkeypatch.setitem(sys.modules, "routes", routes_module)
monkeypatch.setitem(sys.modules, "routes.settings", settings_module)
monkeypatch.setitem(sys.modules, "routes.llama", llama_module)
monkeypatch.setitem(sys.modules, "routes.prompts", prompts_module)
import studio.backend.main as backend_main

80
studio/backend/tests/test_external_provider_proxy_env.py Normal file
View file

@ -0,0 +1,80 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from pathlib import Path
import importlib.util
import sys
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_EXTERNAL_PROVIDER_PATH = (
Path(__file__).resolve().parent.parent / "core/inference/external_provider.py"
)
def _load_external_provider_module():
spec = importlib.util.spec_from_file_location(
"external_provider_under_test",
_EXTERNAL_PROVIDER_PATH,
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def test_shared_http_client_ignores_unsupported_proxy_scheme(monkeypatch):
ep_mod = _load_external_provider_module()
calls = []
class FakeAsyncClient:
def __init__(self, **kwargs):
calls.append(kwargs)
if kwargs.get("trust_env") is not False:
raise ValueError("Unknown scheme for proxy URL URL('socks4://127.0.0.1:12345')")
monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient)
client = ep_mod._create_shared_http_client()
assert isinstance(client, FakeAsyncClient)
assert calls == [{}, {"trust_env": False}]
def test_shared_http_client_ignores_missing_socksio(monkeypatch):
ep_mod = _load_external_provider_module()
calls = []
class FakeAsyncClient:
def __init__(self, **kwargs):
calls.append(kwargs)
if kwargs.get("trust_env") is not False:
raise ImportError(
"Using SOCKS proxy, but the 'socksio' package is not installed. "
"Make sure to install httpx using `pip install httpx[socks]`."
)
monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient)
client = ep_mod._create_shared_http_client()
assert isinstance(client, FakeAsyncClient)
assert calls == [{}, {"trust_env": False}]
def test_shared_http_client_reraises_other_value_errors(monkeypatch):
ep_mod = _load_external_provider_module()
class FakeAsyncClient:
def __init__(self, **kwargs):
raise ValueError("different httpx setup error")
monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient)
try:
ep_mod._create_shared_http_client()
except ValueError as exc:
assert str(exc) == "different httpx setup error"
else:
raise AssertionError("expected ValueError")

33
studio/backend/tests/test_inference_orchestrator_crash_message.py Normal file
View file

@ -0,0 +1,33 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from pathlib import Path
from types import SimpleNamespace
import importlib.util
import sys
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
def test_subprocess_crash_message_includes_signal_and_oom_hint():
spec = importlib.util.spec_from_file_location(
"inference_orchestrator_under_test",
Path(__file__).resolve().parent.parent / "core/inference/orchestrator.py",
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
orchestrator = module.InferenceOrchestrator.__new__(module.InferenceOrchestrator)
orchestrator._proc = SimpleNamespace(pid = 1234, exitcode = -9)
msg = orchestrator._subprocess_crash_message("wait")
assert msg.startswith("The inference worker stopped unexpectedly while loading the model.")
assert "memory pressure" in msg
assert "smaller model" in msg
assert "Details:" in msg
assert "signal=SIGKILL" in msg
assert "exitcode=-9" in msg

173
studio/backend/tests/test_install_resolve_prebuilt.py Normal file
View file

@ -0,0 +1,173 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode.
These back the in-app update for source-build (markerless) installs: the backend
asks the installer whether an official prebuilt exists for this host without
downloading. Network and host detection are stubbed; no GPU or internet needed.
"""
from __future__ import annotations
import importlib
import json
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
ilp = importlib.import_module("install_llama_prebuilt")
if not hasattr(ilp, "published_repo_for_host") or not hasattr(
ilp, "resolve_simple_install_release_plans"
):
pytest.skip("PR symbols not present - check branch", allow_module_level = True)
FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp
UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp
def _host(**kw):
base = dict(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = False,
is_macos = False,
is_x86_64 = False,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
macos_version = None,
)
base.update(kw)
return ilp.HostInfo(**base)
def test_published_repo_for_host():
# CPU-only Linux (x64 and arm64) -> ggml-org upstream.
assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM
assert (
ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64"))
== UPSTREAM
)
# GPU Linux -> fork.
assert (
ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True))
== FORK
)
assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK
# CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle).
assert (
ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True))
== UPSTREAM
)
# GPU Windows -> fork.
assert (
ilp.published_repo_for_host(
_host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True)
)
== FORK
)
# macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS).
assert (
ilp.published_repo_for_host(
_host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64")
)
== FORK
)
# Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling).
assert (
ilp.published_repo_for_host(
_host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True
)
== FORK
)
# The tooling hint is Linux-only: Windows CPU stays on ggml-org.
assert (
ilp.published_repo_for_host(
_host(system = "Windows", is_windows = True, is_x86_64 = True),
linux_amd_tooling_present = True,
)
== UPSTREAM
)
def _run_resolve(monkeypatch, capsys, plans_or_exc):
monkeypatch.setattr(
ilp,
"detect_host",
lambda: _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"),
)
def _resolver(tag, host, repo, published_release_tag):
if isinstance(plans_or_exc, Exception):
raise plans_or_exc
return ("b9585", plans_or_exc)
monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
monkeypatch.setattr(
sys,
"argv",
["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"],
)
rc = ilp.main()
assert rc == ilp.EXIT_SUCCESS
return json.loads(capsys.readouterr().out.strip().splitlines()[-1])
def test_resolve_prebuilt_available(monkeypatch, capsys):
plan = SimpleNamespace(
release_tag = "b9585",
llama_tag = "b9585",
attempts = [
SimpleNamespace(name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64")
],
)
out = _run_resolve(monkeypatch, capsys, [plan])
assert out["prebuilt_available"] is True
assert out["repo"] == FORK
assert out["release_tag"] == "b9585"
assert out["asset"] == "llama-b9585-bin-macos-arm64.tar.gz"
assert out["install_kind"] == "macos-arm64"
def test_resolve_prebuilt_unavailable(monkeypatch, capsys):
out = _run_resolve(monkeypatch, capsys, ilp.PrebuiltFallback("no macOS asset"))
assert out["prebuilt_available"] is False
assert out["repo"] == FORK
def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys):
# CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the
# fork so a HIP source build is not offered an upstream CPU prebuilt.
monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True))
monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo")
seen = {}
def _resolver(tag, host, repo, published_release_tag):
seen["repo"] = repo
raise ilp.PrebuiltFallback("no asset")
monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
monkeypatch.setattr(
sys,
"argv",
["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"],
)
assert ilp.main() == ilp.EXIT_SUCCESS
out = json.loads(capsys.readouterr().out.strip().splitlines()[-1])
assert seen["repo"] == FORK
assert out["repo"] == FORK

236
studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
View file

@ -1,10 +1,11 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validates that the installer resolves lemonade ROCm prebuilt assets.
"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets.
Uses a faked HostInfo so no AMD GPU is needed. The lemonade GitHub API calls
are stubbed so the suite runs offline and isn't subject to rate limits.
Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade
GitHub API are stubbed out so the suite runs without internet access and is
not subject to rate limits.
"""
from __future__ import annotations
@ -31,14 +32,19 @@ if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
@pytest.fixture(autouse = True)
def _clear_lemonade_release_cache():
"""Prevent cross-test pollution of the lemonade release lru_cache when
tests vary the fetch_json mock return value."""
"""Prevent cross-test pollution of the lemonade release lru_cache and
selection-log dedup set when tests vary the fetch_json mock return value."""
_cache = getattr(_mod, "_fetch_lemonade_release_cached", None)
_logged: set | None = getattr(_mod, "_lemonade_selection_logged", None)
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
if _logged is not None:
_logged.clear()
yield
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
if _logged is not None:
_logged.clear()
_STUB_TAG = "b1262"
@ -89,7 +95,9 @@ def _lookup_family(gfx: str) -> str | None:
return None
# ---------------------------------------------------------------------------
# GPU family mapping
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
@ -111,7 +119,9 @@ def test_unknown_gpu_not_in_families():
assert _lookup_family("gfx999") is None
# ---------------------------------------------------------------------------
# Asset resolution - hits real lemonade GitHub API
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
@ -141,59 +151,71 @@ def test_unknown_gpu_falls_through_to_upstream():
assert result is None
# Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts.
# This is the path setup.sh invokes (via --simple-policy), so the lemonade
# integration is useless if it isn't wired in here.
# ---------------------------------------------------------------------------
# The Linux attempt builder must plan a lemonade ROCm attempt for AMD-only hosts.
# This is the path setup.sh actually invokes (fork hosts now select from the
# manifest), so the lemonade integration is useless if it isn't wired in here.
# ---------------------------------------------------------------------------
direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None)
_linux_published_attempts = getattr(_mod, "_linux_published_attempts", None)
direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None)
PublishedLlamaArtifact = _mod.PublishedLlamaArtifact
PublishedReleaseBundle = _mod.PublishedReleaseBundle
def _stub_unsloth_release(release_tag: str = "b9022") -> dict:
# Minimal payload parse_direct_linux_release_bundle accepts. It needs at
# least one `app-{label}-linux-x64*.tar.gz` asset to recognise the bundle;
# we ship a bare CPU one so the planner has a baseline non-ROCm fallback.
asset_name = f"app-{release_tag}-linux-x64.tar.gz"
return {
"tag_name": release_tag,
"name": release_tag,
"assets": [
{
"name": asset_name,
"browser_download_url": f"https://example.invalid/{asset_name}",
},
],
}
def _rocm_bundle(gfx_family: str, mapped_targets: list[str]) -> "PublishedReleaseBundle":
"""A fork manifest bundle exposing a per-gfx linux-rocm artifact, so
published_rocm_choice_for_host can match the host before the lemonade
fallback is appended."""
asset_name = f"app-b9457-linux-x64-rocm-{gfx_family}.tar.gz"
artifact = PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = "linux-rocm",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
rank = 1000,
gfx_target = gfx_family,
mapped_targets = mapped_targets,
)
return PublishedReleaseBundle(
repo = "unslothai/llama.cpp",
release_tag = "v1.0",
upstream_tag = "b9457",
assets = {asset_name: f"https://example.invalid/{asset_name}"},
artifacts = [artifact],
)
@pytest.mark.skipif(
direct_linux_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
_linux_published_attempts is None,
reason = "Linux attempt builder not present on this branch",
)
def test_simple_policy_plans_lemonade_for_rocm_host():
def test_linux_attempts_include_fork_rocm_and_lemonade_for_rocm_host():
host = _make_rocm_host("gfx1151")
bundle = _rocm_bundle("gfx1151", ["gfx1151"])
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_linux_release_plan(
_stub_unsloth_release(),
host,
"unslothai/llama.cpp",
"latest",
)
assert plan is not None, "ROCm host should not be skipped by simple-policy planner"
kinds = [a.install_kind for a in plan.attempts]
assert (
"linux-rocm" in kinds
), f"simple-policy planner did not include a lemonade ROCm attempt; got {kinds}"
rocm_attempt = next(a for a in plan.attempts if a.install_kind == "linux-rocm")
assert rocm_attempt.source_label == "lemonade"
assert "gfx1151" in rocm_attempt.name
attempts = _linux_published_attempts(host, bundle, "latest")
kinds = [a.install_kind for a in attempts]
assert "linux-rocm" in kinds, f"builder did not include any linux-rocm attempt; got {kinds}"
sources = {a.source_label for a in attempts if a.install_kind == "linux-rocm"}
# The fork's own per-gfx bundle is preferred, with the lemonade prebuilt as
# the fallback -- both must be present for a covered ROCm host.
assert "published" in sources, f"fork ROCm bundle missing; got {sources}"
assert "lemonade" in sources, f"lemonade ROCm fallback missing; got {sources}"
lemonade_attempt = next(a for a in attempts if a.source_label == "lemonade")
assert "gfx1151" in lemonade_attempt.name
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
reason = "direct release planners not present on this branch",
)
def test_simple_policy_plans_lemonade_for_windows_hip_host():
def test_direct_upstream_plan_includes_lemonade_for_windows_hip_host():
host = _make_rocm_host("gfx1151", windows = True)
release = {
"tag_name": "b9022",
@ -204,16 +226,14 @@ def test_simple_policy_plans_lemonade_for_windows_hip_host():
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt"
kinds = [a.install_kind for a in plan.attempts]
assert (
"windows-hip" in kinds
), f"simple-policy planner did not include a lemonade HIP attempt; got {kinds}"
assert "windows-hip" in kinds, f"planner did not include a lemonade HIP attempt; got {kinds}"
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
reason = "direct release planners not present on this branch",
)
def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable():
def test_windows_hip_falls_back_to_upstream_when_lemonade_unavailable():
"""If lemonade returns None (e.g. gfx999 or transient API failure), the planner
must still include the upstream HIP asset rather than silently downgrading to CPU."""
host = _make_rocm_host("gfx999", windows = True)
@ -247,8 +267,9 @@ def test_lemonade_release_api_url_pinned_tag():
def test_lemonade_release_api_url_encodes_tag():
"""Slashes / hashes in the tag must be URL-encoded so the URL can't be
reshaped (defence in depth -- tags should already be sanitised upstream)."""
"""Unexpected slashes / hashes in the tag must be URL-encoded so the URL
cannot be reshaped (defence in depth -- tags should already be sanitised
upstream)."""
url = _mod._lemonade_release_api_for("b1260/../latest")
assert "/releases/tags/b1260%2F..%2Flatest" in url
assert "//latest" not in url.split("/releases/tags/", 1)[1]
@ -263,9 +284,9 @@ def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch):
def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
"""If the GitHub API response contained an off-host download URL, the
resolver must refuse it (lemonade assets aren't in the approved-hash
manifest)."""
"""If the GitHub API response somehow contained an off-host download URL,
the resolver must refuse to use it (lemonade assets are not in the
approved-hash manifest)."""
bad_release = {
"tag_name": _STUB_TAG,
"assets": [
@ -289,7 +310,7 @@ def test_lemonade_resolver_rejects_http_scheme():
def test_lemonade_resolver_accepts_github_cdn():
# Real GitHub release CDN URLs carry the /github-production-release-asset- prefix
# Real GitHub release CDN URLs carry the /github-production-release-asset- prefix.
assert _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x",
"lemonade-sdk/llamacpp-rocm",
@ -297,7 +318,7 @@ def test_lemonade_resolver_accepts_github_cdn():
def test_lemonade_resolver_rejects_arbitrary_cdn_path():
# A CDN URL without the release-asset path prefix must be rejected
# A CDN URL without the release-asset path prefix must be rejected.
assert not _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/abc/def",
"lemonade-sdk/llamacpp-rocm",
@ -339,7 +360,7 @@ def test_lemonade_runtime_patterns_include_hip_runtime():
Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp,
...) whose names change across ROCm releases. A broad ``lib*.so*`` glob
avoids enumerating every transitive dependency by name.
avoids having to enumerate every transitive dependency by name.
"""
from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice
@ -353,7 +374,7 @@ def test_lemonade_runtime_patterns_include_hip_runtime():
)
pats = runtime_patterns_for_choice(choice)
# The broad glob must be present so every .so in the lemonade bundle
# (including future transitive deps) gets overlaid.
# (including transitive deps added in future ROCm releases) gets overlaid.
assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}"
@ -365,9 +386,9 @@ _pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None)
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch):
"""AMD HIP honours CUDA_VISIBLE_DEVICES like HIP_VISIBLE_DEVICES; on a
gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
# Two GPUs; rocminfo reports each token twice (as in real tool output).
"""AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES;
on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
# Two GPUs; rocminfo reports each token twice (as in the real tool output).
probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
@ -396,7 +417,7 @@ def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
"""Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must
return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the
two gfx1100 entries into one and making index 2 out of range."""
# rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
# Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
# Each GPU gets its own Agent section with a few token mentions.
probe_out = (
"***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n"
@ -407,3 +428,96 @@ def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2")
assert _pick_rocm_gfx_target(probe_out) == "gfx1151"
# ---------------------------------------------------------------------------
# Fork release scan: Windows ROCm resolves lemonade by the requested tag
# ---------------------------------------------------------------------------
_resolve_release_asset_choice = getattr(_mod, "resolve_release_asset_choice", None)
_ApprovedReleaseChecksums = getattr(_mod, "ApprovedReleaseChecksums", None)
@pytest.mark.skipif(
_resolve_release_asset_choice is None or _ApprovedReleaseChecksums is None,
reason = "fork release planner not present on this branch",
)
def test_fork_scan_windows_rocm_resolves_lemonade_by_requested_tag():
"""The fork release scan pins llama_tag to per-release upstream tags
(b9457, ...) that lemonade's own tag series never contains, so the
lemonade lookup must use the requested tag ("latest") instead. Pinning
lemonade to the per-release tag 404s on every scanned release and a
Windows ROCm host ends in a rate-limited fatal instead of the lemonade
prebuilt."""
host = _make_rocm_host("gfx1151", windows = True)
# No windows-rocm artifact in the bundle, matching current fork releases.
bundle = _rocm_bundle("gfx1151", ["gfx1151"])
checksums = _ApprovedReleaseChecksums(
repo = "unslothai/llama.cpp",
release_tag = "v1.0",
upstream_tag = "b9457",
artifacts = {},
)
seen_urls: list[str] = []
def _fake_fetch(api_url, *args, **kwargs):
seen_urls.append(api_url)
if "lemonade-sdk" in api_url:
if api_url.endswith("/releases/latest"):
return _stub_lemonade_release()
raise RuntimeError(f"unexpected pinned lemonade fetch: {api_url}")
# ggml-org asset listing for the upstream HIP/CPU filename fallbacks.
return {"tag_name": "b9457", "assets": []}
with patch.object(_mod, "fetch_json", side_effect = _fake_fetch):
attempts = _resolve_release_asset_choice(
host,
"b9457", # concrete per-release upstream tag from the scan loop
bundle,
checksums,
requested_tag = "latest",
)
lemonade = [a for a in attempts if a.source_label == "lemonade"]
assert lemonade, f"lemonade attempt missing for Windows ROCm host; got {attempts}"
assert "gfx1151" in lemonade[0].name
assert any(
u.endswith("/releases/latest") for u in seen_urls
), f"lemonade was never resolved via /releases/latest; fetches: {seen_urls}"
assert not any(
"lemonade-sdk" in u and "/releases/tags/" in u for u in seen_urls
), f"lemonade lookup was pinned to the fork release tag: {seen_urls}"
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "direct release planners not present on this branch",
)
def test_direct_upstream_plan_includes_lemonade_for_linux_rocm_host():
"""A Linux ROCm host on the ggml-org direct path (e.g. a --published-repo
override) must plan lemonade before the CPU tarball, mirroring the Windows
branch. The lemonade planning previously lived in the removed
--simple-policy dispatcher, so without this leg such hosts silently
install the CPU build."""
host = _make_rocm_host("gfx1151")
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [
{
"name": "llama-b9022-bin-ubuntu-x64.tar.gz",
"browser_download_url": (
"https://github.com/ggml-org/llama.cpp/releases/download/"
"b9022/llama-b9022-bin-ubuntu-x64.tar.gz"
),
}
],
}
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None, "Linux ROCm host should produce a direct plan"
kinds = [a.install_kind for a in plan.attempts]
sources = [a.source_label for a in plan.attempts]
assert "linux-rocm" in kinds, f"lemonade ROCm attempt missing; got {kinds}"
assert sources[0] == "lemonade", f"lemonade must be the first attempt; got {sources}"
assert "gfx1151" in plan.attempts[0].name

134
studio/backend/tests/test_llama_cpp_freshness.py
View file

@ -21,12 +21,25 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
class _NoopLogger:
"""structlog-style logger: every method swallows positional + kwargs.
A stdlib logging.Logger rejects structlog's keyword fields (e.g.
``logger.warning(msg, error=...)``), which leaked into the update module's
error path and failed only when this file's stub loaded first.
"""
def __getattr__(self, _name):
return lambda *a, **k: None
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
_loggers_stub.get_logger = lambda *a, **k: _NoopLogger()
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
_structlog_stub.get_logger = lambda *a, **k: _NoopLogger()
sys.modules.setdefault("structlog", _structlog_stub)
import pytest
@ -51,6 +64,11 @@ def _write_marker(install_dir: Path, **overrides) -> Path:
.replace("+00:00", "Z"),
}
payload.update(overrides)
# The installer always writes `tag` and `release_tag` from the same release
# (a normalized base vs the full release tag), so keep the pair consistent
# when a test overrides only `tag`.
if "tag" in overrides and "release_tag" not in overrides:
payload["release_tag"] = overrides["tag"]
install_dir.mkdir(parents = True, exist_ok = True)
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(payload))
return install_dir / "UNSLOTH_PREBUILT_INFO.json"
@ -303,3 +321,115 @@ def test_format_stale_warning_singular_day():
msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1})
assert "1 day" in msg
assert "1 days" not in msg
# parse_base_build / is_behind.
def test_parse_base_build():
assert fr.parse_base_build("b9596") == 9596
assert fr.parse_base_build(" b9596 ") == 9596
assert fr.parse_base_build("b9596-mix-e6f2453") == 9596 # mix suffix doesn't defeat it
assert fr.parse_base_build("9596") is None
assert fr.parse_base_build("master-abc") is None
assert fr.parse_base_build("") is None
assert fr.parse_base_build(None) is None
@pytest.mark.parametrize(
"installed, latest, expected",
[
(
"b9596-mix-e6f2453",
"b9596-mix-e6f2453",
False,
), # already on the mix latest -> not behind
("b9596", "b9594", False), # latest is an older build -> downgrade guard
("b9596", "b9594-mix-xxx", False), # older mix latest -> still guarded
("b9500", "b9596-mix-e6f2453", True), # newer base -> behind
("b9596-mix-aaa", "b9596-mix-bbb", True), # new mix at same base -> behind
("b9596", "b9596-mix-bbb", True), # clean -> mix at same base -> behind
("b9596-mix-aaa", "b9596", False), # bare base never supersedes a mix install
("b9596", "b9596", False), # identical -> not behind
(" b9596 ", "b9596", False), # whitespace-only diff -> not behind
("master-abc", "master-def", True), # non-bNNNN both -> plain inequality
("master-abc", "master-abc", False),
(None, "b9596", False),
("b9596", None, False),
],
)
def test_is_behind(installed, latest, expected):
assert fr.is_behind(installed, latest) is expected
def test_check_prebuilt_freshness_not_behind_on_mix_latest(monkeypatch, tmp_path):
# Installed the mix latest: marker base tag b9596, full release_tag with sha,
# GitHub latest is that same full tag. Must not report behind (sticky bug).
install_dir = tmp_path / "llama.cpp"
_write_marker(install_dir, tag = "b9596", release_tag = "b9596-mix-e6f2453")
bin_path = _fake_binary(install_dir, layout = "root")
monkeypatch.setattr(
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453"
)
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["behind"] is False
assert info["stale"] is False
def test_check_prebuilt_freshness_downgrade_guard(monkeypatch, tmp_path):
# A lagging latest (older build than installed) must never read as behind/stale.
install_dir = tmp_path / "llama.cpp"
_write_marker(
install_dir,
tag = "b9585",
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 30))
.isoformat()
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["behind"] is False
assert info["stale"] is False
def test_fetch_latest_release_tag_uses_publish_time(monkeypatch):
# Resolves newest by published_at (like the installer), skips drafts/prereleases,
# and does NOT just take GitHub's first/`/releases/latest` item.
import urllib.request
class _Resp:
def __init__(self, payload):
self._p = json.dumps(payload).encode()
def read(self):
return self._p
def __enter__(self):
return self
def __exit__(self, *a):
return False
payload = [
{
"tag_name": "b9518",
"draft": False,
"prerelease": False,
"published_at": "2026-06-04T21:11:19Z",
},
{
"tag_name": "b9596-mix-e6f2453",
"draft": False,
"prerelease": False,
"published_at": "2026-06-11T22:50:41Z",
},
{
"tag_name": "b9999-draft",
"draft": True,
"prerelease": False,
"published_at": "2026-06-12T00:00:00Z",
},
]
monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload))
assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453"

187
studio/backend/tests/test_llama_cpp_mmproj_fallback.py Normal file
View file

@ -0,0 +1,187 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the llama-server mmproj text-only fallback.
A GGUF vision model is launched with ``--mmproj <projector>``. When the
installed llama.cpp prebuilt is older than the model's projector format,
llama-server aborts at startup with ``clip.cpp:NNNN: Unknown projector
type`` (exit -6). load_model now retries once WITHOUT ``--mmproj`` so the
base model still loads text-only, warns the user to update llama.cpp, and
marks the session non-vision. These tests pin the two decision helpers:
``_is_projector_incompatibility`` (when to retry) and ``_strip_mmproj_args``
(how the retry argv is built). Unrelated failures must NOT trigger a retry.
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Match the stubbing pattern in sibling tests so the module imports in a
# lightweight env without fastapi.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
if not hasattr(sys.modules["structlog"], "get_logger"):
sys.modules["structlog"].get_logger = _structlog_stub.get_logger
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
_detect = LlamaCppBackend._is_projector_incompatibility
_strip = LlamaCppBackend._strip_mmproj_args
# Real abort captured loading gemma-4 on a 3-day-old prebuilt (build b9496).
_GEMMA4_OLD_LLAMACPP_OUT = (
"srv load_model: loading model 'gemma-4-E2B-it-UD-Q4_K_XL.gguf'\n"
"/build_work/src/llama.cpp-b9496/tools/mtmd/clip.cpp:4391: "
"Unknown projector type\n"
"libggml-base.so.0(ggml_abort+0x152)\n"
"libmtmd.so.0(clip_n_mmproj_embd)\n"
)
# Unrelated failures that must keep their own handling (no projector retry).
_OOM_OUT = (
"ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on "
"device 0: cudaMalloc failed: out of memory"
)
_BAD_ARCH_OUT = "llama_model_load: error loading model: unknown model architecture: 'qwen_image'"
_PORT_OUT = "srv start: failed to bind: address already in use"
_MISSING_OUT = "error: failed to open GGUF file: no such file or directory"
# A healthy startup log that merely mentions the projector must not match.
_HEALTHY_VISION_OUT = (
"Using mmproj for vision: /cache/mmproj-F16.gguf\n"
"clip_model_loader: loaded meta data with 20 key-value pairs\n"
"srv update_slots: all slots are idle"
)
class TestProjectorIncompatibilityDetector:
def test_gemma4_on_old_llamacpp_triggers_retry(self):
# Headline case: a 3-day-old llama.cpp aborts on Gemma-4's projector.
assert _detect(_GEMMA4_OLD_LLAMACPP_OUT) is True
@pytest.mark.parametrize(
"out",
[
"clip.cpp:4391: Unknown projector type",
"error: unsupported projector type for this model",
"llama_mmproj: unsupported mmproj file version",
"clip.cpp: projector type 'gemma4' is not supported",
],
)
def test_projector_format_errors_match(self, out):
assert _detect(out) is True
def test_case_insensitive(self):
assert _detect("UNKNOWN PROJECTOR TYPE") is True
@pytest.mark.parametrize(
"out",
[
_OOM_OUT,
_BAD_ARCH_OUT,
_PORT_OUT,
_MISSING_OUT,
_HEALTHY_VISION_OUT,
"",
# bare multimodal words without a failure term must not match
"loading clip model",
"mmproj file resolved from cache",
],
)
def test_unrelated_failures_do_not_retry(self, out):
assert _detect(out) is False
# A realistic vision launch argv (mirrors the live "Starting llama-server"
# command), projector pair at the end.
_VISION_CMD = [
"/home/u/.unsloth/llama.cpp/build/bin/llama-server",
"-m",
"/cache/gemma-4-E2B-it-UD-Q4_K_XL.gguf",
"--port",
"55473",
"-c",
"131072",
"--parallel",
"1",
"--flash-attn",
"on",
"--no-context-shift",
"-ngl",
"-1",
"--threads",
"-1",
"--jinja",
"--spec-default",
"--mmproj",
"/cache/mmproj-F16.gguf",
]
class TestStripMmprojArgs:
def test_removes_mmproj_pair(self):
stripped = _strip(_VISION_CMD)
assert "--mmproj" not in stripped
assert "/cache/mmproj-F16.gguf" not in stripped
def test_preserves_every_text_flag(self):
stripped = _strip(_VISION_CMD)
for flag in (
"-m",
"/cache/gemma-4-E2B-it-UD-Q4_K_XL.gguf",
"--port",
"55473",
"-c",
"131072",
"-ngl",
"-1",
"--jinja",
"--spec-default",
"--flash-attn",
"on",
):
assert flag in stripped
# Exactly the two projector tokens are dropped.
assert len(stripped) == len(_VISION_CMD) - 2
def test_strips_mmproj_in_the_middle(self):
cmd = ["llama-server", "--mmproj", "/p/mm.gguf", "-c", "4096", "--jinja"]
assert _strip(cmd) == ["llama-server", "-c", "4096", "--jinja"]
def test_noop_when_no_mmproj(self):
cmd = ["llama-server", "-m", "/p/model.gguf", "-c", "4096", "--jinja"]
assert _strip(cmd) == cmd
def test_returns_new_list(self):
cmd = ["llama-server", "--mmproj", "/p/mm.gguf"]
out = _strip(cmd)
assert out is not cmd
assert cmd[-1] == "/p/mm.gguf" # input untouched
class TestRetryContract:
"""The two helpers compose into the load_model retry decision."""
def test_gemma4_failure_yields_valid_text_only_command(self):
# Old-llama.cpp projector abort -> retry, and the retry argv is a
# valid text-only launch (model + context kept, projector gone).
assert _detect(_GEMMA4_OLD_LLAMACPP_OUT) is True
retry_cmd = _strip(_VISION_CMD)
assert "--mmproj" not in retry_cmd
assert "-m" in retry_cmd and "--jinja" in retry_cmd
def test_oom_does_not_retry_text_only(self):
# An OOM with --mmproj present must NOT be treated as a projector
# problem: load_model errors out instead of dropping vision.
assert _detect(_OOM_OUT) is False

284
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -1056,8 +1056,8 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF"
("mtp", False, _MTP_MODEL, "draft-mtp", "3", False),
# ── mtp forced on sub-3B: engage anyway ──
("mtp", True, _SUB_3B_MTP_MODEL, "draft-mtp", "2", False),
# ── mtp forced on non-MTP: engage anyway ──
("mtp", True, _NON_MTP_MODEL, "draft-mtp", "2", False),
# ── mtp forced on non-MTP: default back (no head/drafter) ──
("mtp", True, _NON_MTP_MODEL, None, None, False),
# ── ngram forced: ngram-mod alone on BOTH platforms ──
("ngram", True, _MTP_MODEL, "ngram-mod", None, True),
("ngram", False, _MTP_MODEL, "ngram-mod", None, True),
@ -1066,6 +1066,8 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF"
("mtp+ngram", True, _MTP_MODEL, "ngram-mod,draft-mtp", "2", True),
("mtp+ngram", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True),
("mtp+ngram", True, _SUB_3B_MTP_MODEL, "ngram-mod,draft-mtp", "2", True),
# ── mtp+ngram forced on non-MTP: keep ngram, drop draft-mtp ──
("mtp+ngram", True, _NON_MTP_MODEL, "ngram-mod", None, True),
# ── off: nothing emitted ──
("off", True, _MTP_MODEL, None, None, False),
("off", False, _MTP_MODEL, None, None, False),
@ -1178,3 +1180,281 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
# choice is still reflected in _requested_spec_mode.
assert backend.requested_spec_mode == "mtp"
assert backend.speculative_type is None
def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch):
# Forcing MTP on a model with no head/drafter must NOT emit draft-mtp:
# llama-server aborts on it ("failed to measure MTP context memory")
# rather than no-op'ing. Default back to --spec-default instead.
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "mtp",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _NON_MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert "--spec-type" not in flags
assert "--spec-default" in flags
assert backend.speculative_type == "default"
assert backend.requested_spec_mode == "mtp"
def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch):
# mtp+ngram on a non-MTP model drops the doomed draft-mtp chain but keeps
# the ngram half, which needs no head.
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "mtp+ngram",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _NON_MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "ngram-mod"
assert backend.speculative_type == "ngram-mod"
assert backend.requested_spec_mode == "mtp+ngram"
# ── Full named-repo resolver matrix (the shipping Studio families) ─────
#
# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and
# gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay
# resolver-only (122B-A10B / 397B-A17B). Expectations are derived from the
# same signals load_model uses -- _extract_model_size_b (active>effective>
# total, so E2B->2, A3B->3, A10B->10, A17B->17), _is_mtp_model_name, and the
# separate-drafter flag -- so each row mirrors what the loader emits on a
# B200 (GPU default, n=2). gemma carries no -MTP marker; its MTP comes from
# the root mtp-*.gguf drafter, modelled here by passing mtp_draft_path.
#
# auto_spec: "draft-mtp" = head/drafter engaged (>=3B MTP, or any size with a
# separate drafter); "ngram-mod" = embedded sub-3B drop (zero-VRAM); None =
# non-MTP -> llama-server --spec-default.
_GEMMA_DRAFTER = "/snap/mtp-gemma-4-it.gguf" # stand-in separate drafter
_REAL_REPO_MATRIX = [
# repo, drafter, auto_spec, auto_ngram_knobs
("unsloth/Qwen3.5-0.8B-MTP-GGUF", None, "ngram-mod", True),
("unsloth/Qwen3.5-2B-MTP-GGUF", None, "ngram-mod", True),
("unsloth/Qwen3.5-4B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-9B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-27B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-35B-A3B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-122B-A10B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-397B-A17B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-0.8B-GGUF", None, None, False),
("unsloth/Qwen3.5-2B-GGUF", None, None, False),
("unsloth/Qwen3.5-4B-GGUF", None, None, False),
("unsloth/Qwen3.5-9B-GGUF", None, None, False),
# E2B is 2B but ships a separate drafter -> exempt from the sub-3B drop.
("unsloth/gemma-4-E2B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-E4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-12b-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-26B-A4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-31B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-E2B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-E4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-12b-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-26B-A4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-31B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
]
def _resolve_real(monkeypatch, repo, drafter, mode):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = mode,
spec_draft_n_max = None,
extra_args = None,
model_identifier = repo,
model_path = None,
gpus = True, # B200 default
binary = "/fake/llama-server",
mtp_draft_path = drafter,
)
return backend, flags, _flags_dict(flags)
@pytest.mark.parametrize(
"repo, drafter, auto_spec, auto_ngram_knobs",
_REAL_REPO_MATRIX,
ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX],
)
def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs):
# Auto is the default mode the dropdown ships with.
backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "auto")
if auto_spec is None:
# Non-MTP: no draft-mtp, hand off to llama-server's own default.
assert "--spec-type" not in parsed
assert "--spec-default" in flags
assert backend.speculative_type == "default"
elif auto_spec == "draft-mtp":
assert parsed.get("--spec-type") == "draft-mtp"
assert parsed.get("--spec-draft-n-max") == "2"
assert backend.speculative_type == "draft-mtp"
# gemma ships a separate drafter; Qwen bakes the head into the GGUF.
assert (
(parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed)
)
else: # ngram-mod (sub-3B MTP drop)
assert parsed.get("--spec-type") == "ngram-mod"
assert "--model-draft" not in parsed # draft head dropped
assert backend.speculative_type == "ngram-mod"
if auto_ngram_knobs:
assert "--spec-ngram-mod-n-match" in parsed
assert backend.requested_spec_mode == "auto"
@pytest.mark.parametrize(
"repo, drafter",
[(r[0], r[1]) for r in _REAL_REPO_MATRIX],
ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX],
)
def test_real_repo_off_emits_nothing(monkeypatch, repo, drafter):
# Off must suppress speculative decoding for every family.
backend, flags, _ = _resolve_real(monkeypatch, repo, drafter, "off")
assert flags == []
assert backend.speculative_type is None
assert backend.requested_spec_mode == "off"
@pytest.mark.parametrize(
"repo, drafter",
[(r[0], r[1]) for r in _REAL_REPO_MATRIX],
ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX],
)
def test_real_repo_forced_mtp_never_aborts(monkeypatch, repo, drafter):
# Forcing MTP on the dropdown: real MTP models (name marker or separate
# drafter) engage draft-mtp even below 3B; non-MTP models default back to
# --spec-default instead of emitting a draft-mtp llama-server will abort on.
backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "mtp")
is_real_mtp = _is_mtp_model_name(repo) or bool(drafter)
if is_real_mtp:
assert parsed.get("--spec-type") == "draft-mtp"
assert backend.speculative_type == "draft-mtp"
assert (
(parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed)
)
else:
assert "--spec-type" not in parsed
assert "--spec-default" in flags
assert backend.speculative_type == "default"
assert backend.requested_spec_mode == "mtp"
# ── Sub-3B separate-drafter exemption (Gemma) ─────────────────────────
#
# The sub-3B MTP drop is an embedded-head cost (Qwen). A separate drafter
# (Gemma's root mtp-*.gguf) is a cheap standalone model that wins below 3B
# (B200 Q4_K_XL: gemma-4-E2B draft-mtp n=2 = 1.21x vs OFF), so it is exempt.
def test_sub3b_gemma_separate_drafter_engages_mtp(monkeypatch):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = "unsloth/gemma-4-E2B-it-GGUF", # 2B
model_path = None,
gpus = True,
binary = "/fake/llama-server",
mtp_draft_path = "/snap/mtp-gemma-4-E2B-it.gguf", # separate drafter
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "draft-mtp"
assert parsed.get("--model-draft") == "/snap/mtp-gemma-4-E2B-it.gguf"
assert "--spec-ngram-mod-n-match" not in parsed
assert backend.speculative_type == "draft-mtp"
def test_sub3b_qwen_embedded_head_still_drops_to_ngram(monkeypatch):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", # 2B, embedded head
model_path = None,
gpus = True,
binary = "/fake/llama-server",
mtp_draft_path = None, # no separate drafter
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "ngram-mod"
assert "--model-draft" not in parsed
assert backend.speculative_type == "ngram-mod"
def test_auto_mode_drops_mtp_exempts_separate_drafter():
from core.inference.llama_cpp import _auto_mode_drops_mtp
assert _auto_mode_drops_mtp("auto", 2.0) is True
assert _auto_mode_drops_mtp("auto", 2.0, has_separate_drafter = True) is False
assert _auto_mode_drops_mtp("auto", 4.0) is False
assert _auto_mode_drops_mtp("mtp", 2.0) is False # forced engages regardless
# ── spec_fallback_reason (drives the "update llama.cpp" UI hint) ───────
def test_spec_fallback_reason_set_when_binary_lacks_mtp(monkeypatch):
# Outdated llama-server with no mtp token: a forced MTP request can't emit
# draft-mtp, so record the reason for the UI update affordance.
backend = _resolver_backend(monkeypatch, mtp_token = None)
backend._build_speculative_flags(
speculative_type = "mtp",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert backend.spec_fallback_reason == "binary_no_mtp"
def test_spec_fallback_reason_none_when_mtp_engages(monkeypatch):
backend = _resolver_backend(monkeypatch)
backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert backend.speculative_type == "draft-mtp"
assert backend.spec_fallback_reason is None
def test_spec_fallback_reason_reset_on_off(monkeypatch):
# A subsequent off load must clear a stale reason.
backend = _resolver_backend(monkeypatch, mtp_token = None)
backend._build_speculative_flags(
speculative_type = "mtp",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert backend.spec_fallback_reason == "binary_no_mtp"
backend._build_speculative_flags(
speculative_type = "off",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert backend.spec_fallback_reason is None

254
studio/backend/tests/test_llama_cpp_props_readback.py Normal file
View file

@ -0,0 +1,254 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the post-launch /props context readback.
llama-server's memory-fit step or --parallel slot split can allocate less
context than the requested -c while Studio keeps advertising the requested
value; clients sized to it then die on exceed_context_size_error 400s.
``_reconcile_effective_ctx_with_server`` must adopt the server's real
``default_generation_settings.n_ctx`` whenever it is smaller.
Stubbed httpx; no subprocess, GPU, or network. Cross-platform.
"""
from __future__ import annotations
import json
import sys
import types as _types
from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Stub heavy/unavailable deps before importing the module under test.
# Mirrors test_llama_cpp_context_fit.py.
# ---------------------------------------------------------------------------
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Prefer the real modules so importing this file first cannot poison later
# test modules with stubs; only stub what the environment genuinely lacks.
try:
import loggers # noqa: F401
except ImportError:
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
try:
import structlog # noqa: F401
except ImportError:
sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
try:
import httpx # noqa: F401
except ImportError:
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"WriteError",
"HTTPError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
class _FakeTimeout:
def __init__(self, *a, **kw):
pass
_httpx_stub.Timeout = _FakeTimeout
_httpx_stub.Client = type(
"Client",
(),
{
"__init__": lambda self, **kw: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *a: None,
},
)
_httpx_stub.get = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("unstubbed httpx.get"))
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend
import core.inference.llama_cpp as llama_cpp_mod
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _FakeResponse:
def __init__(
self,
status_code = 200,
body = None,
):
self.status_code = status_code
self._body = body or {}
def json(self):
return self._body
def _make_backend(effective_ctx = 98304, port = 51234):
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._port = port
inst._effective_context_length = effective_ctx
inst._context_length = 262144
return inst
def _stub_props(
monkeypatch,
status_code = 200,
body = None,
exc = None,
):
def fake_get(url, timeout = None):
assert url.endswith("/props")
if exc is not None:
raise exc
return _FakeResponse(status_code, body)
monkeypatch.setattr(llama_cpp_mod.httpx, "get", fake_get, raising = False)
# ---------------------------------------------------------------------------
# _query_server_n_ctx parsing
# ---------------------------------------------------------------------------
def test_query_n_ctx_reads_default_generation_settings(monkeypatch):
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 67584}},
)
assert _make_backend()._query_server_n_ctx() == 67584
def test_query_n_ctx_non_200_returns_none(monkeypatch):
_stub_props(monkeypatch, status_code = 503)
assert _make_backend()._query_server_n_ctx() is None
def test_query_n_ctx_missing_key_returns_none(monkeypatch):
_stub_props(monkeypatch, body = {"default_generation_settings": {}})
assert _make_backend()._query_server_n_ctx() is None
def test_query_n_ctx_swallows_transport_errors(monkeypatch):
_stub_props(monkeypatch, exc = RuntimeError("connection refused"))
assert _make_backend()._query_server_n_ctx() is None
# ---------------------------------------------------------------------------
# _reconcile_effective_ctx_with_server decisions
# ---------------------------------------------------------------------------
def test_fit_shrunk_ctx_overwrites_advertised_value(monkeypatch):
"""The Nick repro: requested/advertised 98304, server really at 67584."""
inst = _make_backend(effective_ctx = 98304)
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 67584}},
)
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 67584
assert inst.context_length == 67584
def test_matching_ctx_is_left_alone(monkeypatch):
inst = _make_backend(effective_ctx = 98304)
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 98304}},
)
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 98304
def test_larger_server_ctx_does_not_inflate_advertised_value(monkeypatch):
"""Never advertise more than the user asked for, even if the server could."""
inst = _make_backend(effective_ctx = 32768)
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 65536}},
)
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 32768
def test_unset_effective_ctx_adopts_server_value(monkeypatch):
inst = _make_backend(effective_ctx = None)
inst._context_length = None
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 40960}},
)
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 40960
def test_props_failure_keeps_studio_value(monkeypatch):
"""A flaky /props must never wipe the computed context."""
inst = _make_backend(effective_ctx = 98304)
_stub_props(monkeypatch, exc = RuntimeError("boom"))
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 98304
# ---------------------------------------------------------------------------
# _ctx_integrity_flags: keep the per-request window equal to the advertised ctx
# ---------------------------------------------------------------------------
_CAPS_ALL = {"supports_kv_unified": True, "supports_fit_ctx": True}
_CAPS_NONE = {"supports_kv_unified": False, "supports_fit_ctx": False}
def test_kv_unified_added_for_multi_slot():
"""Explicit --parallel N disables llama-server's auto-slots kv-unified
default, splitting -c into per-slot windows of -c/N; Studio must restore
the shared pool so one request can use the full advertised context."""
flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL)
assert "--kv-unified" in flags
def test_kv_unified_skipped_for_single_slot_or_old_build():
assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags(
1, False, 98304, 98304, _CAPS_ALL
)
assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags(
4, False, 98304, 98304, _CAPS_NONE
)
def test_fit_ctx_floors_explicit_request_under_fit():
flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL)
assert flags[flags.index("--fit-ctx") + 1] == "98304"
def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support():
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
1, False, 98304, 98304, _CAPS_ALL
)
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL)
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
1, True, 98304, 98304, _CAPS_NONE
)
def test_probe_missing_binary_reports_new_capabilities_false():
info = LlamaCppBackend.probe_server_capabilities(binary = "/nonexistent/llama-server")
assert info["found"] is False
assert info["supports_kv_unified"] is False
assert info["supports_fit_ctx"] is False

843
studio/backend/tests/test_llama_cpp_update.py Normal file
View file

@ -0,0 +1,843 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hermetic tests for the in-app llama.cpp update orchestration.
No network, no real install: the GitHub release lookup and the installer
subprocess are both monkeypatched. Verifies detection (update_available) and
the apply flow (job lifecycle, installer invocation, post-swap re-read).
"""
from __future__ import annotations
import json
import sys
import time
from pathlib import Path
from types import ModuleType
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import utils.llama_cpp_freshness as freshness # noqa: E402
import utils.llama_cpp_update as upd # noqa: E402
MARKER = "UNSLOTH_PREBUILT_INFO.json"
class _FakeInstallerPopen:
"""Stands in for the streamed installer process in _run_update."""
def __init__(
self,
cmd,
*,
returncode = 0,
lines = None,
on_start = None,
captured_kwargs = None,
**kwargs,
):
if captured_kwargs is not None:
captured_kwargs.update(kwargs)
if on_start is not None:
on_start(list(cmd))
self.returncode = returncode
self.stdout = iter(lines or [])
def wait(self):
return self.returncode
def kill(self):
pass
def _patch_installer_popen(
monkeypatch,
*,
returncode = 0,
lines = None,
on_start = None,
captured_kwargs = None,
):
monkeypatch.setattr(
upd.subprocess,
"Popen",
lambda cmd, **kw: _FakeInstallerPopen(
cmd,
returncode = returncode,
lines = lines,
on_start = on_start,
captured_kwargs = captured_kwargs,
**kw,
),
)
def _write_install(
dir_: Path,
tag: str,
repo: str = "unslothai/llama.cpp",
asset: str | None = None,
release_tag: str | None = None,
) -> str:
"""Create a fake prebuilt install tree and return the llama-server path.
``asset`` is the bundle filename recorded in the marker; omit it to model an
older marker that predates asset-based ROCm forwarding (backward compat).
``release_tag`` is the full release tag (e.g. a ``b9596-mix-<sha>`` mix
build); defaults to ``tag`` for a plain prebuilt."""
bin_dir = dir_ / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
binary = bin_dir / "llama-server"
binary.write_text("#!/bin/sh\necho stub\n")
marker = {
"tag": tag,
"release_tag": release_tag or tag,
"published_repo": repo,
"installed_at_utc": "2020-01-01T00:00:00Z",
"bundle_profile": "cuda13-newer",
"runtime_line": "cuda13",
}
if asset is not None:
marker["asset"] = asset
(dir_ / MARKER).write_text(json.dumps(marker))
return str(binary)
@pytest.fixture(autouse = True)
def _clean_state(monkeypatch, tmp_path):
freshness.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
# Isolate the freshness disk cache so the suite never writes the real
# ~/.unsloth cache (the default when storage_roots can't be imported).
monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".freshness_cache")
# Deterministic markerless paths: no host-pinned binary, no custom dir.
monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
# Never hit the network in these tests.
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
yield
freshness.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
def _no_prebuilt(monkeypatch):
"""Stub the host prebuilt probe to 'none available' (no source-build offer)."""
monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: None)
def _prebuilt(
monkeypatch,
*,
repo = "unslothai/llama.cpp",
release_tag = "b9585",
llama_tag = None,
asset = None,
):
"""Stub the host prebuilt probe to report an available prebuilt."""
payload = {
"prebuilt_available": True,
"repo": repo,
"release_tag": release_tag,
"llama_tag": llama_tag or release_tag,
"asset": asset or f"llama-{release_tag}-bin-macos-arm64.tar.gz",
"install_kind": "macos-arm64",
}
monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: payload)
def test_status_no_marker_no_prebuilt(monkeypatch, tmp_path):
# No marker AND no prebuilt available for the host -> unsupported (the genuine
# source-build-with-nothing-to-offer case).
binary = tmp_path / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # no marker file alongside
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_no_prebuilt(monkeypatch)
st = upd.get_update_status()
assert st["supported"] is False
assert st["update_available"] is False
assert st["installed_tag"] is None
def test_status_source_build_offers_prebuilt(monkeypatch, tmp_path):
# Markerless source build with a prebuilt now available for the host: surface
# the update. Unknown installed version (source build) is treated as behind.
binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(monkeypatch, release_tag = "b9585")
monkeypatch.setattr(upd, "_installed_build_number", lambda b: None)
st = upd.get_update_status()
assert st["supported"] is True
assert st["update_available"] is True
assert st["source_build"] is True
assert st["latest_tag"] == "b9585"
assert st["published_repo"] == "unslothai/llama.cpp"
def test_status_source_build_compares_llama_tag(monkeypatch, tmp_path):
# release_tag may be a fork wrapper (v1.0); compare/display the upstream
# llama_tag (b9457) so a source build is not wrongly judged newer.
binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(monkeypatch, release_tag = "v1.0", llama_tag = "b9457")
monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9000)
st = upd.get_update_status()
assert st["latest_tag"] == "b9457" # not the wrapper tag
assert st["update_available"] is True # 9000 < 9457
def test_status_source_build_pinned_binary_not_offered(monkeypatch, tmp_path):
# LLAMA_SERVER_PATH pins a custom binary outside any llama.cpp dir; an apply
# could not take effect, so the button must not surface.
binary = tmp_path / "custom" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary))
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(monkeypatch)
st = upd.get_update_status()
assert st["supported"] is False
assert st["update_available"] is False
def test_llama_install_root_pinned_returns_none(monkeypatch, tmp_path):
binary = tmp_path / "custom" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary))
assert upd._llama_install_root(str(binary)) is None
def test_status_source_build_suppressed_when_newer(monkeypatch, tmp_path):
# A source build already newer than the latest prebuilt is not nagged.
binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(monkeypatch, release_tag = "b9518")
monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600)
st = upd.get_update_status()
assert st["supported"] is True
assert st["update_available"] is False
assert st["installed_tag"] == "b9600"
def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path):
# While the updater swaps the tree, status polls must not exec the binary
# being replaced (on Windows that exec can fail the installer's os.replace);
# the 3s poller only consumes job progress.
binary = tmp_path / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
probes = {"resolve": 0, "version": 0}
def _count_resolve(*, force_refresh = False):
probes["resolve"] += 1
return None
def _count_version(b):
probes["version"] += 1
return None
monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", _count_resolve)
monkeypatch.setattr(upd, "_installed_build_number", _count_version)
with upd._job_lock:
upd._job["state"] = upd._JOB_RUNNING
st = upd.get_update_status()
assert st["job"]["state"] == "running"
assert probes == {"resolve": 0, "version": 0}
def test_status_update_available(monkeypatch, tmp_path):
binary = _write_install(tmp_path, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
st = upd.get_update_status(force_refresh = True)
assert st["supported"] is True
assert st["installed_tag"] == "b9493"
assert st["latest_tag"] == "b9518"
assert st["update_available"] is True
def test_status_up_to_date(monkeypatch, tmp_path):
binary = _write_install(tmp_path, "b9518")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
st = upd.get_update_status(force_refresh = True)
assert st["installed_tag"] == "b9518"
assert st["latest_tag"] == "b9518"
assert st["update_available"] is False
def test_start_update_no_marker_no_prebuilt_refuses(monkeypatch, tmp_path):
binary = tmp_path / "llama-server"
binary.write_text("stub") # no marker
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
_no_prebuilt(monkeypatch)
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "no_prebuilt_available"
def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path):
# Markerless install + available prebuilt: install in place into the resolved
# root, with the asset-derived ROCm forwarding and the resolved repo.
install_dir = tmp_path / "llama.cpp"
binary = install_dir / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # no marker
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
_prebuilt(
monkeypatch, repo = "unslothai/llama.cpp", asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz"
)
captured = {}
class _Proc:
returncode = 0
stdout = "installed"
stderr = ""
def _fake_run(cmd, **kwargs):
cmd = list(cmd)
assert "--version" in cmd # only status polls still use run()
return _Proc()
def _on_start(cmd):
captured["cmd"] = cmd
_write_install(install_dir, "b9585") # installer writes the marker
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
_patch_installer_popen(monkeypatch, on_start = _on_start)
res = upd.start_update()
assert res["started"] is True, res
deadline = time.time() + 10
while time.time() < deadline:
if upd.get_update_status()["job"]["state"] in ("success", "error"):
break
time.sleep(0.05)
cmd = captured["cmd"]
assert "--install-dir" in cmd and str(install_dir) in cmd
assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd
assert "--llama-tag" in cmd and "latest" in cmd
assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x"
assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd
def test_start_update_happy_path(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
captured = {}
class _Proc:
returncode = 0
stdout = "installed"
stderr = ""
def _on_start(cmd):
captured["cmd"] = cmd
# Simulate the installer writing a new marker with the latest tag.
_write_install(install_dir, "b9518")
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
lines = [
"[llama-prebuilt] resolving release\n",
"Downloading llama.zip: 35.0% (12.0 MiB/35.0 MiB) at 9.0 MiB/s\n",
"Downloading llama.zip: 80.0% (28.0 MiB/35.0 MiB) at 9.0 MiB/s\n",
],
on_start = _on_start,
captured_kwargs = popen_kwargs,
)
res = upd.start_update()
assert res["started"] is True
assert res["job"]["from_tag"] == "b9493"
assert res["job"]["progress"] == 0.0
# Wait for the background worker.
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
assert job["to_tag"] == "b9518"
# Installer was invoked with the resolved install dir + latest + repo.
assert "--install-dir" in captured["cmd"]
assert str(install_dir) in captured["cmd"]
assert "--llama-tag" in captured["cmd"] and "latest" in captured["cmd"]
assert "unslothai/llama.cpp" in captured["cmd"]
# Progress lines were parsed and success pins progress at 1.0.
assert job["progress"] == 1.0
# The worker asks the installer for fine-grained progress milestones.
assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
_patch_installer_popen(monkeypatch, returncode = 2, lines = ["boom: network error\n"])
res = upd.start_update()
assert res["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "error"
assert "boom" in (job["error"] or "")
# --- installer-argument construction (mirrors the post-#5963 setup scripts) ---
def test_rocm_install_args_lemonade_gfx():
# Lemonade HIP app bundle: gfx family lives in the asset name.
assert upd._rocm_install_args("app-b9585-linux-x64-rocm-gfx110X.tar.gz") == [
"--rocm-gfx",
"gfx110x",
]
assert upd._rocm_install_args("app-b9585-windows-x64-rocm-gfx1150.zip") == [
"--rocm-gfx",
"gfx1150",
]
def test_rocm_install_args_fork_version_bundle():
# Fork ROCm bundles encode a ROCm version, not a gfx -> forward --has-rocm.
assert upd._rocm_install_args("llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz") == ["--has-rocm"]
def test_rocm_install_args_windows_hip():
assert upd._rocm_install_args("llama-b9334-bin-win-hip-radeon-x64.zip") == ["--has-rocm"]
def test_rocm_install_args_non_rocm_and_missing():
assert upd._rocm_install_args("llama-b9334-bin-ubuntu-x64.tar.gz") == []
assert upd._rocm_install_args("app-b9585-linux-x64-cuda13.tar.gz") == []
assert upd._rocm_install_args(None) == []
def _capture_install_cmd(
monkeypatch,
tmp_path,
*,
tag = "b9493",
repo = "unslothai/llama.cpp",
asset = None,
latest = "b9518",
) -> list:
"""Run start_update() with the installer subprocess stubbed; return the argv."""
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, tag, repo = repo, asset = asset)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest)
captured = {}
class _Proc:
returncode = 0
stdout = "installed"
stderr = ""
def _fake_run(cmd, **kwargs):
cmd = list(cmd)
assert "--version" in cmd # only status polls still use run()
return _Proc()
def _on_start(cmd):
captured["cmd"] = cmd
_write_install(install_dir, latest, repo = repo, asset = asset)
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
_patch_installer_popen(monkeypatch, on_start = _on_start)
res = upd.start_update()
assert res["started"] is True, res
deadline = time.time() + 10
while time.time() < deadline:
if upd.get_update_status()["job"]["state"] in ("success", "error"):
break
time.sleep(0.05)
return captured.get("cmd", [])
def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path):
cmd = _capture_install_cmd(
monkeypatch, tmp_path, asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz"
)
assert "--rocm-gfx" in cmd
assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x"
assert "--has-rocm" not in cmd
assert "--cpu-fallback" not in cmd
assert "--simple-policy" not in cmd
assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd
def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path):
cmd = _capture_install_cmd(
monkeypatch, tmp_path, asset = "llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz"
)
assert "--has-rocm" in cmd
assert "--rocm-gfx" not in cmd
def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path):
# CPU installs come from ggml-org. Re-running into the same install-dir/repo
# reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU
# detection) is reserved for setup.sh's arm64 rescue and must not appear here.
cmd = _capture_install_cmd(
monkeypatch,
tmp_path,
repo = "ggml-org/llama.cpp",
asset = "llama-b9334-bin-ubuntu-x64.tar.gz",
)
assert "--cpu-fallback" not in cmd
assert "--rocm-gfx" not in cmd
assert "--has-rocm" not in cmd
assert "--simple-policy" not in cmd
assert "--published-repo" in cmd and "ggml-org/llama.cpp" in cmd
def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tmp_path):
# Marker without an asset field (older install): no ROCm flags, no crash, and
# never the obsolete --simple-policy that #5963 removed from setup.
cmd = _capture_install_cmd(monkeypatch, tmp_path, asset = None)
assert "--simple-policy" not in cmd
assert "--rocm-gfx" not in cmd
assert "--has-rocm" not in cmd
assert "--cpu-fallback" not in cmd
# --- refusal + maintenance-state coordination ---
def test_start_update_already_running_refuses(monkeypatch, tmp_path):
binary = _write_install(tmp_path / "llama.cpp", "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
with upd._job_lock:
upd._job.update(state = upd._JOB_RUNNING)
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "already_running"
def test_start_update_installer_missing_refuses(monkeypatch, tmp_path):
binary = _write_install(tmp_path / "llama.cpp", "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: None)
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "installer_missing"
class _FakeBackend:
"""Minimal stand-in for LlamaCppBackend's update-coordination surface."""
def __init__(self):
import threading
self._serial_load_lock = threading.Lock()
self._llama_update_in_progress = False
self.is_active = True
self.unloaded = False
def unload_model(self):
self.unloaded = True
def _inject_backend(monkeypatch, backend):
routes_pkg = ModuleType("routes")
routes_pkg.__path__ = []
inference_mod = ModuleType("routes.inference")
inference_mod.get_llama_cpp_backend = lambda: backend
monkeypatch.setitem(sys.modules, "routes", routes_pkg)
monkeypatch.setitem(sys.modules, "routes.inference", inference_mod)
def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
backend = _FakeBackend()
_inject_backend(monkeypatch, backend)
seen = {}
def _on_start(cmd):
# The maintenance flag must be set while the installer runs.
seen["flag_during_install"] = backend._llama_update_in_progress
_write_install(install_dir, "b9518")
_patch_installer_popen(monkeypatch, on_start = _on_start)
res = upd.start_update()
assert res["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
if upd.get_update_status()["job"]["state"] in ("success", "error"):
break
time.sleep(0.05)
assert backend.unloaded is True
assert seen.get("flag_during_install") is True
# Cleared in the finally so model loads work again after the swap.
assert backend._llama_update_in_progress is False
def test_update_clears_maintenance_flag_on_installer_failure(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
backend = _FakeBackend()
_inject_backend(monkeypatch, backend)
_patch_installer_popen(monkeypatch, returncode = 1, lines = ["boom\n"])
res = upd.start_update()
assert res["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
if upd.get_update_status()["job"]["state"] in ("success", "error"):
break
time.sleep(0.05)
assert upd.get_update_status()["job"]["state"] == "error"
assert backend._llama_update_in_progress is False
def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
def _raise():
raise RuntimeError("no backend")
inference_mod = ModuleType("routes.inference")
inference_mod.get_llama_cpp_backend = lambda: _raise()
routes_pkg = ModuleType("routes")
routes_pkg.__path__ = []
monkeypatch.setitem(sys.modules, "routes", routes_pkg)
monkeypatch.setitem(sys.modules, "routes.inference", inference_mod)
_patch_installer_popen(monkeypatch, on_start = lambda cmd: _write_install(install_dir, "b9518"))
res = upd.start_update()
assert res["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
# --- markerless helper units ---
def test_resolve_prebuilt_parses_and_caches(monkeypatch, tmp_path):
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
calls = {"n": 0}
class _Proc:
returncode = 0
# stderr noise plus the JSON line on stdout (installer logs to stderr).
stdout = (
'{"prebuilt_available": true, "repo": "unslothai/llama.cpp", "release_tag": "b9585"}'
)
stderr = "[llama-prebuilt] some log\n"
def _fake_run(cmd, **kwargs):
calls["n"] += 1
assert "--resolve-prebuilt" in cmd
return _Proc()
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
res = upd._resolve_prebuilt_for_host()
assert res["prebuilt_available"] is True and res["release_tag"] == "b9585"
# Second call is memoized (no second subprocess).
upd._resolve_prebuilt_for_host()
assert calls["n"] == 1
def test_resolve_prebuilt_fails_open(monkeypatch, tmp_path):
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
def _boom(cmd, **kwargs):
raise OSError("subprocess failed")
monkeypatch.setattr(upd.subprocess, "run", _boom)
assert upd._resolve_prebuilt_for_host() is None
# Failures are not cached: a later success is observed.
class _Proc:
returncode = 0
stdout = '{"prebuilt_available": false}'
stderr = ""
monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc())
assert upd._resolve_prebuilt_for_host() == {"prebuilt_available": False}
def test_installed_build_number(monkeypatch):
def _ver(text):
class _Proc:
returncode = 0
stdout = ""
stderr = text
monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc())
return upd._installed_build_number("/bin/llama-server")
assert _ver("version: 9585 (abc1234)\nbuilt with clang\n") == 9585
assert _ver("version: 1 (deadbee)\n") is None # source build without tags
assert _ver("no version here") is None
assert upd._installed_build_number(None) is None
def test_llama_install_root_finds_llama_cpp_ancestor(monkeypatch, tmp_path):
root = tmp_path / "llama.cpp"
binary = root / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
assert upd._llama_install_root(str(binary)) == root
def test_llama_install_root_unmanaged_path_returns_none(monkeypatch, tmp_path):
# A binary on PATH (no marker, no env pin, no llama.cpp ancestor) is foreign:
# installing elsewhere would not replace it, so report no manageable root.
binary = tmp_path / "usr" / "local" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
assert upd._llama_install_root(str(binary)) is None
def test_llama_install_root_unsloth_env_dir(monkeypatch, tmp_path):
# UNSLOTH_LLAMA_CPP_PATH dir holding the active binary is the managed root.
root = tmp_path / "vendor" / "llama"
binary = root / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(root))
assert upd._llama_install_root(str(binary)) == root
def test_llama_install_root_ignores_inactive_env_root(monkeypatch, tmp_path):
# UNSLOTH_LLAMA_CPP_PATH set but the active binary is not under it: do not
# target the stale env root, resolve from the binary's own llama.cpp tree.
inactive = tmp_path / "custom-empty"
inactive.mkdir()
active = tmp_path / "llama.cpp"
binary = active / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(inactive))
assert upd._llama_install_root(str(binary)) == active
def test_llama_install_root_refuses_pinned_checkout_under_llama_cpp(monkeypatch, tmp_path):
# The LLAMA_SERVER_PATH pin guard must run before the ancestor scan, or a
# user's own llama.cpp checkout could be handed to the installer.
root = tmp_path / "my-project" / "llama.cpp"
binary = root / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary))
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
assert upd._llama_install_root(str(binary)) is None
def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path):
# A direct POST on a source build already newer than the prebuilt must not
# downgrade it; start_update mirrors the detection suppression.
install_dir = tmp_path / "llama.cpp"
binary = install_dir / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # no marker
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
_prebuilt(monkeypatch, release_tag = "b9518")
monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600)
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"
# --- mix-tag detection + apply guard (the reported banner bug) ---
def test_status_not_offered_on_mix_latest(monkeypatch, tmp_path):
# Installed the mix latest; GitHub latest is that same full tag -> no banner.
binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453"
)
st = upd.get_update_status()
assert st["update_available"] is False
assert st["installed_tag"] == "b9596"
assert st["latest_tag"] == "b9596-mix-e6f2453"
def test_status_not_offered_when_latest_lags(monkeypatch, tmp_path):
# A lagging latest (older build than installed) must never be offered.
binary = _write_install(tmp_path / "llama.cpp", "b9585")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
st = upd.get_update_status()
assert st["update_available"] is False
def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path):
# A direct POST / stale banner must not reinstall when already on the latest.
binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453"
)
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"

108
studio/backend/tests/test_llama_cpp_wait_for_health.py
View file

@ -145,3 +145,111 @@ class TestWaitForHealthResilience:
monkeypatch.setattr(httpx, "get", should_not_be_called)
assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False
assert called["n"] == 0
class TestCrashLogTail:
"""The "exited with code X" log must keep the TAIL of the output.
Crash diagnostics (abort reason, ROCm/CUDA error text) print last,
after the long startup banner; head truncation has cut off exactly
the diagnostic line in field reports (gfx1151 fit-step abort)."""
@staticmethod
def _capture_error_logs(monkeypatch) -> list:
"""Capture module-logger .error() messages directly -- immune to
whatever logging/structlog config sibling test modules installed."""
import core.inference.llama_cpp as _llama_mod
records: list = []
fake_logger = mock.Mock()
fake_logger.error = mock.Mock(side_effect = lambda msg, *a, **k: records.append(msg))
monkeypatch.setattr(_llama_mod, "logger", fake_logger)
return records
def test_crash_log_keeps_tail_not_head(self, monkeypatch):
records = self._capture_error_logs(monkeypatch)
b = _make_backend()
b._process.poll.return_value = 1
b._process.returncode = 1
# >2000 chars of banner, diagnostic on the final line.
banner = [f"load_model: tensor blk.{i} buffer ROCm0" for i in range(80)]
diagnostic = "ggml-cuda.cu:103: ROCm error: out of memory"
b._stdout_lines = banner + [diagnostic]
assert b._wait_for_health(timeout = 1.0, interval = 0.01) is False
crash_logs = [m for m in records if "exited with code" in m]
assert crash_logs, "crash must produce an exited-with-code log"
assert diagnostic in crash_logs[-1]
assert "Output (tail)" in crash_logs[-1]
# The head of the banner must be the part sacrificed to truncation.
assert "blk.0 buffer" not in crash_logs[-1]
def test_crash_log_mentions_log_file_when_present(self, monkeypatch):
records = self._capture_error_logs(monkeypatch)
b = _make_backend()
b._process.poll.return_value = 1
b._process.returncode = 1
b._stdout_lines = ["boom"]
b._llama_log_path = Path("C:/logs/llama-123-port-1234.log")
assert b._wait_for_health(timeout = 1.0, interval = 0.01) is False
crash_logs = [m for m in records if "exited with code" in m]
assert crash_logs and "llama-123-port-1234.log" in crash_logs[-1]
class TestRetryLogFilenameUnique:
"""The --fit off retry can respawn within the same epoch second; the log
filename must carry the attempt index or the second open ("w") truncates
the crash log the retry warning just referenced (found by simulation:
frozen time.time -> single file, crash evidence gone)."""
def test_log_name_includes_attempt_index(self):
src = (
Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py"
).read_text(encoding = "utf-8")
assert "-try{_spawn_attempt}.log" in src
class TestFitOffRetryEligible:
"""Gate for the one-shot --fit off startup-crash retry.
Retry only when Studio's own VRAM math placed the model and nothing
on the command line chose the fit mode explicitly."""
def test_eligible_for_plain_ngl_launch(self):
cmd = ["llama-server", "-m", "x.gguf", "-ngl", "-1", "--jinja"]
assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is True
def test_not_eligible_when_use_fit(self):
cmd = ["llama-server", "-m", "x.gguf", "--fit", "on"]
assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = True) is False
@pytest.mark.parametrize(
"fit_args",
[
["--fit", "on"],
["--fit", "off"],
["-fit", "off"],
["--fit=on"],
["-fit=off"],
],
)
def test_not_eligible_with_explicit_fit_flag(self, fit_args):
cmd = ["llama-server", "-m", "x.gguf", *fit_args]
assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is False
@pytest.mark.parametrize(
"tuning_args",
[
["--fit-ctx", "8192"],
["--fit-target", "1024"],
["-fitc", "4096"],
["-fitt", "512"],
["--fit-ctx=8192"],
],
)
def test_fit_tuning_flags_do_not_block_retry(self, tuning_args):
cmd = ["llama-server", "-m", "x.gguf", *tuning_args]
assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is True

111
studio/backend/tests/test_llama_route.py Normal file
View file

@ -0,0 +1,111 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""routes/llama.py: the source_build field is exposed and the handlers run the
(now subprocess-touching) detection off the event loop via a worker thread.
The route file is loaded standalone with a stubbed auth dependency so the test
does not pull the whole routes package (matplotlib-heavy training router) and
works in a minimal env.
"""
from __future__ import annotations
import asyncio
import importlib.util
import sys
import threading
import types
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
pytest.importorskip("fastapi")
def _load_route():
# Prefer the real auth module; stub it only in minimal envs where its
# deps are absent. Stubs are popped after the load so they never leak
# into sys.modules for the rest of the suite.
stubbed = []
try:
import auth.authentication # noqa: F401
except Exception:
auth_pkg = types.ModuleType("auth")
auth_pkg.__path__ = []
auth_mod = types.ModuleType("auth.authentication")
auth_mod.get_current_subject = lambda: "test"
for name, stub in (("auth", auth_pkg), ("auth.authentication", auth_mod)):
if name not in sys.modules:
sys.modules[name] = stub
stubbed.append(name)
try:
spec = importlib.util.spec_from_file_location(
"llama_route_under_test", str(_BACKEND / "routes" / "llama.py")
)
mod = importlib.util.module_from_spec(spec)
sys.modules["llama_route_under_test"] = mod # so pydantic resolves forward refs
spec.loader.exec_module(mod)
return mod
finally:
for name in stubbed:
sys.modules.pop(name, None)
rl = _load_route()
def test_status_response_exposes_source_build():
payload = {
"supported": True,
"update_available": True,
"stale": False,
"installed_tag": None,
"latest_tag": "b9585",
"published_repo": "unslothai/llama.cpp",
"installed_at_utc": None,
"age_days": None,
"source_build": True,
"job": {"state": "idle"},
}
model = rl.LlamaUpdateStatusResponse(**payload)
assert model.model_dump()["source_build"] is True
# Extra/unknown keys must not crash the response model.
rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1})
def test_status_handler_runs_off_event_loop(monkeypatch):
seen = {}
def fake_status(force_refresh = False):
seen["thread"] = threading.current_thread()
return {
"supported": True,
"update_available": True,
"source_build": True,
"latest_tag": "b9585",
"job": {"state": "idle"},
}
monkeypatch.setattr(rl, "get_update_status", fake_status)
out = asyncio.run(rl.llama_update_status(force_refresh = False, current_subject = "t"))
assert out.source_build is True
# Detection ran in a worker thread, not the event-loop thread.
assert seen["thread"] is not threading.main_thread()
def test_update_handler_runs_off_event_loop(monkeypatch):
seen = {}
def fake_start():
seen["thread"] = threading.current_thread()
return {"started": True, "reason": None, "job": {"state": "running"}}
monkeypatch.setattr(rl, "start_update", fake_start)
out = asyncio.run(rl.llama_update(current_subject = "t"))
assert out.started is True
assert seen["thread"] is not threading.main_thread()

25
studio/backend/tests/test_llama_server_args.py
View file

@ -507,3 +507,28 @@ def test_strip_shadowing_flags_defaults_strip_everything():
["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"]
)
assert out == []
def test_strip_shadowing_flags_drops_model_draft_with_spec():
# --model-draft (and aliases) are Studio-managed since the separate
# MTP drafter support: an inherited copy must not last-wins-override
# the auto-detected drafter.
out = strip_shadowing_flags(
["--model-draft", "/old/mtp.gguf", "-md", "/old2.gguf", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = True,
strip_template = False,
)
assert out == ["--top-k", "20"]
def test_strip_shadowing_flags_keeps_model_draft_without_spec():
out = strip_shadowing_flags(
["--model-draft", "/custom/mtp.gguf"],
strip_context = True,
strip_cache = False,
strip_spec = False,
strip_template = False,
)
assert out == ["--model-draft", "/custom/mtp.gguf"]

341
studio/backend/tests/test_mcp_config_import.py Normal file
View file

@ -0,0 +1,341 @@
"""Tests for MCP config-file import (issue #5936).
Covers the round-trip-safe command join/split inverse (join_stdio_command ↔
parse_stdio_command, on both posix and win32 using the issue's Windows
fixtures), the pure config parser (parse_mcp_config), and the POST /import
route (stdio gate on/off, url dedup, one bad entry not sinking the batch).
Run from studio/backend: python -m pytest tests/test_mcp_config_import.py -q
"""
import sys
import pytest
from core.inference import mcp_client
from core.inference.mcp_config_import import parse_mcp_config
from storage import mcp_servers_db
def _reset_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
def _enable(monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
def _disable(monkeypatch):
monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
# ── 1. join_stdio_command ↔ parse_stdio_command round-trip ──────────
@pytest.mark.parametrize(
"parts",
[
["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
["python", "-m", "mod", "--name", "a b"],
["uvx", "some-server", "--flag"],
["/usr/local/bin/my-server"],
["mcp-server-sqlite"],
],
)
def test_join_parse_roundtrip_posix(monkeypatch, parts):
monkeypatch.setattr(sys, "platform", "linux")
joined = mcp_client.join_stdio_command(parts)
assert mcp_client.parse_stdio_command(joined) == parts
@pytest.mark.parametrize(
"parts",
[
# Issue #5936's literal Windows examples: absolute .exe with a path, and
# backslash drive/dir args must survive the join→split round-trip intact.
[
"C:\\Users\\user\\Documents\\Office-Word-MCP-Server\\.venv\\Scripts\\python.exe",
"C:\\Users\\user\\Documents\\Office-Word-MCP-Server\\word_mcp_server.py",
],
[
"node",
"C:\\Users\\user\\Documents\\DesktopCommanderMCP\\dist\\index.js",
"--no-onboarding",
],
[
"node",
"C:\\Users\\user\\AppData\\Roaming\\npm\\node_modules\\@modelcontextprotocol\\server-filesystem\\dist\\index.js",
"D:\\",
"O:\\",
],
# A command path with spaces is the case that actually needs quoting.
["C:\\Program Files\\node\\node.exe", "server.js"],
["C:\\Program Files\\Foo\\", "server.js"],
["C:\\Program Files\\Foo\\", '{"foo":"bar"}'],
["'C:\\Program Files\\node\\node.exe'", "server.js"],
["node", "O'Reilly"],
["node", "C:\\Users\\O'Reilly\\server.js"],
["node", "'draft'"],
["node", "'open", "close'"],
["node", ""],
],
)
def test_join_parse_roundtrip_win32(monkeypatch, parts):
monkeypatch.setattr(sys, "platform", "win32")
joined = mcp_client.join_stdio_command(parts)
assert mcp_client.parse_stdio_command(joined) == parts
def test_parse_rejects_manual_single_quoted_windows_executable(monkeypatch):
monkeypatch.setattr(sys, "platform", "win32")
command = "'C:\\Program Files\\node\\node.exe' server.js"
with pytest.raises(ValueError):
mcp_client.parse_stdio_command(command)
def test_parse_windows_apostrophes_as_literals(monkeypatch):
monkeypatch.setattr(sys, "platform", "win32")
assert mcp_client.parse_stdio_command("node O'Reilly") == ["node", "O'Reilly"]
assert mcp_client.parse_stdio_command("node C:\\Users\\O'Reilly\\server.js") == [
"node",
"C:\\Users\\O'Reilly\\server.js",
]
assert mcp_client.parse_stdio_command("node 'draft'") == ["node", "'draft'"]
assert mcp_client.parse_stdio_command("node 'open close'") == ["node", "'open", "close'"]
def test_parse_rejects_unterminated_windows_double_quote(monkeypatch):
monkeypatch.setattr(sys, "platform", "win32")
with pytest.raises(ValueError):
mcp_client.parse_stdio_command('node "C:\\path with spaces')
# ── 2. parse_mcp_config ─────────────────────────────────────────────
def test_parse_stdio_entry():
cfg = {
"mcpServers": {
"fs": {
"command": "npx",
"args": ["-y", "server", "/tmp"],
"env": {"K": "v"},
}
}
}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert len(entries) == 1
entry = entries[0]
assert entry.display_name == "fs"
assert entry.is_stdio is True
assert entry.headers == {"K": "v"}
assert mcp_client.parse_stdio_command(entry.url) == ["npx", "-y", "server", "/tmp"]
def test_parse_remote_entry():
cfg = {
"mcpServers": {
"remote": {
"url": "https://example.com/mcp",
"headers": {"Authorization": "Bearer x"},
}
}
}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert entries[0].url == "https://example.com/mcp"
assert entries[0].is_stdio is False
assert entries[0].headers == {"Authorization": "Bearer x"}
def test_parse_preserves_disabled_and_oauth():
cfg = {
"servers": {
"remote": {
"type": "http",
"url": "https://example.com/mcp",
"oauth": {"clientId": "client"},
"disabled": True,
}
}
}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert entries[0].is_enabled is False
assert entries[0].use_oauth is True
def test_parse_accepts_cline_streamable_http_alias():
cfg = {
"mcpServers": {
"remote": {
"type": "streamableHttp",
"url": "https://example.com/mcp",
}
}
}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert entries[0].url == "https://example.com/mcp"
assert entries[0].is_stdio is False
@pytest.mark.parametrize(
"server",
[
{"command": "node", "args": ["server.js"], "cwd": "/tmp/server"},
{"command": "node", "args": ["server.js"], "envFile": ".env"},
{"command": "node", "args": ["server.js"], "env": {"API_KEY": "${input:api-key}"}},
{"command": "node", "args": ["${workspaceFolder}/server.js"]},
{"command": "node", "args": ["server.js"], "env": {"HTTP_PROXY": None}},
{"command": "node", "args": ["server.js"], "sandboxEnabled": True},
{"url": "https://example.com/mcp", "headers": {"Authorization": "Bearer ${input:token}"}},
{"url": "https://example.com/mcp", "headers": {"Authorization": None}},
{"type": "http", "url": "https://example.com/sse"},
{"type": "http", "url": "https://example.com/sse "},
{"type": "streamableHttp", "url": "https://example.com/sse"},
{"url": "https://example.com/mcp", "timeout": 120},
{"url": "https://example.com/mcp", "timeoutMs": 120000},
{"url": "https://example.com/mcp", "timeoutSeconds": 120},
{"type": "sse", "url": "https://example.com/custom"},
],
)
def test_parse_rejects_unrepresentable_imports(server):
entries, errors = parse_mcp_config({"servers": {"bad": server}})
assert entries == []
assert len(errors) == 1
def test_servers_alias_key():
# VS Code uses "servers" instead of "mcpServers".
cfg = {"servers": {"fs": {"command": "node", "args": ["x.js"]}}}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert len(entries) == 1
def test_env_and_args_values_coerced_to_str():
cfg = {"mcpServers": {"fs": {"command": "node", "args": [8080], "env": {"PORT": 8080}}}}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert entries[0].headers == {"PORT": "8080"}
assert mcp_client.parse_stdio_command(entries[0].url) == ["node", "8080"]
def test_args_optional():
cfg = {"mcpServers": {"sqlite": {"command": "mcp-server-sqlite"}}}
entries, errors = parse_mcp_config(cfg)
assert errors == []
assert entries[0].url == "mcp-server-sqlite"
assert entries[0].headers is None
def test_bad_entry_does_not_sink_batch():
cfg = {
"mcpServers": {
"good": {"command": "node", "args": ["x.js"]},
"both": {"command": "node", "url": "https://x/mcp"},
"neither": {"name": "oops"},
"bad_args": {"command": "node", "args": "x.js"},
"bad_env": {"command": "node", "env": ["NOT", "A", "DICT"]},
}
}
entries, errors = parse_mcp_config(cfg)
assert {e.display_name for e in entries} == {"good"}
assert len(errors) == 4
def test_not_a_dict():
entries, errors = parse_mcp_config([])
assert entries == []
assert len(errors) == 1
def test_missing_servers_key():
entries, errors = parse_mcp_config({"foo": {}})
assert entries == []
assert len(errors) == 1
def test_servers_alias_error_names_actual_key():
entries, errors = parse_mcp_config({"servers": []})
assert entries == []
assert errors == ["'servers' must be an object mapping name -> server."]
# ── 3. POST /import route ───────────────────────────────────────────
def test_import_route_creates_and_dedups(tmp_path, monkeypatch):
import asyncio
from models.mcp_servers import McpServerImportRequest
import routes.mcp_servers as routes_mcp
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
cfg = {
"mcpServers": {
"fs": {
"command": "npx",
"args": ["-y", "server", "/tmp"],
"env": {"API_KEY": "sk"},
},
"remote": {"url": "https://example.com/mcp"},
"oauth": {
"type": "http",
"url": "https://auth.example.com/mcp",
"oauth": {"clientId": "client"},
},
"disabled": {
"url": "https://disabled.example.com/mcp",
"disabled": True,
},
}
}
res = asyncio.run(
routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u")
)
assert res.errors == []
assert res.skipped == []
assert {c.display_name for c in res.created} == {"fs", "remote", "oauth", "disabled"}
fs = next(c for c in res.created if c.display_name == "fs")
assert fs.headers == {"API_KEY": "sk"}
assert fs.use_oauth is False
assert fs.is_enabled is True
oauth = next(c for c in res.created if c.display_name == "oauth")
assert oauth.use_oauth is True
disabled = next(c for c in res.created if c.display_name == "disabled")
assert disabled.is_enabled is False
# Re-importing the same config skips both by url.
res2 = asyncio.run(
routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u")
)
assert res2.created == []
assert set(res2.skipped) == {"fs", "remote", "oauth", "disabled"}
def test_import_route_gates_stdio_when_disabled(tmp_path, monkeypatch):
import asyncio
from models.mcp_servers import McpServerImportRequest
import routes.mcp_servers as routes_mcp
_reset_db(tmp_path, monkeypatch)
_disable(monkeypatch)
cfg = {
"mcpServers": {
"fs": {"command": "npx", "args": ["server"]},
"remote": {"url": "https://example.com/mcp"},
}
}
res = asyncio.run(
routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u")
)
# Remote still imports; the stdio entry is rejected per-entry (gate off).
assert {c.display_name for c in res.created} == {"remote"}
assert any("fs" in err for err in res.errors)
assert len(mcp_servers_db.list_servers()) == 1

124
studio/backend/tests/test_mlx_training_worker_config.py
View file

@ -63,6 +63,10 @@ _worker = _load_worker_module()
_normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer
_normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler
_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size
_mlx_vlm_resized_image_layout = _worker._mlx_vlm_resized_image_layout
_copy_mlx_vlm_image_processor = _worker._copy_mlx_vlm_image_processor
_resize_mlx_vlm_image = _worker._resize_mlx_vlm_image
_adapt_for_mlx_vlm = _worker._adapt_for_mlx_vlm
def test_mlx_studio_optimizer_aliases_are_explicit():
@ -90,3 +94,123 @@ def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer():
# Half-pixel cases must match the Torch collator (not banker's round).
assert _mlx_vlm_max_resized_size(333, 1000, 500) == (167, 500)
assert _mlx_vlm_max_resized_size(1000, 333, 500) == (500, 167)
def test_mlx_vlm_resize_keeps_default_numpy_layout_hwc():
Image = pytest.importorskip("PIL.Image")
image = Image.new("RGB", (320, 200), color = (10, 20, 30))
resized = _resize_mlx_vlm_image(image, 128)
assert resized.shape == (80, 128, 3)
assert resized.flags.c_contiguous
def test_mlx_vlm_resize_uses_requested_chw_numpy_layout():
Image = pytest.importorskip("PIL.Image")
image = Image.new("RGB", (320, 200), color = (10, 20, 30))
resized = _resize_mlx_vlm_image(image, 128, image_layout = "chw")
assert resized.shape == (3, 80, 128)
assert resized.flags.c_contiguous
def test_mlx_vlm_resized_image_layout_probes_processor_contract():
class ChwOnlyImageProcessor:
def __call__(self, images = None):
image = images[0]
if image.shape[0] == 3:
return {"pixel_values": image}
raise ValueError("expected CHW")
class HwcImageProcessor:
def __call__(self, images = None):
image = images[0]
if image.shape[-1] == 3:
return {"pixel_values": image}
raise ValueError("expected HWC")
assert (
_mlx_vlm_resized_image_layout(
types.SimpleNamespace(image_processor = ChwOnlyImageProcessor())
)
== "chw"
)
assert (
_mlx_vlm_resized_image_layout(types.SimpleNamespace(image_processor = HwcImageProcessor()))
is None
)
def test_mlx_vlm_layout_probe_copies_image_processor():
class StatefulImageProcessor:
def __init__(self):
self.calls = 0
def __call__(self, images = None):
self.calls += 1
image = images[0]
if image.shape[0] == 3:
return {"pixel_values": image}
raise ValueError("expected CHW")
image_processor = StatefulImageProcessor()
layout = _mlx_vlm_resized_image_layout(types.SimpleNamespace(image_processor = image_processor))
assert layout == "chw"
assert image_processor.calls == 0
def test_mlx_vlm_image_processor_copy_refuses_uncopyable_processors():
class UncopyableImageProcessor:
def __copy__(self):
raise RuntimeError("no copy")
def __deepcopy__(self, _memo):
raise RuntimeError("no deepcopy")
image_processor = UncopyableImageProcessor()
assert _copy_mlx_vlm_image_processor(image_processor) is None
def test_mlx_vlm_layout_probe_skips_uncopyable_processors():
class UncopyableImageProcessor:
def __copy__(self):
raise RuntimeError("no copy")
def __deepcopy__(self, _memo):
raise RuntimeError("no deepcopy")
def __call__(self, images = None):
raise AssertionError("live processor should not be probed")
assert (
_mlx_vlm_resized_image_layout(
types.SimpleNamespace(image_processor = UncopyableImageProcessor())
)
is None
)
def test_mlx_vlm_adapter_applies_chw_layout_to_message_images():
Image = pytest.importorskip("PIL.Image")
image = Image.new("RGB", (320, 200), color = (10, 20, 30))
item = {
"messages": [
{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": "Describe it."},
],
}
]
}
adapted = _adapt_for_mlx_vlm([item], resize = 128, image_layout = "chw")
assert adapted[0]["image"].shape == (3, 80, 128)
assert adapted[0]["messages"][0]["content"][0] == {"type": "image"}

276
studio/backend/tests/test_mtp_drafter_companion.py Normal file
View file

@ -0,0 +1,276 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Separate-file MTP drafter (Gemma 4) contracts.
Pins: the drafter-path predicate and its two layering mirrors, Gemma
effective-size extraction, companion classification in variant plans
(including resume from pre-fix manifests where the drafter leaked into a
quant's main files), and local drafter detection / self-pairing rejection.
"""
from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from hub.utils.download_manifest import ExpectedFile
from hub.utils.gguf import is_mtp_drafter_path
from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files
from utils.models.model_config import (
_is_mtp_drafter,
detect_gguf_model,
detect_mtp_file,
extract_model_size_b,
)
# ── Predicate + layering mirrors ─────────────────────────────────────
DRAFTER_CASES = [
("mtp-gemma-4-12b-it.gguf", True),
("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True),
("foo/MTP/bar.gguf", True),
("gemma-4-12b-it-Q8_0.gguf", False),
# Baked-in Qwen MTP repos: the head is inside the main GGUF, the file
# IS the model -- must never be classified as a companion.
("Qwen3.6-27B-MTP-Q4_K_M.gguf", False),
("prompt-mtp-test.gguf", False),
("smtp/model.gguf", False),
("mtp-readme.txt", False),
]
@pytest.mark.parametrize("path,expected", DRAFTER_CASES)
def test_drafter_predicate_and_mirrors_agree(path, expected):
from core.inference.llama_cpp import _is_companion_gguf_path
assert is_mtp_drafter_path(path) is expected
assert _is_mtp_drafter(path) is expected
# The core mirror bundles mmproj; none of these inputs are mmproj, so
# it must agree with the canonical predicate.
assert _is_companion_gguf_path(path) is expected
# ── Gemma effective-size extraction ──────────────────────────────────
@pytest.mark.parametrize(
"model_id,size_b",
[
("unsloth/gemma-4-E2B-it-GGUF", 2.0),
("unsloth/gemma-4-E4B-it", 4.0),
("unsloth/gemma-3n-E4B-it", 4.0),
# MoE active params beat effective and total notation.
("unsloth/Qwen3.5-35B-A3B", 3.0),
("unsloth/gemma-4-12b-it-GGUF", 12.0),
("unsloth/Qwen3.5-9B-MTP-GGUF", 9.0),
("no-size-here", None),
],
)
def test_extract_model_size_b(model_id, size_b):
assert extract_model_size_b(model_id) == size_b
# ── Variant plan companion classification ────────────────────────────
def _sib(name: str, size: int, sha: str):
return SimpleNamespace(rfilename = name, size = size, lfs = {"sha256": sha})
GEMMA_SIBLINGS = [
_sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"),
_sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"),
_sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"),
_sib("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", 100, "mtp-sub-q8"),
_sib("MTP/gemma-4-12b-it-BF16-MTP.gguf", 200, "mtp-sub-bf16"),
_sib("mmproj-F16.gguf", 500, "mmproj"),
]
def test_variant_plans_carry_drafter_as_companion():
plans = build_gguf_variant_plans(GEMMA_SIBLINGS)
# No phantom quants from the drafter's Q8_0 label or the MTP/ copies.
assert set(plans) == {"q4_k_m", "q8_0"}
for plan in plans.values():
assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames
assert not any("MTP/" in name for name in plan.target_filenames)
assert "drafter" in plan.companion_hashes
assert "drafter" not in plan.main_hashes
assert plan.mmproj_filenames == frozenset({"mmproj-F16.gguf"})
q4 = plans["q4_k_m"]
assert q4.main_filenames == frozenset({"gemma-4-12b-it-Q4_K_M.gguf"})
assert q4.main_size_bytes == 4_000
# Download size = main + mmproj + drafter.
assert q4.download_size_bytes == 4_600
def test_baked_in_repo_plans_unchanged():
plans = build_gguf_variant_plans([_sib("Qwen3.6-27B-MTP-Q4_K_M.gguf", 4_000, "q4")])
assert plans["q4_k_m"].target_filenames == ("Qwen3.6-27B-MTP-Q4_K_M.gguf",)
def test_old_manifest_resume_reclassifies_drafter():
# Pre-fix manifests could leak the drafter into a quant's expected
# files; resume must classify it as a companion, not a main shard.
old = [
ExpectedFile(path = "gemma-4-12b-it-Q8_0.gguf", size = 8_000, sha256 = "main-q8"),
ExpectedFile(path = "mtp-gemma-4-12b-it.gguf", size = 100, sha256 = "drafter"),
]
plan = plan_from_expected_files("Q8_0", old)
assert plan.main_hashes == frozenset({"main-q8"})
assert plan.companion_hashes == frozenset({"drafter"})
assert plan.mmproj_filenames == frozenset()
# ── Local detection / self-pairing ───────────────────────────────────
def test_detect_mtp_file_finds_root_sibling(tmp_path):
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"x")
(tmp_path / "mtp-model.gguf").write_bytes(b"x")
(tmp_path / "MTP").mkdir()
(tmp_path / "MTP" / "model-Q8_0-MTP.gguf").write_bytes(b"x")
found = detect_mtp_file(str(tmp_path / "model-Q4_K_M.gguf"))
assert found is not None
assert found.endswith("mtp-model.gguf")
def test_detect_mtp_file_none_without_sibling(tmp_path):
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"x")
assert detect_mtp_file(str(tmp_path / "model-Q4_K_M.gguf")) is None
def test_detect_gguf_model_rejects_drafter_file(tmp_path):
drafter = tmp_path / "mtp-model.gguf"
drafter.write_bytes(b"x")
assert detect_gguf_model(str(drafter)) is None
def test_detect_gguf_model_dir_skips_companions(tmp_path):
main = tmp_path / "model-Q4_K_M.gguf"
main.write_bytes(b"xxxx")
# Companions are larger so a size-sorted pick would wrongly win.
(tmp_path / "mtp-model.gguf").write_bytes(b"x" * 64)
(tmp_path / "mmproj-F16.gguf").write_bytes(b"x" * 128)
assert detect_gguf_model(str(tmp_path)) == str(main.resolve())
def test_detect_mtp_file_pairs_by_weight_name(tmp_path):
# Multi-model folder: each weight must get its own drafter, never the
# first-sorted foreign one.
(tmp_path / "gemma-4-12b-it-Q4_K_M.gguf").write_bytes(b"x")
(tmp_path / "gemma-4-31B-it-Q4_K_M.gguf").write_bytes(b"x")
(tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x")
(tmp_path / "mtp-gemma-4-31B-it.gguf").write_bytes(b"x")
found = detect_mtp_file(str(tmp_path / "gemma-4-31B-it-Q4_K_M.gguf"))
assert found is not None and found.endswith("mtp-gemma-4-31B-it.gguf")
def test_detect_mtp_file_skips_foreign_drafter(tmp_path):
(tmp_path / "qwen3-8b-Q4_K_M.gguf").write_bytes(b"x")
(tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x")
assert detect_mtp_file(str(tmp_path / "qwen3-8b-Q4_K_M.gguf")) is None
def test_detect_mtp_file_qat_prefix_layout(tmp_path):
# unsloth's qat repo: drafter stem omits the -qat suffix but prefixes
# the weight name (mtp-gemma-4-12B-it.gguf / gemma-4-12B-it-qat-Q4_0.gguf).
(tmp_path / "gemma-4-12B-it-qat-Q4_0.gguf").write_bytes(b"x")
(tmp_path / "mtp-gemma-4-12B-it.gguf").write_bytes(b"x")
found = detect_mtp_file(str(tmp_path / "gemma-4-12B-it-qat-Q4_0.gguf"))
assert found is not None and found.endswith("mtp-gemma-4-12B-it.gguf")
def test_detect_mtp_file_search_root(tmp_path):
# Weight in a quant subdir, drafter at the granted directory root.
sub = tmp_path / "Q4_K_M"
sub.mkdir()
(sub / "gemma-4-12b-it-Q4_K_M.gguf").write_bytes(b"x")
(tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x")
found = detect_mtp_file(str(sub / "gemma-4-12b-it-Q4_K_M.gguf"), search_root = str(tmp_path))
assert found is not None and found.endswith("mtp-gemma-4-12b-it.gguf")
# ── Reload dedup includes the drafter ────────────────────────────────
def _loaded_backend(weight, drafter_path):
from core.inference.llama_cpp import LlamaCppBackend
b = LlamaCppBackend()
# Shape matches atexit cleanup expectations (terminate/wait/kill).
b._process = SimpleNamespace(
poll = lambda: None,
terminate = lambda: None,
wait = lambda timeout = None: 0,
kill = lambda: None,
)
b._healthy = True
b._model_identifier = "local-gemma"
b._gguf_path = str(weight)
b._hf_variant = None
b._requested_n_ctx = 4096
b._cache_type_kv = None
b._requested_spec_mode = "auto"
b._speculative_type = "draft-mtp" if drafter_path else "default"
b._spec_draft_n_max = None
b._chat_template_override = None
b._extra_args = None
b._mtp_draft_path = drafter_path
return b
def _target_state_kwargs(weight, mtp_draft_path):
return dict(
model_identifier = "local-gemma",
hf_variant = None,
n_ctx = 4096,
cache_type_kv = None,
speculative_type = "auto",
spec_draft_n_max = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
gguf_path = str(weight),
mtp_draft_path = mtp_draft_path,
)
def test_already_in_target_state_bounces_on_new_drafter(tmp_path):
weight = tmp_path / "gemma-4-12b-it-Q4_K_M.gguf"
weight.write_bytes(b"x")
drafter = tmp_path / "mtp-gemma-4-12b-it.gguf"
drafter.write_bytes(b"x")
# Loaded without a drafter; one now exists on disk -> must reload.
b = _loaded_backend(weight, None)
assert not b._already_in_target_state(**_target_state_kwargs(weight, str(drafter)))
# Same drafter as launched -> still deduped.
b = _loaded_backend(weight, str(drafter))
assert b._already_in_target_state(**_target_state_kwargs(weight, str(drafter)))
def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path):
# Direct selection of an MTP/ copy: the basename alone has no mtp-
# prefix, so rejection relies on the parent dir name.
sub = tmp_path / "MTP"
sub.mkdir()
copy = sub / "gemma-4-12b-it-BF16-MTP.gguf"
copy.write_bytes(b"x")
assert detect_gguf_model(str(copy)) is None
# Selecting the MTP dir itself must not surface the copies as models.
assert detect_gguf_model(str(sub)) is None

39
studio/backend/tests/test_native_context_length.py
View file

@ -271,6 +271,7 @@ class TestPydanticModels:
def test_load_response_has_field(self):
"""Field exists in LoadResponse.model_fields."""
assert "native_context_length" in LoadResponse.model_fields
assert "context_length" in LoadResponse.model_fields
def test_load_response_defaults_none(self):
"""Omitting native_context_length defaults to None."""
@ -319,6 +320,7 @@ class TestPydanticModels:
def test_status_response_has_field(self):
"""Field exists in InferenceStatusResponse.model_fields."""
assert "native_context_length" in InferenceStatusResponse.model_fields
assert "context_length" in InferenceStatusResponse.model_fields
def test_status_response_has_chat_template_field(self):
"""Status includes chat_template so the UI can rehydrate after refresh."""
@ -347,6 +349,18 @@ class TestPydanticModels:
roundtripped = LoadResponse.model_validate_json(resp.model_dump_json())
assert roundtripped.native_context_length == 131072
def test_context_length_roundtrip(self):
"""Runtime context_length serializes for non-GGUF/hub models."""
resp = LoadResponse(
status = "loaded",
model = "test",
display_name = "Test",
inference = {},
context_length = 8192,
)
roundtripped = LoadResponse.model_validate_json(resp.model_dump_json())
assert roundtripped.context_length == 8192
# =====================================================================
# D. TestRouteCompleteness -- source-level verification
@ -408,6 +422,16 @@ class TestRouteCompleteness:
"native_context_length" not in block
), f"Non-GGUF LoadResponse should not set native_context_length:\n{block[:200]}"
def test_non_gguf_load_responses_set_runtime_context_length(self):
"""Non-GGUF LoadResponse blocks report runtime context_length."""
blocks = self._find_construction_blocks("LoadResponse")
non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b]
assert non_gguf, "Expected at least one non-GGUF LoadResponse block"
for block in non_gguf:
assert (
"context_length" in block
), f"Non-GGUF LoadResponse should set context_length:\n{block[:200]}"
def test_status_path(self):
"""InferenceStatusResponse construction with llama_backend has the field."""
blocks = self._find_construction_blocks("InferenceStatusResponse")
@ -420,6 +444,21 @@ class TestRouteCompleteness:
found
), "No InferenceStatusResponse block with llama_backend has native_context_length"
def test_non_gguf_status_path_reports_runtime_context_length(self):
"""Non-GGUF InferenceStatusResponse reports context_length from model_info."""
blocks = self._find_construction_blocks("InferenceStatusResponse")
found = False
for block in blocks:
if "is_gguf = False" in block and "context_length" in block:
found = True
break
assert found, "No non-GGUF InferenceStatusResponse block with context_length"
def test_openai_models_listing_reports_context_length(self):
"""/v1/models includes context_length when the backend knows it."""
assert 'entry["context_length"]' in self._source
assert 'model_info.get("context_length")' in self._source
# =====================================================================
# E. TestEdgeCases

103
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -166,10 +166,11 @@ class TestChatMessageToolRoles:
with pytest.raises(ValidationError):
ChatMessage(role = "user", content = [])
def test_tool_empty_content_rejected(self):
with pytest.raises(ValidationError) as exc_info:
ChatMessage(role = "tool", tool_call_id = "call_1", content = "")
assert "content" in str(exc_info.value)
def test_tool_empty_content_accepted(self):
# Empty tool output (mkdir, git add, ...) is routine in agentic loops;
# OpenAI and llama-server both accept it, so Studio must not 400.
msg = ChatMessage(role = "tool", tool_call_id = "call_1", content = "")
assert msg.content == ""
def test_assistant_without_content_or_tool_calls_tolerated(self):
# Stop-button leaves an empty assistant turn; tolerate for replay.
@ -654,6 +655,100 @@ class TestBuildPassthroughPayloadToolChoice:
]
# =====================================================================
# Passthrough reasoning kwargs — enable_thinking / reasoning_effort /
# preserve_thinking must reach llama-server via chat_template_kwargs,
# gated on template capabilities like the non-passthrough paths.
# =====================================================================
def _reasoning_backend(
supports_reasoning = True,
reasoning_style = "enable_thinking",
reasoning_always_on = False,
supports_preserve_thinking = False,
):
"""Bare LlamaCppBackend with just the reasoning capability flags set,
so _build_openai_passthrough_body exercises the real
_request_reasoning_kwargs gating."""
from core.inference.llama_cpp import LlamaCppBackend
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._supports_reasoning = supports_reasoning
backend._reasoning_style = reasoning_style
backend._reasoning_always_on = reasoning_always_on
backend._supports_preserve_thinking = supports_preserve_thinking
return backend
class TestPassthroughReasoningKwargs:
def _payload(self, **fields):
return ChatCompletionRequest(
model = "default",
messages = [{"role": "user", "content": "hi"}],
**fields,
)
def test_enable_thinking_forwarded(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = False),
backend_ctx = 4096,
llama_backend = _reasoning_backend(),
)
assert body["chat_template_kwargs"] == {"enable_thinking": False}
def test_preserve_thinking_forwarded_when_template_supports_it(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = True, preserve_thinking = True),
backend_ctx = 4096,
llama_backend = _reasoning_backend(supports_preserve_thinking = True),
)
assert body["chat_template_kwargs"] == {
"enable_thinking": True,
"preserve_thinking": True,
}
def test_preserve_thinking_dropped_when_template_lacks_it(self):
body = _build_openai_passthrough_body(
self._payload(preserve_thinking = True),
backend_ctx = 4096,
llama_backend = _reasoning_backend(supports_preserve_thinking = False),
)
assert "chat_template_kwargs" not in body
def test_reasoning_effort_forwarded_for_effort_style_models(self):
body = _build_openai_passthrough_body(
self._payload(reasoning_effort = "high"),
backend_ctx = 4096,
llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"),
)
assert body["chat_template_kwargs"] == {"reasoning_effort": "high"}
def test_enable_thinking_maps_to_effort_for_effort_style_models(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = False),
backend_ctx = 4096,
llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"),
)
assert body["chat_template_kwargs"] == {"reasoning_effort": "low"}
def test_always_on_reasoning_skips_thinking_kwargs(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = False),
backend_ctx = 4096,
llama_backend = _reasoning_backend(reasoning_always_on = True),
)
assert "chat_template_kwargs" not in body
def test_no_reasoning_fields_omits_chat_template_kwargs(self):
body = _build_openai_passthrough_body(
self._payload(),
backend_ctx = 4096,
llama_backend = _reasoning_backend(supports_preserve_thinking = True),
)
assert "chat_template_kwargs" not in body
# =====================================================================
# OpenAI API compatibility helpers — verified spec edge cases
# =====================================================================

113
studio/backend/tests/test_responses_tool_passthrough.py
View file

@ -58,6 +58,7 @@ from routes.inference import (
_build_chat_request,
_chat_tool_calls_to_responses_output,
_normalise_responses_input,
_responses_tool_output_text,
_responses_stream,
_translate_responses_tool_choice_to_chat,
_translate_responses_tools_to_chat,
@ -393,6 +394,100 @@ class TestNormaliseResponsesInputWithTools:
# Content is serialised so llama-server sees a string.
assert json.loads(msgs[0].content) == [{"type": "output_text", "text": "ok"}]
def test_empty_function_call_output_gets_no_output_sentinel(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": "",
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].role == "tool"
assert msgs[0].tool_call_id == "call_1"
assert msgs[0].content == "(no output)"
ChatMessage(**msgs[0].model_dump(exclude_none = True))
def test_whitespace_function_call_output_gets_no_output_sentinel(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": " \n\t",
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].content == "(no output)"
def test_empty_content_array_output_gets_no_output_sentinel(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [],
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].content == "(no output)"
def test_image_content_array_tool_output_is_serialised(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": "iVBORw0KGgo=",
},
}
],
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].role == "tool"
assert json.loads(msgs[0].content)[0]["type"] == "image"
def test_image_payload_outside_output_gets_no_output_sentinel(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": "",
"content": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": "iVBORw0KGgo=",
},
}
],
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].role == "tool"
assert msgs[0].tool_call_id == "call_1"
assert msgs[0].content == "(no output)"
def test_tool_output_serializer_preserves_non_empty_text(self):
assert _responses_tool_output_text("done") == "done"
assert _responses_tool_output_text(" done ") == " done "
# =====================================================================
# Response mapping — tool_calls → function_call output items
@ -533,6 +628,10 @@ class TestResponsesStreamAdapter:
is_vision = False,
context_length = 4096,
base_url = "http://llama.test",
# Non-reasoning template: the real backend returns None here.
_request_reasoning_kwargs = (
lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None
),
),
)
@ -806,3 +905,17 @@ class TestTranslatedMessagesValidate:
# Building a fresh ChatMessage from the dump round-trips the
# role-shape validator — the passthrough's key invariant.
ChatMessage(**m.model_dump(exclude_none = True))
def test_empty_tool_output_round_trips_through_chat_message_validator(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_empty",
"output": "",
},
],
)
msgs = _normalise_responses_input(payload)
for m in msgs:
ChatMessage(**m.model_dump(exclude_none = True))

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

@ -13,6 +13,7 @@ wrong headroom factor on a 128 GiB unified-memory pool.
from __future__ import annotations
from pathlib import Path
from types import SimpleNamespace
import pytest
@ -28,6 +29,46 @@ def _props(**kwargs) -> SimpleNamespace:
return SimpleNamespace(**kwargs)
# ── Path 0: props.is_integrated (driver's own unified-memory answer) ─────────
class TestIsIntegratedSignal:
"""hipDeviceProp_t.integrated wins when truthy; 0/absent never downgrades.
Same universal gate PR #5988's UMA safetensors fast-load uses -- keeps
Studio's two unified-memory consumers on one signal."""
def test_integrated_upgrades_unknown_apu(self) -> None:
# gfx1103 Phoenix iGPU: outside the hardcoded arch set, but the
# driver says integrated -> unified.
props = _props(gcnArchName = "gfx1103", name = "Radeon 780M", is_integrated = 1)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1103"
assert is_unified is True
def test_integrated_wins_without_any_arch(self) -> None:
props = _props(name = "Some Future APU", is_integrated = 1)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert is_unified is True
def test_zero_does_not_downgrade_known_apu(self) -> None:
# A wheel that zeroes the field must not flip Strix Halo to discrete.
props = _props(gcnArchName = "gfx1151", name = "x", is_integrated = 0)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert is_unified is True
def test_absent_keeps_existing_behavior(self) -> None:
props = _props(gcnArchName = "gfx1201", name = "RX 9070 XT")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert is_unified is False
def test_discrete_with_zero_stays_discrete(self) -> None:
props = _props(gcnArchName = "gfx1100", name = "RX 7900 XTX", is_integrated = 0)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert is_unified is False
# ── Path 1: canonical gcnArchName ────────────────────────────────────────────
@ -168,3 +209,35 @@ class TestDeviceNameFallback:
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert is_unified is False
# ── Fraction selection (source-pinned) ───────────────────────────────────────
_WORKER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py"
class TestMemFractionSelection:
"""Pin the per-platform fraction policy in worker.py section 1g.
On native Windows, torch.cuda.mem_get_info's total is the WDDM budget
the driver grants HIP -- the OS share of RAM is already outside it, so
a 0.80 cap double-taxes (field report: 48.49 GiB budget -> '38.79 GiB
allowed' OOM denying a 47.29 GiB load that fit in free memory). 1.0
removes the double-tax; current AMD Windows wheels enforce only
sub-1.0 fractions, so it behaves like torch's uncapped default with
WDDM arbitrating residency (measured on gfx1151)."""
def test_unified_win32_uses_budget_exact_fraction(self) -> None:
source = _WORKER_PY.read_text(encoding = "utf-8")
assert '1.0 if sys.platform == "win32" else 0.80' in source
def test_discrete_keeps_090(self) -> None:
source = _WORKER_PY.read_text(encoding = "utf-8")
assert "_mem_fraction = 0.90" in source
def test_win32_unified_logs_vgm_hint(self) -> None:
"""Users must learn the WDDM budget is raisable (BIOS UMA / AMD
Software Variable Graphics Memory) instead of assuming a bug."""
source = _WORKER_PY.read_text(encoding = "utf-8")
assert "Variable Graphics Memory" in source

99
studio/backend/tests/test_server_disk_logging.py Normal file
View file

@ -0,0 +1,99 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the server session log + native-crash capture in run.py.
Field regression: Studio "terminates without a warning" -- a native crash in
the GPU runtime kills the process with no Python traceback, and a desktop-
shortcut console closes before anything can be read. The server must tee its
console output to disk and aim faulthandler at the same file so even hard
crashes leave evidence.
"""
from __future__ import annotations
import io
import sys
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import run as run_mod # noqa: E402
class TestTeeStream:
def test_writes_reach_both_and_return_original(self):
console, log = io.StringIO(), io.StringIO()
tee = run_mod._TeeStream(console, log)
n = tee.write("hello")
assert console.getvalue() == "hello" == log.getvalue()
assert n == 5 # delegate's return value, console contract unchanged
def test_log_failure_never_breaks_console(self):
class Broken:
def write(self, data):
raise OSError("disk full")
def flush(self):
raise OSError("disk full")
console = io.StringIO()
tee = run_mod._TeeStream(console, Broken())
assert tee.write("still works") == len("still works")
tee.flush() # must not raise
assert console.getvalue() == "still works"
def test_attribute_proxy(self):
console, log = io.StringIO(), io.StringIO()
tee = run_mod._TeeStream(console, log)
# isatty / encoding probes must see the original stream's answers.
assert tee.isatty() == console.isatty()
class TestSetupServerDiskLogging:
def test_opt_out_env(self, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_NO_FILE_LOG", "1")
assert run_mod._setup_server_disk_logging() is None
def test_creates_log_and_enables_faulthandler(self, monkeypatch, tmp_path):
import faulthandler
monkeypatch.delenv("UNSLOTH_STUDIO_NO_FILE_LOG", raising = False)
monkeypatch.delenv("PYTHONFAULTHANDLER", raising = False)
# Both resolution paths (utils.paths.studio_root and the env
# fallback) honor UNSLOTH_STUDIO_HOME, so this redirects the log dir.
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
orig_out, orig_err = sys.stdout, sys.stderr
was_enabled = faulthandler.is_enabled()
try:
log_path = run_mod._setup_server_disk_logging()
assert log_path is not None
assert Path(log_path).is_file()
assert "logs" in str(log_path)
# faulthandler armed at the file; children inherit the env switch.
assert faulthandler.is_enabled()
import os
assert os.environ.get("PYTHONFAULTHANDLER") == "1"
print("tee-capture-marker")
sys.stdout.flush()
assert "tee-capture-marker" in Path(log_path).read_text(
encoding = "utf-8", errors = "replace"
)
finally:
sys.stdout, sys.stderr = orig_out, orig_err
if not was_enabled:
faulthandler.disable()
def test_run_server_wires_logging_before_main_import(self):
src = (Path(_BACKEND_DIR) / "run.py").read_text(encoding = "utf-8")
call_idx = src.index("_setup_server_disk_logging()", src.index("def run_server"))
main_import_idx = src.index("from main import app", src.index("def run_server"))
assert call_idx < main_import_idx, (
"disk logging must be armed before importing main so import-time "
"failures leave evidence on disk"
)

53
studio/backend/tests/test_tool_message_empty_content.py Normal file
View file

@ -0,0 +1,53 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Empty ``role="tool"`` content must be accepted on the OpenAI-compat surface.
Agentic clients send ``content: ""`` when a command produced no output;
OpenAI and llama-server both accept it. Studio used to 400, which standard
clients treat as non-retryable and kill the session. The validator must
normalize empty/missing tool content to ``""`` instead of raising.
"""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from models.inference import ChatMessage
def test_tool_message_empty_string_content_is_accepted():
msg = ChatMessage(role = "tool", content = "", tool_call_id = "call_1")
assert msg.content == ""
def test_tool_message_none_content_normalizes_to_empty_string():
msg = ChatMessage(role = "tool", content = None, tool_call_id = "call_1")
assert msg.content == ""
def test_tool_message_empty_list_content_normalizes_to_empty_string():
msg = ChatMessage(role = "tool", content = [], tool_call_id = "call_1")
assert msg.content == ""
def test_tool_message_real_content_is_preserved():
msg = ChatMessage(role = "tool", content = "ok", tool_call_id = "call_1")
assert msg.content == "ok"
def test_user_message_still_requires_content():
with pytest.raises(ValueError):
ChatMessage(role = "user", content = None)
def test_assistant_empty_content_still_collapses_to_none():
msg = ChatMessage(role = "assistant", content = "")
assert msg.content is None

110
studio/backend/tests/test_training_nan_loss_handling.py Normal file
View file

@ -0,0 +1,110 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Pin Studio's behavior when a training event reports non-finite (NaN/Inf) loss.
The training event handler used to filter NaN/Inf to None silently while
leaving the previous finite loss in progress.loss — so the API kept reporting
the stale value as if everything were fine. We now drop the stale value:
clients see loss=None at the affected step and a one-shot warning is logged.
Training continues; the run is not marked failed.
"""
from __future__ import annotations
import math
import os
import sys
import pytest
_BACKEND = os.path.join(os.path.dirname(__file__), "..")
if _BACKEND not in sys.path:
sys.path.insert(0, _BACKEND)
from core.training.training import TrainingBackend
def _make_backend() -> TrainingBackend:
return TrainingBackend()
def _progress_event(
step: int,
loss: float,
lr: float = 1e-4,
) -> dict:
return {
"type": "progress",
"step": step,
"loss": loss,
"learning_rate": lr,
"epoch": 0.0,
"total_steps": 100,
}
class TestNonfiniteLossSoftHandling:
def test_finite_loss_updates_progress_normally(self):
b = _make_backend()
b._handle_event(_progress_event(step = 1, loss = 0.97))
assert b._progress.loss == pytest.approx(0.97)
assert b._progress.error is None
assert b._should_stop is False
assert getattr(b._progress, "_nonfinite_loss_warned", False) is False
def test_nan_loss_clears_progress_loss(self):
b = _make_backend()
b._handle_event(_progress_event(step = 1, loss = 0.97))
assert b._progress.loss == pytest.approx(0.97)
b._handle_event(_progress_event(step = 2, loss = float("nan")))
# Stale finite loss must NOT leak through
assert b._progress.loss is None
# Run is not marked failed
assert b._progress.error is None
assert b._should_stop is False
# Warning flag is set so we don't re-log on every subsequent NaN step
assert b._progress._nonfinite_loss_warned is True
def test_inf_loss_clears_progress_loss(self):
b = _make_backend()
b._handle_event(_progress_event(step = 1, loss = float("inf")))
assert b._progress.loss is None
assert b._progress.error is None
assert b._should_stop is False
assert b._progress._nonfinite_loss_warned is True
def test_negative_inf_loss_clears_progress_loss(self):
b = _make_backend()
b._handle_event(_progress_event(step = 1, loss = float("-inf")))
assert b._progress.loss is None
assert b._progress.error is None
assert b._should_stop is False
assert b._progress._nonfinite_loss_warned is True
def test_repeated_nan_only_warns_once(self):
"""Subsequent NaN events must not re-fire the warning flag setter.
The flag should already be True after the first NaN."""
b = _make_backend()
b._handle_event(_progress_event(step = 1, loss = 0.97))
b._handle_event(_progress_event(step = 2, loss = float("nan")))
assert b._progress._nonfinite_loss_warned is True
# Further NaN steps don't change anything we care about
b._handle_event(_progress_event(step = 3, loss = float("nan")))
b._handle_event(_progress_event(step = 4, loss = float("nan")))
assert b._progress._nonfinite_loss_warned is True
assert b._progress.loss is None
assert b._progress.error is None
assert b._should_stop is False
def test_recovery_updates_loss_when_finite_again(self):
"""If a NaN step is followed by a finite step, progress.loss must
reflect the new finite value (not stay stuck at None)."""
b = _make_backend()
b._handle_event(_progress_event(step = 1, loss = 0.97))
b._handle_event(_progress_event(step = 2, loss = float("nan")))
assert b._progress.loss is None
b._handle_event(_progress_event(step = 3, loss = 0.85))
assert b._progress.loss == pytest.approx(0.85)
# Warning flag stays set (we don't reset it on recovery)
assert b._progress._nonfinite_loss_warned is True

129
studio/backend/tests/test_training_progress_stream_nan.py Normal file
View file

@ -0,0 +1,129 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The SSE progress stream must follow the live progress step during
non-finite-loss stretches (loss reported as null) instead of replaying the
last finite step/loss pair from the metric histories, which skip NaN steps."""
import asyncio
import json
import sys
import types
import pytest
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules["structlog"] = types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
)
import routes.training as rt
class _Progress:
def __init__(self):
self.step = 5
self.total_steps = 10
self.loss = None # cleared by the NaN honesty fix in core training
self.learning_rate = 8e-5
self.epoch = 0.1
self.grad_norm = None
self.num_tokens = None
self.eval_loss = None
self.elapsed_seconds = None
self.eta_seconds = None
class _FakeBackend:
"""Finite history stops at step 2; live progress is at step 5 with NaN
(loss=None). Active for a few polls, then done."""
def __init__(self, active_polls = 2):
self.current_job_id = "job-1"
self.step_history = [1, 2]
self.loss_history = [2.0, 1.5]
self.lr_history = [1e-4, 9e-5]
self.eval_enabled = False
self._active_calls = 0
self._active_polls = active_polls
self.trainer = types.SimpleNamespace(training_progress = _Progress())
def is_training_active(self):
self._active_calls += 1
return self._active_calls <= self._active_polls
class _FakeRequest:
headers = {}
def _collect_events(response, timeout = 15):
async def _drain():
chunks = []
async for chunk in response.body_iterator:
chunks.append(chunk)
return "".join(c.decode() if isinstance(c, bytes) else c for c in chunks)
return asyncio.run(asyncio.wait_for(_drain(), timeout))
def _progress_payloads(raw):
payloads = []
for block in raw.split("\n\n"):
lines = block.strip().splitlines()
data = next((l[6:] for l in lines if l.startswith("data: ")), None)
if data:
payloads.append(json.loads(data))
return payloads
def test_stream_reports_live_step_with_null_loss_during_nan(monkeypatch):
backend = _FakeBackend(active_polls = 2)
monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))
raw = _collect_events(response)
payloads = _progress_payloads(raw)
assert payloads, f"no SSE payloads parsed from: {raw!r}"
live = [p for p in payloads if p.get("step") == 5]
assert live, (
"stream never advanced to the live progress step during the NaN "
f"stretch; steps seen: {[p.get('step') for p in payloads]}"
)
assert live[0]["loss"] is None
# The stale finite pair must not be re-emitted as the latest progress.
stale = [p for p in payloads if p.get("step") == 2 and p.get("loss") == 1.5]
assert not stale
def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch):
# Fresh connection after the run already ended during a NaN stretch: the
# immediate complete event must not replay the stale finite pair either.
backend = _FakeBackend(active_polls = 0)
monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))
payloads = _progress_payloads(_collect_events(response))
final = payloads[-1]
assert final["step"] == 5
assert final["loss"] is None
def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch):
backend = _FakeBackend(active_polls = 2)
# Live progress agrees with the history tail: normal finite behavior.
backend.trainer.training_progress.step = 2
backend.trainer.training_progress.loss = 1.5
monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))
payloads = _progress_payloads(_collect_events(response))
finite = [p for p in payloads if p.get("step") == 2]
assert finite and finite[0]["loss"] == 1.5

224
studio/backend/tests/test_transformers_version.py
View file

@ -30,8 +30,11 @@ sys.modules.setdefault("loggers", _loggers_stub)
from utils.transformers_version import (
_resolve_base_model,
_check_tokenizer_config_needs_v5,
_check_config_needs_510,
_check_config_needs_550,
_config_json_cache,
_tokenizer_class_cache,
_config_needs_510_cache,
_config_needs_550_cache,
needs_transformers_5,
get_transformers_tier,
@ -200,6 +203,7 @@ class TestCheckConfigNeeds550:
"""Tests for _check_config_needs_550() local config.json checks."""
def setup_method(self):
_config_json_cache.clear()
_config_needs_550_cache.clear()
def test_gemma4_architecture(self, tmp_path: Path):
@ -253,6 +257,106 @@ class TestCheckConfigNeeds550:
mock_urlopen.assert_not_called()
# ---------------------------------------------------------------------------
# _check_config_needs_510 — config.json architecture/model_type check
# ---------------------------------------------------------------------------
class TestCheckConfigNeeds510:
"""Tests for _check_config_needs_510() local config.json checks."""
def setup_method(self):
_config_json_cache.clear()
_config_needs_510_cache.clear()
def test_gemma4_unified_architecture(self, tmp_path: Path):
"""config.json with Gemma4UnifiedForConditionalGeneration should return True."""
cfg = {
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
"model_type": "gemma4_unified",
}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is True
def test_gemma4_unified_model_type_only(self, tmp_path: Path):
"""config.json with model_type=gemma4_unified should return True."""
cfg = {"model_type": "gemma4_unified"}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is True
def test_gemma4_unified_assistant_architecture(self, tmp_path: Path):
"""Assistant Gemma 4 Unified configs should return True."""
cfg = {
"architectures": ["Gemma4UnifiedAssistantForCausalLM"],
"model_type": "gemma4_unified_assistant",
}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is True
def test_gemma4_unified_assistant_model_type_only(self, tmp_path: Path):
"""Assistant Gemma 4 Unified model_type should return True."""
cfg = {"model_type": "gemma4_unified_assistant"}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is True
def test_gemma4_assistant_architecture(self, tmp_path: Path):
"""Assistant Gemma 4 configs should return True."""
cfg = {
"architectures": ["Gemma4AssistantForCausalLM"],
"model_type": "gemma4_assistant",
}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is True
def test_gemma4_assistant_model_type_only(self, tmp_path: Path):
"""Assistant Gemma 4 model_type should return True."""
cfg = {"model_type": "gemma4_assistant"}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is True
def test_gemma4_non_unified_returns_false(self, tmp_path: Path):
"""Older Gemma 4 config should stay on the 550 tier."""
cfg = {
"architectures": ["Gemma4ForConditionalGeneration"],
"model_type": "gemma4",
}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert _check_config_needs_510(str(tmp_path)) is False
def test_no_config_json(self, tmp_path: Path):
"""Missing config.json should return False (fail-open)."""
# Patch network call to avoid real fetch
with patch("urllib.request.urlopen") as mock_urlopen:
mock_urlopen.side_effect = Exception("no network")
assert _check_config_needs_510(str(tmp_path)) is False
def test_result_is_cached(self, tmp_path: Path):
"""Subsequent calls should use the cache."""
cfg = {"architectures": ["Gemma4UnifiedForConditionalGeneration"]}
(tmp_path / "config.json").write_text(json.dumps(cfg))
key = str(tmp_path)
_check_config_needs_510(key)
assert key in _config_needs_510_cache
assert _config_needs_510_cache[key] is True
def test_local_file_skips_network(self, tmp_path: Path):
"""When local config.json exists, no network request should be made."""
cfg = {"architectures": ["LlamaForCausalLM"]}
(tmp_path / "config.json").write_text(json.dumps(cfg))
with patch("urllib.request.urlopen") as mock_urlopen:
_check_config_needs_510(str(tmp_path))
mock_urlopen.assert_not_called()
# ---------------------------------------------------------------------------
# get_transformers_tier — tier detection
# ---------------------------------------------------------------------------
@ -263,11 +367,19 @@ class TestGetTransformersTier:
def setup_method(self):
_tokenizer_class_cache.clear()
_config_json_cache.clear()
_config_needs_510_cache.clear()
_config_needs_550_cache.clear()
def test_gemma4_substring_returns_550(self):
assert get_transformers_tier("google/gemma-4-E2B-it") == "550"
def test_gemma4_12b_substring_returns_510(self):
assert get_transformers_tier("unsloth/gemma-4-12b-it") == "510"
def test_gemma4_assistant_substring_returns_510(self):
assert get_transformers_tier("google/gemma-4-E2B-it-assistant") == "510"
def test_gemma4_alt_substring_returns_550(self):
assert get_transformers_tier("unsloth/gemma4-E4B-it") == "550"
@ -281,17 +393,92 @@ class TestGetTransformersTier:
assert get_transformers_tier(str(tmp_path)) == "550"
def test_gemma4_unified_config_json_returns_510(self, tmp_path: Path):
"""Local checkpoint with Gemma4 Unified architecture → 510."""
cfg = {
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
"model_type": "gemma4_unified",
}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert get_transformers_tier(str(tmp_path)) == "510"
def test_gemma4_assistant_config_json_returns_510(self, tmp_path: Path):
"""Local checkpoint with Gemma4 Assistant architecture → 510."""
cfg = {
"architectures": ["Gemma4AssistantForCausalLM"],
"model_type": "gemma4_assistant",
}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert get_transformers_tier(str(tmp_path)) == "510"
def test_local_config_json_short_circuits_path_substrings(self, tmp_path: Path):
"""Local config.json should prevent false matches from parent directory names."""
model_dir = tmp_path / "gemma-4-12b-experiment" / "llama-checkpoint"
model_dir.mkdir(parents = True)
(model_dir / "config.json").write_text(
json.dumps(
{
"architectures": ["LlamaForCausalLM"],
"model_type": "llama",
}
)
)
(model_dir / "tokenizer_config.json").write_text(
json.dumps({"tokenizer_class": "LlamaTokenizerFast"})
)
with patch("urllib.request.urlopen") as mock_urlopen:
assert get_transformers_tier(str(model_dir)) == "default"
mock_urlopen.assert_not_called()
def test_remote_config_json_is_fetched_once_for_config_tiers(self):
"""510 and 550 slow-path checks should share one config.json fetch."""
class _Response:
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return False
def read(self):
return json.dumps(
{
"architectures": ["Gemma4ForConditionalGeneration"],
"model_type": "gemma4",
}
).encode()
with patch("urllib.request.urlopen", return_value = _Response()) as mock_urlopen:
assert get_transformers_tier("org/no-fast-substring-model") == "550"
assert mock_urlopen.call_count == 1
def test_qwen35_returns_530(self):
with patch(
"utils.transformers_version._check_config_needs_550",
return_value = False,
with (
patch(
"utils.transformers_version._check_config_needs_550",
return_value = False,
),
patch(
"utils.transformers_version._check_config_needs_510",
return_value = False,
),
):
assert get_transformers_tier("Qwen/Qwen3.5-9B") == "530"
def test_ministral_returns_530(self):
with patch(
"utils.transformers_version._check_config_needs_550",
return_value = False,
with (
patch(
"utils.transformers_version._check_config_needs_550",
return_value = False,
),
patch(
"utils.transformers_version._check_config_needs_510",
return_value = False,
),
):
assert get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530"
@ -301,6 +488,10 @@ class TestGetTransformersTier:
"utils.transformers_version._check_config_needs_550",
return_value = False,
),
patch(
"utils.transformers_version._check_config_needs_510",
return_value = False,
),
patch(
"utils.transformers_version._check_tokenizer_config_needs_v5",
return_value = False,
@ -309,15 +500,22 @@ class TestGetTransformersTier:
assert get_transformers_tier("meta-llama/Llama-3-8B") == "default"
def test_550_checked_before_530(self):
"""5.5.0 is checked first — a model matching both gets 550."""
"""5.5.0 is checked before 5.3.0 - a model matching both gets 550."""
assert get_transformers_tier("gemma-4-model") == "550"
def test_needs_transformers_5_compat(self):
"""needs_transformers_5 should return True for both 530 and 550 models."""
"""needs_transformers_5 should return True for 510, 530, and 550 models."""
assert needs_transformers_5("unsloth/gemma-4-12b-it") is True
assert needs_transformers_5("google/gemma-4-E2B-it") is True
with patch(
"utils.transformers_version._check_config_needs_550",
return_value = False,
with (
patch(
"utils.transformers_version._check_config_needs_550",
return_value = False,
),
patch(
"utils.transformers_version._check_config_needs_510",
return_value = False,
),
):
assert needs_transformers_5("Qwen/Qwen3.5-9B") is True
with (
@ -325,6 +523,10 @@ class TestGetTransformersTier:
"utils.transformers_version._check_config_needs_550",
return_value = False,
),
patch(
"utils.transformers_version._check_config_needs_510",
return_value = False,
),
patch(
"utils.transformers_version._check_tokenizer_config_needs_v5",
return_value = False,

51
studio/backend/utils/datasets/cache_safe.py Normal file
View file

@ -0,0 +1,51 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Permission-safe wrapper around datasets.load_dataset.
A shared HF datasets cache can contain subtrees owned by another user (for
example populated by an earlier root-run job). datasets then raises
"[Errno 13] Permission denied: ..._builder.lock" while locking the cached
builder, killing the training run even though the dataset itself is fine.
Retry such loads in a Studio-owned cache so the run proceeds; the worst case
is one rebuild of the dataset in the fallback location.
"""
import logging
import os
from utils.paths.storage_roots import cache_root
logger = logging.getLogger(__name__)
def studio_datasets_cache() -> str:
path = cache_root() / "hf-datasets"
path.mkdir(parents = True, exist_ok = True)
return str(path)
def load_dataset_cache_safe(*args, **kwargs):
"""datasets.load_dataset, retried in a Studio-owned cache on EACCES."""
from datasets import load_dataset
try:
return load_dataset(*args, **kwargs)
except PermissionError as error:
fallback = studio_datasets_cache()
logger.warning(
"HF datasets cache is not writable (%s); rebuilding in %s",
error,
fallback,
)
kwargs["cache_dir"] = fallback
# Nested builders consult the env var while the load runs; restore it
# after so other datasets keep trying the shared cache first.
old_env = os.environ.get("HF_DATASETS_CACHE")
os.environ["HF_DATASETS_CACHE"] = fallback
try:
return load_dataset(*args, **kwargs)
finally:
if old_env is None:
os.environ.pop("HF_DATASETS_CACHE", None)
else:
os.environ["HF_DATASETS_CACHE"] = old_env

68
studio/backend/utils/hardware/amd.py
View file

@ -12,6 +12,7 @@ import math
import os
import platform
import re
import shutil
import subprocess
import sys
from typing import Any, Optional
@ -33,24 +34,85 @@ _amd_smi_consecutive_failures = 0
_amd_smi_disabled = False
def _hip_sdk_present() -> bool:
"""True if a HIP SDK is detectable (hipinfo on PATH or under HIP_PATH/
ROCM_PATH), meaning amd-smi has a working runtime and runs un-elevated."""
if shutil.which("hipinfo"):
return True
for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
root = os.environ.get(var)
if root and os.path.exists(os.path.join(root, "bin", "hipinfo.exe")):
return True
return False
def _amd_smi_allowed() -> bool:
"""Whether it is safe to spawn amd-smi here.
On Windows without a working HIP runtime, amd-smi elevates a child at
runtime -- popping a UAC/DiskPart prompt that RunAsInvoker can't suppress
(its manifest is asInvoker). So only call it on Windows with a HIP SDK
present or UNSLOTH_ENABLE_AMD_SMI=1. Linux amd-smi never elevates.
"""
if platform.system() != "Windows":
return True
flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower()
if flag in ("1", "true", "yes", "on"):
return True
if flag in ("0", "false", "no", "off"):
return False
return _hip_sdk_present()
def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]:
"""Run amd-smi with the given args and return parsed JSON, or None."""
global _amd_smi_consecutive_failures, _amd_smi_disabled
if _amd_smi_disabled:
return None
if not _amd_smi_allowed():
# Permanently skip amd-smi on Windows w/o a HIP SDK: every call would
# pop a UAC/DiskPart prompt (see _amd_smi_allowed). VRAM polling is then
# unavailable, but that beats the prompt. Opt back in with
# UNSLOTH_ENABLE_AMD_SMI=1.
if not _amd_smi_disabled:
logger.info(
"amd-smi disabled on Windows (no HIP SDK detected) to avoid a "
"UAC/DiskPart elevation prompt; GPU VRAM polling unavailable. "
"Set UNSLOTH_ENABLE_AMD_SMI=1 to force amd-smi."
)
_amd_smi_disabled = True
return None
if shutil.which("amd-smi") is None:
# amd-smi does not exist on Windows (neither Adrenalin nor the HIP SDK
# ship a CLI) and can be absent on minimal Linux installs. Disable the
# poller in one step instead of burning the 3-strike circuit breaker
# on guaranteed FileNotFoundError spawns. Studio's VRAM display falls
# back to torch mem_get_info.
if not _amd_smi_disabled:
logger.info(
"amd-smi not found on PATH; GPU utilization polling via "
"amd-smi unavailable (VRAM falls back to torch mem_get_info)."
)
_amd_smi_disabled = True
return None
_amd_env = child_env_without_native_path_secret()
if platform.system() == "Windows":
# RunAsInvoker belt-and-suspenders for any manifest-elevating helper;
# the real guard is _amd_smi_allowed() above. Mirrors install scripts.
_amd_env = {**_amd_env, "__COMPAT_LAYER": "RunAsInvoker"}
try:
result = subprocess.run(
["amd-smi", *args, "--json"],
capture_output = True,
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
env = _amd_env,
**windows_hidden_subprocess_kwargs(),
)
except (OSError, subprocess.TimeoutExpired) as e:
if isinstance(e, FileNotFoundError):
# amd-smi ships with Adrenalin, not the HIP SDK; absence is expected
# on HIP SDK-only Windows setups.
# Raced a PATH change after the which() check above; absence is
# expected on Windows (no AMD product ships an amd-smi CLI there).
logger.debug("amd-smi not found (not in PATH): %s", e)
else:
logger.warning("amd-smi query failed: %s", e)

8
studio/backend/utils/hardware/hardware.py
View file

@ -1818,7 +1818,13 @@ def apply_gpu_ids(gpu_ids) -> None:
)
if _is_rocm:
os.environ["HIP_VISIBLE_DEVICES"] = value
os.environ["ROCR_VISIBLE_DEVICES"] = value
# ROCR_VISIBLE_DEVICES operates at the HSA agent level and uses
# different indexing semantics to HIP_VISIBLE_DEVICES. Setting it
# to a physical GPU index breaks multi-GPU ROCm systems where the
# parent already set ROCR_VISIBLE_DEVICES (e.g. "0,1"): narrowing
# to "1" causes torch.cuda.is_available() to return False in the
# worker subprocess. HIP_VISIBLE_DEVICES is sufficient for GPU
# selection on ROCm -- leave ROCR_VISIBLE_DEVICES inherited.
_visible_gpu_count = None
if _is_rocm:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)

87
studio/backend/utils/llama_cpp_freshness.py
View file

@ -13,6 +13,7 @@ from __future__ import annotations
import json
import os
import re
import time
from datetime import datetime, timezone
from pathlib import Path
@ -104,11 +105,18 @@ def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None:
def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
"""GitHub API call. None on any failure (offline, rate-limited, etc)."""
"""Newest published release tag for `repo`, by publish time.
Resolves "latest" the way install_llama_prebuilt.py does (newest
non-draft/non-prerelease by ``published_at``), NOT via GitHub's
``/releases/latest`` pointer. That pointer sorts by commit date and can lag
behind the build the installer actually installs, so detection and apply
disagreed -- the cause of the downgrade/sticky banner. None on any failure
(offline, rate-limited, etc)."""
import urllib.error
import urllib.request
url = f"https://api.github.com/repos/{repo}/releases/latest"
url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
headers = {
"Accept": "application/vnd.github+json",
"User-Agent": "unsloth-studio-freshness-check",
@ -128,8 +136,21 @@ def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
) as exc:
logger.debug("freshness fetch failed", repo = repo, error = str(exc))
return None
tag = data.get("tag_name")
return tag if isinstance(tag, str) and tag else None
if not isinstance(data, list):
return None
published = [
r
for r in data
if isinstance(r, dict)
and not r.get("draft")
and not r.get("prerelease")
and isinstance(r.get("tag_name"), str)
and r.get("tag_name")
]
if not published:
return None
newest = max(published, key = lambda r: r.get("published_at") or "")
return newest["tag_name"]
def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]:
@ -172,19 +193,58 @@ def _parse_installed_at(value: object) -> Optional[datetime]:
return dt
def parse_base_build(tag: object) -> Optional[int]:
"""Numeric base build from a release tag. Handles both a plain ``bNNNN`` and
a mix-build tag like ``b9596-mix-<sha>`` (anchored at the start, so the mix
suffix doesn't defeat it). None for anything not starting with ``bNNNN``."""
if not isinstance(tag, str):
return None
m = re.match(r"b(\d+)", tag.strip())
return int(m.group(1)) if m else None
def is_behind(installed: Optional[str], latest: Optional[str]) -> bool:
"""Whether `installed` is genuinely behind `latest`, comparing the FULL
release identity (so a mix build can legitimately be the latest) with a
base-build guard so a lagging GitHub /releases/latest can never read as an
update or a downgrade.
- identical tags -> not behind (clears the sticky banner post-update)
- higher base build on `latest` -> behind; lower -> NOT behind (downgrade guard)
- same base build: a different/new mix -> behind, but a bare ``bNNNN`` never
supersedes a mix build (extra PRs) at that base -> not behind
- non-bNNNN tags -> behind (plain inequality, since they already differ)
"""
if not installed or not latest:
return False
installed, latest = installed.strip(), latest.strip()
if installed == latest:
return False
ib, lb = parse_base_build(installed), parse_base_build(latest)
if ib is None or lb is None:
return True
if lb != ib:
return lb > ib
# Same base build, different tags: offer a mix (latest carries a suffix), but
# never offer a bare base over a mix install at the same base.
return latest != f"b{lb}"
def check_prebuilt_freshness(
binary_path: Optional[str],
*,
threshold_days: int = STALENESS_THRESHOLD_DAYS,
now: Optional[datetime] = None,
) -> dict:
"""Returns {has_marker, stale, installed_tag, latest_tag,
"""Returns {has_marker, stale, behind, installed_tag, latest_tag,
installed_at_utc, age_days, published_repo, threshold_days}.
stale = True iff installed != latest AND age >= threshold.
Fails open on missing data (stale stays False)."""
behind = installed genuinely older than latest (see is_behind).
stale = behind AND age >= threshold.
Fails open on missing data (behind/stale stay False)."""
out: dict = {
"has_marker": False,
"stale": False,
"behind": False,
"installed_tag": None,
"latest_tag": None,
"installed_at_utc": None,
@ -196,16 +256,25 @@ def check_prebuilt_freshness(
if not marker:
return out
out["has_marker"] = True
# Display prefers the normalized base ("tag"); comparison below prefers the
# full "release_tag" -- deliberately opposite fallbacks.
out["installed_tag"] = marker.get("tag") or marker.get("release_tag")
out["installed_at_utc"] = marker.get("installed_at_utc")
out["published_repo"] = marker.get("published_repo")
# The marker records both a normalized base tag ("tag", e.g. b9596) and the
# full release tag ("release_tag", e.g. b9596-mix-<sha>). Compare against the
# FULL identity, since GitHub /releases/latest returns the full tag_name --
# comparing the normalized base against the full latest is what produced the
# permanent "downgrade" banner on every mix release.
installed_full = marker.get("release_tag") or marker.get("tag")
repo = out["published_repo"]
if not repo or not out["installed_tag"]:
if not repo or not installed_full:
return out
latest = latest_published_release(repo)
out["latest_tag"] = latest
if not latest or latest == out["installed_tag"]:
out["behind"] = is_behind(installed_full, latest)
if not out["behind"]:
return out
installed_at = _parse_installed_at(out["installed_at_utc"])

560
studio/backend/utils/llama_cpp_update.py Normal file
View file

@ -0,0 +1,560 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""In-app llama.cpp prebuilt update.
Builds on utils.llama_cpp_freshness (which detects whether a newer prebuilt
release exists) and adds the *apply* half: run install_llama_prebuilt.py to
download the newest bundle for this host and atomically swap it in place, so
the next model load uses it.
Design notes:
- Detection is delegated to check_prebuilt_freshness(). We surface an
``update_available`` flag (installed_tag != latest_tag) which is laxer than
freshness' ``stale`` (which additionally requires the install to be >= 3 days
old). The UI shows the "Update llama.cpp" affordance on update_available.
- The install is slow (download + extract + validate), so it runs on a daemon
thread; callers poll get_update_status() for the job state.
- Everything fails open: a missing marker / offline GitHub / source build just
reports update_available=False and never blocks the app.
"""
from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Optional
import structlog
from utils.llama_cpp_freshness import (
_INSTALL_MARKER_NAME,
check_prebuilt_freshness,
latest_published_release,
read_install_marker,
reset_caches,
)
logger = structlog.get_logger(__name__)
DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp"
_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate
# Background job state. Single in-flight update at a time, guarded by _job_lock.
_JOB_IDLE = "idle"
_JOB_RUNNING = "running"
_JOB_SUCCESS = "success"
_JOB_ERROR = "error"
_job_lock = threading.Lock()
_job: dict = {
"state": _JOB_IDLE,
"message": "",
"from_tag": None,
"to_tag": None,
"error": None,
"progress": None,
"started_at": None,
"finished_at": None,
}
# Matches the installer's download progress lines, e.g.
# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s".
_PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(")
# The download dominates the update; extract/validate fill the last slice.
_DOWNLOAD_PROGRESS_CEILING = 0.95
def _utcnow() -> str:
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
def _find_binary() -> Optional[str]:
"""Locate the active llama-server binary via the inference backend's own
resolver, so update targets exactly what Studio runs. Lazy import keeps the
heavy inference module off this module's import path."""
try:
from core.inference.llama_cpp import LlamaCppBackend
return LlamaCppBackend._find_llama_server_binary()
except Exception as exc: # pragma: no cover - defensive
logger.debug("llama update: binary discovery failed", error = str(exc))
return None
def _install_dir_for(binary_path: Optional[str]) -> Optional[Path]:
"""The directory holding UNSLOTH_PREBUILT_INFO.json -- i.e. the install root
install_llama_prebuilt.py wrote and the one we re-install into. Walks up from
the binary the same way read_install_marker() does."""
if not binary_path:
return None
p = Path(binary_path)
for parent in p.parents[:5]:
if (parent / _INSTALL_MARKER_NAME).is_file():
return parent
return None
def _installer_script() -> Optional[Path]:
"""Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then
searches up from this file for both ``<root>/install_llama_prebuilt.py`` and
``<root>/studio/install_llama_prebuilt.py`` so it works in the dev tree and
in an installed Studio layout."""
env = os.environ.get("UNSLOTH_LLAMA_INSTALLER")
if env and Path(env).is_file():
return Path(env)
here = Path(__file__).resolve()
for up in here.parents:
for cand in (up / "install_llama_prebuilt.py", up / "studio" / "install_llama_prebuilt.py"):
if cand.is_file():
return cand
return None
# Markerless (source-build) installs have no UNSLOTH_PREBUILT_INFO.json, so we
# ask the installer whether an official prebuilt now exists for this host. Memo
# is 24h; only successful answers are cached so a network blip retries.
_RESOLVE_TTL_SECONDS = 24 * 60 * 60
_resolve_memo: dict = {}
def _resolve_prebuilt_for_host(*, force_refresh: bool = False) -> Optional[dict]:
"""Run install_llama_prebuilt.py --resolve-prebuilt (no download) and return
{prebuilt_available, repo, release_tag, llama_tag, asset, install_kind} or
None. Fail-open: any error -> None so a source build never blocks the app."""
now = time.time()
if not force_refresh and _resolve_memo:
if now - _resolve_memo.get("at", 0.0) < _RESOLVE_TTL_SECONDS:
return _resolve_memo.get("value")
script = _installer_script()
if script is None:
return None
value: Optional[dict] = None
try:
proc = subprocess.run(
[
sys.executable,
str(script),
"--resolve-prebuilt",
"latest",
"--output-format",
"json",
],
capture_output = True,
text = True,
timeout = 60,
)
out = (proc.stdout or "").strip()
if proc.returncode == 0 and out:
parsed = json.loads(out.splitlines()[-1])
if isinstance(parsed, dict):
value = parsed
except Exception as exc: # pragma: no cover - subprocess/json defensive
logger.debug("llama update: resolve-prebuilt failed", error = str(exc))
value = None
if value is not None: # cache real answers; let failures retry next poll
_resolve_memo.update(at = now, value = value)
return value
def _installed_build_number(binary: Optional[str]) -> Optional[int]:
"""Best-effort build number from ``llama-server --version`` (e.g.
'version: 9585 (abc)'). None when unparseable or <= 1: a source build with
no git tags reports 'version: 1', which we treat as unknown (offer update)."""
if not binary:
return None
try:
proc = subprocess.run([binary, "--version"], capture_output = True, text = True, timeout = 20)
except Exception: # pragma: no cover - defensive
return None
m = re.search(r"version:\s*(\d+)", (proc.stderr or "") + (proc.stdout or ""))
if not m:
return None
n = int(m.group(1))
return n if n > 1 else None
def _is_under(path: Path, root: Path) -> bool:
try:
p, r = path.resolve(), root.resolve()
except (OSError, ValueError):
p, r = path, root
return p == r or r in p.parents
def _llama_install_root(binary: Optional[str]) -> Optional[Path]:
"""The Studio-managed llama.cpp root the active binary lives under, or None
when the binary is unmanaged. Installing anywhere the active binary is not
would not replace what _find_llama_server_binary runs (which prefers a pinned
LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we
refuse rather than silently install into an inactive or foreign tree."""
marked = _install_dir_for(binary)
if marked is not None:
return marked
if not binary:
return None
# LLAMA_SERVER_PATH is an explicit user pin that always wins in discovery;
# never auto-replace its tree (even a user's own llama.cpp checkout).
if os.environ.get("LLAMA_SERVER_PATH"):
return None
p = Path(binary)
env = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
if env and _is_under(p, Path(env)):
return Path(env)
for parent in p.parents:
if parent.name == "llama.cpp":
return parent
# PATH / system / custom install: not a managed tree, so do not offer.
return None
def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
"""Update status for a markerless (source-build) install: offer the official
prebuilt when one exists for this host and is newer than the installed
binary. None -> caller falls through to the no-marker default (unsupported)."""
res = _resolve_prebuilt_for_host(force_refresh = force_refresh)
if not res or not res.get("prebuilt_available"):
return None
# llama_tag is the upstream build (bNNNN, what --version reports); release_tag
# can be a fork wrapper tag, so compare/display against llama_tag.
latest = res.get("llama_tag") or res.get("release_tag")
if not latest:
return None
# No resolvable install root (e.g. a pinned LLAMA_SERVER_PATH we cannot
# manage) means an apply would not take effect, so do not offer.
if _llama_install_root(binary) is None:
return None
installed_build = _installed_build_number(binary)
m = re.search(r"(\d+)", latest)
latest_build = int(m.group(1)) if m else None
# Suppress only when the source build is reliably newer/equal; unknown
# version (the involuntary source-build case) is treated as behind.
update_available = (
installed_build is None or latest_build is None or installed_build < latest_build
)
with _job_lock:
job = dict(_job)
return {
"supported": True,
"update_available": update_available,
"stale": False,
"installed_tag": (f"b{installed_build}" if installed_build else None),
"latest_tag": latest,
"published_repo": res.get("repo"),
"installed_at_utc": None,
"age_days": None,
"source_build": True,
"job": job,
}
def get_update_status(*, force_refresh: bool = False) -> dict:
"""Report whether a newer prebuilt exists plus the current job state.
force_refresh bypasses the 24h release cache for an explicit "check now".
"""
binary = _find_binary()
marker = read_install_marker(binary)
with _job_lock:
job_running = _job["state"] == _JOB_RUNNING
# No marker = source build / custom path. Offer the official prebuilt if one
# now exists for this host (this is why macOS source builds showed no button).
# Skipped while the updater swaps the tree: each 3s poll would exec the
# half-replaced binary (on Windows that exec can make the installer's
# os.replace fail) and the poller only consumes job progress.
if marker is None and binary is not None and not job_running:
src = _source_build_status(binary, force_refresh = force_refresh)
if src is not None:
return src
repo = (marker or {}).get("published_repo") or DEFAULT_PUBLISHED_REPO
if force_refresh and repo:
# Prime the cache so the freshness read below sees the newest tag.
try:
latest_published_release(repo, force_refresh = True)
except Exception as exc: # pragma: no cover - network defensive
logger.debug("llama update: force refresh failed", error = str(exc))
freshness = check_prebuilt_freshness(binary)
installed = freshness.get("installed_tag")
latest = freshness.get("latest_tag")
# `behind` compares the full release identity with a base-build guard, so a
# lagging /releases/latest or a mix-tagged latest can't show a false update
# (see llama_cpp_freshness.is_behind).
update_available = bool(freshness.get("has_marker") and freshness.get("behind"))
with _job_lock:
job = dict(_job)
return {
"supported": bool(freshness.get("has_marker")),
"update_available": update_available,
"stale": bool(freshness.get("stale")),
"installed_tag": installed,
"latest_tag": latest,
"published_repo": freshness.get("published_repo") or repo,
"installed_at_utc": freshness.get("installed_at_utc"),
"age_days": freshness.get("age_days"),
"source_build": False,
"job": job,
}
def _rocm_install_args(asset: Optional[str]) -> list[str]:
"""Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh.
The installer probe can miss the gfx arch on amd-smi-only hosts; lemonade
bundles carry the family in the name (rocm-gfx110X), fork bundles only rocm/hip."""
if not asset:
return []
low = asset.lower()
if "rocm" not in low and "hip" not in low:
return []
gfx = re.search(r"-gfx[0-9a-z]+", low)
if gfx:
# _normalize_forwarded_gfx accepts the family form (gfx110x -> gfx110X).
return ["--rocm-gfx", gfx.group(0).lstrip("-")]
return ["--has-rocm"]
def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path) -> None:
"""Worker: put the backend into a maintenance state, run the installer for
the latest prebuilt, then refresh caches so the next load uses the new build."""
backend = None
model_was_active = False
try:
# Maintenance state so no load starts a server from the half-swapped binary
# (and the old binary is freed for the swap). Fails open without a backend.
try:
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
except Exception as exc:
logger.debug(
"llama update: backend unavailable, skipping load coordination", error = str(exc)
)
backend = None
if backend is not None:
try:
with backend._serial_load_lock:
backend._llama_update_in_progress = True
# is_active covers the loading/unhealthy window is_loaded misses
# (a live process also locks the exe on Windows during the swap).
if getattr(backend, "is_active", False):
model_was_active = True
backend.unload_model()
except Exception as exc:
logger.debug("llama update: load coordination failed", error = str(exc))
cmd = [
sys.executable,
str(script),
"--install-dir",
str(install_dir),
"--llama-tag",
"latest",
"--published-repo",
repo,
]
cmd.extend(_rocm_install_args(asset))
logger.info("llama update: installing", cmd = " ".join(cmd))
# Stream the installer output so download percent lines feed
# job["progress"]; finer milestones via UNSLOTH_PROGRESS_PERCENT_STEP.
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
env = env,
)
timed_out = threading.Event()
def _kill_on_timeout() -> None:
timed_out.set()
proc.kill()
watchdog = threading.Timer(_INSTALL_TIMEOUT_SECONDS, _kill_on_timeout)
watchdog.daemon = True
watchdog.start()
tail_lines: list[str] = []
try:
assert proc.stdout is not None
for line in proc.stdout:
tail_lines.append(line)
if len(tail_lines) > 80:
del tail_lines[0]
m = _PROGRESS_LINE_RE.search(line)
if m is None:
continue
fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING
with _job_lock:
_job["progress"] = max(_job.get("progress") or 0.0, fraction)
returncode = proc.wait()
finally:
watchdog.cancel()
if timed_out.is_set():
raise RuntimeError(f"installer timed out after {_INSTALL_TIMEOUT_SECONDS}s")
if returncode != 0:
tail = "".join(tail_lines).strip()[-1500:]
raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}")
# New UNSLOTH_PREBUILT_INFO.json is on disk; drop in-memory caches and
# re-prime the 24h disk freshness cache with the true newest, so the
# banner can't linger on a stale same-base value after the swap.
reset_caches()
try:
latest_published_release(repo, force_refresh = True)
except Exception as exc: # pragma: no cover - network defensive
logger.debug("llama update: post-install freshness refresh failed", error = str(exc))
new_marker = read_install_marker(_find_binary())
new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag")
with _job_lock:
_job.update(
state = _JOB_SUCCESS,
message = (
f"Updated llama.cpp to {new_tag}."
+ (" Reload your model to use it." if model_was_active else "")
),
to_tag = new_tag,
error = None,
progress = 1.0,
finished_at = _utcnow(),
)
logger.info("llama update: success", to_tag = new_tag)
except Exception as exc:
logger.warning("llama update: failed", error = str(exc))
with _job_lock:
_job.update(
state = _JOB_ERROR,
message = "llama.cpp update failed.",
error = str(exc),
finished_at = _utcnow(),
)
finally:
# Lift the maintenance state so model loads work again, success or not.
if backend is not None:
try:
backend._llama_update_in_progress = False
except Exception: # pragma: no cover - defensive
pass
def start_update() -> dict:
"""Kick off a background update. Idempotent: a second call while one is
running returns the in-flight job rather than starting another."""
binary = _find_binary()
marker = read_install_marker(binary)
script = _installer_script()
if script is None:
return {
"started": False,
"reason": "installer_missing",
"message": "install_llama_prebuilt.py could not be located.",
"job": get_update_status()["job"],
}
# A job already in flight wins over any freshness re-check below (and skips
# its network call). The final lock block re-checks to close the TOCTOU.
with _job_lock:
if _job["state"] == _JOB_RUNNING:
return {"started": False, "reason": "already_running", "job": dict(_job)}
if marker:
# Mirror the detection guard: a direct POST or a stale banner must not
# start an install when the latest is not actually newer (force a fresh
# check so a stale 24h cache can't wrongly block a real update either).
status = get_update_status(force_refresh = True)
if not status.get("update_available"):
return {
"started": False,
"reason": "up_to_date",
"message": "The installed llama.cpp build is already at the latest prebuilt.",
"job": status["job"],
}
install_dir = _install_dir_for(binary)
repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO
from_tag = marker.get("tag") or marker.get("release_tag")
asset = marker.get("asset")
else:
# Source build / custom path: only proceed when the same detection logic
# would offer the update (prebuilt exists, install is behind, root is
# manageable), so a direct POST cannot downgrade a newer source build.
src = _source_build_status(binary, force_refresh = True) if binary else None
if src is None:
return {
"started": False,
"reason": "no_prebuilt_available",
"message": (
"No official llama.cpp prebuilt is available for this host, "
"so the source build cannot be swapped automatically."
),
"job": get_update_status()["job"],
}
if not src.get("update_available"):
return {
"started": False,
"reason": "up_to_date",
"message": "The installed llama.cpp build is already at or newer than the latest prebuilt.",
"job": get_update_status()["job"],
}
res = _resolve_prebuilt_for_host()
install_dir = _llama_install_root(binary)
repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO
from_tag = None
asset = (res or {}).get("asset")
if install_dir is None:
return {
"started": False,
"reason": "no_install_dir",
"message": "Could not determine the llama.cpp install directory.",
"job": get_update_status()["job"],
}
with _job_lock:
if _job["state"] == _JOB_RUNNING:
return {"started": False, "reason": "already_running", "job": dict(_job)}
_job.update(
state = _JOB_RUNNING,
message = "Downloading and installing the latest llama.cpp prebuilt...",
from_tag = from_tag,
to_tag = None,
error = None,
progress = 0.0,
started_at = _utcnow(),
finished_at = None,
)
job_snapshot = dict(_job)
thread = threading.Thread(
target = _run_update,
args = (install_dir, repo, asset, script),
name = "llama-cpp-update",
daemon = True,
)
thread.start()
return {"started": True, "reason": None, "job": job_snapshot}
def _reset_job_for_tests() -> None:
"""Test-only: return the job tracker to idle."""
with _job_lock:
_job.update(
state = _JOB_IDLE,
message = "",
from_tag = None,
to_tag = None,
error = None,
progress = None,
started_at = None,
finished_at = None,
)

122
studio/backend/utils/models/model_config.py
View file

@ -58,25 +58,28 @@ import re as _re
_MODEL_SIZE_RE = _re.compile(r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
# MoE active-parameter pattern: "A3B", "A3.5B", etc.
_ACTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
# Gemma 3n/4 effective-parameter pattern: "E2B", "E4B" -- the runtime
# footprint (MatFormer + per-layer embeddings), which is the size that
# matters for size-gated policies like sub-3B speculative-decoding fallback.
_EFFECTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])e(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
def extract_model_size_b(model_id: str) -> float | None:
"""Extract model size in billions from a model identifier.
Prefers MoE active-parameter notation (e.g. ``A3B`` in
``Qwen3.5-35B-A3B``) over total params. Handles ``B`` (billions)
and ``M`` (millions) suffixes.
``Qwen3.5-35B-A3B``), then Gemma effective-parameter notation
(e.g. ``E2B``), over total params. Handles ``B`` (billions) and
``M`` (millions) suffixes.
"""
mid = (model_id or "").lower()
active = _ACTIVE_SIZE_RE.search(mid)
if active:
val = float(active.group(1))
return val / 1000.0 if active.group(2).lower() == "m" else val
size = _MODEL_SIZE_RE.search(mid)
if not size:
return None
val = float(size.group(1))
return val / 1000.0 if size.group(2).lower() == "m" else val
# First match wins, in priority order: active > effective > total.
for pattern in (_ACTIVE_SIZE_RE, _EFFECTIVE_SIZE_RE, _MODEL_SIZE_RE):
m = pattern.search(mid)
if m:
val = float(m.group(1))
return val / 1000.0 if m.group(2).lower() == "m" else val
return None
# Maps equivalent model names to their canonical YAML config file.
@ -510,7 +513,7 @@ _VLM_MODEL_TYPES = {
_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"}
# Pre-computed .venv_t5 paths and backend dir for subprocess version switching.
# Vision check uses 5.5.0 (newest, recognizes all architectures).
# Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures.
from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402
_VENV_T5_DIR = str(_studio_root() / ".venv_t5_550")
@ -930,6 +933,24 @@ def _is_mmproj(filename: str) -> bool:
return "mmproj" in filename.lower()
def _is_mtp_drafter(path: str) -> bool:
"""True for a separate-file MTP drafter (speculative head), a companion
to the main model rather than a selectable quant: the repo-root
``mtp-*.gguf`` or the ``MTP/`` subdir copies (Gemma 4).
Mirrors hub.utils.gguf.is_mtp_drafter_path (utils cannot import hub).
Must be excluded everywhere mmproj is, or the drafter leaks into variant
menus (a phantom quant) and quant-matched file lookups -- e.g. a ``Q8_0``
request must not resolve to ``MTP/...-Q8_0-MTP.gguf``, which sorts ahead
of the real weight.
"""
p = path.lower()
if not p.endswith(".gguf"):
return False
name = p.rsplit("/", 1)[-1]
return name.startswith("mtp-") or "/mtp/" in f"/{p}"
# Family tokens for #5347's filename fallback. Lowercase; order irrelevant.
_MODEL_FAMILY_TOKENS: tuple[str, ...] = (
"qwen",
@ -1137,6 +1158,48 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
return str(best[1])
def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
"""Find the separate MTP drafter (``mtp-*.gguf``) for a local GGUF model.
The drafter that pairs with the main weights sits at the repo/snapshot
root (Gemma 4); the weight itself may be at the root or in a quant subdir,
so scan the weight's directory and ``search_root``. Matches by the
``mtp-`` filename prefix unsloth uses for ``-hf`` auto-discovery -- the
same signal as the HF download path. Repos that bake the head into the
main GGUF (Qwen) have no such sibling, so this returns None.
Pairs by name so a multi-model folder can't attach a foreign drafter:
unsloth names the drafter ``mtp-<model>.gguf`` where ``<model>`` prefixes
the weight filename across all Gemma 4 repos (e.g.
``mtp-gemma-4-12B-it.gguf`` next to ``gemma-4-12B-it-qat-Q4_0.gguf``).
An unmatched drafter is skipped (fail-safe: no MTP).
"""
p = Path(path)
weight_name = p.name.lower() if p.suffix.lower() == ".gguf" else None
start_dir = p.parent if p.is_file() else p
dirs = [start_dir]
if search_root is not None:
dirs.append(Path(search_root))
for d in dirs:
try:
entries = sorted(d.iterdir())
except OSError:
continue
for f in entries:
name = f.name.lower()
if not (name.startswith("mtp-") and name.endswith(".gguf")):
continue
stem = name[len("mtp-") : -len(".gguf")]
if not stem or (weight_name is not None and not weight_name.startswith(stem)):
continue
try:
if f.is_file():
return str(f.resolve())
except OSError:
continue
return None
def detect_gguf_model(path: str) -> Optional[str]:
"""Check if a local path is or contains a GGUF model file.
@ -1148,7 +1211,12 @@ def detect_gguf_model(path: str) -> Optional[str]:
# Case 1: direct .gguf file
if p.suffix.lower() == ".gguf":
if _is_mmproj(p.name):
# Companions are not models: rejecting a drafter here also keeps
# detect_mtp_file from pairing the same file with itself
# (-m drafter --model-draft drafter). Include the immediate parent
# dir so the MTP/ subdir copies are caught -- the basename alone
# (...-MTP.gguf) doesn't match the predicate's mtp- prefix.
if _is_mmproj(p.name) or _is_mtp_drafter(f"{p.parent.name}/{p.name}"):
return None
# Extension is authoritative: don't gate on is_file()/exists(), which
# can fail in the Windows lock window after llama-server is killed.
@ -1160,10 +1228,14 @@ def detect_gguf_model(path: str) -> Optional[str]:
return str(p.absolute()) # absolute() keeps symlink names readable
# Directory named "*.gguf": fall through to the dir scan below.
# Case 2: directory containing .gguf files (skip mmproj)
# Case 2: directory containing .gguf files (skip mmproj / MTP drafter)
if p.is_dir():
gguf_files = sorted(
(f for f in _iter_gguf_files(p) if not _is_mmproj(f.name)),
(
f
for f in _iter_gguf_files(p)
if not _is_mmproj(f.name) and not _is_mtp_drafter(f"{f.parent.name}/{f.name}")
),
key = lambda f: f.stat().st_size,
reverse = True,
)
@ -1393,6 +1465,9 @@ def list_gguf_variants(
if "mmproj" in fname.lower():
has_vision = True
continue
# MTP drafters are speculative-decoding companions, not quants.
if _is_mtp_drafter(fname):
continue
quant = _extract_quant_label(fname)
quant_totals[quant] = quant_totals.get(quant, 0) + size
@ -1468,6 +1543,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
# Use the relative path so ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf``
# get distinct quant labels instead of collapsing on basename.
rel = f.relative_to(p).as_posix()
if _is_mtp_drafter(rel):
continue
quant = _extract_quant_label(rel)
quant_totals[quant] = quant_totals.get(quant, 0) + size
if quant not in quant_first_file:
@ -1504,7 +1581,9 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
matches = sorted(
f
for f in _iter_gguf_files(p, recursive = True)
if not _is_mmproj(f.name) and _extract_quant_label(f.relative_to(p).as_posix()) == variant
if not _is_mmproj(f.name)
and not _is_mtp_drafter(f.relative_to(p).as_posix())
and _extract_quant_label(f.relative_to(p).as_posix()) == variant
)
if matches:
return str(matches[0].resolve())
@ -1519,9 +1598,9 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
"""
for snap in _iter_hf_cache_snapshots(repo_id):
rel_files = [
f.relative_to(snap).as_posix()
rel
for f in _iter_gguf_files(snap, recursive = True)
if not _is_mmproj(f.name)
if not _is_mtp_drafter(rel := f.relative_to(snap).as_posix()) and not _is_mmproj(f.name)
]
if rel_files:
return _pick_best_gguf(rel_files)
@ -2068,6 +2147,7 @@ class ModelConfig:
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
gguf_mtp_file: Optional[str] = None # Full path to the separate MTP drafter (local mode)
gguf_hf_repo: Optional[str] = (
None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
)
@ -2207,6 +2287,11 @@ class ModelConfig:
elif base_is_vision:
logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
# Separate MTP drafter sibling (Gemma 4), mirroring mmproj.
mtp_file = detect_mtp_file(gguf_file, search_root = path)
if mtp_file:
logger.info(f"Detected MTP drafter: {mtp_file}")
return cls(
identifier = identifier,
display_name = display_name,
@ -2218,6 +2303,7 @@ class ModelConfig:
is_gguf = True,
gguf_file = gguf_file,
gguf_mmproj_file = mmproj_file,
gguf_mtp_file = mtp_file,
)
else:
# Does the HF repo contain GGUF files?

5
studio/backend/utils/paths/storage_roots.py
View file

@ -60,6 +60,11 @@ def cache_root() -> Path:
return studio_root() / "cache"
def studio_bin_root() -> Path:
"""Dir for Studio-managed executables (the `unsloth` shim, downloaded tools like cloudflared)."""
return studio_root() / "bin"
def assets_root() -> Path:
return studio_root() / "assets"

35
studio/backend/utils/studio_version.py
View file

@ -15,6 +15,7 @@ _DEV_VERSION = "dev"
_GIT_TIMEOUT_SECONDS = 1.0
_STUDIO_TAG_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$")
_GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$")
_GIT_BRANCH_RE = re.compile(r"^[0-9A-Za-z._/-]+$")
_MAX_VERSION_LENGTH = 64
@ -71,6 +72,35 @@ def _exact_git_studio_tag(repo_root: Path) -> str | None:
return tag if is_valid_studio_release_version(tag) else None
def _git_branch(repo_root: Path) -> str | None:
try:
result = subprocess.run(
["git", "rev-parse", "--abbrev-ref", "HEAD"],
cwd = repo_root,
check = False,
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = _GIT_TIMEOUT_SECONDS,
)
except (OSError, subprocess.TimeoutExpired):
return None
if result.returncode != 0:
return None
branch = result.stdout.strip()
# "HEAD" means detached, e.g. a tag or commit checkout.
if (
not branch
or branch == "HEAD"
or len(branch) > _MAX_VERSION_LENGTH
or _GIT_BRANCH_RE.fullmatch(branch) is None
):
return None
return branch
def get_studio_version(repo_root: Path | None = None) -> str:
"""Return the installed Studio release tag for display, or ``dev``.
@ -81,7 +111,10 @@ def get_studio_version(repo_root: Path | None = None) -> str:
if _is_source_checkout(resolved_repo_root):
git_tag = _exact_git_studio_tag(resolved_repo_root)
return git_tag if git_tag is not None else _DEV_VERSION
if git_tag is not None:
return git_tag
branch = _git_branch(resolved_repo_root)
return f"GitHub {branch}" if branch is not None else _DEV_VERSION
stamped_version = _studio_release_build.STUDIO_RELEASE_VERSION
if is_valid_studio_release_version(stamped_version):

302
studio/backend/utils/transformers_version.py
View file

@ -3,12 +3,27 @@
"""Automatic transformers version switching.
Some newer architectures need transformers>=5.3.0 (.venv_t5_530/); Gemma 4
needs >=5.5.0 (.venv_t5_550/). Everything else uses the default 4.57.x. A
custom-named LoRA adapter's base model is resolved from adapter_config.json.
Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE,
tiny_qwen3_moe) require transformers>=5.3.0, while Gemma 4 models require a
newer 5.x sidecar. Everything else needs the default 4.57.x that ships with
Unsloth.
Training/inference run in subprocesses that activate the right version via
sys.path; export (in-process) uses ensure_transformers_version() for the swap.
Two separate target directories are maintained:
- .venv_t5_530/ — transformers 5.3.0 (Ministral-3, GLM, Qwen3 MoE, etc.)
- .venv_t5_550/ — transformers 5.5.0 (Gemma 4)
- .venv_t5_510/ — transformers 5.10.2 (Gemma 4 Unified / 12B)
When loading a LoRA adapter with a custom name, we resolve the base model from
``adapter_config.json`` and check *that* against the model list.
Strategy:
Training and inference run in subprocesses that activate the correct version
via sys.path (prepending the appropriate .venv_t5_*/ directory). See:
- core/training/worker.py
- core/inference/worker.py
For export (still in-process), ensure_transformers_version() does a lightweight
sys.path swap using the same directories pre-installed by setup.sh.
"""
import importlib
@ -53,13 +68,32 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
"lfm2.5-vl-450m", # LiquidAI/LFM2.5-VL-450M
)
# Lowercase substrings for models that require transformers 5.5.0 (checked first).
# Lowercase substrings for models that require transformers 5.10.x (checked first).
TRANSFORMERS_510_MODEL_SUBSTRINGS: tuple[str, ...] = (
"gemma-4-12b", # Gemma 4 Unified 12B
"gemma4-12b",
)
# Lowercase substrings for models that require the Gemma 4 transformers 5.5 sidecar.
TRANSFORMERS_550_MODEL_SUBSTRINGS: tuple[str, ...] = (
"gemma-4", # Gemma-4 (E2B-it, E4B-it, 31B-it, 26B-A4B-it)
"gemma4", # Gemma-4 alternate naming
"qwen3.6",
)
# Architecture classes / model_type values that require transformers 5.10.x.
# Checked via config.json (local or HuggingFace).
_TRANSFORMERS_510_ARCHITECTURES: set[str] = {
"Gemma4UnifiedForConditionalGeneration",
"Gemma4AssistantForCausalLM",
"Gemma4UnifiedAssistantForCausalLM",
}
_TRANSFORMERS_510_MODEL_TYPES: set[str] = {
"gemma4_unified",
"gemma4_assistant",
"gemma4_unified_assistant",
}
# Architecture classes / model_type values that require transformers 5.5.0.
# Checked via config.json (local or HuggingFace).
_TRANSFORMERS_550_ARCHITECTURES: set[str] = {
@ -78,21 +112,26 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
_tokenizer_class_cache: dict[str, bool] = {}
# Cache for dynamic config.json lookups (architecture/model_type checks).
_config_json_cache: dict[str, dict | None] = {}
_config_needs_510_cache: dict[str, bool] = {}
_config_needs_550_cache: dict[str, bool] = {}
# Versions
TRANSFORMERS_510_VERSION = "5.10.2"
TRANSFORMERS_550_VERSION = "5.5.0"
TRANSFORMERS_530_VERSION = "5.3.0"
TRANSFORMERS_DEFAULT_VERSION = "4.57.6"
# Backwards-compat alias — points to 5.5.0 (the highest 5.x tier).
# Consumers should prefer TRANSFORMERS_530_VERSION / TRANSFORMERS_550_VERSION.
TRANSFORMERS_5_VERSION = TRANSFORMERS_550_VERSION
# Backwards-compat alias — points to the highest 5.x tier.
# Consumers should prefer TRANSFORMERS_510_VERSION / TRANSFORMERS_550_VERSION /
# TRANSFORMERS_530_VERSION.
TRANSFORMERS_5_VERSION = TRANSFORMERS_510_VERSION
# Pre-installed directories — created by setup.sh / setup.ps1.
from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402
_VENV_T5_530_DIR = str(_studio_root() / ".venv_t5_530")
_VENV_T5_550_DIR = str(_studio_root() / ".venv_t5_550")
_VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510")
# Backwards-compat alias
_VENV_T5_DIR = _VENV_T5_550_DIR
@ -108,15 +147,34 @@ def activate_transformers_for_subprocess(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
tier = get_transformers_tier(resolved)
if tier == "550":
if tier == "510":
if not _ensure_venv_t5_510_exists():
raise RuntimeError(
f"Cannot activate transformers {TRANSFORMERS_510_VERSION}: "
f".venv_t5_510 missing at {_VENV_T5_510_DIR}"
)
if _VENV_T5_510_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_510_DIR)
logger.info(
"Activated transformers %s from %s",
TRANSFORMERS_510_VERSION,
_VENV_T5_510_DIR,
)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_510_DIR + (os.pathsep + _pp if _pp else "")
elif tier == "550":
if not _ensure_venv_t5_550_exists():
raise RuntimeError(
f"Cannot activate transformers 5.5.0: "
f"Cannot activate transformers {TRANSFORMERS_550_VERSION}: "
f".venv_t5_550 missing at {_VENV_T5_550_DIR}"
)
if _VENV_T5_550_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_550_DIR)
logger.info("Activated transformers 5.5.0 from %s", _VENV_T5_550_DIR)
logger.info(
"Activated transformers %s from %s",
TRANSFORMERS_550_VERSION,
_VENV_T5_550_DIR,
)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_550_DIR + (os.pathsep + _pp if _pp else "")
elif tier == "530":
@ -260,6 +318,67 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
return False
def _load_config_json(model_name: str) -> dict | None:
"""Return parsed ``config.json`` for *model_name*, checking local files first."""
if model_name in _config_json_cache:
return _config_json_cache[model_name]
local_cfg = Path(model_name) / "config.json"
if local_cfg.is_file():
try:
with open(local_cfg) as f:
cfg = json.load(f)
_config_json_cache[model_name] = cfg
return cfg
except Exception as exc:
logger.debug("Could not read %s: %s", local_cfg, exc)
_config_json_cache[model_name] = None
return None
if _env_offline():
_config_json_cache[model_name] = None
return None
import urllib.request
url = f"https://huggingface.co/{model_name}/raw/main/config.json"
try:
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = 10) as resp:
cfg = json.loads(resp.read().decode())
_config_json_cache[model_name] = cfg
return cfg
except Exception as exc:
logger.debug("Could not fetch config.json for '%s': %s", model_name, exc)
_config_json_cache[model_name] = None
return None
def _config_matches_tier(cfg: dict, architectures: set[str], model_types: set[str]) -> bool:
archs = cfg.get("architectures", [])
if any(a in architectures for a in archs):
return True
if cfg.get("model_type") in model_types:
return True
return False
def _config_needs_550(cfg: dict) -> bool:
return _config_matches_tier(
cfg,
_TRANSFORMERS_550_ARCHITECTURES,
_TRANSFORMERS_550_MODEL_TYPES,
)
def _config_needs_510(cfg: dict) -> bool:
return _config_matches_tier(
cfg,
_TRANSFORMERS_510_ARCHITECTURES,
_TRANSFORMERS_510_MODEL_TYPES,
)
def _check_config_needs_550(model_name: str) -> bool:
"""True if ``config.json`` has architectures/model_type needing transformers
5.5.0 (e.g. Gemma 4).
@ -270,81 +389,88 @@ def _check_config_needs_550(model_name: str) -> bool:
if model_name in _config_needs_550_cache:
return _config_needs_550_cache[model_name]
def _check_cfg(cfg: dict) -> bool:
archs = cfg.get("architectures", [])
if any(a in _TRANSFORMERS_550_ARCHITECTURES for a in archs):
return True
if cfg.get("model_type") in _TRANSFORMERS_550_MODEL_TYPES:
return True
return False
# --- Check local config.json first ------------------------------------
local_path = Path(model_name)
local_cfg = local_path / "config.json"
if local_cfg.is_file():
try:
with open(local_cfg) as f:
cfg = json.load(f)
result = _check_cfg(cfg)
if result:
logger.info(
"Local config.json check: %s needs transformers 5.5.0 "
"(architectures=%s, model_type=%s)",
model_name,
cfg.get("architectures", []),
cfg.get("model_type"),
)
_config_needs_550_cache[model_name] = result
return result
except Exception as exc:
logger.debug("Could not read %s: %s", local_cfg, exc)
# Offline: skip the 10s urllib fetch (fail-open to lower tier).
if _env_offline():
cfg = _load_config_json(model_name)
if cfg is None:
_config_needs_550_cache[model_name] = False
return False
# --- Fall back to fetching from HuggingFace ---------------------------
import urllib.request
result = _config_needs_550(cfg)
if result:
logger.info(
"config.json check: %s needs transformers %s (architectures=%s, model_type=%s)",
model_name,
TRANSFORMERS_550_VERSION,
cfg.get("architectures", []),
cfg.get("model_type"),
)
_config_needs_550_cache[model_name] = result
return result
url = f"https://huggingface.co/{model_name}/raw/main/config.json"
try:
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = 10) as resp:
cfg = json.loads(resp.read().decode())
result = _check_cfg(cfg)
if result:
logger.info(
"Dynamic config.json check: %s needs transformers 5.5.0 "
"(architectures=%s, model_type=%s)",
model_name,
cfg.get("architectures", []),
cfg.get("model_type"),
)
_config_needs_550_cache[model_name] = result
return result
except Exception as exc:
logger.debug("Could not fetch config.json for '%s': %s", model_name, exc)
_config_needs_550_cache[model_name] = False
def _check_config_needs_510(model_name: str) -> bool:
"""Check ``config.json`` for Gemma 4 Unified / 12B architectures."""
if model_name in _config_needs_510_cache:
return _config_needs_510_cache[model_name]
cfg = _load_config_json(model_name)
if cfg is None:
_config_needs_510_cache[model_name] = False
return False
result = _config_needs_510(cfg)
if result:
logger.info(
"config.json check: %s needs transformers %s (architectures=%s, model_type=%s)",
model_name,
TRANSFORMERS_510_VERSION,
cfg.get("architectures", []),
cfg.get("model_type"),
)
_config_needs_510_cache[model_name] = result
return result
def get_transformers_tier(model_name: str) -> str:
"""Return the transformers tier required for *model_name*.
``"550"`` for transformers 5.5.0 (e.g. Gemma 4), ``"530"`` for 5.3.0
(e.g. Ministral-3, Qwen3 MoE), or ``"default"`` for everything else (4.57.x).
The 5.5.0 check runs first, then 5.3.0.
Returns ``"510"`` for models needing transformers 5.10.x (Gemma 4 Unified),
``"550"`` for models needing transformers 5.5.0 (Gemma 4),
``"530"`` for models needing transformers 5.3.0 (e.g. Ministral-3, Qwen3 MoE),
or ``"default"`` for everything else (4.57.x).
Higher 5.x tiers run first.
"""
lowered = model_name.lower()
# Local checkpoint names can contain architecture substrings in their
# directory names (for example a pytest temp dir). If config.json exists,
# trust it before using name heuristics.
local_cfg = Path(model_name) / "config.json"
if local_cfg.is_file():
cfg = _load_config_json(model_name)
if cfg is not None and _config_needs_510(cfg):
return "510"
if cfg is not None and _config_needs_550(cfg):
return "550"
if cfg is not None:
local_tc = Path(model_name) / "tokenizer_config.json"
if local_tc.is_file() and _check_tokenizer_config_needs_v5(model_name):
return "530"
return "default"
# --- Fast substring checks (no I/O) ------------------------------------
if "assistant" in lowered and ("gemma-4" in lowered or "gemma4" in lowered):
return "510"
if any(sub in lowered for sub in TRANSFORMERS_510_MODEL_SUBSTRINGS):
return "510"
if any(sub in lowered for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS):
return "550"
if any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS):
return "530"
# --- Slow config fallbacks (local file first, then network) -----------
# --- Slow config fallbacks (network for HF IDs) ------------------------
if _check_config_needs_510(model_name):
return "510"
if _check_config_needs_550(model_name):
return "550"
if _check_tokenizer_config_needs_v5(model_name):
@ -419,6 +545,13 @@ _VENV_T5_530_PACKAGES = (
"tiktoken",
)
_VENV_T5_510_PACKAGES = (
f"transformers=={TRANSFORMERS_510_VERSION}",
"huggingface_hub==1.8.0",
"hf_xet==1.4.2",
"tiktoken",
)
_VENV_T5_550_PACKAGES = (
f"transformers=={TRANSFORMERS_550_VERSION}",
"huggingface_hub==1.8.0",
@ -475,7 +608,7 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool:
def _venv_t5_is_valid() -> bool:
"""Backwards-compat: check the 5.5.0 venv."""
"""Backwards-compat: check the Gemma 4 sidecar venv."""
return _venv_dir_is_valid(_VENV_T5_550_DIR, _VENV_T5_550_PACKAGES)
@ -553,11 +686,24 @@ def _ensure_venv_t5_530_exists() -> bool:
def _ensure_venv_t5_550_exists() -> bool:
"""Ensure .venv_t5_550/ exists with transformers 5.5.0."""
return _ensure_venv_dir(_VENV_T5_550_DIR, _VENV_T5_550_PACKAGES, "transformers 5.5.0")
return _ensure_venv_dir(
_VENV_T5_550_DIR,
_VENV_T5_550_PACKAGES,
f"transformers {TRANSFORMERS_550_VERSION}",
)
def _ensure_venv_t5_510_exists() -> bool:
"""Ensure .venv_t5_510/ exists with transformers 5.10.x."""
return _ensure_venv_dir(
_VENV_T5_510_DIR,
_VENV_T5_510_PACKAGES,
f"transformers {TRANSFORMERS_510_VERSION}",
)
def _ensure_venv_t5_exists() -> bool:
"""Backwards-compat: ensure the 5.5.0 venv exists."""
"""Backwards-compat: ensure the Gemma 4 5.5 sidecar venv exists."""
return _ensure_venv_t5_550_exists()
@ -577,7 +723,7 @@ def _activate_venv(venv_dir: str, label: str) -> None:
def _deactivate_5x() -> None:
"""Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport."""
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR):
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR):
while d in sys.path:
sys.path.remove(d)
logger.info("Removed venv_t5 dirs from sys.path")
@ -593,7 +739,9 @@ def _deactivate_5x() -> None:
def ensure_transformers_version(model_name: str) -> None:
"""Ensure the correct ``transformers`` version is active for *model_name*.
Uses sys.path with .venv_t5_530/ or .venv_t5_550/ (pre-installed by setup.sh):
Uses sys.path with .venv_t5_510/, .venv_t5_550/, or .venv_t5_530/
(pre-installed by setup.sh):
• Need 5.10.x → prepend .venv_t5_510/ to sys.path, purge modules.
• Need 5.5.0 → prepend .venv_t5_550/ to sys.path, purge modules.
• Need 5.3.0 → prepend .venv_t5_530/ to sys.path, purge modules.
• Need 4.x → remove all .venv_t5_*/ from sys.path, purge modules.
@ -608,7 +756,11 @@ def ensure_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
tier = get_transformers_tier(resolved)
if tier == "550":
if tier == "510":
target_version = TRANSFORMERS_510_VERSION
venv_dir = _VENV_T5_510_DIR
ensure_fn = _ensure_venv_t5_510_exists
elif tier == "550":
target_version = TRANSFORMERS_550_VERSION
venv_dir = _VENV_T5_550_DIR
ensure_fn = _ensure_venv_t5_550_exists
@ -643,7 +795,7 @@ def ensure_transformers_version(model_name: str) -> None:
model_name,
)
return
# Different 5.x → must switch (e.g. 5.3.0 loaded but need 5.5.0).
# Different 5.x -> need to switch (e.g. 5.3.0 loaded but need 5.10.x).
in_memory_major = int(in_memory.split(".")[0])
if in_memory_major == target_major and venv_dir is None:
# Both are default (4.x) — close enough.

7
studio/frontend/.npmrc
View file

@ -6,11 +6,8 @@
# upstream removal. npm interprets the bare integer as DAYS; do not
# append `d`, npm 11.x will parse `7d` as a Date string and abort.
min-release-age=7
# Defensive alias: `minimum-release-age` takes minutes (10080 = 7 days).
# Some npm versions / wrappers consult one key but not the other; setting
# both means a single setting-name parse change upstream cannot silently
# disable the cooldown. The two keys MUST agree; do not let them drift.
minimum-release-age=10080
# Do not re-add the old `minimum-release-age` alias: npm >=11.16 warns on
# unknown project configs and npm 12 stops accepting them.
# Belt-and-braces: refuse to write back loose `^x.y.z` ranges into
# package.json when a maintainer runs `npm install <pkg>` locally. This
# does NOT rewrite already-present ranges (those need an explicit

6
studio/frontend/package-lock.json generated
View file

@ -10264,9 +10264,9 @@
}
},
"node_modules/hono": {
"version": "4.12.18",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.18.tgz",
"integrity": "sha512-RWzP96k/yv0PQfyXnWjs6zot20TqfpfsNXhOnev8d1InAxubW93L11/oNUc3tQqn2G0bSdAOBpX+2uDFHV7kdQ==",
"version": "4.12.21",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.21.tgz",
"integrity": "sha512-uV63apnb0kyPtAUwoWgaGh9HyIFcv8lgmzPZSiTBQAFOFGIzka5EZ1dZocmGnn0XdX0+XTqJ6Tqv7selMuGLRQ==",
"license": "MIT",
"engines": {
"node": ">=16.9.0"

7
studio/frontend/package.json
View file

@ -86,7 +86,7 @@
"@tanstack/router-core": "1.169.2",
"@tanstack/history": "1.161.6",
"mermaid": "11.15.0",
"hono": "4.12.18",
"hono": "4.12.21",
"qs": "6.15.2",
"ip-address": "10.1.1",
"brace-expansion@5.0.5": "5.0.6"
@ -108,5 +108,10 @@
"typescript": "~5.9.3",
"typescript-eslint": "^8.55.0",
"vite": "^8.0.1"
},
"allowScripts": {
"@biomejs/biome@1.9.4": true,
"msw@2.14.3": true,
"fsevents": true
}
}

18
studio/frontend/src/app/provider.tsx
View file

@ -11,6 +11,7 @@ import {
import { Toaster } from "@/components/ui/sonner";
import { TooltipProvider } from "@/components/ui/tooltip";
import { WebUpdateBanner } from "@/components/web/update-banner";
import { LlamaUpdateBanner } from "@/components/llama-update-banner";
import { DownloadManagerPanel } from "@/features/hub/download-manager";
import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth";
import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain";
@ -259,6 +260,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
{children}
<DownloadManagerPanel />
<WebUpdateBanner enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)} />
<LlamaUpdateBanner enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)} />
</>
);
}
@ -294,7 +296,18 @@ function TauriWrapper({ children }: { children: ReactNode }) {
/>
);
if (!shouldUseCustomWindowTitlebar()) return content;
if (!shouldUseCustomWindowTitlebar()) {
// macOS desktop uses the native titlebar and returns here before the
// custom-titlebar branch, so mount the updater banner on this path too.
return (
<>
{content}
<LlamaUpdateBanner
enabled={showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname)}
/>
</>
);
}
const showSidebarSurface =
showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname);
@ -305,6 +318,9 @@ function TauriWrapper({ children }: { children: ReactNode }) {
<div className="min-h-0 flex-1 overflow-hidden">
{content}
</div>
<LlamaUpdateBanner
enabled={showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname)}
/>
</div>
);
}

7
studio/frontend/src/app/router.tsx
View file

@ -3,6 +3,7 @@
import { Link, createRouter, useRouterState } from "@tanstack/react-router";
import { Button } from "@/components/ui/button";
import { MascotImg } from "@/components/mascot-img";
import { useT } from "@/i18n";
import { Route as rootRoute } from "./routes/__root";
import { Route as dataRecipesRoute } from "./routes/data-recipes";
@ -41,11 +42,7 @@ function DefaultNotFound() {
return (
<div className="flex flex-1 flex-col items-center justify-center gap-4 p-8 text-center">
<img
src="/Sloth%20emojis/sloth%20shy%20large.png"
alt="Sloth mascot"
className="size-24"
/>
<MascotImg src="Sloth emojis/sloth shy large.png" className="size-24" />
<div className="flex flex-col items-center gap-1">
<h1 className="font-heading font-semibold text-2xl tracking-tight">
{t("shell.notFound.title")}

9
studio/frontend/src/app/routes/__root.tsx
View file

@ -140,9 +140,14 @@ function RootLayout() {
<SidebarInset className={isChatRoute ? "overflow-hidden" : "overflow-y-auto"}>
<Navbar />
<div
className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"} ${isChatRoute ? "" : "pt-14 md:pt-0"}`}
className={`relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"} ${isChatRoute ? "" : "pt-14 md:pt-0"}`}
>
<AnimatePresence initial={false} mode="wait">
{/* Use mode="popLayout" instead of "wait" to prevent UI freezes when
switching from heavy pages (like Export with many checkpoints).
"popLayout" allows the new route to mount immediately while the
old one animates out, avoiding blocking on expensive exit renders.
See issue #5850. */}
<AnimatePresence initial={false} mode="popLayout">
<motion.div
key={pathname}
initial={{ opacity: 0 }}

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