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

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-06-11 09:22:24 -07:00
commit 4d2afc62a0
96 changed files with 9609 additions and 522 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -333,19 +333,16 @@ jobs:
run: |
python -m pytest -v --tb=short tests/test_callback_signature_drift.py
- name: batched left-padding generation guard (HARD GATE)
# Guards _fast_prepare_inputs_for_generation against the bug class of
# issues #1066 / #3699: position_ids taken from cache_position (which
# counts left-pad tokens) or the 2D attention mask truncated to its
# last column. Both shipped in cc4c5d77 and were fixed by #2216 and
# #4100; nothing tested this path, so each regression reached users.
# Layer 1 in the file is stdlib-ast-only (survives unsloth import
# breakage), layer 2 calls the real function on CPU via the
# tests/conftest.py CUDA spoof. Validated to fail on the pre-#2216
# and pre-#4100 code states; staging proof on GPU-less runners:
# danielhanchen/unsloth-staging-2 PR 170 (green, gate passed in all combos) / PR 172 (red, gate failed in all combos).
- 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
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/

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

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

@ -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
@ -829,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
@ -841,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

82
install.ps1
View file

@ -983,10 +983,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
@ -1406,14 +1409,58 @@ shell.Run cmd, 0, False
}
}
# ── 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) {
@ -1423,8 +1470,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 {}
}
}
@ -1694,7 +1740,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.
@ -1830,7 +1876,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.2" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -1844,7 +1890,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.2" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -1891,7 +1937,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.2" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
@ -1903,7 +1949,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.2" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -1931,7 +1977,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.2" --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)
@ -2046,6 +2092,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
@ -2056,6 +2107,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

100
install.sh
View file

@ -1717,8 +1717,15 @@ _ensure_rocm_probe_env() {
# 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
@ -1726,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 ──
@ -1763,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
@ -1847,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' \
@ -2098,6 +2155,21 @@ _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
@ -2333,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.2" 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.
@ -2346,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.2" unsloth-zoo
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2550,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.2" 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
@ -2568,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.2" 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..."
@ -2600,7 +2672,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.2" --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..."

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.

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

@ -514,9 +514,18 @@ _CTX_FIT_VRAM_FRACTION = 0.90
_MTP_VRAM_RESERVE_FRAC = 0.05
def _auto_mode_drops_mtp(req_mode: Optional[str], size_b: Optional[float]) -> bool:
"""Auto mode drops MTP below _MTP_MIN_SIZE_B (draft-mtp regresses there);
forced mtp / mtp+ngram engage regardless of size."""
def _auto_mode_drops_mtp(
req_mode: Optional[str],
size_b: Optional[float],
*,
has_separate_drafter: bool = False,
) -> bool:
"""Auto mode drops MTP below _MTP_MIN_SIZE_B for an embedded draft head
(its per-token cost regresses there); a separate drafter (Gemma) is a tiny
standalone model that still speeds up below 3B, so it never drops. Forced
mtp / mtp+ngram engage regardless of size."""
if has_separate_drafter:
return False
return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B
@ -667,6 +676,11 @@ class LlamaCppBackend:
# Separate MTP drafter launched with the current model; reload-dedup
# key so a drafter that appears next to the weights forces a reload.
self._mtp_draft_path: Optional[str] = None
# Why MTP was disabled on the last load that asked for it (auto on an
# MTP model, or forced mtp / mtp+ngram), else None. Drives the "update
# llama.cpp" hint in the UI. "binary_no_mtp" / "binary_outdated" ->
# a newer prebuilt would help; "runtime_error" -> it may not.
self._spec_fallback_reason: Optional[str] = None
self._hf_variant: Optional[str] = None
self._is_vision: bool = False
self._healthy = False
@ -785,6 +799,11 @@ class LlamaCppBackend:
def mtp_draft_path(self) -> Optional[str]:
return self._mtp_draft_path
@property
def spec_fallback_reason(self) -> Optional[str]:
"""Why MTP was disabled on the last MTP-requesting load, else None."""
return self._spec_fallback_reason
@property
def extra_args(self) -> Optional[List[str]]:
"""Extra llama-server flags from the last load (a copy). None =
@ -1134,6 +1153,8 @@ class LlamaCppBackend:
"ngram_mod_flavor": None,
"supports_ngram_mod": False,
"spec_draft_n_max_flag": None,
"supports_kv_unified": False,
"supports_fit_ctx": False,
}
try:
mtime = int(Path(bin_path).stat().st_mtime)
@ -1147,6 +1168,8 @@ class LlamaCppBackend:
mtp_token: Optional[str] = None
ngram_mod_flavor: Optional[str] = None
spec_draft_n_max_flag: Optional[str] = None
supports_kv_unified = False
supports_fit_ctx = False
try:
result = subprocess.run(
[bin_path, "--help"],
@ -1234,6 +1257,9 @@ class LlamaCppBackend:
spec_draft_n_max_flag = "--spec-draft-n-max"
elif _is_real("--draft-max"):
spec_draft_n_max_flag = "--draft-max"
supports_kv_unified = _is_real("--kv-unified")
supports_fit_ctx = _is_real("--fit-ctx")
except (OSError, subprocess.SubprocessError) as exc:
logger.debug(f"llama-server --help probe failed: {exc}")
@ -1244,6 +1270,8 @@ class LlamaCppBackend:
"ngram_mod_flavor": ngram_mod_flavor,
"supports_ngram_mod": ngram_mod_flavor is not None,
"spec_draft_n_max_flag": spec_draft_n_max_flag,
"supports_kv_unified": supports_kv_unified,
"supports_fit_ctx": supports_fit_ctx,
}
cls._capability_cache[cache_key] = info
return info
@ -2910,19 +2938,14 @@ class LlamaCppBackend:
# Auto-download the separate MTP drafter (e.g. Gemma) when
# the requested spec mode can use it. Repos with the head
# baked into the main GGUF (Qwen) have no mtp- sibling and
# this no-ops. Skipped when the user disabled MTP, drives
# --spec-type manually via extra_args, or in auto mode on a
# sub-3B model (e.g. Gemma E2B) where the resolver drops
# MTP anyway -- no point fetching a drafter it never uses.
# Forced mtp / mtp+ngram still download (user override).
# this no-ops, so the size gate stays out of it: a separate
# drafter speeds up even sub-3B (Gemma E2B), and the resolver
# below decides the final emission. Skipped only when the
# user disabled MTP or drives --spec-type manually.
_spec_canon = _canonicalize_spec_mode(speculative_type) or "auto"
_auto_drops_mtp = _auto_mode_drops_mtp(
_spec_canon, _extract_model_size_b(model_identifier)
)
if (
not mtp_draft_path
and _spec_canon in ("auto", "mtp", "mtp+ngram")
and not _auto_drops_mtp
and not _extra_args_set_spec_type(extra_args)
):
mtp_draft_path = self._download_mtp(
@ -3002,8 +3025,12 @@ class LlamaCppBackend:
_mtp_canonical = _canonicalize_spec_mode(speculative_type)
_mtp_effective = _mtp_canonical or "auto"
_mtp_size_for_fit = _extract_model_size_b(model_identifier)
# Sub-3B drops MTP only for an embedded head; a separate
# drafter (Gemma) engages and needs its VRAM reserved.
_mtp_sub_3b_for_fit = (
_mtp_size_for_fit is not None and _mtp_size_for_fit < _MTP_MIN_SIZE_B
_mtp_size_for_fit is not None
and _mtp_size_for_fit < _MTP_MIN_SIZE_B
and not bool(mtp_draft_path)
)
_mtp_will_engage = bool(
not _extra_args_set_spec_type(extra_args)
@ -3214,6 +3241,16 @@ class LlamaCppBackend:
# Fits on selected GPU(s) -- offload all layers
cmd.extend(["-ngl", "-1"])
cmd.extend(
self._ctx_integrity_flags(
n_parallel,
use_fit,
requested_ctx,
effective_ctx,
self.probe_server_capabilities(binary),
)
)
# -1 = llama.cpp auto-detect (physical cores). Pass explicitly
# so we don't inherit llama-server's internal default, which
# has varied (hardware concurrency incl. hyperthreads on some
@ -3442,7 +3479,7 @@ class LlamaCppBackend:
# Pin to selected GPU(s). On ROCm, narrowing only
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
# HIP/ROCR set, so set those too.
# set, so set HIP_VISIBLE_DEVICES too.
if gpu_indices is not None:
pinned = ",".join(str(i) for i in gpu_indices)
env["CUDA_VISIBLE_DEVICES"] = pinned
@ -3450,7 +3487,19 @@ class LlamaCppBackend:
import torch as _torch
if getattr(_torch.version, "hip", None) is not None:
env["HIP_VISIBLE_DEVICES"] = pinned
env["ROCR_VISIBLE_DEVICES"] = pinned
# Do NOT also set ROCR_VISIBLE_DEVICES to the same
# value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr
# layer and HIP_VISIBLE_DEVICES at the HIP layer, so
# setting both with the same physical indices applies
# the mask twice: ROCR reduces the visible set and
# re-indexes it from 0, then HIP indexes into the
# already-reduced set. A single non-zero pin (e.g.
# "1") then points out of range at the HIP layer, HIP
# enumerates 0 devices, and llama.cpp falls back to
# CPU ("ggml_cuda_init: no ROCm-capable device is
# detected"). The HIP mask alone narrows correctly;
# clear any inherited ROCR mask so it can't double up.
env.pop("ROCR_VISIBLE_DEVICES", None)
except Exception as e:
logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
@ -3567,6 +3616,7 @@ class LlamaCppBackend:
self._effective_context_length = (
effective_ctx if effective_ctx > 0 else self._context_length
)
self._reconcile_effective_ctx_with_server()
self._max_context_length = (
max_available_ctx if max_available_ctx > 0 else self._effective_context_length
)
@ -3586,8 +3636,13 @@ class LlamaCppBackend:
# failing (unknown arch / draft or context build); an
# unrelated crash (e.g. OOM) gets a neutral message.
_lo = "\n".join(self._stdout_lines).lower()
# Only an unknown architecture proves the prebuilt predates
# this MTP model (an update fixes it). The memory/context
# build failures are generic (VRAM / ctx pressure), where an
# update may not help, so classify those as runtime_error.
_arch_unsupported = "unknown model architecture" in _lo
if (
"unknown model architecture" in _lo
_arch_unsupported
or "failed to measure draft model memory" in _lo
or "failed to measure mtp context memory" in _lo
or "failed to create llama_context" in _lo
@ -3597,10 +3652,14 @@ class LlamaCppBackend:
"speculative decoding -- run `unsloth studio "
"update` for MTP"
)
self._spec_fallback_reason = (
"binary_outdated" if _arch_unsupported else "runtime_error"
)
else:
_retry_reason = (
"retrying without speculative decoding in case MTP is the cause"
)
self._spec_fallback_reason = "runtime_error"
_drafter = (
Path(launch_mtp_draft_path).name
if launch_mtp_draft_path
@ -3771,23 +3830,32 @@ class LlamaCppBackend:
https://github.com/ggml-org/llama.cpp/pull/18471
MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide
Sub-3B dense MTP regresses vs spec-off: the draft head's per-token
cost exceeds the acceptance savings at this scale. Q4_K_XL clean
bench (each prompt once after an unrelated warmup) on B200 + x86 CPU:
Sub-3B dense MTP regresses vs spec-off when the head is baked into the
main GGUF (Qwen): the draft head's per-token cost exceeds the
acceptance savings at this scale. Q4_K_XL clean bench (each prompt once
after an unrelated warmup) on B200 + x86 CPU:
0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x
2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x
0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x
2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x
4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
A separate drafter (Gemma's root mtp-*.gguf) is a different, cheaper
mechanism that wins even below 3B, so it is exempt from the sub-3B drop
(``mtp_draft_path`` set -> not too small). B200 Q4_K_XL bench, draft-mtp
n=2 vs OFF: gemma-4-E2B (2B) = 1.21x, accept ~0.65 (vs ngram = 1.00x);
gemma-4-E4B (4B) and 12B engage as usual.
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
diverse content); forced MTP on a model with no head/drafter defaults
back (mtp -> spec-default, mtp+ngram -> ngram-mod) since llama-server
aborts otherwise; sub-3B real-MTP engages with a warning.
diverse content) for an embedded sub-3B head; forced MTP on a model
with no head/drafter defaults back (mtp -> spec-default, mtp+ngram ->
ngram-mod) since llama-server aborts otherwise; a drafter the binary
cannot build (older prebuilt, or a CUDA kernel limit) aborts the spawn
and the load retries once without speculative decoding.
"""
flags: List[str] = []
# Reset; emit branches re-set on the resolved emission.
self._spec_draft_n_max = None
self._speculative_type = None
self._spec_fallback_reason = None
# Canonical UI-facing requested mode (legacy values mapped via
# _canonicalize_spec_mode).
@ -3801,7 +3869,11 @@ class LlamaCppBackend:
)
user_owns_spec_type = _extra_args_set_spec_type(extra_args)
_mtp_size_b = _extract_model_size_b(model_identifier)
_mtp_too_small = _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B
# The sub-3B regression is an embedded-head cost; a separate drafter
# (Gemma) is a cheap standalone model that wins below 3B, so exempt it.
_mtp_too_small = (
_mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path)
)
if user_owns_spec_type:
# User --spec-type wins outright; suppress auto-emit to avoid a
@ -3833,6 +3905,7 @@ class LlamaCppBackend:
"run `unsloth studio update`. Loading without "
"speculative decoding."
)
self._spec_fallback_reason = "binary_no_mtp"
return False
draft_n_max = _resolved_draft_n_max()
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
@ -4098,6 +4171,7 @@ class LlamaCppBackend:
self._gguf_path = None
self._hf_repo = None
self._mtp_draft_path = None
self._spec_fallback_reason = None
self._hf_variant = None
self._is_vision = False
self._is_audio = False
@ -4426,6 +4500,64 @@ class LlamaCppBackend:
logger.error(f"llama-server health check timed out after {timeout}s")
return False
@staticmethod
def _ctx_integrity_flags(
n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict
) -> list[str]:
"""Flags that keep the per-request window equal to the advertised ctx.
Explicit ``--parallel`` disables llama-server's auto-slots
``--kv-unified`` default, silently splitting ``-c`` into per-slot
windows of ``-c / N``; restore the shared pool so one request can use
the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step
at an explicitly requested ctx (default floor is 4096) so it offloads
or fails instead of silently shrinking the window.
"""
flags: list[str] = []
if n_parallel > 1 and caps.get("supports_kv_unified"):
flags.append("--kv-unified")
if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"):
flags.extend(["--fit-ctx", str(effective_ctx)])
return flags
def _query_server_n_ctx(self) -> Optional[int]:
"""Per-slot context llama-server actually allocated, from ``/props``.
The memory-fit step or ``--parallel`` slot split can leave this below
the requested ``-c``; requests are validated against this value.
"""
url = f"http://127.0.0.1:{self._port}/props"
try:
resp = httpx.get(url, timeout = 5.0)
if resp.status_code != 200:
return None
settings = resp.json().get("default_generation_settings") or {}
n_ctx = settings.get("n_ctx")
return int(n_ctx) if n_ctx else None
except Exception:
return None
def _reconcile_effective_ctx_with_server(self) -> None:
"""Adopt the server's real ``n_ctx`` when it is below Studio's value.
Keeps ``context_length`` (load response, status route, passthrough
``max_tokens`` ceiling) honest; clients sized to the requested value
would otherwise hit ``exceed_context_size_error`` 400s early.
"""
actual_n_ctx = self._query_server_n_ctx()
if not actual_n_ctx or actual_n_ctx <= 0:
return
if self._effective_context_length and actual_n_ctx < self._effective_context_length:
logger.warning(
"llama-server allocated a smaller per-request context than "
f"requested ({self._effective_context_length} -> {actual_n_ctx}; "
"memory fit or --parallel slot split); clients must treat "
f"{actual_n_ctx} as the real context window."
)
self._effective_context_length = actual_n_ctx
elif not self._effective_context_length:
self._effective_context_length = actual_n_ctx
# ── Message building (OpenAI format) ──────────────────────────
@staticmethod

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

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)

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

@ -1389,6 +1389,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"
@ -1852,7 +1858,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]:
@ -2143,14 +2149,14 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# 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. Values seeded by
# the installer are redetectable defaults, while caller overrides
# remain authoritative.
# 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
@ -2163,6 +2169,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),
@ -2174,15 +2181,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 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,
)
# 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

15
studio/backend/main.py
View file

@ -83,9 +83,9 @@ if sys.platform == "win32":
# ── 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.
# 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 (
@ -95,7 +95,6 @@ if sys.platform == "win32":
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:
@ -118,11 +117,13 @@ 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:
# 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"

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

@ -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:
@ -681,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/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:

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

@ -242,6 +242,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
@ -417,7 +580,6 @@ try:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_auto_mode_drops_mtp,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
@ -431,7 +593,6 @@ try:
from utils.inference import load_inference_config
from utils.models.model_config import (
detect_mtp_file,
extract_model_size_b,
load_model_defaults,
)
from utils.native_path_leases import (
@ -450,7 +611,6 @@ except ImportError:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_auto_mode_drops_mtp,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
@ -464,7 +624,6 @@ except ImportError:
from utils.inference import load_inference_config
from utils.models.model_config import (
detect_mtp_file,
extract_model_size_b,
load_model_defaults,
)
from utils.native_path_leases import (
@ -1025,14 +1184,15 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC
else:
if list(request.llama_extra_args) != backend_extra:
return False
# A drafter that appeared next to the loaded weights since the last load
# changes the launch command (--model-draft) when the mode can use it;
# without this, a duplicate /load is deduped and MTP can't engage short
# of an unload. Runs last: it stats the filesystem (two dir scans against
# the resolved weight path -- covers local dirs and HF cache snapshots
# alike), so every pure-memory comparison above short-circuits first.
# Skipped when auto drops MTP anyway (sub-3B) or the user owns
# --spec-type, where a drafter changes nothing. Resolve both sides: the
# 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:
@ -1041,10 +1201,7 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC
if request.llama_extra_args is not None
else llama_backend.extra_args
)
size_b = extract_model_size_b(llama_backend.model_identifier or "")
if not _auto_mode_drops_mtp(req_mode, size_b) and not _extra_args_set_spec_type(
effective_extras
):
if not _extra_args_set_spec_type(effective_extras):
detected = detect_mtp_file(llama_backend.gguf_path)
stored = llama_backend.mtp_draft_path
try:
@ -1912,6 +2069,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,
@ -4461,26 +4619,35 @@ 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",
}
# Extension fields: the real per-request window (post /props readback)
# so clients can budget/compact against the enforced limit.
if llama_backend.context_length:
entry["context_length"] = llama_backend.context_length
if llama_backend.max_context_length:
entry["max_context_length"] = llama_backend.max_context_length
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",
}
entry = {
"id": backend.active_model_name,
"object": "model",
"created": _created,
"owned_by": "local",
}
_sf_ctx = getattr(backend, "context_length", None) or getattr(
backend, "max_seq_length", None
)
if _sf_ctx:
entry["context_length"] = _sf_ctx
models.append(entry)
return models
@ -4730,6 +4897,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.
@ -4789,10 +4967,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",
@ -5037,7 +5214,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"
@ -6743,7 +6922,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
@ -6754,12 +6937,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,
@ -6794,7 +6984,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)
@ -6812,27 +7004,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(
@ -6845,6 +7042,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:
@ -6939,22 +7144,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

11
studio/backend/routes/llama.py
View file

@ -13,6 +13,7 @@ never blocks on a missing marker / offline GitHub.
from __future__ import annotations
import asyncio
from typing import Optional
from fastapi import APIRouter, Depends, Query
@ -48,6 +49,9 @@ class LlamaUpdateStatusResponse(BaseModel):
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)
@ -65,11 +69,14 @@ async def llama_update_status(
),
current_subject: str = Depends(get_current_subject),
) -> LlamaUpdateStatusResponse:
return LlamaUpdateStatusResponse(**get_update_status(force_refresh = force_refresh))
# 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:
return LlamaUpdateActionResponse(**start_update())
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(

52
studio/backend/run.py
View file

@ -414,9 +414,26 @@ 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.
"""
if not _cloudflare_url:
return
from startup_banner import stdout_supports_color
accent = "\033[38;5;150;1m"
reset = "\033[0m"
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.
@ -584,6 +601,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")
@ -594,6 +618,10 @@ _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
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
@ -769,6 +797,7 @@ def run_server(
silent: bool = False,
api_only: bool = False,
llama_parallel_slots: int = 1,
cloudflare: bool = True,
):
"""
Start the FastAPI server.
@ -973,6 +1002,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)
@ -1011,6 +1055,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
@ -1037,6 +1088,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)

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

@ -0,0 +1,425 @@
# 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
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

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

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")

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

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

@ -1219,3 +1219,242 @@ def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch):
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

325
studio/backend/tests/test_llama_cpp_update.py
View file

@ -60,24 +60,151 @@ def _write_install(
def _clean_state(monkeypatch):
freshness.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
# 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 test_status_no_marker(monkeypatch, tmp_path):
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)
@ -99,13 +226,62 @@ def test_status_up_to_date(monkeypatch, tmp_path):
assert st["update_available"] is False
def test_start_update_no_marker_refuses(monkeypatch, tmp_path):
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_marker"
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)
# Status polls probe `llama-server --version`; keep the installer argv.
if "--version" in cmd:
return _Proc()
captured["cmd"] = cmd
_write_install(install_dir, "b9585") # installer writes the marker
return _Proc()
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
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):
@ -123,6 +299,10 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
stderr = ""
def _fake_run(cmd, **kwargs):
cmd = list(cmd)
# Status polls probe `llama-server --version`; keep the installer argv.
if "--version" in cmd:
return _Proc()
captured["cmd"] = cmd
# Simulate the installer writing a new marker with the latest tag.
_write_install(install_dir, "b9518")
@ -229,7 +409,11 @@ def _capture_install_cmd(
stderr = ""
def _fake_run(cmd, **kwargs):
captured["cmd"] = list(cmd)
cmd = list(cmd)
# Status polls probe `llama-server --version`; keep the installer argv.
if "--version" in cmd:
return _Proc()
captured["cmd"] = cmd
_write_install(install_dir, latest, repo = repo, asset = asset)
return _Proc()
@ -442,3 +626,136 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path):
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"

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()

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

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

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

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)

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

@ -21,6 +21,7 @@ Design notes:
from __future__ import annotations
import json
import os
import re
import subprocess
@ -108,6 +109,143 @@ def _installer_script() -> Optional[Path]:
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.
@ -115,6 +253,20 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
"""
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:
@ -143,6 +295,7 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
"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,
}
@ -254,19 +407,7 @@ 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()
install_dir = _install_dir_for(binary)
marker = read_install_marker(binary)
if install_dir is None or not marker:
return {
"started": False,
"reason": "no_prebuilt_marker",
"message": (
"This llama.cpp install was not provisioned from an Unsloth "
"prebuilt (source build or custom path); in-app update is "
"unavailable."
),
"job": get_update_status()["job"],
}
script = _installer_script()
if script is None:
return {
@ -275,9 +416,47 @@ def start_update() -> dict:
"message": "install_llama_prebuilt.py could not be located.",
"job": get_update_status()["job"],
}
repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO
from_tag = marker.get("tag") or marker.get("release_tag")
asset = marker.get("asset")
if marker:
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:

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"

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"

2
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"

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

4
studio/frontend/src/features/auth/change-password-page.tsx
View file

@ -10,11 +10,11 @@ export function ChangePasswordPage() {
<div className="relative flex min-h-[calc(100dvh-var(--studio-titlebar-height,0px))] items-center justify-center overflow-hidden bg-background px-4 py-8 sm:px-6 sm:py-10 md:px-10">
<LightRays
count={6}
color="rgba(34, 197, 94, 0.25)"
color="rgba(34, 197, 94, 0.1)"
blur={34}
speed={15}
length="70vh"
style={{ opacity: 0.4 }}
className="opacity-35 dark:opacity-15"
/>
<Card className="relative z-10 w-full max-w-sm px-5 py-6 shadow-border ring-1 ring-border sm:px-6 sm:py-8">
<AuthForm mode="change-password" />

4
studio/frontend/src/features/auth/login-page.tsx
View file

@ -10,11 +10,11 @@ export function LoginPage() {
<div className="relative flex min-h-[calc(100dvh-var(--studio-titlebar-height,0px))] items-center justify-center overflow-hidden bg-background px-4 py-8 sm:px-6 sm:py-10 md:px-10">
<LightRays
count={6}
color="rgba(34, 197, 94, 0.25)"
color="rgba(34, 197, 94, 0.1)"
blur={34}
speed={15}
length="70vh"
style={{ opacity: 0.4 }}
className="opacity-35 dark:opacity-15"
/>
<Card className="relative z-10 w-full max-w-sm px-5 py-6 shadow-border ring-0 sm:px-6 sm:py-8">
<AuthForm mode="login" />

15
studio/frontend/src/features/chat/api/mcp-servers-api.ts
View file

@ -21,6 +21,12 @@ export interface McpServerProbeResult {
error: string | null;
}
export interface McpServerImportResult {
created: McpServerConfig[];
skipped: string[];
errors: string[];
}
function parseErrorText(status: number, body: unknown): string {
if (body && typeof body === "object") {
const { detail, message } = body as { detail?: unknown; message?: unknown };
@ -114,3 +120,12 @@ export function testMcpServer(payload: {
},
});
}
// Bulk-import servers from a standard mcpServers JSON config (Claude Desktop,
// Cursor, Cline, VS Code). The backend skips duplicates and reports per-entry
// errors instead of failing the whole batch.
export function importMcpServers(
config: unknown,
): Promise<McpServerImportResult> {
return mcpRequest("/import", { method: "POST", body: { config } });
}

88
studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx
View file

@ -1,11 +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
import { useCallback, useEffect, useState } from "react";
import { type ChangeEvent, useCallback, useEffect, useRef, useState } from "react";
import { toast } from "sonner";
import { Delete02Icon, Edit03Icon, PlusSignIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { RefreshCwIcon } from "lucide-react";
import { RefreshCwIcon, UploadIcon } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
@ -23,6 +23,7 @@ import {
type McpServerConfig,
createMcpServer,
deleteMcpServer,
importMcpServers,
listMcpServers,
refreshMcpServerTools,
testMcpServer,
@ -194,7 +195,9 @@ export function ChatMcpServersDialog({
const [form, setForm] = useState<FormState>(EMPTY_FORM);
const [saving, setSaving] = useState(false);
const [testing, setTesting] = useState(false);
const [importing, setImporting] = useState(false);
const [refreshingId, setRefreshingId] = useState<string | null>(null);
const fileInputRef = useRef<HTMLInputElement>(null);
const refresh = useCallback(async () => {
setLoading(true);
@ -315,6 +318,49 @@ export function ChatMcpServersDialog({
}
}
async function onImportFile(e: ChangeEvent<HTMLInputElement>) {
const file = e.target.files?.[0];
e.target.value = ""; // let the user re-pick the same file later
if (!file) return;
let config: unknown;
try {
config = JSON.parse(await file.text());
} catch {
toast.error("Invalid JSON file");
return;
}
setImporting(true);
try {
const result = await importMcpServers(config);
const parts = [`${result.created.length} added`];
if (result.skipped.length) parts.push(`${result.skipped.length} skipped`);
if (result.errors.length) {
parts.push(
`${result.errors.length} error${result.errors.length === 1 ? "" : "s"}`,
);
}
const summary = parts.join(", ");
if (result.errors.length) {
toast.warning(summary, {
description: (
<div className="whitespace-pre-line">
{result.errors.slice(0, 5).join("\n")}
</div>
),
});
} else {
toast.success(summary);
}
await refresh();
} catch (err) {
toast.error("Import failed", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
setImporting(false);
}
}
async function removeServer(server: McpServerConfig) {
const ok = window.confirm(`Delete MCP server "${server.display_name}"?`);
if (!ok) return;
@ -385,9 +431,35 @@ export function ChatMcpServersDialog({
Register remote (HTTP) or local (stdio command) MCP servers.
</DialogDescription>
</DialogHeader>
<input
ref={fileInputRef}
type="file"
accept="application/json,.json"
className="hidden"
onChange={onImportFile}
/>
{showForm ? (
<div className="flex flex-col gap-4">
{view.kind === "create" && (
<div className="flex items-center justify-between gap-3 rounded-md border border-dashed px-3 py-2">
<span className="text-xs text-muted-foreground">
Import servers from a config file.
</span>
<Button
type="button"
size="sm"
variant="outline"
className="shrink-0"
onClick={() => fileInputRef.current?.click()}
disabled={importing}
title="Import servers from a mcpServers JSON config (Claude Desktop, Cursor, VS Code…)"
>
{importing ? <Spinner /> : <UploadIcon size={14} />}
Import config
</Button>
</div>
)}
<div className="grid gap-2">
<Label htmlFor="mcp-display-name">Display name</Label>
<Input
@ -467,7 +539,17 @@ export function ChatMcpServersDialog({
</div>
) : (
<div className="flex min-w-0 flex-col gap-3">
<div className="flex justify-end">
<div className="flex justify-end gap-2">
<Button
size="sm"
variant="outline"
onClick={() => fileInputRef.current?.click()}
disabled={importing}
title="Import servers from a mcpServers JSON config (Claude Desktop, Cursor, VS Code…)"
>
{importing ? <Spinner /> : <UploadIcon size={14} />}
Import config
</Button>
<Button size="sm" onClick={startCreate}>
<HugeiconsIcon icon={PlusSignIcon} size={14} />
Add server

50
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -50,6 +50,7 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip";
import { useIsMobile } from "@/hooks/use-mobile";
import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check";
import { cn } from "@/lib/utils";
import {
ArrowDown01Icon,
@ -62,7 +63,7 @@ import { HugeiconsIcon } from "@hugeicons/react";
import { ChevronDown, ExternalLink } from "lucide-react";
import { Tooltip as TooltipPrimitive } from "radix-ui";
import { Fragment, type ReactNode } from "react";
import { useEffect, useMemo, useRef, useState } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { toast } from "@/lib/toast";
import { OpenAICodeExecSection } from "./components/openai-code-exec-section";
import {
@ -487,6 +488,27 @@ export function ChatSettingsPanel({
const loadedSpeculativeType = useChatRuntimeStore(
(s) => s.loadedSpeculativeType,
);
const specFallbackReason = useChatRuntimeStore((s) => s.specFallbackReason);
// "binary_no_mtp" / "binary_outdated" mean a newer prebuilt would re-enable
// MTP; "runtime_error" means the current build cannot run it (no update push).
const mtpUpdatable =
specFallbackReason === "binary_no_mtp" ||
specFallbackReason === "binary_outdated";
const {
status: llamaUpdateStatus,
applying: llamaUpdating,
apply: applyLlamaUpdate,
} = useLlamaUpdateCheck({ enabled: mtpUpdatable });
const handleMtpUpdate = useCallback(async () => {
const result = await applyLlamaUpdate();
if (result.ok) {
toast.success(
`llama.cpp updated to ${result.tag ?? "the latest build"}. Reload your model to enable MTP.`,
);
} else {
toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`);
}
}, [applyLlamaUpdate]);
const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax);
const setSpecDraftNMax = useChatRuntimeStore((s) => s.setSpecDraftNMax);
const loadedSpecDraftNMax = useChatRuntimeStore(
@ -923,6 +945,32 @@ export function ChatSettingsPanel({
</Select>
</div>
</div>
{specFallbackReason &&
(speculativeType === "auto" ||
speculativeType === "mtp" ||
speculativeType === "mtp+ngram") && (
<div className="rounded-lg bg-amber-500/[0.08] px-3 py-2 text-[12px] leading-[1.4] text-nav-fg/80">
<p>
{specFallbackReason === "runtime_error"
? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding."
: "MTP is not available in the installed llama.cpp build, so this model is running without it." +
(llamaUpdateStatus?.update_available
? " Update llama.cpp to enable it."
: "")}
</p>
{mtpUpdatable && llamaUpdateStatus?.update_available && (
<Button
size="sm"
className="corner-squircle mt-2 h-7 text-[12px]"
onClick={handleMtpUpdate}
disabled={llamaUpdating}
data-test-id="mtp-update-button"
>
{llamaUpdating ? "Updating..." : "Update llama.cpp"}
</Button>
)}
</div>
)}
{(speculativeType === "mtp" ||
speculativeType === "mtp+ngram") && (
<div className="flex items-center justify-between gap-3">

4
studio/frontend/src/features/chat/components/chat-search-dialog.tsx
View file

@ -64,10 +64,10 @@ export function ChatSearchDialog() {
<CommandDialog
open={isOpen}
onOpenChange={setOpen}
className="chat-search-surface corner-squircle top-[25%] w-[635px] max-w-[calc(100%-2rem)] gap-0 p-0 ring-0 sm:max-w-[635px]"
className="chat-search-surface rounded-3xl! top-1/2 -translate-y-1/2 w-[635px] max-w-[calc(100%-2rem)] gap-0 p-0 ring-0 sm:max-w-[635px]"
overlayClassName="bg-transparent"
>
<Command className="rounded-4xl p-0" filter={chatSearchFilter}>
<Command className="rounded-3xl p-0" filter={chatSearchFilter}>
<div className="flex items-center gap-3 border-b border-border/40 px-4 py-3">
<HugeiconsIcon
icon={SearchIcon}

1
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -412,6 +412,7 @@ export function useChatModelRuntime() {
statusRes.requires_trust_remote_code ?? false,
defaultChatTemplate: nextDefaultChatTemplate,
loadedIsMultimodal: isMultimodalResponse(statusRes),
specFallbackReason: statusRes.spec_fallback_reason ?? null,
...(prevState.loadedSpeculativeType === null && {
speculativeType: currentSpecType,
loadedSpeculativeType: currentSpecType,

7
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -416,6 +416,11 @@ type ChatRuntimeStore = {
loadedKvCacheDtype: string | null;
speculativeType: string | null;
loadedSpeculativeType: string | null;
/**
* Why MTP was disabled on the loaded model despite being requested, or null.
* Mirrors InferenceStatusResponse.spec_fallback_reason.
*/
specFallbackReason: string | null;
/** User --spec-draft-n-max override (null = platform default). */
specDraftNMax: number | null;
loadedSpecDraftNMax: number | null;
@ -765,6 +770,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
loadedKvCacheDtype: null,
speculativeType: "auto",
loadedSpeculativeType: null,
specFallbackReason: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
@ -977,6 +983,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
loadedKvCacheDtype: null,
speculativeType: "auto",
loadedSpeculativeType: null,
specFallbackReason: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,

6
studio/frontend/src/features/chat/types/api.ts
View file

@ -174,6 +174,12 @@ export interface InferenceStatusResponse {
/** Canonical UI-facing mode currently active. See LoadModelRequest. */
speculative_type?: string | null;
spec_draft_n_max?: number | null;
/**
* Why MTP was disabled on the loaded model despite being requested.
* "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable
* it; "runtime_error" -> the current build could not run it. Null otherwise.
*/
spec_fallback_reason?: string | null;
}
export interface AudioGenerationResponse {

5
studio/frontend/src/features/training/stores/training-runtime-store.ts
View file

@ -274,7 +274,10 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
jobId: payload.job_id || state.jobId,
currentStep: step,
totalSteps: Math.max(payload.total_steps, state.totalSteps),
currentLoss: currentLoss ?? state.currentLoss,
// A null loss at a new step means the backend reported a non-finite
// loss; clear the display instead of keeping the stale value.
currentLoss:
currentLoss ?? (step > state.currentStep ? null : state.currentLoss),
currentLearningRate: currentLearningRate ?? state.currentLearningRate,
progressPercent: payload.progress_percent,
currentEpoch: payload.epoch ?? state.currentEpoch,

3
studio/frontend/src/features/training/types/runtime.ts
View file

@ -90,7 +90,8 @@ export interface TrainingRuntimeState {
currentStep: number;
totalSteps: number;
currentEpoch: number;
currentLoss: number;
// null while the latest reported loss is non-finite
currentLoss: number | null;
currentLearningRate: number;
progressPercent: number;
elapsedSeconds: number | null;

10
studio/frontend/src/index.css
View file

@ -897,8 +897,7 @@
/* Chat search box: borderless, soft Gemini-style elevation. */
.chat-search-surface {
border: none;
/* Pin to the dark --radius so the corners are the same (less round) in
both themes; rounded-4xl here and on the inner Command follow this. */
/* Pin to the dark --radius so rounded-3xl corners stay consistent. */
--radius: 0.625rem;
/* Match the chat box: soft elevation. */
box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
@ -1751,10 +1750,11 @@
margin: 0 !important;
}
/* Boost shadow on dark surfaces; mirrors .shadow-border / .menu-soft-surface pattern. */
/* Composer shadow on the dark background color so toasts
do not merge into card-colored surfaces behind them. */
.dark [data-sonner-toast][data-styled='true'] {
background-color: var(--card) !important;
box-shadow: none !important;
background-color: var(--background) !important;
box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16) !important;
}
/* Selectable toast text; non-selectable toast buttons. */

185
studio/install_llama_prebuilt.py
View file

@ -177,6 +177,7 @@ VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf"
INSTALL_LOCK_TIMEOUT_SECONDS = 300
INSTALL_STAGING_ROOT_NAME = ".staging"
GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"}
HF_AUTH_HOSTS = {"huggingface.co", "www.huggingface.co"}
RETRYABLE_HTTP_STATUS = {408, 429, 500, 502, 503, 504}
HTTP_FETCH_ATTEMPTS = 4
HTTP_FETCH_BASE_DELAY_SECONDS = 0.75
@ -484,6 +485,10 @@ def should_send_github_auth(url: str | None) -> bool:
return parsed_hostname(url) in GITHUB_AUTH_HOSTS
def should_send_hf_auth(url: str | None) -> bool:
return parsed_hostname(url) in HF_AUTH_HOSTS
def auth_headers(url: str | None = None) -> dict[str, str]:
headers = {
"User-Agent": "unsloth-studio-llama-prebuilt",
@ -491,9 +496,35 @@ def auth_headers(url: str | None = None) -> dict[str, str]:
token = os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
if token and should_send_github_auth(url):
headers["Authorization"] = f"Bearer {token}"
return headers
# Anonymous huggingface.co fetches share a per-IP rate limit that CI
# fleets exhaust (HTTP 429), sinking the prebuilt path into a source
# build. Authenticate when a token is available.
hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
if hf_token and should_send_hf_auth(url):
headers["Authorization"] = f"Bearer {hf_token}"
return headers
class _CrossHostAuthStrippingRedirectHandler(urllib.request.HTTPRedirectHandler):
"""Drop Authorization when a redirect leaves the original host.
huggingface.co redirects file downloads to CDN hosts whose signed URLs
can reject foreign Authorization headers; urllib forwards headers to
redirect targets by default (requests/huggingface_hub strip them).
"""
def redirect_request(self, req, fp, code, msg, headers, newurl):
new_request = super().redirect_request(req, fp, code, msg, headers, newurl)
if new_request is not None and parsed_hostname(newurl) != parsed_hostname(req.full_url):
new_request.headers.pop("Authorization", None)
new_request.unredirected_hdrs.pop("Authorization", None)
return new_request
_URL_OPENER = urllib.request.build_opener(_CrossHostAuthStrippingRedirectHandler())
def github_api_headers(url: str | None = None) -> dict[str, str]:
return {
"Accept": "application/vnd.github+json",
@ -936,7 +967,7 @@ def download_bytes(
for attempt in range(1, attempts + 1):
try:
request = urllib.request.Request(url, headers = headers or auth_headers(url))
with urllib.request.urlopen(request, timeout = timeout) as response:
with _URL_OPENER.open(request, timeout = timeout) as response:
total_bytes: int | None = None
content_length = response.headers.get("Content-Length")
if content_length and content_length.isdigit():
@ -1015,7 +1046,7 @@ def download_file(url: str, destination: Path) -> None:
delete = False,
) as handle:
tmp_path = Path(handle.name)
with urllib.request.urlopen(request, timeout = 120) as response:
with _URL_OPENER.open(request, timeout = 120) as response:
total_bytes: int | None = None
content_length = response.headers.get("Content-Length")
if content_length and content_length.isdigit():
@ -1095,6 +1126,20 @@ def commit_source_archive_urls(repo: str, source_commit: str) -> list[str]:
]
def release_asset_download_url(
repo: str | None, release_tag: str | None, asset_name: str | None
) -> str | None:
"""Direct download URL for a release asset, or None if any part is missing.
A mix build's merged commit is never pushed, so its source tree is only
reachable as this asset (codeload would 404 on the merge commit)."""
if not repo or not release_tag or not asset_name:
return None
return (
f"https://github.com/{repo}/releases/download/"
f"{urllib.parse.quote(release_tag, safe = '')}/{urllib.parse.quote(asset_name, safe = '')}"
)
def github_release_assets(repo: str, tag: str) -> dict[str, str]:
payload = fetch_json(
f"https://api.github.com/repos/{repo}/releases/tags/{urllib.parse.quote(tag, safe = '')}"
@ -2890,7 +2935,27 @@ def detect_host() -> HostInfo:
except Exception:
pass
# Linux /proc/driver/nvidia/gpus fallback: the NVIDIA driver exposes one
# subdir per GPU here regardless of nvidia-smi state, so a host whose
# nvidia-smi is absent from PATH, wedged, or failing is still recognised as
# NVIDIA. Mirrors the fallback added to install.sh / install_python_stack.py
# in PR 6174 so the prebuilt installer does not misroute such hosts to ROCm
# or CPU. driver_cuda_version / compute_caps stay unset here; downstream
# CUDA asset selection treats unknown SMs as "prefer portable" and an
# unknown driver runtime line as "no published CUDA match" (returns None,
# no crash), so planning falls back to a source build with GGML_CUDA=ON.
if is_linux and not has_physical_nvidia:
try:
proc_gpu_dir = "/proc/driver/nvidia/gpus"
if os.path.isdir(proc_gpu_dir) and os.listdir(proc_gpu_dir):
has_physical_nvidia = True
has_usable_nvidia = visible_device_tokens != []
except OSError:
pass
# Detect AMD ROCm (HIP) -- require actual GPU, not just tools installed
# NVIDIA takes precedence: when an NVIDIA GPU is usable, skip ROCm probing
# entirely so co-installed ROCm tools cannot misroute the host (PR 6174).
def _amd_smi_has_gpu(stdout: str) -> bool:
"""Check for 'GPU: <number>' data rows, not just a table header."""
@ -2898,7 +2963,7 @@ def detect_host() -> HostInfo:
has_rocm = False
rocm_gfx_target: str | None = None
if is_linux:
if is_linux and not has_usable_nvidia:
# 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
# rocminfo can live only under /opt/rocm/bin (the profile.d PATH
@ -2937,7 +3002,7 @@ def detect_host() -> HostInfo:
has_rocm = True
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
elif is_windows:
elif is_windows and not has_usable_nvidia:
# Windows: prefer active probes that validate GPU presence.
# hipinfo / amd-smi are often NOT on PATH -- the HIP SDK installer
# sets HIP_PATH / ROCM_PATH but does not always add the bin dir to
@ -3045,6 +3110,21 @@ def _apply_host_overrides(
return host
def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str:
"""The release repo setup.sh / setup.ps1 pick for this host: macOS always the
fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows
-> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA
or ROCm) -> the fork. linux_amd_tooling_present mirrors setup.sh routing Linux
hosts that expose AMD tooling (rocminfo/amd-smi/hipconfig/hipinfo) to the fork
even when the probe cannot confirm an active GPU. Mirrors the shell routing."""
if host.is_macos:
return DEFAULT_PUBLISHED_REPO
has_gpu = (
host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present)
)
return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO
def pick_windows_cuda_runtime(host: HostInfo) -> str | None:
if not host.driver_cuda_version:
return None
@ -4460,13 +4540,17 @@ def hydrate_source_tree(
expected_sha256: str | None,
source_label: str | None = None,
exact_source: bool = False,
asset_url: str | None = None,
) -> None:
archive_path = work_dir / f"llama.cpp-source-{source_ref}.tar.gz"
source_urls = (
repo_urls = (
commit_source_archive_urls(source_repo, source_ref)
if exact_source
else upstream_source_archive_urls(source_ref)
)
# Prefer the published release asset (the only copy of a mix build's merged
# tree); fall back to codeload/archive for vanilla builds whose commit is real.
source_urls = ([asset_url] if asset_url else []) + repo_urls
label = source_label or f"llama.cpp source tree for {source_ref}"
extract_dir = Path(tempfile.mkdtemp(prefix = "source-extract-", dir = work_dir))
@ -4817,6 +4901,35 @@ def confirm_install_tree(install_dir: Path, host: HostInfo) -> None:
raise RuntimeError("activated install was missing expected files: " + ", ".join(missing))
def activate_staged_dir(staging_dir: Path, dst: Path) -> None:
"""Move a freshly extracted ``staging_dir`` onto ``dst``.
``os.replace`` is attempted first as the fast path. On Windows ARM64 the
antivirus scanner can transiently hold a freshly extracted DLL open at the
moment ``MoveFileEx`` runs, surfacing as ``[WinError 5] Access is denied``;
a file-by-file copy bypasses the rename entirely.
This fallback is intentionally limited to staging trees we just extracted.
It must not be used to move an existing/active install aside: there an
``os.replace`` failure means the directory is genuinely in use, and a
silent copy + ``rmtree`` could partially delete a live install.
Only busy/lock errors (``is_busy_lock_error``) trigger the copy; anything
else (disk full, cross-device, missing path) re-raises so it cannot leave
a partially copied install behind. A copy is preferred over retrying the
rename because antivirus scans of large DLLs can outlast any reasonable
retry window.
"""
try:
os.replace(staging_dir, dst)
except OSError as exc:
if not is_busy_lock_error(exc):
raise
log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree")
shutil.copytree(staging_dir, dst, dirs_exist_ok = True)
remove_tree(staging_dir)
def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) -> None:
rollback_dir: Path | None = None
failed_dir: Path | None = None
@ -4828,7 +4941,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo)
log(f"moved existing install to rollback path {rollback_dir.name}")
log(f"activating staged install {staging_dir} -> {install_dir}")
os.replace(staging_dir, install_dir)
activate_staged_dir(staging_dir, install_dir)
log(f"activated staged install at {install_dir}")
log(f"confirming activated install tree at {install_dir}")
confirm_install_tree(install_dir, host)
@ -6369,6 +6482,15 @@ def validate_prebuilt_choice(
source_repo, source_ref, source_archive, exact_source = preferred_source_archive(
approved_checksums, llama_tag
)
# For an exact (mix) source the merge commit lives only in the release asset,
# not in any repo, so fetch the asset directly; codeload stays the fallback.
asset_url = (
release_asset_download_url(
approved_checksums.repo, approved_checksums.release_tag, source_archive.asset_name
)
if exact_source and source_archive is not None
else None
)
if exact_source:
log(f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}")
else:
@ -6385,6 +6507,7 @@ def validate_prebuilt_choice(
else f"llama.cpp source tree for {llama_tag}"
),
exact_source = exact_source,
asset_url = asset_url,
)
log(f"overlaying prebuilt bundle {choice.name} into {install_dir}")
server_path, quantize_path = install_from_archives(choice, host, install_dir, work_dir)
@ -6690,6 +6813,16 @@ def parse_args() -> argparse.Namespace:
const = "latest",
help = ("Resolve the source-build fallback plan."),
)
resolve_group.add_argument(
"--resolve-prebuilt",
nargs = "?",
const = "latest",
help = (
"Report whether an official prebuilt exists for this host without "
"downloading. Picks the host's published repo when --published-repo "
"is left at the default. Use --output-format json."
),
)
parser.add_argument(
"--output-format",
choices = ("plain", "json"),
@ -6774,6 +6907,46 @@ def main() -> int:
)
return EXIT_SUCCESS
if args.resolve_prebuilt is not None:
# Host-aware "is a prebuilt available" probe, no download. A default repo
# means "pick the repo for this host"; PrebuiltFallback == source build.
host = _apply_host_overrides(
detect_host(),
override_has_rocm = args.has_rocm,
override_rocm_gfx = args.rocm_gfx,
force_cpu = args.cpu_fallback,
)
# setup.sh routes Linux hosts with AMD tooling to the fork even when no GPU
# is probed; mirror that so a HIP source build is not offered a CPU prebuilt.
amd_tooling = host.is_linux and any(
shutil.which(t) for t in ("rocminfo", "amd-smi", "hipconfig", "hipinfo")
)
repo = (
published_repo_for_host(host, linux_amd_tooling_present = amd_tooling)
if args.published_repo == DEFAULT_PUBLISHED_REPO
else args.published_repo
)
try:
_requested, plans = resolve_simple_install_release_plans(
args.resolve_prebuilt, host, repo, args.published_release_tag or ""
)
choice = plans[0].attempts[0] if plans and plans[0].attempts else None
if choice is None:
payload = {"prebuilt_available": False, "repo": repo}
else:
payload = {
"prebuilt_available": True,
"repo": repo,
"release_tag": plans[0].release_tag,
"llama_tag": plans[0].llama_tag,
"asset": choice.name,
"install_kind": choice.install_kind,
}
except PrebuiltFallback:
payload = {"prebuilt_available": False, "repo": repo}
emit_resolver_output(payload, output_format = args.output_format)
return EXIT_SUCCESS
if not args.install_dir:
raise SystemExit(
"install_llama_prebuilt.py: --install-dir is required unless --resolve-llama-tag, --resolve-install-tag, or --resolve-source-build is used"

270
studio/install_python_stack.py
View file

@ -90,6 +90,18 @@ _PYTORCH_WHL_BASE = (
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
).rstrip("/")
# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are
# pinned to the torch<2.11 family rather than left bare: the install uses an
# exclusive --index-url (no PyPI fallback), so a bare name could resolve a
# torchvision built against a different torch major (e.g. 0.27 for torch 2.12)
# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm
# spec above, which targets the same torch family.
_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
)
# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/).
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs.
_ROCM_WINDOWS_INDEX_BASE = (
@ -557,7 +569,15 @@ def _persist_bnb_rocm_version(version: str) -> bool:
def _has_rocm_gpu() -> bool:
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed).
Always returns False when an NVIDIA GPU is present -- NVIDIA takes
priority on mixed hosts and prevents every detection path below
(rocminfo, amd-smi, KFD sysfs) from producing a false positive even
if ROCm tools are installed alongside the NVIDIA driver.
"""
if _has_usable_nvidia_gpu():
return False
for cmd, check_fn in (
# rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit).
# gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like
@ -598,6 +618,13 @@ def _has_rocm_gpu() -> bool:
# runtime-only detection. On minimal package-managed installs (no
# rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via
# /sys/class/kfd so `studio update` can still detect and repair.
#
# Guard: reject any KFD node whose properties file reports a non-AMD
# vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs
# can register KFD topology nodes with a non-zero gpu_id; those nodes
# have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002).
# Without this check the fallback returns True on NVIDIA-only systems,
# causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware.
if sys.platform != "win32":
try:
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
@ -609,29 +636,69 @@ def _has_rocm_gpu() -> bool:
gpu_id = fh.read().strip()
except OSError:
continue
if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node
return True
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
continue
# Require AMD vendor_id 4098 (0x1002) in the properties file.
# KFD properties files exist on every kernel that exposes
# /sys/class/kfd, so absence of the file means we cannot
# confirm AMD ownership -- skip the node rather than risk a
# false positive (e.g. NVIDIA open driver KFD nodes that
# lack a properties file on some kernel versions).
props_path = os.path.join(kfd_nodes, entry, "properties")
try:
with open(props_path) as fh:
props = fh.read()
except OSError:
continue # can't confirm vendor -- skip
if not re.search(r"\bvendor_id\s+4098\b", props):
continue
return True
except OSError:
pass
return False
def _has_usable_nvidia_gpu() -> bool:
"""Return True only when nvidia-smi exists AND reports at least one GPU."""
"""Return True when an NVIDIA GPU is present and usable.
Primary probe: nvidia-smi -L (subprocess).
Fallback: /proc/driver/nvidia/gpus/ sysfs (Linux only) -- handles the
case where nvidia-smi is present but the subprocess fails (PATH gap,
timeout, driver initialisation race). If either probe confirms an
NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked.
CUDA_VISIBLE_DEVICES set to "" or "-1" hides every NVIDIA device (mixed
AMD+NVIDIA hosts steering work to the AMD card); neither probe honours
that env var, so check it first and report the GPU as not usable. Unset
means all devices visible.
"""
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd is not None and cvd.strip() in ("", "-1"):
return False
exe = shutil.which("nvidia-smi")
if not exe:
return False
try:
result = subprocess.run(
[exe, "-L"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
except Exception:
return False
return result.returncode == 0 and "GPU " in result.stdout
if exe:
try:
result = subprocess.run(
[exe, "-L"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
if result.returncode == 0 and "GPU " in result.stdout:
return True
except Exception:
pass
# Fallback: the NVIDIA driver exposes one subdirectory per GPU under
# /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state.
if sys.platform != "win32":
try:
gpu_dir = "/proc/driver/nvidia/gpus"
if os.path.isdir(gpu_dir) and os.listdir(gpu_dir):
return True
except OSError:
pass
return False
def _detect_amd_gfx_codes() -> list[str]:
@ -739,6 +806,139 @@ def _install_bnb_windows_rocm() -> bool:
return True
def _detect_cuda_torch_index_url() -> str:
"""Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver.
Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update`
repairs to the same wheel family a fresh `curl | sh` install would pick.
Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the
legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings.
Defaults to cu126 when nvidia-smi is missing or the version is unreadable
(e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
"""
exe = shutil.which("nvidia-smi")
if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
exe = "/usr/bin/nvidia-smi"
tag = "cu126" # default when the driver CUDA version cannot be read
if exe:
try:
result = subprocess.run(
[exe],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
if result.returncode == 0:
m = re.search(r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", result.stdout)
if m:
major, minor = int(m.group(1)), int(m.group(2))
if major >= 13:
tag = "cu130"
elif major == 12 and minor >= 8:
tag = "cu128"
elif major == 12 and minor >= 6:
tag = "cu126"
elif major >= 12:
tag = "cu124"
elif major >= 11:
tag = "cu118"
else:
tag = "cpu" # ancient driver: no usable CUDA wheels
except Exception:
pass
return f"{_PYTORCH_WHL_BASE}/{tag}"
def _ensure_cuda_torch() -> None:
"""Repair a venv whose torch is a ROCm build on an NVIDIA host.
Counterpart to _ensure_rocm_torch. A venv poisoned by the pre-fix KFD
gpu_id false positive (ROCm torch installed on an NVIDIA-only machine)
keeps that broken torch on `studio update`, because a torch+rocm wheel
satisfies the version constraint and nothing force-reinstalls it. This
detects that exact case and reinstalls CUDA torch.
Only repairs when torch actually links against HIP/ROCm. Healthy CUDA
torch and deliberate CPU-only torch are left untouched.
"""
# Respect an explicit backend choice from install.sh: only "" (standalone
# `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu"
# (or any unrecognised value) are deliberate and must not be overridden.
if _TORCH_BACKEND not in ("", "cuda"):
return
# No CUDA torch on macOS; Windows venv/torch lifecycle is owned by
# install.ps1 (and the KFD poisoning bug is Linux-only), so skip both.
if IS_MACOS or IS_WINDOWS or NO_TORCH:
return
# Never undo a deliberate ROCm install (setup.ps1 sets this marker).
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
return
# CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for
# example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card);
# never force CUDA wheels over that choice.
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if _cvd is not None and _cvd.strip() in ("", "-1"):
return
# Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu()
# covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent.
if not _has_usable_nvidia_gpu():
return
# Classify the installed torch: "hip" (ROCm build -- the poisoning
# signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A
# non-zero exit means torch is missing or un-importable; the base install
# step handles that, so leave it alone.
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip = getattr(torch.version, 'hip', '') or ''; "
"cuda = getattr(torch.version, 'cuda', '') or ''; "
"ver = getattr(torch, '__version__', '').lower(); "
"print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
except (OSError, subprocess.TimeoutExpired):
return
if probe.returncode != 0:
return
# Take the last non-empty stdout line: stray output from sitecustomize or
# an import hook must not mask the marker (fail-closed either way).
_marker_lines = [
line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
]
if not _marker_lines or _marker_lines[-1] != "hip":
return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
index_url = _detect_cuda_torch_index_url()
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
print(
f" torch is a ROCm build on an NVIDIA host -- reinstalling "
f"CUDA torch from {index_url}\n"
f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch "
f"on a mixed AMD+NVIDIA host)"
)
pip_install(
"CUDA torch repair",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
def _ensure_rocm_torch() -> None:
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
@ -749,6 +949,13 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family
# ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh
# already selected a non-ROCm backend -- this is the authoritative signal
# and avoids re-running GPU detection in a subprocess that may see a
# different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
if _TORCH_BACKEND in ("cuda", "cpu"):
return
# setup.ps1 sets this after installing AMD wheels; skip the probe only when
# torch is actually importable as ROCm. If the venv was wiped between runs,
# the stale env-var would suppress a needed reinstall.
@ -1088,6 +1295,29 @@ def _infer_no_torch() -> bool:
NO_TORCH = _infer_no_torch()
# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so
# that this script knows which torch variant was selected without re-running
# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
# (standalone `unsloth studio update` runs, where we re-detect normally).
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
def _torch_step_label(suffix: str) -> str:
"""Return a progress label like 'torch check (cuda)' using the known backend.
Falls back to GPU detection when UNSLOTH_TORCH_BACKEND is not set (e.g.
standalone `unsloth studio update` runs that bypass install.sh).
"""
backend = _TORCH_BACKEND
if not backend:
if _has_usable_nvidia_gpu():
backend = "cuda"
elif _has_rocm_gpu():
backend = "rocm"
else:
backend = "cpu"
return f"torch {suffix} ({backend})"
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal in-place one-line progress bar.
@ -1770,7 +2000,8 @@ def install_python_stack() -> int:
# venv got CPU-only torch (common when pip resolves torch from PyPI).
# Must follow base packages so torch is present for inspection.
if not IS_MACOS and not NO_TORCH:
_progress("ROCm torch check")
_progress(_torch_step_label("check"))
_ensure_cuda_torch()
_ensure_rocm_torch()
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
@ -1955,7 +2186,8 @@ def install_python_stack() -> int:
# Running the repair last ensures ROCm torch is in place at runtime,
# whichever intermediate step clobbered it.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress("ROCm torch (final)")
_progress(_torch_step_label("final"))
_ensure_cuda_torch()
_ensure_rocm_torch()
# 14. Final check (silent; third-party conflicts are expected)

202
studio/setup.ps1
View file

@ -299,7 +299,10 @@ function Get-CudaComputeCapability {
if (-not $smiExe) { return $null }
try {
$raw = & $smiExe --query-gpu=compute_cap --format=csv,noheader 2>$null
# Bounded: a wedged nvidia-smi must not hang setup after the initial
# -L probe succeeded (the helper merges stderr after stdout, so the
# first line is still the compute_cap value).
$raw = Invoke-NvidiaSmiBounded $smiExe @('--query-gpu=compute_cap', '--format=csv,noheader')
if ($LASTEXITCODE -ne 0 -or -not $raw) { return $null }
# nvidia-smi may return multiple GPUs; take the first one
@ -363,10 +366,10 @@ function Get-PytorchCudaTag {
if (-not $smiExe) { return "cu126" }
try {
# 2>&1 | Out-String merges stderr into stdout then converts to a single
# string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 --
# ErrorRecord objects leak into $output and break the -match.
$output = & $smiExe 2>&1 | Out-String
# Bounded: a wedged nvidia-smi must not hang setup. The helper merges
# stderr into the returned string, matching the old 2>&1 | Out-String
# shape (plain 2>$null leaks ErrorRecord objects in PS 5.1).
$output = Invoke-NvidiaSmiBounded $smiExe
# 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.
@ -667,16 +670,58 @@ try {
# ============================================
# 1a. GPU detection
# ============================================
# ── 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 below) so detection
# cannot hang setup. 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+:')
}
$HasNvidiaSmi = $false
$NvidiaSmiExe = $null # Absolute path -- survives Refresh-Environment
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 {}
# Fallback: nvidia-smi may not be on PATH even though a GPU + driver exist.
@ -689,8 +734,7 @@ if (-not $HasNvidiaSmi) {
foreach ($p in $nvSmiDefaults) {
if (Test-Path $p) {
try {
& $p *> $null
if ($LASTEXITCODE -eq 0) {
if (Test-NvidiaSmiHasGpu $p) {
$HasNvidiaSmi = $true
$NvidiaSmiExe = $p
Write-Host " Found nvidia-smi at $(Split-Path $p -Parent)" -ForegroundColor Gray
@ -1151,7 +1195,16 @@ function Resolve-CudaToolkit {
$DriverMaxCuda = $null
try {
$smiOut = & $NvidiaSmiExe 2>&1 | Out-String
# Bounded: source-build toolkit resolution must not hang on a wedged smi.
# test_resolve_cuda_toolkit.ps1 extracts this function alone into a child
# pwsh (no Invoke-NvidiaSmiBounded in scope) and stubs nvidia-smi with a
# .ps1 script, so fall back to direct invocation when the bounded runner
# is unavailable; production setup.ps1 always has it defined.
$smiOut = if (Get-Command Invoke-NvidiaSmiBounded -ErrorAction SilentlyContinue) {
Invoke-NvidiaSmiBounded $NvidiaSmiExe
} else {
& $NvidiaSmiExe 2>&1 | Out-String
}
# Newer drivers report "CUDA UMD Version: X.Y" instead of "CUDA Version: X.Y"; accept both.
if ($smiOut -match "CUDA(?: UMD)? Version:\s+([\d]+)\.([\d]+)") {
$DriverMaxCuda = "$($Matches[1]).$($Matches[2])"
@ -1499,10 +1552,48 @@ if ($IsPipInstall) {
}
}
# 1g. Python (>= 3.11 and < 3.14). Prefer the Studio venv that install.ps1
# just created, then py.exe so a 3.14 ahead of 3.13 on PATH does not trip the gate.
# Conda CPython ships modified DLL search paths that break torch's c10.dll
# loading on Windows; a venv made from conda Python inherits its base_prefix,
# so check the executable path AND sys.base_prefix.
$CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)'
function Test-IsConda {
param([string]$Exe)
if ($Exe -match $CondaSkipPattern) { return $true }
try {
$basePrefix = (& $Exe -c "import sys; print(sys.base_prefix)" 2>$null | Out-String).Trim()
if ($basePrefix -match $CondaSkipPattern) { return $true }
} catch { }
return $false
}
# 1g. Python (>= 3.11 and < 3.14). Prefer the interpreter install.ps1 already
# resolved and built the venv with (UNSLOTH_SETUP_PYTHON), or the existing
# venv python, before re-probing a system where a 3.14 or a WindowsApps stub
# ahead on PATH would trip the gate. setup.ps1 only updates packages in that
# venv, so the handoff is safe to reuse once validated.
function Resolve-ReusedSetupPython {
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_SETUP_PYTHON) -and
(Test-Path -LiteralPath $env:UNSLOTH_SETUP_PYTHON)) {
return $env:UNSLOTH_SETUP_PYTHON
}
# Standalone `unsloth studio setup/update` (install.ps1 did not run): derive
# the venv python from the studio root, mirroring the resolver below.
$root = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() }
elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() }
else { Join-Path $env:USERPROFILE ".unsloth\studio" }
if ($root -eq "~") {
# Join-Path with an empty child throws on Windows PowerShell 5.1.
$root = $env:USERPROFILE
} elseif ($root -like "~/*" -or $root -like "~\*") {
$root = Join-Path $env:USERPROFILE $root.Substring(1).TrimStart('/', '\')
}
$venvPy = Join-Path $root "unsloth_studio\Scripts\python.exe"
if (Test-Path -LiteralPath $venvPy) { return $venvPy }
return $null
}
$ReusedSetupPython = Resolve-ReusedSetupPython
$HasPython = $null -ne (Get-Command python -ErrorAction SilentlyContinue)
$PyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue
$PythonOk = $false
$DetectedPyVer = $null
@ -1531,28 +1622,24 @@ function Add-PythonDirToProcessPath {
} catch { }
}
$_prereqStudioHome = $null
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) {
$_prereqStudioHome = $env:UNSLOTH_STUDIO_HOME.Trim()
} elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) {
$_prereqStudioHome = $env:STUDIO_HOME.Trim()
} else {
$_prereqStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
}
if ($_prereqStudioHome -eq "~" -or $_prereqStudioHome -like "~/*" -or $_prereqStudioHome -like "~\*") {
$_prereqStudioHome = (Join-Path $env:USERPROFILE $_prereqStudioHome.Substring(1).TrimStart('/','\'))
}
$_prereqVenvPython = Join-Path $_prereqStudioHome "unsloth_studio\Scripts\python.exe"
if (Test-Path -LiteralPath $_prereqVenvPython) {
$_venvPyVer = Get-CompatiblePythonVersion $_prereqVenvPython
if ($_venvPyVer) {
$DetectedPyVer = $_venvPyVer
Add-PythonDirToProcessPath $_prereqVenvPython
# Reuse the install.ps1 / venv interpreter before any system probe.
if ($ReusedSetupPython) {
$_reusedVer = Get-CompatiblePythonVersion $ReusedSetupPython
if ($_reusedVer -and -not (Test-IsConda $ReusedSetupPython)) {
$DetectedPyVer = $_reusedVer
Add-PythonDirToProcessPath $ReusedSetupPython
$PythonOk = $true
}
}
if (-not $PythonOk -and $PyLauncher) {
# Fall back to every py.exe on PATH (all-users and per-user launchers can both
# register). -All is required: Windows PowerShell 5.1 returns only the first
# launcher without it, and the PowerShell 7 multi-match array breaks the call
# operator if used directly.
$PyLaunchers = if ($PythonOk) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) }
foreach ($PyLauncher in $PyLaunchers) {
if ($PyLauncher.Source -match $CondaSkipPattern) { continue }
foreach ($minor in @("3.13", "3.12", "3.11")) {
try {
$out = & $PyLauncher.Source "-$minor" --version 2>&1 | Out-String
@ -1572,6 +1659,7 @@ if (-not $PythonOk -and $PyLauncher) {
}
} catch { }
}
if ($PythonOk) { break }
}
if (-not $PythonOk -and $HasPython) {
@ -1804,36 +1892,33 @@ if (Test-Path $OxcValidatorDir) {
Write-Host ""
substep "setting up Python environment..."
# Find Python -- skip Anaconda/Miniconda distributions.
# Conda-bundled CPython ships modified DLL search paths that break
# torch's c10.dll loading on Windows. Standalone CPython (python.org,
# winget, uv) does not have this issue.
# Uses Get-Command -All to look past conda entries that shadow a valid
# standalone Python further down PATH, and probes py.exe (the Python
# Launcher) which reliably finds python.org installs.
#
# NOTE: A venv created from conda Python inherits conda's base_prefix
# even though the venv path itself does not contain "conda". We check
# both the executable path AND sys.base_prefix to catch this case.
$CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)'
# Find Python -- skip Anaconda/Miniconda distributions ($CondaSkipPattern and
# Test-IsConda are defined above the 1g gate). Standalone CPython (python.org,
# winget, uv) does not have conda's torch c10.dll loading issue.
$PythonCmd = $null
# Helper: check if a Python executable is conda-based by inspecting
# both the path and sys.base_prefix (catches venvs created from conda).
function Test-IsConda {
param([string]$Exe)
if ($Exe -match $CondaSkipPattern) { return $true }
# 0. Reuse the interpreter install.ps1 already resolved and built the venv with
# (UNSLOTH_SETUP_PYTHON, or the existing venv python) before probing the
# system -- it is already validated as supported and non-conda.
if ($ReusedSetupPython) {
try {
$basePrefix = (& $Exe -c "import sys; print(sys.base_prefix)" 2>$null | Out-String).Trim()
if ($basePrefix -match $CondaSkipPattern) { return $true }
$out = & $ReusedSetupPython --version 2>&1 | Out-String
if ($out -match 'Python 3\.(\d+)') {
$pyMinor = [int]$Matches[1]
if ($pyMinor -ge 11 -and $pyMinor -le 13 -and -not (Test-IsConda $ReusedSetupPython)) {
$PythonCmd = $ReusedSetupPython
}
}
} catch { }
return $false
}
# 1. Try the Python Launcher (py.exe) first -- most reliable on Windows.
# py.exe is installed by python.org and resolves to standalone CPython.
$pyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue
if ($pyLauncher -and $pyLauncher.Source -notmatch $CondaSkipPattern) {
# Enumerate every launcher with -All (Windows PowerShell 5.1 returns only
# the first match without it) and search each for a supported, non-conda
# interpreter.
$PyLaunchersResolve = if ($PythonCmd) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) }
foreach ($pyLauncher in $PyLaunchersResolve) {
if ($pyLauncher.Source -match $CondaSkipPattern) { continue }
foreach ($minor in @("3.13", "3.12", "3.11")) {
try {
$out = & $pyLauncher.Source "-$minor" --version 2>&1 | Out-String
@ -1851,6 +1936,7 @@ if ($pyLauncher -and $pyLauncher.Source -notmatch $CondaSkipPattern) {
}
} catch { }
}
if ($PythonCmd) { break }
}
# 2. Fall back to scanning python3.x / python3 / python on PATH.

166
studio/setup.sh
View file

@ -155,9 +155,60 @@ _nvcc_meets_llama_minimum() {
echo "$_raw"
}
# Run a GPU probe under a 10s timeout when `timeout` is available so a wedged
# NVIDIA driver cannot hang setup; fall back to a bare call where it is not.
_setup_run_smi() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 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.
_setup_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_setup_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_setup_cvd_trim" ] || [ "$_setup_cvd_trim" = "-1" ]
}
# Returns 0 when an NVIDIA GPU is present and usable. Primary probe is
# `nvidia-smi -L` (timeout-bounded). Fallback is /proc/driver/nvidia/gpus,
# which the driver populates per GPU regardless of nvidia-smi state -- handles
# PATH gaps and driver init races. Mirrors install.sh _has_usable_nvidia_gpu
# (PR 6174) so setup routes the same way as the torch installer. A GPU hidden
# via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so the AMD probes still run
# and a mixed host steered to its AMD card keeps the ROCm route.
_setup_has_usable_nvidia_gpu() {
if _setup_cvd_hides_nvidia; then
return 1
fi
_setup_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_setup_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_setup_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_setup_nvsmi" ]; then
if _setup_run_smi "$_setup_nvsmi" -L 2>/dev/null \
| awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
_cuda_driver_max_version() {
command -v nvidia-smi >/dev/null 2>&1 || return 0
nvidia-smi 2>/dev/null \
_setup_run_smi nvidia-smi 2>/dev/null \
| sed -nE 's/.*CUDA( UMD)? Version:[[:space:]]*([0-9]+)\.([0-9]+).*/\2.\3/p' \
| head -1 || true
}
@ -815,25 +866,42 @@ _setup_amd_detected=false
_setup_nvidia_usable=false
_setup_gfx_all=""
_setup_mkt=""
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
_setup_amd_detected=true
_setup_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_setup_mkt=$(rocminfo 2>/dev/null | awk -F': ' \
'/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
elif command -v amd-smi >/dev/null 2>&1 && \
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
_setup_amd_detected=true
_setup_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
[ -z "$_setup_gfx_all" ] && \
_setup_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_setup_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
'/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
# NVIDIA priority: classify NVIDIA first and skip the AMD probes entirely on
# a usable-NVIDIA host (mirrors _has_rocm_gpu in install_python_stack.py).
# This also keeps a wedged rocminfo/amd-smi from hanging setup before the
# host is classified; the AMD probes themselves run under _setup_run_smi.
if _setup_has_usable_nvidia_gpu; then
_setup_nvidia_usable=true
fi
if [ "$_setup_nvidia_usable" != true ]; then
if command -v rocminfo >/dev/null 2>&1 && \
_setup_run_smi rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
_setup_amd_detected=true
_setup_gfx_all=$(_setup_run_smi rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_setup_mkt=$(_setup_run_smi rocminfo 2>/dev/null | awk -F': ' \
'/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
elif command -v amd-smi >/dev/null 2>&1 && \
_setup_run_smi amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
_setup_amd_detected=true
_setup_gfx_all=$(_setup_run_smi amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
[ -z "$_setup_gfx_all" ] && \
_setup_gfx_all=$(_setup_run_smi amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_setup_mkt=$(_setup_run_smi amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
'/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
elif [ -e /dev/kfd ] && \
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
# KFD sysfs fallback, AMD vendor_id 4098 only (mirrors install.sh
# _has_amd_rocm_gpu): covers AMD hosts where rocminfo/amd-smi are
# missing but the kernel exposes the GPU, so the source-build gate
# below does not drop them to a CPU llama.cpp build. No gfx arch is
# available from this path; name-based inference handles it.
_setup_amd_detected=true
fi
fi
if command -v nvidia-smi >/dev/null 2>&1 && \
nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
_setup_nvidia_usable=true
if [ "$_setup_nvidia_usable" = true ]; then
step "gpu" "NVIDIA GPU detected"
elif [ "$_setup_amd_detected" = true ]; then
_setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
@ -918,15 +986,15 @@ _HOST_MACHINE="$(uname -m 2>/dev/null || true)"
# use unslothai.
_LINUX_HAS_GPU=false
# Route to the fork only for a usable GPU. NVIDIA counts only when a device is
# actually enumerated (_setup_nvidia_usable, from the nvidia-smi -L probe above)
# AND not hidden via CUDA_VISIBLE_DEVICES=-1 -- mirroring install_llama_prebuilt.py's
# has_usable_nvidia. Mere nvidia-smi presence (CPU-only CUDA-toolkit containers,
# broken drivers) or a hidden GPU therefore takes the ggml-org CPU prebuilt
# instead of a slow source build. AMD is deliberately left on tooling presence,
# not usability: an unusable NVIDIA host has a good CPU prebuilt to fall back to,
# whereas tightening AMD would regress ROCm hosts exposing only hipconfig/hipinfo
# into an unnecessary CPU build.
if [ "$_setup_nvidia_usable" = true ] && [ "${CUDA_VISIBLE_DEVICES:-}" != "-1" ]; then
# actually enumerated and not hidden via CUDA_VISIBLE_DEVICES=""/-1
# (_setup_nvidia_usable, from _setup_has_usable_nvidia_gpu above) -- mirroring
# install_llama_prebuilt.py's has_usable_nvidia. Mere nvidia-smi presence
# (CPU-only CUDA-toolkit containers, broken drivers) or a hidden GPU therefore
# takes the ggml-org CPU prebuilt instead of a slow source build. AMD is
# deliberately left on tooling presence, not usability: an unusable NVIDIA host
# has a good CPU prebuilt to fall back to, whereas tightening AMD would regress
# ROCm hosts exposing only hipconfig/hipinfo into an unnecessary CPU build.
if [ "$_setup_nvidia_usable" = true ]; then
_LINUX_HAS_GPU=true
else
for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do
@ -1271,23 +1339,35 @@ else
GPU_BACKEND=""
NVCC_PATH=""
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
GPU_BACKEND="cuda"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
GPU_BACKEND="cuda"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
GPU_BACKEND="cuda"
# Gate the CUDA toolkit search on an actually-usable NVIDIA GPU
# (_setup_nvidia_usable, computed in the GPU summary block above;
# already false when hidden via CUDA_VISIBLE_DEVICES=""/-1).
# A CUDA toolkit alone (CPU-only build container, leftover packages)
# is not proof of a GPU: building with -DGGML_CUDA=ON there yields a
# binary that fails at runtime, so fall through to the CPU build.
if [ "$_setup_nvidia_usable" = true ]; then
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
GPU_BACKEND="cuda"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
GPU_BACKEND="cuda"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
GPU_BACKEND="cuda"
fi
fi
# Check for ROCm (AMD) only if CUDA was not already selected
# Check for ROCm (AMD) only if CUDA was not already selected, and
# only when an AMD GPU was actually detected (_setup_amd_detected).
# hipcc presence alone (HIP SDK, no GPU) must not select a HIP build.
# NVIDIA-usable hosts never build HIP (defense in depth: the AMD
# probes above are already skipped when NVIDIA is usable).
ROCM_HIPCC=""
if [ -z "$GPU_BACKEND" ]; then
if [ -z "$GPU_BACKEND" ] && [ "$_setup_nvidia_usable" != true ] && [ "$_setup_amd_detected" = true ]; then
if command -v hipcc &>/dev/null; then
ROCM_HIPCC="$(command -v hipcc)"
GPU_BACKEND="rocm"
@ -1349,7 +1429,7 @@ else
CUDA_ARCHS=""
if command -v nvidia-smi &>/dev/null; then
_raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
_raw_caps=$(_setup_run_smi nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
while IFS= read -r _cap; do
_cap=$(echo "$_cap" | tr -d '[:space:]')
if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
@ -1455,7 +1535,7 @@ else
CMAKE_ARGS="$CMAKE_ARGS -DGPU_TARGETS=${GPU_TARGETS}"
_BUILD_DESC="building (ROCm, ${GPU_TARGETS//;/+})"
fi
elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
elif [ -d /usr/local/cuda ] || _setup_run_smi nvidia-smi &>/dev/null; then
_BUILD_DESC="building (CPU, CUDA driver found but nvcc missing)"
elif [ -d /opt/rocm ] || command -v rocm-smi &>/dev/null; then
_BUILD_DESC="building (CPU, ROCm driver found but hipcc missing)"

51
tests/sh/test_get_torch_index_url.sh
View file

@ -13,6 +13,10 @@ FAIL=0
_FUNC_FILE=$(mktemp)
_FAKE_SMI_DIR=$(mktemp -d)
{
sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_cvd_hides_nvidia()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_has_amd_rocm_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
@ -107,7 +111,7 @@ MOCK
# Build a minimal tools directory with symlinks to essential commands
# (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi.
_TOOLS_DIR=$(mktemp -d)
for _cmd in uname grep sed head sh bash cat awk printf; do
for _cmd in uname grep sed head sh bash cat awk printf tr; do
_real=$(command -v "$_cmd" 2>/dev/null || true)
[ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd"
done
@ -116,12 +120,19 @@ done
# $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test
run_func() {
_mock_dir="$1"
# Default: strip CUDA_VISIBLE_DEVICES so the host environment cannot leak
# in; a second argument sets it explicitly (hidden-GPU scenarios).
if [ "$#" -ge 2 ]; then
_cvd_setup="export CUDA_VISIBLE_DEVICES='$2'"
else
_cvd_setup="unset CUDA_VISIBLE_DEVICES"
fi
if [ "$_mock_dir" = "none" ]; then
# Minimal PATH with only basic tools, no nvidia-smi anywhere
PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
PATH="$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
else
# Put mock nvidia-smi dir first, then basic tools
PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
PATH="$_mock_dir:$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
fi
}
@ -332,6 +343,40 @@ _result=$(run_func "$_dir")
assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 34) CUDA_VISIBLE_DEVICES="" hides the NVIDIA GPU -> cpu (no AMD present)
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" "")
assert_eq "CVD='' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 35) CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU -> cpu (no AMD present)
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" "-1")
assert_eq "CVD=-1 hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 36) Mixed AMD+NVIDIA host with NVIDIA hidden -> ROCm route is restored
_cuda_dir=$(make_mock_smi "12.6")
_amd_dir=$(make_mock_amd_smi "6.4")
_combined_dir=$(mktemp -d)
ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi"
ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi"
_result=$(run_func "$_combined_dir" "-1")
assert_eq "CUDA+ROCm with CVD=-1 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result"
rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir"
# 37) CUDA_VISIBLE_DEVICES=0 (a visible device) must NOT hide the GPU
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" "0")
assert_eq "CVD=0 keeps NVIDIA -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 38) Whitespace-padded "-1" still hides the GPU
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" " -1 ")
assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"

248
tests/studio/install/test_cuda_repair.py Normal file
View file

@ -0,0 +1,248 @@
"""Tests for CUDA torch repair on poisoned NVIDIA venvs.
Verifies _ensure_cuda_torch (studio/install_python_stack.py) reinstalls CUDA
torch when a venv on an NVIDIA host carries a ROCm torch build (the pre-fix KFD
gpu_id false positive), without touching healthy CUDA, deliberate CPU wheels,
ROCm hosts, macOS, or Windows. All tests use mocks -- no GPU required.
"""
import importlib.util
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
# ── Load module under test (mirrors test_rocm_support.py) ────────────────────
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH)
assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
sys.modules[_STACK_SPEC.name] = stack_mod
_STACK_SPEC.loader.exec_module(stack_mod)
_ensure_cuda_torch = stack_mod._ensure_cuda_torch
_detect_cuda_torch_index_url = stack_mod._detect_cuda_torch_index_url
# ── Helpers ──────────────────────────────────────────────────────────────────
def _make_run(
torch_state = "hip",
cuda_version = "12.8",
torch_rc = 0,
smi_rc = 0,
):
"""Build a subprocess.run side_effect.
The torch-classify probe runs sys.executable and reads bytes stdout; the
nvidia-smi version probe runs the smi path with text=True. Distinguish by
the executable.
"""
def _run(cmd, *args, **kwargs):
result = MagicMock()
exe = str(cmd[0]) if cmd else ""
if exe == sys.executable:
result.returncode = torch_rc
result.stdout = (torch_state + "\n").encode()
return result
# nvidia-smi version probe (text = True)
result.returncode = smi_rc
out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n"
result.stdout = out if kwargs.get("text") else out.encode()
return result
return _run
def _run_cuda_repair(
*,
backend = "",
nvidia = True,
torch_state = "hip",
cuda_version = "12.8",
torch_rc = 0,
smi_rc = 0,
is_macos = False,
is_windows = False,
no_torch = False,
rocm_marker = False,
smi_path = "/usr/bin/nvidia-smi",
cvd = None,
):
"""Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
cvd controls CUDA_VISIBLE_DEVICES: None removes it from the environment
(the host machine may export one), any string sets it explicitly.
"""
env = {}
if rocm_marker:
env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
if cvd is not None:
env["CUDA_VISIBLE_DEVICES"] = cvd
def _which(name, *a, **k):
if name == "nvidia-smi":
return smi_path
return None
with (
patch.object(stack_mod, "_TORCH_BACKEND", backend),
patch.object(stack_mod, "IS_MACOS", is_macos),
patch.object(stack_mod, "IS_WINDOWS", is_windows),
patch.object(stack_mod, "NO_TORCH", no_torch),
patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia),
patch.object(stack_mod.shutil, "which", side_effect = _which),
patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)),
patch.object(stack_mod, "pip_install") as mock_pip,
patch.object(
stack_mod.subprocess,
"run",
side_effect = _make_run(torch_state, cuda_version, torch_rc, smi_rc),
),
patch.dict(stack_mod.os.environ, env, clear = False),
):
if not rocm_marker:
stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None)
if cvd is None:
stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
_ensure_cuda_torch()
return mock_pip
def _index_url(mock_pip) -> str:
"""Return the --index-url value from the recorded pip_install call."""
args = [str(a) for a in mock_pip.call_args.args]
return args[args.index("--index-url") + 1]
# ── Repair fires only on the poisoning signature ─────────────────────────────
class TestCudaRepairFires:
def test_hip_build_on_nvidia_triggers_repair(self):
mock_pip = _run_cuda_repair(torch_state = "hip", cuda_version = "12.8")
assert mock_pip.call_count == 1
call_args = [str(a) for a in mock_pip.call_args.args]
assert "--force-reinstall" in call_args
assert "--no-cache-dir" in call_args
assert "cu128" in _index_url(mock_pip)
assert mock_pip.call_args.kwargs["constrain"] is False
def test_rocm_in_version_string_triggers_repair(self):
# AMD SDK / Radeon wheels may not set torch.version.hip but encode
# rocm in __version__; the probe prints "hip" for both.
mock_pip = _run_cuda_repair(torch_state = "hip")
assert mock_pip.call_count == 1
# ── No-op cases ──────────────────────────────────────────────────────────────
class TestCudaRepairSkips:
def test_healthy_cuda_torch_no_repair(self):
mock_pip = _run_cuda_repair(torch_state = "cuda")
mock_pip.assert_not_called()
def test_deliberate_cpu_wheel_no_repair(self):
mock_pip = _run_cuda_repair(torch_state = "cpu")
mock_pip.assert_not_called()
def test_backend_rocm_skips(self):
mock_pip = _run_cuda_repair(backend = "rocm", torch_state = "hip")
mock_pip.assert_not_called()
def test_backend_cpu_skips(self):
mock_pip = _run_cuda_repair(backend = "cpu", torch_state = "hip")
mock_pip.assert_not_called()
def test_unknown_backend_skips(self):
mock_pip = _run_cuda_repair(backend = "auto", torch_state = "hip")
mock_pip.assert_not_called()
def test_no_nvidia_gpu_skips(self):
mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip")
mock_pip.assert_not_called()
def test_torch_missing_skips(self):
# Non-zero probe exit = torch missing / un-importable.
mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1)
mock_pip.assert_not_called()
def test_macos_skips(self):
mock_pip = _run_cuda_repair(is_macos = True, torch_state = "hip")
mock_pip.assert_not_called()
def test_windows_skips(self):
mock_pip = _run_cuda_repair(is_windows = True, torch_state = "hip")
mock_pip.assert_not_called()
def test_no_torch_mode_skips(self):
mock_pip = _run_cuda_repair(no_torch = True, torch_state = "hip")
mock_pip.assert_not_called()
def test_rocm_install_marker_skips(self):
mock_pip = _run_cuda_repair(rocm_marker = True, torch_state = "hip")
mock_pip.assert_not_called()
def test_cvd_minus_one_skips(self):
# CUDA_VISIBLE_DEVICES=-1 deliberately hides the NVIDIA GPU (mixed
# AMD+NVIDIA host running ROCm torch on the AMD card).
mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip")
mock_pip.assert_not_called()
def test_cvd_empty_skips(self):
mock_pip = _run_cuda_repair(cvd = "", torch_state = "hip")
mock_pip.assert_not_called()
def test_cvd_explicit_device_still_repairs(self):
mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip")
assert mock_pip.call_count == 1
# ── CUDA index ladder ────────────────────────────────────────────────────────
class TestCudaIndexResolution:
def test_cuda_128_selects_cu128(self):
assert "cu128" in _index_url(_run_cuda_repair(cuda_version = "12.8"))
def test_cuda_130_selects_cu130(self):
assert "cu130" in _index_url(_run_cuda_repair(cuda_version = "13.0"))
def test_cuda_126_selects_cu126(self):
assert "cu126" in _index_url(_run_cuda_repair(cuda_version = "12.6"))
def test_cuda_124_selects_cu124(self):
assert "cu124" in _index_url(_run_cuda_repair(cuda_version = "12.4"))
def test_cuda_118_selects_cu118(self):
assert "cu118" in _index_url(_run_cuda_repair(cuda_version = "11.8"))
def test_unreadable_version_defaults_cu126(self):
# nvidia-smi runs but prints no CUDA version line (or fails).
mock_pip = _run_cuda_repair(cuda_version = "", smi_rc = 1)
assert "cu126" in _index_url(mock_pip)
def test_proc_fallback_no_smi_defaults_cu126(self):
# NVIDIA usable via /proc fallback, nvidia-smi absent entirely.
mock_pip = _run_cuda_repair(smi_path = None)
assert "cu126" in _index_url(mock_pip)
def test_detect_index_url_uses_pytorch_base(self):
with (
patch.object(stack_mod.shutil, "which", return_value = None),
patch.object(stack_mod.os.path, "isfile", return_value = False),
):
url = _detect_cuda_torch_index_url()
assert url == f"{stack_mod._PYTORCH_WHL_BASE}/cu126"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-q"]))

480
tests/studio/install/test_gpu_detection_followups.py Normal file
View file

@ -0,0 +1,480 @@
"""Tests for the GPU-detection follow-ups to PR 6174.
PR 6174 made NVIDIA take precedence and added a /proc/driver/nvidia/gpus
fallback in install.sh and studio/install_python_stack.py. These tests cover the
same hardening ported to the llama.cpp prebuilt installer
(studio/install_llama_prebuilt.py) and the Studio shell setup (studio/setup.sh):
* detect_host() recognises NVIDIA via /proc/driver/nvidia/gpus when nvidia-smi
is unavailable, and skips ROCm probing when NVIDIA is usable.
* setup.sh routes through a timeout-bounded NVIDIA probe with a /proc fallback
and only selects a CUDA/ROCm source build when the matching GPU is detected.
All tests use mocks or source-level assertions -- no GPU, network, or real
nvidia-smi/rocminfo invocation.
"""
import importlib.util
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
# Load studio/install_llama_prebuilt.py the same way the sibling suite does.
_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
_SPEC = importlib.util.spec_from_file_location(
"studio_install_llama_prebuilt_followups", _MODULE_PATH
)
assert _SPEC is not None and _SPEC.loader is not None
prebuilt_mod = importlib.util.module_from_spec(_SPEC)
sys.modules[_SPEC.name] = prebuilt_mod
_SPEC.loader.exec_module(prebuilt_mod)
detect_host = prebuilt_mod.detect_host
_apply_host_overrides = prebuilt_mod._apply_host_overrides
SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
def _make_run_capture(rocminfo_stdout: str = ""):
"""Return a fake run_capture: rocminfo reports rocminfo_stdout, everything
else (nvidia-smi, amd-smi) returns empty so only the patched probes matter."""
def _run_capture(cmd, *args, **kwargs):
exe = str(cmd[0]) if cmd else ""
result = MagicMock()
if exe.endswith("rocminfo"):
result.returncode = 0
result.stdout = rocminfo_stdout
else:
result.returncode = 1
result.stdout = ""
result.stderr = ""
return result
return _run_capture
def _run_detect_host(
*,
machine: str = "x86_64",
system: str = "Linux",
which_map: dict | None = None,
proc_dir_entries: list | None = None,
rocminfo_stdout: str = "",
env: dict | None = None,
):
"""Drive detect_host() against a fully synthetic host."""
which_map = which_map or {}
proc_dir_entries = proc_dir_entries if proc_dir_entries is not None else []
real_isdir = prebuilt_mod.os.path.isdir
real_listdir = prebuilt_mod.os.listdir
proc_path = "/proc/driver/nvidia/gpus"
def fake_isdir(p):
if str(p) == proc_path:
return bool(proc_dir_entries)
return real_isdir(p)
def fake_listdir(p):
if str(p) == proc_path:
if not proc_dir_entries:
raise OSError("no such dir")
return list(proc_dir_entries)
return real_listdir(p)
patches = [
patch.object(prebuilt_mod.platform, "system", return_value = system),
patch.object(prebuilt_mod.platform, "machine", return_value = machine),
patch.object(prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "")),
patch.object(prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n)),
patch.object(prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout)),
patch.object(prebuilt_mod.os.path, "isdir", side_effect = fake_isdir),
patch.object(prebuilt_mod.os, "listdir", side_effect = fake_listdir),
patch.object(prebuilt_mod.os, "access", return_value = False),
patch.dict(prebuilt_mod.os.environ, env or {}, clear = False),
]
for p in patches:
p.start()
try:
# Ensure CUDA_VISIBLE_DEVICES does not leak in from the test host unless
# the scenario sets it explicitly.
if env is None or "CUDA_VISIBLE_DEVICES" not in env:
prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
return detect_host()
finally:
for p in patches:
p.stop()
# ── install_llama_prebuilt.detect_host(): /proc NVIDIA fallback ──────────────
class TestDetectHostProcFallback:
def test_proc_fallback_marks_physical_nvidia_when_smi_absent(self):
"""No nvidia-smi, but /proc/driver/nvidia/gpus is populated -> NVIDIA."""
host = _run_detect_host(
which_map = {}, # nvidia-smi resolves to None
proc_dir_entries = ["0000:01:00.0"],
)
assert host.has_physical_nvidia is True
def test_proc_fallback_has_usable_nvidia_when_devices_visible(self):
"""Default CUDA_VISIBLE_DEVICES (unset) -> visible tokens non-empty -> usable."""
host = _run_detect_host(
which_map = {},
proc_dir_entries = ["0000:01:00.0"],
)
assert host.has_usable_nvidia is True
def test_proc_fallback_not_usable_when_devices_hidden(self):
"""CUDA_VISIBLE_DEVICES='' hides all GPUs -> physical yes, usable no."""
host = _run_detect_host(
which_map = {},
proc_dir_entries = ["0000:01:00.0"],
env = {"CUDA_VISIBLE_DEVICES": ""},
)
assert host.has_physical_nvidia is True
assert host.has_usable_nvidia is False
def test_empty_proc_dir_does_not_mark_nvidia(self):
"""A driver dir that exists but is empty must not assert a GPU."""
host = _run_detect_host(which_map = {}, proc_dir_entries = [])
assert host.has_physical_nvidia is False
def test_proc_fallback_is_linux_only(self):
"""The /proc fallback must not run on Windows (path is Linux-only)."""
host = _run_detect_host(
system = "Windows",
machine = "amd64",
which_map = {},
proc_dir_entries = ["0000:01:00.0"],
)
assert host.has_physical_nvidia is False
# ── install_llama_prebuilt.detect_host(): NVIDIA precedence over ROCm ────────
class TestDetectHostNvidiaPrecedence:
def test_rocm_probe_skipped_when_proc_nvidia_present(self):
"""rocminfo reports gfx1100, but a proc-detected NVIDIA GPU wins."""
host = _run_detect_host(
which_map = {"rocminfo": "/usr/bin/rocminfo"},
proc_dir_entries = ["0000:01:00.0"],
rocminfo_stdout = " Name: gfx1100\n",
)
assert host.has_usable_nvidia is True
assert host.has_rocm is False
def test_rocm_detected_when_no_nvidia(self):
"""With no NVIDIA signal at all, rocminfo gfx1100 -> has_rocm True."""
host = _run_detect_host(
which_map = {"rocminfo": "/usr/bin/rocminfo"},
proc_dir_entries = [],
rocminfo_stdout = " Name: gfx1100\n",
)
assert host.has_usable_nvidia is False
assert host.has_rocm is True
# ── _apply_host_overrides: forwarded --rocm-gfx / --has-rocm still win ───────
class TestOverridesStillWin:
def test_forwarded_gfx_forces_rocm_on_non_nvidia_host(self):
host = _run_detect_host(which_map = {}, proc_dir_entries = [])
assert host.has_rocm is False
overridden = _apply_host_overrides(host, override_rocm_gfx = "gfx1100")
assert overridden.has_rocm is True
assert overridden.rocm_gfx_target == "gfx1100"
def test_override_has_rocm_forces_rocm(self):
host = _run_detect_host(which_map = {}, proc_dir_entries = [])
overridden = _apply_host_overrides(host, override_has_rocm = True)
assert overridden.has_rocm is True
def test_force_cpu_drops_nvidia_attributes(self):
host = _run_detect_host(which_map = {}, proc_dir_entries = ["0000:01:00.0"])
assert host.has_usable_nvidia is True
overridden = _apply_host_overrides(host, force_cpu = True)
assert overridden.has_usable_nvidia is False
assert overridden.has_physical_nvidia is False
assert overridden.has_rocm is False
# ── setup.sh source-level guarantees ────────────────────────────────────────
class TestSetupShHardening:
@pytest.fixture(scope = "class")
def setup_src(self) -> str:
return SETUP_SH.read_text(encoding = "utf-8")
def test_has_usable_nvidia_helper_exists(self, setup_src):
assert "_setup_has_usable_nvidia_gpu()" in setup_src
def test_helper_uses_proc_fallback(self, setup_src):
start = setup_src.find("_setup_has_usable_nvidia_gpu()")
end = setup_src.find("\n}", start)
body = setup_src[start:end]
assert (
"/proc/driver/nvidia/gpus" in body
), "_setup_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus"
def test_gpu_summary_uses_helper(self, setup_src):
assert "if _setup_has_usable_nvidia_gpu; then" in setup_src
def test_timeout_wrapper_exists(self, setup_src):
start = setup_src.find("_setup_run_smi()")
assert start >= 0, "_setup_run_smi timeout wrapper must exist"
end = setup_src.find("\n}", start)
body = setup_src[start:end]
assert "timeout 10" in body
assert "command -v timeout" in body
def test_cuda_source_build_gated_on_usable_nvidia(self, setup_src):
"""The nvcc source-build search must be gated on _setup_nvidia_usable.
The hidden-GPU policy (CUDA_VISIBLE_DEVICES=""/-1) lives inside
_setup_has_usable_nvidia_gpu, so the gate itself only needs the flag.
"""
anchor = setup_src.find('NVCC_PATH=""\n')
assert anchor >= 0
window = setup_src[anchor : anchor + 700]
assert (
'if [ "$_setup_nvidia_usable" = true ]' in window
), "CUDA toolkit search must require a usable NVIDIA GPU, not just nvcc"
def test_nvidia_helper_honours_hidden_cvd(self, setup_src):
"""_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so
CUDA_VISIBLE_DEVICES=""/-1 suppresses NVIDIA before the AMD probes are
gated (mixed hosts steered to the AMD card keep the ROCm route)."""
assert "_setup_cvd_hides_nvidia()" in setup_src
start = setup_src.find("_setup_has_usable_nvidia_gpu() {")
end = setup_src.find("\n}", start)
body = setup_src[start:end]
assert "_setup_cvd_hides_nvidia" in body
def test_rocm_source_build_gated_on_amd_detected(self, setup_src):
"""The hipcc source-build search must be gated on _setup_amd_detected."""
anchor = setup_src.find('ROCM_HIPCC=""')
assert anchor >= 0
window = setup_src[anchor : anchor + 400]
assert (
'[ "$_setup_amd_detected" = true ]' in window
), "ROCm toolkit search must require a detected AMD GPU, not just hipcc"
def test_compute_cap_probe_timeout_wrapped(self, setup_src):
assert "_setup_run_smi nvidia-smi --query-gpu=compute_cap" in setup_src
def test_driver_version_probe_timeout_wrapped(self, setup_src):
start = setup_src.find("_cuda_driver_max_version()")
end = setup_src.find("\n}", start)
body = setup_src[start:end]
assert "_setup_run_smi nvidia-smi" in body
# TEST: install.sh -- UNSLOTH_TORCH_BACKEND classified on the final path segment
class TestBackendExportLeafClassification:
"""A custom UNSLOTH_PYTORCH_MIRROR whose base path contains "rocm" or
"gfx" must not mislabel a cu*/cpu index as ROCm; classification uses the
final path segment of TORCH_INDEX_URL only."""
@pytest.fixture(scope = "class")
def install_src(self) -> str:
return (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
def test_export_block_uses_leaf(self, install_src):
anchor = install_src.find("_torch_index_leaf=")
assert anchor >= 0, "backend export must classify on the final path segment"
window = install_src[anchor : anchor + 500]
assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window
assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window
assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window
def test_leaf_classification_behaviour(self, tmp_path):
import subprocess as sp
script = tmp_path / "leaf.sh"
src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
anchor = src.find("_torch_index_leaf=")
block = src[anchor : src.find("esac", anchor) + 4]
# Drive the extracted block with adversarial mirror URLs.
script.write_text(
"#!/bin/sh\n"
'TORCH_INDEX_URL="$1"\n' + block + "\n"
'printf "%s" "$UNSLOTH_TORCH_BACKEND"\n'
)
cases = {
"https://download.pytorch.org/whl/cu128": "cuda",
"https://download.pytorch.org/whl/cpu": "cpu",
"https://download.pytorch.org/whl/rocm6.4": "rocm",
"https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/": "rocm",
"https://repo.amd.com/rocm/whl/gfx1151/": "rocm",
"https://mirror.local/rocm-cache/cu128": "cuda",
"https://mirror.local/gfx-cache/cpu": "cpu",
}
for url, expected in cases.items():
out = sp.run(
["sh", str(script), url], capture_output = True, text = True, timeout = 30
).stdout.strip()
assert out == expected, f"{url} classified as {out!r}, expected {expected!r}"
# TEST: CUDA_VISIBLE_DEVICES=""/-1 hides NVIDIA in every usable-GPU helper
_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
_STACK_SPEC = importlib.util.spec_from_file_location(
"studio_install_python_stack_followups", _STACK_PATH
)
assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
sys.modules[_STACK_SPEC.name] = stack_mod
_STACK_SPEC.loader.exec_module(stack_mod)
def _stack_nvidia_usable(cvd):
"""Drive install_python_stack._has_usable_nvidia_gpu with a mocked
nvidia-smi that always reports a GPU; cvd = None removes the env var."""
def fake_run(cmd, *args, **kwargs):
result = MagicMock()
result.returncode = 0
result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n"
return result
env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd}
with (
patch.object(
stack_mod.shutil,
"which",
side_effect = lambda n: "/usr/bin/nvidia-smi" if n == "nvidia-smi" else None,
),
patch.object(stack_mod.subprocess, "run", side_effect = fake_run),
patch.dict(stack_mod.os.environ, env, clear = False),
):
if cvd is None:
stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
return stack_mod._has_usable_nvidia_gpu()
class TestHiddenCvdNotUsable:
"""CUDA_VISIBLE_DEVICES set to "" or "-1" deliberately hides every NVIDIA
device (mixed AMD+NVIDIA hosts steering work to the AMD card). All three
_has_usable_nvidia_gpu implementations (install_python_stack.py, install.sh,
setup.sh) must report the GPU as not usable so the AMD/CPU routes run,
matching install_llama_prebuilt.py's has_usable_nvidia."""
def test_python_unset_cvd_is_usable(self):
assert _stack_nvidia_usable(None) is True
def test_python_empty_cvd_not_usable(self):
assert _stack_nvidia_usable("") is False
def test_python_minus_one_not_usable(self):
assert _stack_nvidia_usable("-1") is False
def test_python_padded_minus_one_not_usable(self):
assert _stack_nvidia_usable(" -1 ") is False
def test_python_explicit_device_is_usable(self):
assert _stack_nvidia_usable("0") is True
def test_python_device_list_is_usable(self):
assert _stack_nvidia_usable("0,1") is True
def test_hidden_nvidia_restores_rocm_detection(self):
"""Mixed host, NVIDIA hidden via CVD=-1, rocminfo reports gfx1100:
_has_rocm_gpu must proceed past the NVIDIA guard and return True
(before this fix the guard ignored CVD and blocked ROCm)."""
def fake_run(cmd, *args, **kwargs):
result = MagicMock()
result.returncode = 0
exe = str(cmd[0])
if exe.endswith("rocminfo"):
result.stdout = " Name: gfx1100\n"
else:
result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n"
return result
which_map = {
"rocminfo": "/usr/bin/rocminfo",
"nvidia-smi": "/usr/bin/nvidia-smi",
}
with (
patch.object(stack_mod.shutil, "which", side_effect = which_map.get),
patch.object(stack_mod.subprocess, "run", side_effect = fake_run),
patch.dict(stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False),
):
assert stack_mod._has_rocm_gpu() is True
@staticmethod
def _run_sh_helper(tmp_path, src: str, fn_names: list, cvd):
"""Extract shell functions, run the usable-GPU one against a fake
nvidia-smi, and return "usable"/"not_usable"."""
import os as _os
import subprocess as sp
blocks = []
for name in fn_names:
start = src.find(f"{name}() {{")
assert start >= 0, f"{name} missing"
end = src.find("\n}", start) + 2
blocks.append(src[start:end])
fake_bin = tmp_path / "bin"
fake_bin.mkdir(exist_ok = True)
smi = fake_bin / "nvidia-smi"
smi.write_text("#!/bin/sh\necho 'GPU 0: NVIDIA Fake (UUID: GPU-x)'\n")
smi.chmod(0o755)
script = tmp_path / "probe.sh"
script.write_text(
"#!/bin/sh\n" + "\n".join(blocks) + "\n"
f"if {fn_names[-1]}; then echo usable; else echo not_usable; fi\n"
)
env = dict(_os.environ)
env["PATH"] = f"{fake_bin}:{env['PATH']}"
if cvd is None:
env.pop("CUDA_VISIBLE_DEVICES", None)
else:
env["CUDA_VISIBLE_DEVICES"] = cvd
return sp.run(
["sh", str(script)], capture_output = True, text = True, timeout = 30, env = env
).stdout.strip()
@pytest.mark.parametrize(
"cvd, expected",
[(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")],
)
def test_install_sh_helper_cvd(self, tmp_path, cvd, expected):
src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
out = self._run_sh_helper(
tmp_path,
src,
["_run_bounded", "_cvd_hides_nvidia", "_has_usable_nvidia_gpu"],
cvd,
)
assert out == expected
@pytest.mark.parametrize(
"cvd, expected",
[(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")],
)
def test_setup_sh_helper_cvd(self, tmp_path, cvd, expected):
src = SETUP_SH.read_text(encoding = "utf-8")
out = self._run_sh_helper(
tmp_path,
src,
["_setup_run_smi", "_setup_cvd_hides_nvidia", "_setup_has_usable_nvidia_gpu"],
cvd,
)
assert out == expected

114
tests/studio/install/test_hf_auth.py Normal file
View file

@ -0,0 +1,114 @@
"""Tests for Hugging Face auth on the llama.cpp prebuilt installer's fetches.
Anonymous huggingface.co downloads (tiny GGUF validation model) share a
per-IP rate limit that CI fleets exhaust (HTTP 429), forcing the prebuilt
path into a source build. auth_headers now sends HF_TOKEN to huggingface.co
hosts, and a redirect handler strips Authorization when the download is
redirected to a different host (CDN signed URLs). All tests are offline.
"""
import importlib.util
import sys
import urllib.request
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
_SPEC = importlib.util.spec_from_file_location(
"studio_install_llama_prebuilt_hf_auth", _MODULE_PATH
)
assert _SPEC is not None and _SPEC.loader is not None
mod = importlib.util.module_from_spec(_SPEC)
sys.modules[_SPEC.name] = mod
_SPEC.loader.exec_module(mod)
_TOKEN_VARS = ("GH_TOKEN", "GITHUB_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN")
HF_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases"
def _headers(url, env):
"""auth_headers under a fully controlled token environment."""
with patch.dict(mod.os.environ, env, clear = False):
for var in _TOKEN_VARS:
if var not in env:
mod.os.environ.pop(var, None)
return mod.auth_headers(url)
class TestAuthHeaderRouting:
def test_hf_token_sent_to_huggingface(self):
headers = _headers(HF_URL, {"HF_TOKEN": "hf_x"})
assert headers.get("Authorization") == "Bearer hf_x"
def test_hub_token_fallback(self):
headers = _headers(HF_URL, {"HUGGING_FACE_HUB_TOKEN": "hf_y"})
assert headers.get("Authorization") == "Bearer hf_y"
def test_hf_token_not_sent_to_github(self):
headers = _headers(GH_URL, {"HF_TOKEN": "hf_x"})
assert "Authorization" not in headers
def test_hf_token_not_sent_to_other_hosts(self):
headers = _headers("https://cdn-lfs.huggingface.co/x", {"HF_TOKEN": "hf_x"})
assert "Authorization" not in headers
def test_gh_token_not_sent_to_huggingface(self):
headers = _headers(HF_URL, {"GH_TOKEN": "gh_x"})
assert "Authorization" not in headers
def test_gh_token_still_wins_on_github(self):
headers = _headers(GH_URL, {"GH_TOKEN": "gh_x", "HF_TOKEN": "hf_x"})
assert headers.get("Authorization") == "Bearer gh_x"
def test_no_tokens_no_auth(self):
assert "Authorization" not in _headers(HF_URL, {})
def test_validation_model_url_is_hf(self):
assert mod.should_send_hf_auth(mod.TEST_MODEL_URL) is True
class TestCrossHostRedirectStripsAuth:
def _redirect(self, newurl):
req = urllib.request.Request(HF_URL, headers = {"Authorization": "Bearer hf_x"})
handler = mod._CrossHostAuthStrippingRedirectHandler()
return handler.redirect_request(req, None, 302, "Found", {}, newurl)
def test_cross_host_redirect_drops_authorization(self):
new_request = self._redirect("https://cdn-lfs.huggingface.co/signed/blob")
assert new_request is not None
assert "Authorization" not in new_request.headers
assert "Authorization" not in new_request.unredirected_hdrs
def test_same_host_redirect_keeps_authorization(self):
new_request = self._redirect("https://huggingface.co/elsewhere/blob")
assert new_request is not None
assert new_request.headers.get("Authorization") == "Bearer hf_x"
class TestDownloadBytesWiring:
def test_download_bytes_sends_hf_auth(self):
response = MagicMock()
response.__enter__ = lambda s: s
response.__exit__ = lambda s, *a: False
response.headers.get.return_value = None
response.read.side_effect = [b"data", b""]
with (
patch.object(mod._URL_OPENER, "open", return_value = response) as opened,
patch.dict(mod.os.environ, {"HF_TOKEN": "hf_x"}, clear = False),
):
for var in ("GH_TOKEN", "GITHUB_TOKEN"):
mod.os.environ.pop(var, None)
data = mod.download_bytes(HF_URL)
assert data == b"data"
request = opened.call_args.args[0]
assert request.headers.get("Authorization") == "Bearer hf_x"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-q"]))

148
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -1,3 +1,4 @@
import errno
import importlib.util
import io
import json
@ -28,6 +29,7 @@ ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums
hydrate_source_tree = INSTALL_LLAMA_PREBUILT.hydrate_source_tree
validate_prebuilt_choice = INSTALL_LLAMA_PREBUILT.validate_prebuilt_choice
activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree
activate_staged_dir = INSTALL_LLAMA_PREBUILT.activate_staged_dir
create_install_staging_dir = INSTALL_LLAMA_PREBUILT.create_install_staging_dir
sha256_file = INSTALL_LLAMA_PREBUILT.sha256_file
source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name
@ -206,6 +208,110 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents(
assert not (install_dir / f"llama.cpp-{upstream_tag}").exists()
def test_release_asset_download_url():
fn = INSTALL_LLAMA_PREBUILT.release_asset_download_url
assert fn(
"unslothai/llama.cpp", "b9000-mix-abc1234", "llama.cpp-source-commit-deadbeef.tar.gz"
) == (
"https://github.com/unslothai/llama.cpp/releases/download/"
"b9000-mix-abc1234/llama.cpp-source-commit-deadbeef.tar.gz"
)
# Any missing component -> None (no asset url, caller falls back to codeload).
assert fn(None, "b9000", "x.tar.gz") is None
assert fn("unslothai/llama.cpp", None, "x.tar.gz") is None
assert fn("unslothai/llama.cpp", "b9000", None) is None
def _mk_source_tarball(path: Path, tag: str) -> None:
with tarfile.open(path, "w:gz") as archive:
add_bytes_to_tar(
archive, f"llama.cpp-{tag}/CMakeLists.txt", b"cmake_minimum_required(VERSION 3.14)\n"
)
add_bytes_to_tar(
archive,
f"llama.cpp-{tag}/convert_hf_to_gguf.py",
b"#!/usr/bin/env python3\nimport gguf\n",
)
add_bytes_to_tar(archive, f"llama.cpp-{tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n")
def test_hydrate_source_tree_prefers_release_asset_for_mix(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
# A mix build's merge commit is in no repo, so the codeload/archive URLs 404.
# hydrate must fetch the release asset and never touch codeload.
commit = "a" * 40
archive_path = tmp_path / "merged-source.tar.gz"
_mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}")
asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url(
"unslothai/llama.cpp", "b9000-mix-abc1234", f"llama.cpp-source-commit-{commit}.tar.gz"
)
codeload_urls = set(
INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit)
)
seen = []
def fake_download_file(url: str, destination: Path) -> None:
seen.append(url)
if url in codeload_urls:
raise AssertionError("codeload was hit even though the release asset was available")
assert url == asset_url
destination.write_bytes(archive_path.read_bytes())
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
install_dir = tmp_path / "install"
work_dir = tmp_path / "work"
work_dir.mkdir()
hydrate_source_tree(
commit,
install_dir,
work_dir,
source_repo = "unslothai/llama.cpp",
expected_sha256 = sha256_file(archive_path),
exact_source = True,
asset_url = asset_url,
)
assert seen == [asset_url]
assert (install_dir / "CMakeLists.txt").exists()
assert (install_dir / "convert_hf_to_gguf.py").exists()
def test_hydrate_source_tree_falls_back_to_codeload_when_asset_missing(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
# If the release asset 404s, fall back to codeload/archive (vanilla path).
commit = "b" * 40
archive_path = tmp_path / "vanilla-source.tar.gz"
_mk_source_tarball(archive_path, f"commit-{commit[:7]}")
asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url(
"unslothai/llama.cpp", "b9000", f"llama.cpp-source-commit-{commit}.tar.gz"
)
codeload_urls = INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit)
def fake_download_file(url: str, destination: Path) -> None:
if url == asset_url:
raise RuntimeError("404 Not Found")
assert url in codeload_urls
destination.write_bytes(archive_path.read_bytes())
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
install_dir = tmp_path / "install"
work_dir = tmp_path / "work"
work_dir.mkdir()
hydrate_source_tree(
commit,
install_dir,
work_dir,
source_repo = "unslothai/llama.cpp",
expected_sha256 = sha256_file(archive_path),
exact_source = True,
asset_url = asset_url,
)
assert (install_dir / "CMakeLists.txt").exists()
def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
@ -562,6 +668,48 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
assert "removing rollback path" in output
def test_activate_staged_dir_copies_when_replace_hits_busy_lock(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
):
staging_dir = tmp_path / "llama.cpp.staging-test"
(staging_dir / "bin").mkdir(parents = True)
(staging_dir / "bin" / "ggml-base.dll").write_bytes(b"fake dll")
dst = tmp_path / "llama.cpp"
def denied_replace(src, dst_arg):
raise PermissionError(errno.EACCES, "Access is denied", str(src))
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", denied_replace)
activate_staged_dir(staging_dir, dst)
assert (dst / "bin" / "ggml-base.dll").read_bytes() == b"fake dll"
assert not staging_dir.exists()
captured = capsys.readouterr()
assert "falling back to file-by-file copy" in captured.out + captured.err
def test_activate_staged_dir_reraises_non_busy_errors(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
staging_dir = tmp_path / "llama.cpp.staging-test"
staging_dir.mkdir()
(staging_dir / "new.txt").write_text("new install\n")
dst = tmp_path / "llama.cpp"
def out_of_space_replace(src, dst_arg):
raise OSError(errno.ENOSPC, "No space left on device", str(src))
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", out_of_space_replace)
with pytest.raises(OSError, match = "No space left on device"):
activate_staged_dir(staging_dir, dst)
assert not dst.exists()
assert (staging_dir / "new.txt").read_text() == "new install\n"
def test_binary_env_linux_includes_binary_parent_in_ld_library_path(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):

202
tests/studio/install/test_probe_timeouts.py Normal file
View file

@ -0,0 +1,202 @@
"""Tests that NVIDIA probes in the installers are bounded by a timeout.
Covers audit findings 5 and 6: a wedged nvidia-smi must not hang the installer,
and the Windows probe must require a real GPU listing (not just exit code 0).
Source-level assertions verify the guards are present in install.sh / install.ps1
/ setup.ps1; one behavioral shell test confirms the bash helper actually returns
within the timeout when nvidia-smi hangs.
"""
import os
import shutil
import stat
import subprocess
import sys
import tempfile
from pathlib import Path
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
INSTALL_SH = PACKAGE_ROOT / "install.sh"
INSTALL_PS1 = PACKAGE_ROOT / "install.ps1"
SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
def _extract_sh_function_body(source: str, name: str) -> str:
"""Return a shell function body from `source` by brace matching."""
needle = f"{name}() {{"
start = source.find(needle)
if start < 0:
return ""
depth = 0
i = start + len(needle) - 1
n = len(source)
while i < n:
ch = source[i]
if ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return source[start : i + 1]
i += 1
return source[start:]
# ── install.sh: _run_bounded helper and its use at every nvidia-smi call ──
class TestInstallShBoundedProbe:
def _src(self) -> str:
return INSTALL_SH.read_text(encoding = "utf-8")
def test_run_bounded_helper_defined(self):
body = _extract_sh_function_body(self._src(), "_run_bounded")
assert body, "install.sh must define a _run_bounded helper"
assert (
"command -v timeout" in body
), "_run_bounded must check for the `timeout` binary before using it"
assert "timeout 10" in body, "_run_bounded must apply a 10s timeout"
# Must fall back to running unbounded when `timeout` is unavailable
# (e.g. macOS) so semantics are unchanged there.
assert (
"else" in body and '"$@"' in body
), "_run_bounded must run the command unbounded when `timeout` is absent"
def test_nvidia_smi_dash_l_probe_is_bounded(self):
body = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
assert body, "install.sh must define _has_usable_nvidia_gpu"
# The -L probe must go through the bounded runner, not call nvidia-smi raw.
assert (
'_run_bounded "$_nvsmi" -L' in body
), "_has_usable_nvidia_gpu must run nvidia-smi -L through _run_bounded"
# The /proc fallback from PR 6174 must still be present.
assert "/proc/driver/nvidia" in body
def test_cuda_version_parse_is_bounded(self):
body = _extract_sh_function_body(self._src(), "get_torch_index_url")
assert body, "install.sh must define get_torch_index_url"
assert (
"_run_bounded" in body
), "get_torch_index_url CUDA-version parse must run nvidia-smi through _run_bounded"
# The locale must be forced without depending on `env` being on PATH.
assert "LC_ALL=C" in body
# _nvidia_detected gating from PR 6174 must remain.
assert "_nvidia_detected" in body
def test_no_unbounded_nvidia_smi_invocation_remains(self):
"""Every nvidia-smi *execution* in install.sh goes through _run_bounded.
`command -v nvidia-smi` and `-x /usr/bin/nvidia-smi` are resolution
checks, not executions, and are allowed. An execution looks like
`"$_nvsmi" ...` / `$_smi ...` / `nvidia-smi -L`.
"""
body_nvidia = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
body_torch = _extract_sh_function_body(self._src(), "get_torch_index_url")
# In _has_usable_nvidia_gpu the only execution of $_nvsmi must be bounded.
assert '"$_nvsmi" -L' not in body_nvidia.replace(
'_run_bounded "$_nvsmi" -L', ""
), "found an unbounded nvidia-smi -L execution in _has_usable_nvidia_gpu"
# In get_torch_index_url the $_smi execution must be bounded.
assert (
"LC_ALL=C $_smi" not in body_torch
), "found an unbounded LC_ALL=C $_smi execution in get_torch_index_url"
# ── install.ps1 / setup.ps1: bounded, GPU-row-validated Windows probe ──
class TestPowerShellBoundedProbe:
@pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
def test_bounded_helper_present(self, path):
src = path.read_text(encoding = "utf-8")
assert (
"function Invoke-NvidiaSmiBounded" in src
), f"{path.name} must define Invoke-NvidiaSmiBounded"
assert (
"WaitForExit($TimeoutSec * 1000)" in src
), f"{path.name} bounded probe must use WaitForExit with a timeout"
# Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate.
assert (
"$proc.Kill()" in src and "124" in src
), f"{path.name} must kill nvidia-smi and signal a timeout exit code"
@pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
def test_probe_requires_gpu_row(self, path):
src = path.read_text(encoding = "utf-8")
assert (
"function Test-NvidiaSmiHasGpu" in src
), f"{path.name} must define Test-NvidiaSmiHasGpu"
assert "@('-L')" in src, f"{path.name} must probe nvidia-smi with -L"
assert (
"^GPU\\s+\\d+:" in src
), f"{path.name} must require a 'GPU <n>:' data row, not just exit code 0"
@pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
def test_detection_uses_validated_probe(self, path):
src = path.read_text(encoding = "utf-8")
# The exit-code-only pattern must be gone from the detection block.
assert (
"& $nvSmiCmd.Source *> $null" not in src
), f"{path.name} must not use the exit-code-only nvidia-smi probe"
assert (
"Test-NvidiaSmiHasGpu $nvSmiCmd.Source" in src
), f"{path.name} PATH probe must use Test-NvidiaSmiHasGpu"
assert (
"Test-NvidiaSmiHasGpu $p" in src
), f"{path.name} hardcoded-path fallback must use Test-NvidiaSmiHasGpu"
# ── Behavioral: a hanging nvidia-smi must not hang _has_usable_nvidia_gpu ──
def _have_timeout() -> bool:
return shutil.which("timeout") is not None
@pytest.mark.skipif(not _have_timeout(), reason = "`timeout` binary not available")
def test_has_usable_nvidia_gpu_returns_under_timeout():
"""Extract _run_bounded + _has_usable_nvidia_gpu, point them at a fake
nvidia-smi that sleeps 30s, and assert the probe returns well under that.
"""
src = INSTALL_SH.read_text(encoding = "utf-8")
helper = _extract_sh_function_body(src, "_run_bounded")
fn = _extract_sh_function_body(src, "_has_usable_nvidia_gpu")
assert helper and fn
workdir = tempfile.mkdtemp(prefix = "pr6174_timeout_", dir = str(PACKAGE_ROOT.parent))
try:
fake_dir = Path(workdir, "bin")
fake_dir.mkdir()
fake_smi = fake_dir / "nvidia-smi"
fake_smi.write_text("#!/bin/sh\nsleep 30\n")
fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
# Build a minimal PATH that includes the fake nvidia-smi plus the real
# `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins.
real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")}
path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins])
# Force the /proc fallback off so the result depends only on the probe,
# and so a host with real NVIDIA does not mask the timeout behaviour.
script = (
f"{helper}\n{fn}\n"
"if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n"
)
proc = subprocess.run(
["sh", "-c", script],
env = {"PATH": path_env},
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 20, # generous: the internal timeout is 10s, sleep is 30s
)
# The probe must have returned (not hung). On this CI host /proc/driver/
# nvidia/gpus is absent, so a timed-out smi yields NONE; on a real NVIDIA
# host the /proc fallback yields DETECTED. Either way it must not hang.
assert proc.stdout.strip() in {"NONE", "DETECTED"}
finally:
shutil.rmtree(workdir, ignore_errors = True)

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

@ -719,6 +719,143 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
"""UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe."""
with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}):
# Reload _TORCH_BACKEND from the patched environment.
with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"):
_ensure_rocm_torch()
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
"""UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe."""
with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}):
with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"):
_ensure_rocm_torch()
mock_pip.assert_not_called()
# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
class TestHasRocmGpuKfdVendorGuard:
"""Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes.
These tests are source-level: they verify the regex and logic present in
the _has_rocm_gpu implementation rather than running the sysfs traversal
(which requires Linux path conventions).
"""
def _src(self) -> str:
"""Return the source of _has_rocm_gpu from install_python_stack.py."""
import inspect
return inspect.getsource(stack_mod._has_rocm_gpu)
def test_vendor_id_check_present(self):
"""_has_rocm_gpu sysfs fallback must check vendor_id 4098 (AMD 0x1002)."""
src = self._src()
assert "vendor_id" in src, (
"_has_rocm_gpu KFD sysfs fallback must read the properties file "
"to check vendor_id and exclude NVIDIA KFD nodes"
)
assert "4098" in src, (
"_has_rocm_gpu must require AMD vendor_id 4098 (0x1002) in the "
"KFD node properties to avoid false positives on NVIDIA systems"
)
def test_vendor_regex_pattern_anchored(self):
"""The vendor_id regex must use a word boundary to avoid partial matches."""
import re as _re
src = self._src()
# The pattern should have a word boundary before and after the number
# so "vendor_id 41098" doesn't match "vendor_id 4098".
assert (
_re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
), "_has_rocm_gpu vendor_id check should use word boundary anchors"
def test_sysfs_fallback_guarded_by_non_win32(self):
"""KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32')."""
src = self._src()
assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check"
def test_cpu_node_excluded(self):
"""gpu_id == '0' must be excluded (CPU topology nodes)."""
src = self._src()
assert (
'!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src
), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
def test_install_sh_has_vendor_check(self):
"""_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098."""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("_has_amd_rocm_gpu()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
assert "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id"
assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)"
def test_has_rocm_gpu_returns_false_when_nvidia_present(self):
"""_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True.
This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs
produce a false positive, an NVIDIA GPU always wins.
"""
with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True):
with patch("shutil.which", return_value = "/usr/bin/rocminfo"):
# Simulate rocminfo claiming an AMD GPU is present
mock_result = MagicMock()
mock_result.returncode = 0
mock_result.stdout = "Name: gfx1100\n"
with patch("subprocess.run", return_value = mock_result):
assert not stack_mod._has_rocm_gpu(), (
"_has_rocm_gpu must return False when NVIDIA GPU is detected, "
"regardless of what rocminfo reports"
)
def test_install_sh_has_rocm_gpu_nvidia_guard(self):
"""_has_amd_rocm_gpu in install.sh must call _has_usable_nvidia_gpu and return 1 if true."""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("_has_amd_rocm_gpu()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
assert (
"_has_usable_nvidia_gpu" in func_body
), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
assert (
"return 1" in func_body
), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
def test_has_usable_nvidia_gpu_proc_fallback_present(self):
"""`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback."""
import inspect
src = inspect.getsource(stack_mod._has_usable_nvidia_gpu)
assert "/proc/driver/nvidia" in src, (
"_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus when "
"nvidia-smi subprocess fails, to handle PATH gaps and driver init races"
)
def test_install_sh_has_usable_nvidia_gpu_proc_fallback(self):
"""_has_usable_nvidia_gpu in install.sh must also have a /proc/driver/nvidia fallback."""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("_has_usable_nvidia_gpu()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
assert "/proc/driver/nvidia" in func_body, (
"_has_usable_nvidia_gpu in install.sh must fall back to "
"/proc/driver/nvidia/gpus when nvidia-smi fails"
)
# TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping
@ -927,16 +1064,21 @@ class TestInstallShStructure:
source = sh_path.read_text(encoding = "utf-8")
body = _extract_sh_function_body(source, "get_torch_index_url")
nvidia_call = body.find("_has_usable_nvidia_gpu")
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
# Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to
# handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu
# returns true via /proc/driver/nvidia/gpus.
no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]')
if no_nvidia_branch < 0:
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
rocm_call = body.find("_has_amd_rocm_gpu")
assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia-smi"
assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch"
assert (
rocm_call > no_nvidia_branch
), "ROCm detection should sit inside the 'no nvidia-smi' branch"
), "ROCm detection should sit inside the 'no NVIDIA' branch"
assert (
nvidia_call < no_nvidia_branch
), "NVIDIA detection should run before the no-nvidia-smi branch"
), "NVIDIA detection should run before the no-NVIDIA branch"
def test_bitsandbytes_amd_install(self):
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
@ -1018,6 +1160,89 @@ class TestInstallShStructure:
rocm_pos = func_body.find("amd-smi")
assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection"
def test_unsloth_torch_backend_exported_after_get_torch_index_url(self):
"""install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set.
This lets install_python_stack.py skip ROCm torch operations on CUDA
and CPU hosts without re-running GPU detection in a subprocess.
"""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
backend_pos = source.find("UNSLOTH_TORCH_BACKEND")
assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh"
assert (
backend_pos > torch_url_pos
), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
# Verify all three cases are covered
assert '"cuda"' in source[backend_pos : backend_pos + 500]
assert '"rocm"' in source[backend_pos : backend_pos + 500]
assert '"cpu"' in source[backend_pos : backend_pos + 500]
# Must be exported so subprocesses (setup.sh, install_python_stack.py) see it
assert "export UNSLOTH_TORCH_BACKEND" in source
def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self):
"""_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098.
NVIDIA open kernel module (560+) registers KFD nodes with vendor_id
4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0
(true) on NVIDIA-only hosts that have the nvidia-open driver, causing
get_torch_index_url to select a ROCm wheel index.
"""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("_has_amd_rocm_gpu()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
assert (
"vendor_id" in func_body
), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
assert (
"4098" in func_body
), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
def test_kfd_awk_resets_state_per_file(self):
"""KFD sysfs awk must reset gpu/amd state per file (FNR==1).
Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU
agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU
(gpu_id > 0, vendor_id 4318) can produce a false positive: node 0
sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1
before vendor_id 4318 is seen on node 1.
"""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("_has_amd_rocm_gpu()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
assert "FNR==1" in func_body, (
"_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 "
"to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes"
)
def test_get_torch_index_url_uses_nvidia_detected_flag(self):
"""get_torch_index_url must track NVIDIA detection independently of _smi.
When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback
but nvidia-smi is not on PATH, _smi stays empty. Without a separate
_nvidia_detected flag, the function falls into the AMD/CPU branch even
though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA.
"""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("get_torch_index_url()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
assert "_nvidia_detected" in func_body, (
"get_torch_index_url must use a _nvidia_detected flag (separate from "
"_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
)
# The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on
# _smi being empty.
assert (
'_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
# TEST: Live regression on current host (NVIDIA B200 expected)
@ -1389,14 +1614,18 @@ class TestApplyGpuIdsRocmFallback:
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'getattr(_torch.version, "hip", None)' in func_body
def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self):
"""apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm."""
def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self):
"""apply_gpu_ids should set HIP_VISIBLE_DEVICES but leave ROCR_VISIBLE_DEVICES inherited.
ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices.
Overwriting it breaks multi-GPU ROCm systems (see issue #6118).
"""
hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body
assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' in func_body
assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' not in func_body
def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self):
"""torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes."""
@ -2016,7 +2245,22 @@ class TestRuntimeBnbRocmSourceGuards:
"""A failed redetect must not downgrade a persisted suffix to '72'."""
for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH):
source = path.read_text(encoding = "utf-8")
assert 'os.environ.get("BNB_ROCM_VERSION") or "72"' in source, path.name
assert (
'_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source
), path.name
def test_main_requires_found_rocm_dll(self):
"""HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force
a ROCm backend onto a non-ROCm bitsandbytes."""
source = self._MAIN_PATH.read_text(encoding = "utf-8")
assert "if _found_rocm_bnb:" in source
assert "_hip_env" not in source
def test_worker_requires_found_rocm_dll(self):
"""No DLL found: the worker must not write any override or touch the
seeded marker (later import fixes must still see sitecustomize)."""
source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8")
assert "if _found_rocm_bnb:" in source
class TestDetectBnbRocmDllVer:
@ -2218,8 +2462,9 @@ class TestWorkerWindowsRocmPatches:
assert "BNB_ROCM_VERSION" in source
# Detection helper must be used
assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source
# "72" must appear as the safe fallback
assert '"72"' in source or "'72'" in source
# Falls back to the seeded value, never a blind "72" (which would
# force a ROCm backend onto a non-ROCm bitsandbytes wheel)
assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' in source
def test_bnb_rocm_version_set_before_ml_imports(self):
"""BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports."""

15
tests/studio/playwright_extra_ui.py
View file

@ -332,6 +332,21 @@ with sync_playwright() as p:
plus_btn.click(force = True)
page.wait_for_timeout(400)
compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
if compare_item.count() == 0:
# Compare chat moved into the "More" submenu; hover, then click fallback.
more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
if more_trigger.count() > 0:
more_trigger.hover()
page.wait_for_timeout(400)
compare_item = page.get_by_role(
"menuitem", name = re.compile(r"Compare chat", re.I)
).first
if compare_item.count() == 0:
more_trigger.click(force = True)
page.wait_for_timeout(400)
compare_item = page.get_by_role(
"menuitem", name = re.compile(r"Compare chat", re.I)
).first
if compare_item.count() > 0:
compare_item.click(force = True)
compare_opened = True

51
tests/test_import_fixes_drift.py
View file

@ -584,3 +584,54 @@ def test_accelerate_utils_imports_module_present():
"accelerate.utils.imports.is_wandb_available is gone; "
"disable_broken_wandb cannot patch the source module."
)
# ===========================================================================
# bitsandbytes -- ROCm arch / warp-size detection shape
# ===========================================================================
def test_bitsandbytes_rocm_detection_helpers_recognizable():
"""``fix_bitsandbytes_rocm_arch_detection`` swaps bnb's ROCm helpers
only when they shell out via subprocess and never consult torch device
props; a third shape is declined by design, silently restoring Windows
ROCm noise. Fail so the sniff gets updated. Reads source, no import."""
spec = importlib.util.find_spec("bitsandbytes")
if spec is None:
pytest.skip("bitsandbytes not installed -- nothing to drift-check.")
cuda_specs_path = None
for location in spec.submodule_search_locations or []:
candidate = os.path.join(location, "cuda_specs.py")
if os.path.isfile(candidate):
cuda_specs_path = candidate
break
if cuda_specs_path is None:
pytest.skip("bitsandbytes has no cuda_specs.py (pre-ROCm version).")
import ast
with open(cuda_specs_path, "r", encoding = "utf-8") as f:
source = f.read()
helpers = [
node
for node in ast.walk(ast.parse(source))
if isinstance(node, ast.FunctionDef)
and node.name in ("get_rocm_gpu_arch", "get_rocm_warpsize")
]
if not helpers:
pytest.skip("bitsandbytes cuda_specs has no ROCm detection helpers.")
for node in helpers:
segment = ast.get_source_segment(source, node) or ""
recognized = (
"subprocess" in segment
or "get_device_properties" in segment
or "gcnArchName" in segment
)
if not recognized:
pytest.fail(
f"DRIFT DETECTED: bitsandbytes.cuda_specs.{node.name} uses "
"neither subprocess nor torch device properties; "
"fix_bitsandbytes_rocm_arch_detection's shape sniff will "
"decline to patch it and Windows ROCm import-time noise / "
"wrong ROCM_GPU_ARCH may return."
)

312
tests/utils/test_rope_scaling_drift.py Normal file
View file

@ -0,0 +1,312 @@
"""Guard for config.rope_scaling being silently dropped (issue #2405).
Unsloth's replacement rotary classes ignored rope_scaling when constructed
from a config (the modern-transformers path), so Llama-3.1 ran with unscaled
RoPE and collapsed into gibberish past ~32K tokens.
Layers: (1) AST tripwire, stdlib only; (2) CPU checks of the pure helper
_compute_config_rope_inv_freq against transformers' ROPE_INIT_FUNCTIONS;
(3) CUDA checks instantiating the real class (skipped without a real device,
probed by allocating a tensor so import-time CUDA spoofs cannot fool the gate).
Layers 2 and 3 fail on the unfixed code.
"""
import ast
import math
from pathlib import Path
import pytest
import torch
def _has_real_cuda():
try:
torch.zeros(1).to("cuda")
return True
except Exception:
return False
HAS_REAL_CUDA = _has_real_cuda()
requires_cuda = pytest.mark.skipif(
not HAS_REAL_CUDA,
reason = "LlamaRotaryEmbedding builds per-device CUDA caches in __init__",
)
REPO_ROOT = Path(__file__).resolve().parents[2]
LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py"
CLASS_NAME = "LlamaRotaryEmbedding"
# Llama-3.1-style rope_scaling.
LLAMA3_ROPE_SCALING = {
"rope_type": "llama3",
"factor": 8.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
}
ROPE_THETA = 500000.0
HEAD_DIM = 128
MAX_POS = 131072
# --- Layer 1: AST structural tripwire (stdlib only, no unsloth import) ---
def _load_class_init():
tree = ast.parse(LLAMA_PY.read_text())
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == CLASS_NAME:
for sub in node.body:
if isinstance(sub, ast.FunctionDef) and sub.name == "__init__":
return sub
raise AssertionError(
f"{CLASS_NAME}.__init__ not found in {LLAMA_PY}; if it was renamed or "
"moved, update this guard so RoPE scaling stays protected (issue #2405)"
)
def _config_branch(init_fn):
"""The `if config is not None:` block at the top of __init__."""
for node in init_fn.body:
if isinstance(node, ast.If):
test = node.test
is_config_test = (
isinstance(test, ast.Compare)
and isinstance(test.left, ast.Name)
and test.left.id == "config"
)
if is_config_test:
return node
return None
def test_config_path_inspects_rope_scaling():
init_fn = _load_class_init()
branch = _config_branch(init_fn)
assert branch is not None, (
f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` "
"branch; the config constructor path must read config.rope_scaling so "
"scaled models (llama3/linear/longrope) are not silently unscaled "
"(issue #2405)"
)
names = set()
for stmt in branch.body:
for sub in ast.walk(stmt):
if isinstance(sub, ast.Attribute):
names.add(sub.attr)
elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
names.add(sub.value)
assert "rope_scaling" in names, (
f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. "
"When a rotary class is built straight from a config (the path modern "
"transformers takes, since rotary moved to LlamaModel), the llama3 / "
"linear / longrope scaling must still be applied; otherwise long inputs "
"produce repeated-pattern gibberish (issue #2405)."
)
called = {
sub.func.id
for stmt in branch.body
for sub in ast.walk(stmt)
if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)
}
assert "_compute_config_rope_inv_freq" in called, (
f"{CLASS_NAME}.__init__ config path no longer calls "
"_compute_config_rope_inv_freq; the CPU behavioral tests below cover "
"that helper directly, so the constructor must stay wired to it or "
"scaled configs silently lose RoPE scaling again (issue #2405)."
)
# --- Layer 2: CPU behavioral guard (pure helper, no instantiation) ---
def _make_config(rope_scaling):
from transformers import LlamaConfig
return LlamaConfig(
hidden_size = 256,
num_attention_heads = 2,
num_key_value_heads = 2,
head_dim = HEAD_DIM,
rope_theta = ROPE_THETA,
max_position_embeddings = MAX_POS,
rope_scaling = rope_scaling,
)
def _unsloth_rotary(config):
from unsloth.models import llama as llama_mod
return llama_mod.LlamaRotaryEmbedding(config = config)
def _reference_inv_freq(config, rope_type):
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
inv_freq, _attention_factor = ROPE_INIT_FUNCTIONS[rope_type](config, "cpu")
return inv_freq.float().cpu()
def _vanilla_inv_freq():
return 1.0 / (
ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype = torch.int64).float() / HEAD_DIM)
)
def _compute_helper(config, rope_scaling):
from unsloth.models.llama import _compute_config_rope_inv_freq
return _compute_config_rope_inv_freq(config, rope_scaling)
def test_llama3_scaling_applied_to_inv_freq():
config = _make_config(LLAMA3_ROPE_SCALING)
got, attention_scaling = _compute_helper(config, config.rope_scaling)
expected = _reference_inv_freq(config, "llama3")
vanilla = _vanilla_inv_freq()
# Guard against a vacuous test.
assert not torch.allclose(
expected, vanilla, rtol = 1e-4
), "test setup error: llama3-scaled inv_freq should differ from vanilla"
assert got is not None, (
"_compute_config_rope_inv_freq returned None for a llama3 config; the "
"config path is dropping config.rope_scaling, so long-context inference "
"degrades into repeated-pattern gibberish (issue #2405)."
)
got = got.float().cpu()
assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), (
"inv_freq for a llama3 config does not match transformers' llama3 RoPE "
"scaling (issue #2405).\n"
f"got[:6]={got[:6].tolist()}\nexpected[:6]={expected[:6].tolist()}"
)
def test_default_rope_type_matches_vanilla_inv_freq():
config = _make_config(None)
got, attention_scaling = _compute_helper(config, {"rope_type": "default"})
assert got is not None
vanilla = _vanilla_inv_freq()
assert torch.allclose(got.float().cpu(), vanilla, rtol = 1e-4, atol = 1e-6), (
"default rope_type must equal the vanilla inv_freq; "
f"got[:6]={got[:6].tolist()} vanilla[:6]={vanilla[:6].tolist()}"
)
def _cos_at_position(rot, position):
"""cos row at one position, built like _set_cos_sin_cache but CPU-only."""
inv_freq = rot.inv_freq.float().cpu()
t = torch.tensor([position], dtype = torch.float32)
t = rot._apply_time_scaling(t.clone()) if hasattr(rot, "_apply_time_scaling") else t
freqs = torch.outer(t, inv_freq)
emb = torch.cat((freqs, freqs), dim = -1)
return emb.cos().squeeze(0)
# --- Layer 3: CUDA behavioral guard (real instantiation needs a device) ---
@requires_cuda
def test_constructor_applies_llama3_scaling():
config = _make_config(LLAMA3_ROPE_SCALING)
rot = _unsloth_rotary(config)
got = rot.inv_freq.float().cpu()
expected = _reference_inv_freq(config, "llama3")
assert torch.allclose(
got, expected, rtol = 1e-4, atol = 1e-6
), "LlamaRotaryEmbedding built from a llama3 config produced unscaled inv_freq (issue #2405)."
@requires_cuda
def test_constructor_unscaled_config_uses_vanilla_inv_freq():
rot = _unsloth_rotary(_make_config(None))
got = rot.inv_freq.float().cpu()
vanilla = _vanilla_inv_freq()
assert torch.allclose(
got, vanilla, rtol = 1e-4, atol = 1e-6
), "LlamaRotaryEmbedding with no rope_scaling must use the vanilla inv_freq"
@requires_cuda
def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position():
scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
unscaled = _unsloth_rotary(_make_config(None))
pos = 10000
cos_scaled = _cos_at_position(scaled, pos)
cos_unscaled = _cos_at_position(unscaled, pos)
assert not torch.allclose(cos_scaled, cos_unscaled, rtol = 1e-4, atol = 1e-5), (
f"cos values at position {pos} are identical for a llama3-scaled and an "
"unscaled rotary embedding, which means scaling was dropped (issue "
"#2405). With correct llama3 scaling the low-frequency bands shrink by "
"up to 8x and must change the angles at long positions."
)
@requires_cuda
def test_extended_cache_keeps_scaling_after_growth():
scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
# Grow past the initial cache size (mirrors long-context decode).
dummy = torch.zeros(1, dtype = torch.float32)
scaled.extend_rope_embedding(dummy, seq_len = 40960)
config = _make_config(LLAMA3_ROPE_SCALING)
expected = _reference_inv_freq(config, "llama3")
got = scaled.inv_freq.float().cpu()
assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), (
"growing the RoPE cache (extend_rope_embedding) must preserve llama3 "
"scaling of inv_freq; long-context decode loses scaling otherwise "
"(issue #2405)."
)
def test_object_style_rope_scaling_does_not_crash():
# Object-style rope_scaling must be normalized, not .get()'d directly.
from dataclasses import dataclass
from unsloth.models.llama import _compute_config_rope_inv_freq
@dataclass
class FakeRopeScalingConfig:
rope_type: str = "llama3"
factor: float = 8.0
low_freq_factor: float = 1.0
high_freq_factor: float = 4.0
original_max_position_embeddings: int = 8192
config = _make_config(LLAMA3_ROPE_SCALING)
inv_freq, attention_scaling = _compute_config_rope_inv_freq(config, FakeRopeScalingConfig())
assert inv_freq is not None, (
"object-style (non-dict) config.rope_scaling must be normalized, not "
"dropped; otherwise scaled models silently lose RoPE scaling again "
"(issue #2405)."
)
expected = _reference_inv_freq(config, "llama3")
assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6)
def test_object_style_rope_scaling_on_config_delegates_correctly():
# 'linear' has no inline fallback; only the normalized-config retry passes this.
from dataclasses import dataclass
from unsloth.models.llama import _compute_config_rope_inv_freq
@dataclass
class FakeLinearRopeScalingConfig:
rope_type: str = "linear"
factor: float = 4.0
dict_config = _make_config({"rope_type": "linear", "factor": 4.0})
expected = _reference_inv_freq(dict_config, "linear")
object_config = _make_config({"rope_type": "linear", "factor": 4.0})
object_config.rope_scaling = FakeLinearRopeScalingConfig()
inv_freq, attention_scaling = _compute_config_rope_inv_freq(
object_config, object_config.rope_scaling
)
assert inv_freq is not None, (
"linear rope_scaling exposed as a config object was silently dropped; "
"delegation must retry with a config copy carrying the normalized dict "
"(issue #2405)."
)
assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6)

4
unsloth/_gpu_init.py
View file

@ -30,6 +30,7 @@ from .import_fixes import (
disable_broken_causal_conv1d,
disable_broken_vllm,
configure_amdgpu_asic_id_table_path,
fix_bitsandbytes_rocm_arch_detection,
torchvision_compatibility_check,
fix_diffusers_warnings,
fix_huggingface_hub,
@ -67,6 +68,8 @@ except Exception:
# Configure libdrm ids table path early so ROCm can resolve AMD GPU names.
configure_amdgpu_asic_id_table_path()
# Must precede `import unsloth_zoo` below, which imports bnb on ROCm.
fix_bitsandbytes_rocm_arch_detection()
disable_broken_causal_conv1d()
disable_broken_vllm()
fix_message_factory_issue()
@ -75,6 +78,7 @@ torchvision_compatibility_check()
fix_diffusers_warnings()
fix_huggingface_hub()
del configure_amdgpu_asic_id_table_path
del fix_bitsandbytes_rocm_arch_detection
del disable_broken_causal_conv1d
del disable_broken_vllm
del fix_message_factory_issue

18
unsloth/chat_templates.py
View file

@ -2341,13 +2341,15 @@ extra_eos_tokens = None,
You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional.
"""
# Strip only the left
# Strip only the left: trailing whitespace can be part of the repeated example
# (e.g. "{OUTPUT}\n"). Accidental trailing whitespace (#992) is retried on failure.
chat_template = chat_template.lstrip()
assert(tokenizer is not None)
if extra_eos_tokens is None: extra_eos_tokens = []
elif type(extra_eos_tokens) is str: extra_eos_tokens = [extra_eos_tokens,]
original_extra_eos_tokens = list(extra_eos_tokens)
vocab = tokenizer.get_vocab()
for extra_eos in extra_eos_tokens:
@ -2454,6 +2456,20 @@ extra_eos_tokens = None,
f"{left_changed}"
)
except:
# Accidental trailing whitespace (#992) desyncs the two-example detection,
# so retry once without it. Templates that parse as-is are never altered.
rstripped_chat_template = chat_template.rstrip()
if rstripped_chat_template != chat_template:
try:
return construct_chat_template(
tokenizer = tokenizer,
chat_template = rstripped_chat_template,
default_system_message = default_system_message,
extra_eos_tokens = original_extra_eos_tokens,
)
except Exception:
pass
output_pos = chat_template.find("{OUTPUT}")
input_pos = chat_template.find("{INPUT}")
if output_pos == -1 or input_pos == -1:

334
unsloth/import_fixes.py
View file

@ -1823,6 +1823,340 @@ def configure_amdgpu_asic_id_table_path():
return None
# ---------------------------------------------------------------------------
# bitsandbytes Windows ROCm fix: cextension.py runs get_rocm_gpu_arch()
# (bnb >= 0.47) and get_rocm_warpsize() (0.49.x) at import, shelling out to
# rocminfo / hipinfo.exe via PATH. Neither is on PATH on Windows (AMD torch
# wheels put hipInfo.exe in venv Scripts), so every import logs ERROR +
# WARNING, ROCM_GPU_ARCH becomes "unknown", and warp size defaults to 64:
# wrong on RDNA (wave 32), breaking 4-bit blocksizes and
# ALLOW_PREQUANTIZED_MODELS. Upstream fix unmerged (bitsandbytes#1969), so a
# MetaPathFinder swaps both helpers for torch-device-props-first versions
# right after bitsandbytes.cuda_specs executes, before cextension reads
# them. Must run before `import unsloth_zoo` (imports bnb on ROCm).
# ---------------------------------------------------------------------------
_BNB_CUDA_SPECS_MODULE = "bitsandbytes.cuda_specs"
_BNB_ROCM_FIX_FINDER_SENTINEL = "_unsloth_bnb_rocm_fix_finder"
_BNB_ROCM_FIX_FUNCTION_FLAG = "__unsloth_bnb_rocm_fix__"
def _torch_rocm_device_props():
"""Device-0 props on a ROCm torch build with a visible GPU, else None.
Never raises; bnb's own import initializes the device context anyway."""
try:
import torch
if not getattr(getattr(torch, "version", None), "hip", None):
return None
if not torch.cuda.is_available():
return None
return torch.cuda.get_device_properties(0)
except Exception:
return None
def _iter_hipinfo_paths():
"""Yield existing hipInfo.exe paths: PATH, interpreter scripts dir (venv
and conda layouts), then HIP SDK / AMD installer locations."""
import shutil
import sysconfig
candidates = []
try:
resolved = shutil.which("hipinfo.exe")
if resolved:
candidates.append(resolved)
except Exception:
pass
try:
scripts_dir = sysconfig.get_path("scripts")
if scripts_dir:
candidates.append(os.path.join(scripts_dir, "hipInfo.exe"))
except Exception:
pass
executable_dir = os.path.dirname(sys.executable or "")
if executable_dir:
candidates.append(os.path.join(executable_dir, "hipInfo.exe"))
candidates.append(os.path.join(executable_dir, "Scripts", "hipInfo.exe"))
for env_key in ("HIP_PATH", "ROCM_PATH"):
root = os.environ.get(env_key, "").strip()
if root:
candidates.append(os.path.join(root, "bin", "hipInfo.exe"))
rocm_root = os.path.join(os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm")
try:
if os.path.isdir(rocm_root):
for version_dir in sorted(os.listdir(rocm_root), reverse = True):
candidates.append(os.path.join(rocm_root, version_dir, "bin", "hipInfo.exe"))
except Exception:
pass
seen = set()
for candidate in candidates:
try:
key = os.path.normcase(os.path.normpath(candidate))
if key in seen:
continue
seen.add(key)
if os.path.isfile(candidate):
yield candidate
except Exception:
continue
def _run_hipinfo(hipinfo_path):
"""Run hipInfo.exe and return its stdout, or "" on any failure."""
import subprocess
try:
result = subprocess.run(
[hipinfo_path],
capture_output = True,
text = True,
timeout = 15,
creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0),
)
return result.stdout or ""
except Exception as e:
_log_rocm_detection(f"Unsloth: `{hipinfo_path}` failed: {e}")
return ""
def _unsloth_get_rocm_gpu_arch():
"""Replaces bnb's get_rocm_gpu_arch: torch device props first (no
subprocess), then hipInfo.exe by absolute path, then a quiet "unknown"."""
try:
import torch
if not getattr(getattr(torch, "version", None), "hip", None):
return "unknown"
except Exception:
return "unknown"
props = _torch_rocm_device_props()
if props is not None:
try:
# gcnArchName may carry feature flags, e.g. "gfx90a:sramecc+:xnack-"
arch = str(props.gcnArchName).split(":")[0].strip()
if arch.startswith("gfx"):
return arch
except Exception:
pass
for hipinfo_path in _iter_hipinfo_paths():
match = re.search(r"gcnArchName:\s+gfx([a-zA-Z\d]+)", _run_hipinfo(hipinfo_path))
if match:
return "gfx" + match.group(1)
_log_rocm_detection(
"Unsloth: Could not detect the ROCm GPU architecture - bitsandbytes will see `unknown`."
)
return "unknown"
def _unsloth_get_rocm_warpsize():
"""Replaces bnb 0.49.x get_rocm_warpsize: upstream defaults to 64 when
rocminfo is missing, wrong on RDNA (wave 32)."""
try:
import torch
if not getattr(getattr(torch, "version", None), "hip", None):
return 32 # upstream behavior: NVIDIA warp size is always 32
except Exception:
return 64 # upstream behavior: default to 64 on failure
props = _torch_rocm_device_props()
if props is not None:
# torch 2.11 ROCm exposes warp_size; some builds used warpSize.
for attribute_name in ("warp_size", "warpSize"):
warp_size = getattr(props, attribute_name, None)
if isinstance(warp_size, int) and warp_size in (32, 64):
return warp_size
for hipinfo_path in _iter_hipinfo_paths():
match = re.search(r"^\s*warpSize:\s+(\d+)", _run_hipinfo(hipinfo_path), re.MULTILINE)
if match and int(match.group(1)) in (32, 64):
return int(match.group(1))
_log_rocm_detection(
"Unsloth: Could not detect the ROCm warp size - defaulting to 64 "
"(bitsandbytes' own default)."
)
return 64
setattr(_unsloth_get_rocm_gpu_arch, _BNB_ROCM_FIX_FUNCTION_FLAG, True)
setattr(_unsloth_get_rocm_warpsize, _BNB_ROCM_FIX_FUNCTION_FLAG, True)
def _bnb_rocm_helper_is_broken(function):
"""True only for upstream's subprocess-only detectors; co_names works
where getsource fails. Versions consulting torch props are untouched."""
if function is None or not callable(function):
return False
if getattr(function, _BNB_ROCM_FIX_FUNCTION_FLAG, False):
return False # Already ours.
try:
function = inspect.unwrap(function)
except Exception:
pass
code = getattr(function, "__code__", None)
co_names = getattr(code, "co_names", ()) if code is not None else ()
if not co_names:
return False # C function or opaque wrapper -- do not touch.
if "get_device_properties" in co_names or "gcnArchName" in co_names:
return False # Fixed upstream -- no-op.
return "subprocess" in co_names
def _patch_bnb_cuda_specs_module(module):
"""Swap broken ROCm detection helpers on an executed cuda_specs module.
Returns True when the module ends up patched (now or previously)."""
patched = False
for attribute_name, replacement in (
("get_rocm_gpu_arch", _unsloth_get_rocm_gpu_arch),
("get_rocm_warpsize", _unsloth_get_rocm_warpsize),
):
original = getattr(module, attribute_name, None)
if getattr(original, _BNB_ROCM_FIX_FUNCTION_FLAG, False):
patched = True # Already ours.
continue
if not _bnb_rocm_helper_is_broken(original):
continue
setattr(module, attribute_name, replacement)
patched = True
logger.info(
f"Unsloth: Patched bitsandbytes.cuda_specs.{attribute_name} - "
f"avoids PATH-dependent subprocess GPU detection on Windows ROCm."
)
return patched
class _BnbCudaSpecsPatchLoader(importlib.abc.Loader):
__slots__ = ("_loader",)
def __init__(self, loader):
self._loader = loader
def create_module(self, spec):
create_module = getattr(self._loader, "create_module", None)
if create_module is None:
return None
return create_module(spec)
def exec_module(self, module):
self._loader.exec_module(module)
# Patch after the module body ran, before cextension calls it. The
# finder stays on sys.meta_path (same lifecycle as the blockers
# above) so importlib.reload(bitsandbytes.cuda_specs) re-patches.
try:
_patch_bnb_cuda_specs_module(module)
except Exception as e:
_log_rocm_detection(f"Unsloth: bitsandbytes ROCm detection patch failed: {e}")
def __getattr__(self, name):
# Delegate get_source / get_filename etc. so introspection works.
return getattr(self._loader, name)
class _BnbCudaSpecsPatchFinder(importlib.abc.MetaPathFinder):
__slots__ = (_BNB_ROCM_FIX_FINDER_SENTINEL,)
def __init__(self):
setattr(self, _BNB_ROCM_FIX_FINDER_SENTINEL, True)
def find_spec(
self,
fullname,
path = None,
target = None,
):
if fullname != _BNB_CUDA_SPECS_MODULE:
return None
# Delegate to remaining finders (editable installs, frozen apps)
# and wrap the loader that would actually be used.
spec = None
for finder in sys.meta_path:
if finder is self or getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False):
continue
finder_find_spec = getattr(finder, "find_spec", None)
if finder_find_spec is None:
continue
try:
spec = finder_find_spec(fullname, path, target)
except Exception:
spec = None
if spec is not None:
break
if spec is None or spec.loader is None:
return None
if not hasattr(spec.loader, "exec_module"):
return None # Legacy loader -- let the stock machinery handle it.
spec.loader = _BnbCudaSpecsPatchLoader(spec.loader)
return spec
def _repair_imported_bitsandbytes_rocm_constants():
"""bnb imported before unsloth: noise already fired, but fix detectors
and cached constants, incl. by-value ROCM_WARP_SIZE_64 copies."""
cuda_specs = sys.modules.get(_BNB_CUDA_SPECS_MODULE)
if cuda_specs is None:
return
if not _patch_bnb_cuda_specs_module(cuda_specs):
return
try:
arch = cuda_specs.get_rocm_gpu_arch()
except Exception:
arch = "unknown"
warp_size_64 = None
get_rocm_warpsize = getattr(cuda_specs, "get_rocm_warpsize", None)
if callable(get_rocm_warpsize):
try:
warp_size_64 = get_rocm_warpsize() == 64
except Exception:
warp_size_64 = None
for module_name, module in list(sys.modules.items()):
if module is None or module is cuda_specs:
continue
if module_name != "bitsandbytes" and not module_name.startswith("bitsandbytes."):
continue
try:
if arch != "unknown" and getattr(module, "ROCM_GPU_ARCH", None) == "unknown":
module.ROCM_GPU_ARCH = arch
if warp_size_64 is not None and isinstance(
getattr(module, "ROCM_WARP_SIZE_64", None), bool
):
module.ROCM_WARP_SIZE_64 = warp_size_64
except Exception:
continue
logger.info("Unsloth: Repaired bitsandbytes ROCm arch / warp-size constants in place.")
def fix_bitsandbytes_rocm_arch_detection():
"""Fix bnb's import-time ROCm arch / warp-size detection on Windows
(see header above). No-op on non-Windows, non-ROCm, missing or
upstream-fixed bnb. Idempotent. Opt out: UNSLOTH_DISABLE_BNB_ROCM_FIX=1."""
if os.environ.get("UNSLOTH_DISABLE_BNB_ROCM_FIX", "0") == "1":
return
if sys.platform != "win32":
return
if not _is_rocm_torch_build():
return
# Already imported: prevention impossible, repair in place instead.
if _BNB_CUDA_SPECS_MODULE in sys.modules:
try:
_repair_imported_bitsandbytes_rocm_constants()
except Exception:
pass
return
try:
if importlib.util.find_spec("bitsandbytes") is None:
return
except Exception:
return
for finder in sys.meta_path:
if getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False):
return # Already installed -- idempotent.
sys.meta_path.insert(0, _BnbCudaSpecsPatchFinder())
_log_rocm_detection("Unsloth: Installed the bitsandbytes ROCm arch detection patch hook.")
def _is_causal_conv1d_name(module_name: str) -> bool:
return module_name == _CAUSAL_CONV1D_PREFIX or module_name.startswith(
_CAUSAL_CONV1D_PREFIX + "."

20
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2026.6.2"
__version__ = "2026.6.3"
__all__ = [
"SUPPORTS_BFLOAT16",
@ -1264,6 +1264,18 @@ if is_openai_available():
from transformers import AutoTokenizer
from transformers.utils.import_utils import _is_package_available
def _package_available(pkg_name: str) -> bool:
# transformers >= 5.x makes `_is_package_available` always return a
# `(exists, version)` tuple, which is truthy even when the package is
# absent; older versions returned a plain bool. Normalise to a bool so
# callers don't take "package present" branches for missing packages.
result = _is_package_available(pkg_name)
if isinstance(result, tuple):
return bool(result[0])
return bool(result)
SUPPORTS_BFLOAT16 = False
HAS_FLASH_ATTENTION = False
HAS_FLASH_ATTENTION_SOFTCAPPING = False
@ -1274,7 +1286,7 @@ if DEVICE_TYPE == "cuda":
if major_version >= 8:
SUPPORTS_BFLOAT16 = True
if _is_package_available("flash_attn"):
if _package_available("flash_attn"):
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
try:
try:
@ -1319,7 +1331,7 @@ if DEVICE_TYPE == "cuda":
HAS_FLASH_ATTENTION = False
elif DEVICE_TYPE == "hip":
SUPPORTS_BFLOAT16 = True
if _is_package_available("flash_attn"):
if _package_available("flash_attn"):
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
try:
try:
@ -1981,7 +1993,7 @@ def is_bfloat16_supported():
def is_vLLM_available():
return _is_package_available("vllm")
return _package_available("vllm")
# Patches models to add RoPE Scaling

301
unsloth/models/diffusion.py Normal file
View file

@ -0,0 +1,301 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
FastDiffusionModel: a transformers-only slow path for text-diffusion models (e.g. DiffusionGemma).
These models use a block-diffusion sampling loop (custom generate) and a novel backbone, so we skip
Unsloth's autoregressive kernel/compile patching and load the unmodified HF model (outputs stay
bit-identical to transformers), keeping only the safe conveniences: 4bit/8bit loading, PEFT LoRA, the
(model, tokenizer) API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land.
"""
import os
import torch
from transformers import AutoConfig, AutoProcessor, AutoTokenizer
from ._utils import is_bfloat16_supported
from .llama import logger
__all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"]
# transformers model_type strings routed to this slow path
DIFFUSION_MODEL_TYPES = ("diffusion_gemma", "diffusion_gemma4")
# Default LoRA targets: standard nn.Linear modules in the shared Gemma-4 backbone. The 128 MoE experts
# are fused 3D Parameters (gate_up_proj/down_proj), not nn.Linear, so PEFT LoRA cannot target them.
DIFFUSION_LORA_TARGETS = [
"q_proj",
"k_proj",
"v_proj",
"o_proj", # attention
"gate_proj",
"up_proj",
"down_proj", # dense (non-expert) MLP
]
# Vision tower uses a custom Linear with the same suffix names; exclude it so only the text path is wrapped.
DIFFUSION_LORA_EXCLUDE = r".*(vision_tower|embed_vision).*"
def is_diffusion_model_type(model_types):
"""model_types: str or iterable -> True if any is a known diffusion model_type."""
if isinstance(model_types, str):
model_types = (model_types,)
return any(mt in DIFFUSION_MODEL_TYPES for mt in model_types)
def _resolve_diffusion_model_class(config):
"""Resolve the HF model class for a diffusion checkpoint from config.architectures."""
import transformers
archs = getattr(config, "architectures", None) or []
for arch in archs:
cls = getattr(transformers, arch, None)
if cls is not None:
return cls
# Fallbacks across naming revisions.
for name in (
"DiffusionGemmaForBlockDiffusion",
"DiffusionGemma4ModelForBlockDiffusion",
"DiffusionGemma4ForBlockDiffusion",
):
cls = getattr(transformers, name, None)
if cls is not None:
return cls
raise RuntimeError(
f"Unsloth: could not resolve a diffusion model class from architectures={archs}. "
"Ensure you have the transformers build that ships the DiffusionGemma implementation."
)
def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only):
"""Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4``
classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the
type/arch names in-memory, and rebuild."""
try:
return AutoConfig.from_pretrained(
model_name,
token = token,
trust_remote_code = trust_remote_code,
revision = revision,
local_files_only = local_files_only,
)
except ValueError as e:
if "diffusion_gemma" not in str(e):
raise
import json
from transformers.utils import cached_file
cfg_path = cached_file(
model_name,
"config.json",
token = token,
revision = revision,
local_files_only = local_files_only,
)
with open(cfg_path, encoding = "utf-8") as f:
cd = json.load(f)
cd["model_type"] = "diffusion_gemma4"
cd.setdefault("architectures", ["DiffusionGemma4ModelForBlockDiffusion"])
if isinstance(cd.get("text_config"), dict):
cd["text_config"]["model_type"] = "diffusion_gemma4_text"
if isinstance(cd.get("vision_config"), dict):
cd["vision_config"]["model_type"] = "diffusion_gemma4_vision"
from transformers import DiffusionGemma4Config
return DiffusionGemma4Config.from_dict(cd)
class FastDiffusionModel:
"""transformers-only slow path for text-diffusion models."""
@staticmethod
def from_pretrained(
model_name = "google/diffusiongemma-26B-A4B-it",
max_seq_length = None, # API-compat; diffusion uses canvas_length
dtype = None,
load_in_4bit = False,
load_in_8bit = False,
load_in_16bit = False,
full_finetuning = False,
token = None,
device_map = "auto",
trust_remote_code = False,
attn_implementation = "eager", # exact match with the reference golden logits
revision = None,
return_tokenizer = True,
**kwargs,
):
SUPPORTS_BFLOAT16 = is_bfloat16_supported()
if dtype is None:
dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
logger.warning_once("Device does not support bfloat16. Will change to float16.")
dtype = torch.float16
assert dtype in (torch.float16, torch.bfloat16, torch.float32)
# Honor an explicit local_files_only; else fall back to the offline env vars.
local_files_only = kwargs.pop("local_files_only", None)
if local_files_only is None:
local_files_only = (
os.environ.get("HF_HUB_OFFLINE", "0") == "1"
or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1"
)
config = _load_diffusion_config(
model_name,
token,
trust_remote_code,
revision,
local_files_only,
)
model_type = getattr(config, "model_type", None)
if not is_diffusion_model_type(model_type):
raise RuntimeError(
f"Unsloth: FastDiffusionModel only supports diffusion model_types {DIFFUSION_MODEL_TYPES}, "
f"got '{model_type}'. Use FastModel/FastLanguageModel for autoregressive models."
)
model_cls = _resolve_diffusion_model_class(config)
load_kwargs = dict(
dtype = dtype,
device_map = device_map,
token = token,
trust_remote_code = trust_remote_code,
attn_implementation = attn_implementation,
revision = revision,
local_files_only = local_files_only,
)
# Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips
# them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision.
if load_in_4bit or load_in_8bit:
from transformers import BitsAndBytesConfig
if load_in_4bit:
qcfg = BitsAndBytesConfig(
load_in_4bit = True,
bnb_4bit_use_double_quant = True,
bnb_4bit_quant_type = "nf4",
bnb_4bit_compute_dtype = dtype,
llm_int8_skip_modules = [
"lm_head",
"embed_tokens",
"experts",
"self_conditioning",
"router",
],
)
else:
qcfg = BitsAndBytesConfig(load_in_8bit = True)
load_kwargs["quantization_config"] = qcfg
print(f"==(( Unsloth: FastDiffusionModel (slow / transformers-only path) ))==")
print(f" Model: {model_name} | class: {model_cls.__name__} | model_type: {model_type}")
print(
f" dtype: {dtype} | 4bit: {load_in_4bit} | 8bit: {load_in_8bit} | attn: {attn_implementation}"
)
model = model_cls.from_pretrained(model_name, **load_kwargs).eval()
# Mark before any early return so get_peft_model/for_* route to the slow path.
model._unsloth_slow_diffusion = True
if not return_tokenizer:
return model, None
# Prefer the processor (chat template + tokenizer); fall back to a bare tokenizer. Returned as
# "tokenizer" to match the Unsloth (model, tokenizer) contract.
try:
tokenizer = AutoProcessor.from_pretrained(
model_name,
token = token,
trust_remote_code = trust_remote_code,
revision = revision,
local_files_only = local_files_only,
)
except Exception:
tokenizer = AutoTokenizer.from_pretrained(
model_name,
token = token,
trust_remote_code = trust_remote_code,
revision = revision,
local_files_only = local_files_only,
)
return model, tokenizer
@staticmethod
def get_peft_model(
model,
r = 16,
target_modules = None,
lora_alpha = 16,
lora_dropout = 0.0,
bias = "none",
use_gradient_checkpointing = True,
random_state = 3407,
task_type = None,
**kwargs,
):
"""Attach a PEFT LoRA to the diffusion backbone (attention + dense MLP). No fused kernels."""
from peft import LoraConfig, get_peft_model as peft_get_peft_model
if target_modules is None:
target_modules = DIFFUSION_LORA_TARGETS
lora_kwargs = dict(
r = r,
lora_alpha = lora_alpha,
lora_dropout = lora_dropout,
bias = bias,
target_modules = target_modules,
task_type = task_type, # None: diffusion has no standard CAUSAL_LM head
**{k: v for k, v in kwargs.items() if k in ("modules_to_save", "init_lora_weights")},
)
# Exclude the vision tower's custom (non-Linear) modules that share suffix names.
exclude = kwargs.get("exclude_modules", DIFFUSION_LORA_EXCLUDE)
try:
lora_config = LoraConfig(exclude_modules = exclude, **lora_kwargs)
except TypeError:
# Older PEFT without exclude_modules: scope the target to the text decoder by regex.
lora_kwargs["target_modules"] = (
r".*model\.decoder\.layers\.\d+\.(self_attn\.[qkvo]_proj|mlp\.(gate|up|down)_proj)"
)
lora_config = LoraConfig(**lora_kwargs)
if use_gradient_checkpointing:
model.gradient_checkpointing_enable()
if hasattr(model, "enable_input_require_grads"):
model.enable_input_require_grads()
model = peft_get_peft_model(model, lora_config)
model._unsloth_slow_diffusion = True
try:
model.print_trainable_parameters()
except Exception:
pass
return model
@staticmethod
def for_inference(model):
model.eval()
for _, m in model.named_modules():
if hasattr(m, "gradient_checkpointing"):
m.gradient_checkpointing = False
return model
@staticmethod
def for_training(model, use_gradient_checkpointing = True):
model.train()
if use_gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"):
model.gradient_checkpointing_enable()
return model

123
unsloth/models/llama.py
View file

@ -1622,6 +1622,93 @@ def _get_rope_theta(config, default = 10000.0):
return default
def _rope_scaling_as_dict(rope_scaling):
"""Normalize config.rope_scaling (dict or config object) to a dict; {} on failure."""
if isinstance(rope_scaling, dict):
return rope_scaling
for converter in ("to_dict", "dict"):
fn = getattr(rope_scaling, converter, None)
if callable(fn):
try:
d = fn()
if isinstance(d, dict):
return d
except Exception:
pass
try:
return {k: v for k, v in vars(rope_scaling).items() if not k.startswith("_")}
except TypeError:
return {}
def _llama3_inv_freq_from_config(
config,
rope_scaling,
device = "cpu",
):
"""llama3 inv_freq with factors from config; fallback when modeling_rope_utils is missing."""
base = _get_rope_theta(config, default = 10000.0)
dim = getattr(config, "head_dim", None)
if dim is None:
dim = int(config.hidden_size // config.num_attention_heads)
inv_freq = 1.0 / (
base ** (torch.arange(0, dim, 2, dtype = torch.int64, device = device).float() / dim)
)
scale_factor = rope_scaling.get("factor", 8.0)
low_freq_factor = rope_scaling.get("low_freq_factor", 1.0)
high_freq_factor = rope_scaling.get("high_freq_factor", 4.0)
old_context_len = rope_scaling.get("original_max_position_embeddings", 8192)
low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor
assert low_freq_wavelen != high_freq_wavelen
# Vectorized meta-llama bands: high freqs kept, low divided by factor, medium blended.
wavelen = 2 * math.pi / inv_freq
scaled = torch.where(wavelen > low_freq_wavelen, inv_freq / scale_factor, inv_freq)
smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
smoothed = (1 - smooth) * inv_freq / scale_factor + smooth * inv_freq
is_medium = (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen)
return torch.where(is_medium, smoothed, scaled)
def _compute_config_rope_inv_freq(config, rope_scaling):
"""(inv_freq, attention_scaling) per config.rope_scaling via transformers'
ROPE_INIT_FUNCTIONS, with an inline llama3 fallback; (None, 1.0) on failure."""
original_rope_scaling = rope_scaling
rope_scaling = _rope_scaling_as_dict(rope_scaling)
rope_type = rope_scaling.get("rope_type", None) or rope_scaling.get("type", None)
try:
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type]
try:
inv_freq, attention_scaling = rope_init_fn(config, torch.device("cpu"))
except Exception:
# Object-style rope_scaling: retry with a config copy carrying the plain dict.
if isinstance(original_rope_scaling, dict):
raise
import copy as _copy
config_copy = _copy.copy(config)
config_copy.rope_scaling = rope_scaling
inv_freq, attention_scaling = rope_init_fn(config_copy, torch.device("cpu"))
return inv_freq.to(dtype = torch.float32, device = "cpu"), float(attention_scaling)
except Exception as exception:
if rope_type == "llama3":
try:
return _llama3_inv_freq_from_config(config, rope_scaling), 1.0
except Exception:
pass
logger.warning_once(
f"Unsloth: Could not apply RoPE scaling '{rope_type}' from config "
f"({type(exception).__name__}: {exception}); falling back to unscaled RoPE. "
"Long-context generation may degrade."
)
return None, 1.0
# Solves https://github.com/unslothai/unsloth/issues/168
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
@ -1640,6 +1727,12 @@ class LlamaRotaryEmbedding(torch.nn.Module):
config = None, # [TODO] Hack to pass in config - need to remove later
):
super().__init__()
# cos/sin multiplier (1.0 except yarn / longrope); set before any cache build.
self.attention_scaling = 1.0
# Base-class-from-config path (modern transformers): derive inv_freq like
# transformers so config.rope_scaling is not dropped (#2405). Scaled
# subclasses are excluded to avoid double-scaling.
config_inv_freq = None
if config is not None:
# [TODO] Hack to pass in config - need to remove later
base = _get_rope_theta(config, default = base)
@ -1652,6 +1745,13 @@ class LlamaRotaryEmbedding(torch.nn.Module):
device = DEVICE_TYPE_TORCH
max_position_embeddings = config.max_position_embeddings
rope_scaling = getattr(config, "rope_scaling", None)
if rope_scaling is not None and type(self) is LlamaRotaryEmbedding:
config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq(
config,
rope_scaling,
)
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
@ -1660,12 +1760,17 @@ class LlamaRotaryEmbedding(torch.nn.Module):
self.multi_gpu_cos_cached = [None] * DEVICE_COUNT
self.multi_gpu_sin_cached = [None] * DEVICE_COUNT
# Normal Llama-3 RoPE
inv_freq = 1.0 / (
self.base
** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim)
)
inv_freq = self._apply_inv_freq_scaling(inv_freq)
if config_inv_freq is not None:
inv_freq = config_inv_freq # already scaled; skip subclass scaling
else:
# Normal Llama-3 RoPE
inv_freq = 1.0 / (
self.base
** (
torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim
)
)
inv_freq = self._apply_inv_freq_scaling(inv_freq)
self.register_buffer("inv_freq", inv_freq, persistent = False)
# Build here to make `torch.jit.trace` work.
@ -1704,8 +1809,10 @@ class LlamaRotaryEmbedding(torch.nn.Module):
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim = -1)
cos = emb.cos().to(dtype = dtype, device = device, non_blocking = True)
sin = emb.sin().to(dtype = dtype, device = device, non_blocking = True)
# Applied here so attention_scaling survives extend_rope_embedding rebuilds;
# default 1.0 keeps unscaled paths bit-identical.
cos = (emb.cos() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True)
sin = (emb.sin() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True)
self.multi_gpu_cos_cached[device.index] = cos
self.multi_gpu_sin_cached[device.index] = sin
return cos, sin

51
unsloth/models/loader.py
View file

@ -829,6 +829,7 @@ from ..kernels import (
post_patch_loss_function,
)
from .vision import FastBaseModel
from .diffusion import FastDiffusionModel, is_diffusion_model_type
from transformers import (
AutoModelForCausalLM,
)
@ -846,6 +847,25 @@ class FastModel(FastBaseModel):
model = _prepare_model_for_qat(model, qat_scheme)
return model
@staticmethod
def get_peft_model(model, *args, **kwargs):
# Route text-diffusion models (slow path) to the transformers-only PEFT helper.
if getattr(model, "_unsloth_slow_diffusion", False):
return FastDiffusionModel.get_peft_model(model, *args, **kwargs)
return FastBaseModel.get_peft_model(model, *args, **kwargs)
@staticmethod
def for_inference(model):
if getattr(model, "_unsloth_slow_diffusion", False):
return FastDiffusionModel.for_inference(model)
return FastBaseModel.for_inference(model)
@staticmethod
def for_training(model, use_gradient_checkpointing = True):
if getattr(model, "_unsloth_slow_diffusion", False):
return FastDiffusionModel.for_training(model, use_gradient_checkpointing)
return FastBaseModel.for_training(model, use_gradient_checkpointing)
@staticmethod
def from_pretrained(
model_name = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
@ -1065,6 +1085,24 @@ class FastModel(FastBaseModel):
local_files_only = True
kwargs["local_files_only"] = True
# Text-diffusion slow-path dispatch, factored so both the normal route (below) and the
# legacy-config fallback (in the AutoConfig except handler) share one call site.
def _dispatch_diffusion():
return FastDiffusionModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
load_in_8bit = load_in_8bit,
load_in_16bit = load_in_16bit,
full_finetuning = full_finetuning,
token = token,
device_map = device_map,
trust_remote_code = trust_remote_code,
revision = revision,
**kwargs,
)
try:
model_config = AutoConfig.from_pretrained(
model_name,
@ -1078,6 +1116,12 @@ class FastModel(FastBaseModel):
raise
except Exception as error:
autoconfig_error = str(error)
# Legacy text-diffusion configs use model_type "diffusion_gemma", which current
# transformers does not register by name (it ships "diffusion_gemma4"). AutoConfig
# raises before we can dispatch; route straight to the diffusion slow path, whose
# loader aliases the legacy type to the gemma4 classes.
if "diffusion_gemma" in autoconfig_error and is_diffusion_model_type("diffusion_gemma"):
return _dispatch_diffusion()
if "architecture" in autoconfig_error:
if "qwen3_5" in autoconfig_error:
raise ImportError(
@ -1126,6 +1170,13 @@ class FastModel(FastBaseModel):
)
model_types_all = ",".join(model_types) + ","
# ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ----
# These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's
# autoregressive kernel/compile patching and load the unmodified HF model (bit-identical to
# naive transformers), keeping only 4bit/8bit + PEFT LoRA conveniences.
if is_diffusion_model_type(model_types):
return _dispatch_diffusion()
# Save model types and loading method
lowered_model_name = model_name.lower()
string = os.environ.get("UNSLOTH_MODEL_NAME", "") + model_types_all

12
unsloth/models/vision.py
View file

@ -283,6 +283,9 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
input_ids = kwargs["input"]
elif "input_features" in kwargs:
input_ids = kwargs["input_features"]
elif "inputs_embeds" in kwargs:
# canonical HF name for embedding inputs (e.g. multimodal generate)
input_ids = kwargs["inputs_embeds"]
elif "input_embeds" in kwargs:
input_ids = kwargs["input_embeds"]
elif "inputs" in kwargs:
@ -1280,6 +1283,15 @@ class FastBaseModel:
tokenizer.padding_side = "left" # Force inference
if hasattr(tokenizer, "tokenizer"):
tokenizer.tokenizer.padding_side = "left" # Force inference
# Audio feature extractors must stay right padded: left (a text setting,
# forwarded by from_pretrained) shifts Whisper mels and desyncs Gemma 4
# audio token counts (crash on transformers < 5.10).
feature_extractor = getattr(tokenizer, "feature_extractor", None)
if (
feature_extractor is not None
and getattr(feature_extractor, "padding_side", None) == "left"
):
feature_extractor.padding_side = "right"
m = model
while hasattr(m, "model"):
m.max_seq_length = max_seq_length

4
unsloth/save.py
View file

@ -1092,7 +1092,6 @@ def unsloth_save_model(
gc.collect()
# Remove temporary location
import shutil
shutil.rmtree(temporary_location, ignore_errors = True)
@ -1224,7 +1223,6 @@ def install_llama_cpp_old(version = -10):
for i in range(30):
print(f"**[WARNING]** Deleting llama.cpp directory... {30-i} seconds left.")
time.sleep(1)
import shutil
shutil.rmtree("llama.cpp", ignore_errors = True)
@ -2494,7 +2492,6 @@ def unsloth_push_to_hub_gguf(
except Exception as e:
if cleanup_temp:
import shutil
for d in [save_directory, f"{save_directory}_gguf"]:
try:
shutil.rmtree(d)
@ -2681,7 +2678,6 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi
# Clean up temporary directory
if cleanup_temp:
print("Unsloth: Cleaning up temporary files...")
import shutil
for d in [save_directory, f"{save_directory}_gguf"]:
if os.path.exists(d):
try:

22
unsloth_cli/__init__.py
View file

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import os.path as _osp
import os as _os
import sys as _sys
import typer
@ -10,6 +10,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError
from unsloth_cli.commands.train import train
from unsloth_cli.commands.inference import inference
from unsloth_cli.commands.chat import chat
from unsloth_cli.commands.export import export, list_checkpoints
from unsloth_cli.commands.studio import (
run as studio_run,
@ -20,7 +21,7 @@ from unsloth_cli.commands.studio import (
# Canonicalise `-np<N>` only under the `unsloth` console-script;
# third-party scripts that import unsloth_cli keep their argv intact.
_entry_base = _osp.basename(_sys.argv[0]).lower() if _sys.argv else ""
_entry_base = _os.path.basename(_sys.argv[0]).lower() if _sys.argv else ""
if _entry_base in {"unsloth", "unsloth.exe"}:
_expand_attached_np_short()
del _entry_base
@ -53,11 +54,26 @@ def main(
help = "Show version and exit.",
),
):
pass
if (
_sys.platform == "win32"
): # this block catches unsloth running inside of System32 or any subdirs, this WILL cause errors if not prevented.
_cwd = _os.path.normcase(_os.path.normpath(_os.getcwd()))
_system32 = _os.path.normcase(
_os.path.normpath(_os.path.join(_os.environ.get("WINDIR", r"C:\Windows"), "System32"))
)
if _cwd == _system32 or _cwd.startswith(_system32 + _os.sep):
typer.secho(
"Refusing to run Unsloth inside System32 as it will lead to Errors.\n"
"cd to a normal working directory and try again.",
fg = "red",
err = True,
)
raise typer.Exit(code = 1)
app.command()(train)
app.command()(inference)
app.command()(chat)
app.command()(export)
app.command("list-checkpoints")(list_checkpoints)
app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")

417
unsloth_cli/_inference.py Normal file
View file

@ -0,0 +1,417 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Model loading and streaming shared by `inference` and `chat`."""
import os
import re
import sys
from pathlib import Path
from typing import Optional
import typer
_THINK_OPEN = "<think>"
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", re.DOTALL)
def ensure_studio_backend_path() -> None:
backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend")
if backend_dir not in sys.path:
sys.path.insert(0, backend_dir)
def configure_quiet_logging() -> None:
import logging
import structlog
# The CLI never configures structlog, so without this every backend INFO
# line prints. LOG_LEVEL is exported so the worker subprocess inherits it.
level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper()
level = getattr(logging, level_name, logging.WARNING)
structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level))
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
def visible_text(text: str, show_thinking: bool) -> str:
if show_thinking:
return text
text = _THINK_BLOCK.sub("", text)
# Hold back an unclosed trailing <think> so reasoning never leaks mid-stream.
open_idx = text.find(_THINK_OPEN)
if open_idx != -1:
text = text[:open_idx]
max_prefix = min(len(text), len(_THINK_OPEN) - 1)
for size in range(max_prefix, 0, -1):
if _THINK_OPEN.startswith(text[-size:]):
return text[:-size]
return text
def stream_to_stdout(stream, show_thinking: bool) -> str:
# Backends yield the full text-so-far on each step (llama.cpp ends with a
# metadata dict, skipped); print the growing tail, return the raw text.
raw = ""
shown = ""
for chunk in stream:
if not isinstance(chunk, str):
continue
raw = chunk
rendered = visible_text(chunk, show_thinking)
delta = rendered[len(shown) :]
if delta:
sys.stdout.write(delta)
sys.stdout.flush()
shown = rendered
sys.stdout.write("\n")
sys.stdout.flush()
return raw
def stream_markdown(stream, show_thinking: bool, *, console) -> str:
from rich.live import Live
from rich.markdown import Markdown
from rich.text import Text
raw = ""
with Live(console = console, refresh_per_second = 12, vertical_overflow = "visible") as live:
for chunk in stream:
if not isinstance(chunk, str):
continue
raw = chunk
visible = visible_text(chunk, show_thinking)
live.update(Markdown(visible) if visible.strip() else Text(""))
return raw
def collect_stream(stream, show_thinking: bool) -> str:
raw = ""
for chunk in stream:
if isinstance(chunk, str):
raw = chunk
return visible_text(raw, show_thinking)
def render_columns(
left_label: str,
left_text: str,
right_label: str,
right_text: str,
*,
console = None,
) -> None:
from rich import box
from rich.console import Console
from rich.table import Table
table = Table(box = box.MINIMAL, expand = True, padding = (0, 1), pad_edge = False)
table.add_column(left_label, header_style = "bold yellow", ratio = 1, overflow = "fold")
table.add_column(right_label, header_style = "bold magenta", ratio = 1, overflow = "fold")
table.add_row(left_text or "", right_text or "")
(console or Console()).print(table)
class ChatBackend:
"""Uniform stream()/close() over the llama-server and Unsloth backends."""
def __init__(self, kind: str, backend) -> None:
self._kind = kind # "gguf" | "unsloth"
self._backend = backend
def stream(
self,
messages: list,
*,
system_prompt: str,
temperature: float,
top_p: float,
top_k: int,
max_new_tokens: int,
repetition_penalty: float,
enable_thinking: bool,
use_adapter: Optional[bool] = None,
):
if self._kind == "gguf":
# llama-server takes the system prompt as the first message.
msgs = list(messages)
if system_prompt:
msgs = [{"role": "system", "content": system_prompt}, *msgs]
return self._backend.generate_chat_completion(
messages = msgs,
temperature = temperature,
top_p = top_p,
top_k = top_k,
max_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
enable_thinking = enable_thinking,
)
gen_kwargs = dict(
messages = messages,
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
enable_thinking = enable_thinking,
)
if use_adapter is not None:
return self._backend.generate_with_adapter_control(
use_adapter = use_adapter, **gen_kwargs
)
return self._backend.generate_chat_response(**gen_kwargs)
def close(self) -> None:
# Shut the worker down directly: the graceful unload_model waits for
# an ack that compare mode can swallow, hanging exit for minutes.
try:
if self._kind == "gguf":
self._backend.unload_model()
else:
self._backend._shutdown_subprocess(timeout = 2.0)
except Exception:
pass
def resolve_model_config(model: str, *, hf_token: Optional[str]):
ensure_studio_backend_path()
from utils.models import ModelConfig
model_config = ModelConfig.from_identifier(model_id = model, hf_token = hf_token)
if not model_config:
typer.echo("Could not resolve model config", err = True)
raise typer.Exit(code = 1)
return model_config
def _load_gguf_backend(model_config, *, hf_token, max_seq_length):
ensure_studio_backend_path()
from core.inference.llama_cpp import LlamaCppBackend
llama_backend = LlamaCppBackend()
common = dict(
hf_variant = model_config.gguf_variant,
model_identifier = model_config.identifier,
is_vision = model_config.is_vision,
n_ctx = max_seq_length,
)
if model_config.gguf_hf_repo:
loaded = llama_backend.load_model(
hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common
)
else:
loaded = llama_backend.load_model(
gguf_path = model_config.gguf_file,
mmproj_path = model_config.gguf_mmproj_file,
mtp_draft_path = model_config.gguf_mtp_file,
**common,
)
if not loaded:
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
return ChatBackend("gguf", llama_backend)
def load_chat_backend(
model: str,
*,
hf_token: Optional[str],
max_seq_length: int,
load_in_4bit: bool,
model_config = None,
fresh_backend: bool = False,
):
"""Load `model` in-process: GGUF via llama-server, else the orchestrator.
fresh_backend uses a private orchestrator so a second model (compare's
base column) can run alongside the main one.
"""
if model_config is None:
model_config = resolve_model_config(model, hf_token = hf_token)
typer.echo(f"Loading {model}", err = True)
if model_config.is_gguf:
return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length)
if fresh_backend:
ensure_studio_backend_path()
from core.inference import InferenceOrchestrator
backend = InferenceOrchestrator()
else:
ensure_studio_backend_path()
from core.inference import get_inference_backend
backend = get_inference_backend()
if not backend.load_model(
config = model_config,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
hf_token = hf_token,
):
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
return ChatBackend("unsloth", backend)
def find_studio_server(timeout: float = 0.4) -> Optional[str]:
import urllib.request
base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/")
try:
with urllib.request.urlopen(f"{base}/api/health", timeout = timeout):
return base
except Exception:
return None
def _studio_token() -> Optional[str]:
"""Self-issue a JWT: the CLI runs as the same OS user as the server, so it
signs with the same stored secret the server validates against."""
try:
import studio.backend.core # noqa: F401 puts studio/backend on sys.path
from studio.backend.auth import storage
from studio.backend.auth.authentication import create_access_token
row = storage.get_connection().execute("SELECT username FROM auth_user LIMIT 1").fetchone()
return create_access_token(row[0], desktop = True) if row else None
except Exception:
return None
class HttpChatBackend:
"""Chat against a running Studio server over its OpenAI-compatible API.
close() leaves the model loaded on purpose — the next session (or the
UI) starts instantly.
"""
def __init__(self, base_url: str, token: str) -> None:
self._base = base_url
self._token = token
def _request(
self,
method: str,
path: str,
payload = None,
timeout = None,
):
import json
import urllib.request
request = urllib.request.Request(
self._base + path,
data = None if payload is None else json.dumps(payload).encode(),
headers = {
"Authorization": f"Bearer {self._token}",
"Content-Type": "application/json",
},
method = method,
)
return urllib.request.urlopen(request, timeout = timeout)
def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None:
typer.echo(f"Loading {model} on the Studio server", err = True)
try:
self._request(
"POST",
"/api/inference/load",
{
"model_path": model,
"hf_token": hf_token,
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
},
).close()
except Exception as exc:
typer.echo(f"Model load failed: {exc}", err = True)
raise typer.Exit(code = 1)
def stream(
self,
messages: list,
*,
system_prompt: str,
temperature: float,
top_p: float,
top_k: int,
max_new_tokens: int,
repetition_penalty: float,
enable_thinking: bool,
use_adapter: Optional[bool] = None,
):
import json
msgs = list(messages)
if system_prompt:
msgs = [{"role": "system", "content": system_prompt}, *msgs]
resp = self._request(
"POST",
"/v1/chat/completions",
{
"model": "default",
"messages": msgs,
"stream": True,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"max_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
"enable_thinking": enable_thinking,
},
)
def cumulative():
# Accumulate SSE deltas into the full-text-so-far convention the
# stream helpers expect.
text = ""
with resp:
for raw_line in resp:
line = raw_line.decode("utf-8", "replace").strip()
if not line.startswith("data:"):
continue
data = line[len("data:") :].strip()
if data == "[DONE]":
break
try:
parsed = json.loads(data)
except ValueError:
continue
if "error" in parsed:
raise RuntimeError(
f"Server error: {parsed['error'].get('message', 'Unknown server error')}"
)
try:
delta = parsed["choices"][0]["delta"].get("content")
except (KeyError, IndexError):
continue
if not delta:
continue
text += delta
# An emoji can arrive split across two deltas as lone
# surrogate halves: hold back a trailing half, merge pairs.
visible = text
if "\ud800" <= visible[-1] <= "\udbff":
visible = visible[:-1]
yield visible.encode("utf-16", "surrogatepass").decode("utf-16", "replace")
return cumulative()
def close(self) -> None:
pass
def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit):
"""Backend on a running Studio server, or None (caller loads locally)."""
base_url = find_studio_server()
if not base_url:
return None
token = _studio_token()
if not token:
return None
backend = HttpChatBackend(base_url, token)
backend.ensure_loaded(
model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit
)
return backend

340
unsloth_cli/commands/chat.py Normal file
View file

@ -0,0 +1,340 @@
# 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 typing import Optional
import typer
from rich.console import Console
from unsloth_cli._inference import (
collect_stream,
configure_quiet_logging,
connect_studio_server,
ensure_studio_backend_path,
load_chat_backend,
render_columns,
resolve_model_config,
stream_markdown,
visible_text,
)
_HELP = (
"Commands: /exit (quit), /reset (clear history), "
"/think (toggle reasoning), /compare (base vs tuned), /help"
)
def _you_prompt(colors: bool) -> str:
# The prompt must go through input(), not a separate print — readline
# redraws erase anything they didn't draw, eating the label. GNU readline
# wants colors wrapped in \001/\002; libedit (macOS) prints those
# literally, so it gets raw ANSI.
try:
import readline
except ImportError:
return "\n\x1b[1;36mYou: \x1b[0m" if colors else "\nYou: "
libedit = (
"libedit" in (readline.__doc__ or "") or getattr(readline, "backend", "") == "editline"
)
if not colors:
return "\nYou: "
if libedit:
return "\n\x1b[1;36mYou: \x1b[0m"
return "\n\001\x1b[1;36m\002You: \001\x1b[0m\002"
def _compare_blocked_reason(model_config) -> Optional[str]:
if model_config.is_gguf:
return (
"GGUF models can't toggle adapters — load a LoRA fine-tune "
"(transformers backend) to compare base vs tuned."
)
if not model_config.is_lora:
return (
"this isn't a LoRA adapter — compare turns the adapter off for the "
"'base' column, so there's nothing to compare against."
)
return None
def _get_base_load_in_4bit(model_config) -> bool:
"""Determine load_in_4bit for base model based on tuned adapter precision."""
if not model_config.is_lora or not model_config.path:
# Fallback to default if not a LoRA or no path
return True
try:
import json
from pathlib import Path
adapter_cfg_path = Path(model_config.path) / "adapter_config.json"
if not adapter_cfg_path.exists():
return True
with open(adapter_cfg_path) as f:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora":
return False
elif training_method == "qlora":
return True
elif not training_method:
# Fallback: check base model name for -bnb-4bit suffix
if model_config.base_model and "-bnb-4bit" not in model_config.base_model.lower():
return False
return True
return True
except Exception:
return True
def _compare_needs_second_model() -> bool:
# MLX can't toggle the adapter off, so compare loads the base separately.
# detect_hardware() would print into the chat (and import torch), so
# probe its MLX condition quietly: Apple Silicon with mlx installed.
try:
from studio.backend.utils.hardware import hardware as hw
if hw.DEVICE is not None:
return hw.DEVICE == hw.DeviceType.MLX
if not hw.is_apple_silicon():
return False
import mlx.core # noqa: F401
return True
except Exception:
return False
def _pick_trained_model(console) -> str:
ensure_studio_backend_path()
from utils.models import scan_trained_models
trained = scan_trained_models()
if not trained:
typer.echo(
"No trained models found in your outputs folder. "
"Pass a model id or path: `unsloth chat <model>`.",
err = True,
)
raise typer.Exit(code = 1)
console.print("Your trained models (newest first):", style = "bold")
for i, (display_name, _, model_type) in enumerate(trained, 1):
console.print(f" {i}. {display_name} ({model_type})", markup = False)
while True:
try:
raw = input(f"Chat with [1-{len(trained)}, Enter = 1]: ").strip()
except (EOFError, KeyboardInterrupt):
raise typer.Exit(code = 1)
if not raw:
return trained[0][1]
if raw.isdigit() and 1 <= int(raw) <= len(trained):
return trained[int(raw) - 1][1]
console.print(f"Pick a number between 1 and {len(trained)}.", style = "yellow")
def chat(
model: Optional[str] = typer.Argument(
None, help = "HF model id or local path. Omit to pick one of your trained models."
),
hf_token: Optional[str] = typer.Option(
None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed."
),
temperature: float = typer.Option(0.7, "--temperature"),
top_p: float = typer.Option(0.9, "--top-p"),
top_k: int = typer.Option(40, "--top-k"),
max_new_tokens: int = typer.Option(512, "--max-new-tokens"),
repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"),
system_prompt: str = typer.Option(
"", "--system-prompt", help = "Optional system prompt for the conversation."
),
max_seq_length: int = typer.Option(4096, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
think: bool = typer.Option(
False,
"--think/--no-think",
help = "Start with the model's <think> reasoning shown. Toggle live with /think.",
),
compare: bool = typer.Option(
False,
"--compare/--no-compare",
help = "Answer each prompt twice — base vs fine-tuned — side by side. "
"Needs a LoRA adapter. Toggle live with /compare.",
),
verbose: bool = typer.Option(
False, "--verbose", "-v", help = "Show backend and llama-server logs."
),
no_server: bool = typer.Option(
False,
"--no-server",
help = "Load the model in-process even if a Studio server is running.",
),
):
"""Start an interactive chat with a model (loads once, stays warm)."""
if not verbose:
configure_quiet_logging()
console = Console()
err = Console(stderr = True)
if model is None:
model = _pick_trained_model(console)
# Resolve first so --compare can be rejected before the slow load.
model_config = resolve_model_config(model, hf_token = hf_token)
compare_blocked = _compare_blocked_reason(model_config)
if compare and compare_blocked:
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
raise typer.Exit(code = 1)
load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
# Prefer a running Studio server: instant starts, model shared with the UI.
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
server_mode = chat_backend is not None
if server_mode:
console.print(
"(Studio server connected — model stays warm after /exit)",
style = "bright_black",
)
else:
chat_backend = load_chat_backend(model, model_config = model_config, **load_opts)
name = model_config.display_name or model
show_thinking = think
compare_mode = compare
messages = []
# Compare's base column: server mode keeps the tuned model remote and
# loads the base locally; local MLX (no adapter toggle) does the same;
# local CUDA just toggles the adapter on the one loaded model.
dual_compare = compare_blocked is None and (server_mode or _compare_needs_second_model())
base_backend = None
def load_base_for_compare():
nonlocal base_backend
if base_backend is not None:
return True
base_id = model_config.base_model
if not base_id:
console.print(
"(compare unavailable: this adapter doesn't record its base model)",
style = "yellow",
)
return False
console.print(
f"(loading base model {base_id} for compare — keeps two models in memory)",
style = "bright_black",
markup = False,
)
try:
# Use the same precision as the tuned model for fair comparison
base_load_opts = dict(load_opts) # Copy original options
base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config)
base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts)
except Exception as exc:
err.print(f"(base model load failed: {exc})", style = "red", markup = False)
return False
return True
if compare and dual_compare and not load_base_for_compare():
raise typer.Exit(code = 1)
def generate(backend = None, use_adapter = None):
# Reads messages and show_thinking live, so /reset and /think apply.
return (backend or chat_backend).stream(
messages,
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
enable_thinking = show_thinking,
use_adapter = use_adapter,
)
console.print()
console.print(f"Chatting with {name}", style = "bold green", markup = False)
console.print(_HELP, style = "bright_black")
# legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage.
you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows)
assistant_label = "[bold magenta]Assistant:[/bold magenta]"
try:
while True:
try:
user = input(you_prompt).strip()
except (EOFError, KeyboardInterrupt):
console.print()
break
if not user:
continue
if user in ("/exit", "/quit"):
break
if user == "/reset":
messages = []
console.print("(history cleared)", style = "bright_black")
continue
if user == "/think":
show_thinking = not show_thinking
state = "on" if show_thinking else "off"
console.print(f"(thinking {state})", style = "bright_black")
continue
if user == "/compare":
if compare_blocked:
console.print(f"(compare unavailable: {compare_blocked})", style = "yellow")
continue
if not compare_mode and dual_compare and not load_base_for_compare():
continue
compare_mode = not compare_mode
state = "on" if compare_mode else "off"
console.print(f"(compare {state})", style = "bright_black")
continue
if user in ("/help", "/?"):
console.print(_HELP, style = "bright_black")
continue
messages.append({"role": "user", "content": user})
try:
if compare_mode:
console.print("(comparing base vs tuned…)", style = "bright_black")
if dual_compare:
base_text = collect_stream(generate(backend = base_backend), show_thinking)
tuned_text = collect_stream(generate(), show_thinking)
else:
base_text = collect_stream(generate(use_adapter = False), show_thinking)
tuned_text = collect_stream(generate(use_adapter = True), show_thinking)
console.print()
render_columns(
"base", base_text, f"{name} (tuned)", tuned_text, console = console
)
# History continues as the tuned model; base is just the reference.
answer = tuned_text
else:
console.print(assistant_label)
answer = stream_markdown(generate(), show_thinking, console = console)
except KeyboardInterrupt:
# Ctrl-C aborts this answer only; drop the unanswered turn.
console.print("\n(interrupted)", style = "bright_black")
messages.pop()
continue
except Exception as exc:
err.print(f"\n(error: {exc})", style = "red", markup = False)
messages.pop()
continue
messages.append(
{"role": "assistant", "content": visible_text(answer, show_thinking = False)}
)
finally:
chat_backend.close()
if base_backend is not None:
base_backend.close()
err.print("\nBye.", style = "bright_black")

87
unsloth_cli/commands/inference.py
View file

@ -1,11 +1,17 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import sys
from typing import Optional
import typer
from unsloth_cli._inference import (
configure_quiet_logging,
connect_studio_server,
load_chat_backend,
stream_to_stdout,
)
def inference(
model: str = typer.Argument(..., help = "HF model id or local path."),
@ -25,45 +31,46 @@ def inference(
),
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
think: bool = typer.Option(
False,
"--think/--no-think",
help = "Show the model's <think> reasoning. Off by default so reasoning "
"models answer directly instead of spending the token budget thinking.",
),
verbose: bool = typer.Option(
False,
"--verbose",
"-v",
help = "Show backend and llama-server logs (otherwise only the answer).",
),
no_server: bool = typer.Option(
False,
"--no-server",
help = "Load the model in-process even if a Studio server is running.",
),
):
"""Run a single inference using the specified model."""
from studio.backend.core import ModelConfig, get_inference_backend
if not verbose:
configure_quiet_logging()
inference_backend = get_inference_backend()
model_config = ModelConfig.from_ui_selection(
dropdown_value = model, search_value = None, hf_token = hf_token, is_lora = False
)
if not model_config:
typer.echo("Could not resolve model config", err = True)
raise typer.Exit(code = 1)
if not inference_backend.load_model(
config = model_config,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
hf_token = hf_token,
):
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
messages = [{"role": "user", "content": prompt}]
stream = inference_backend.generate_chat_response(
messages = messages,
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
)
typer.echo("Assistant:", nl = True)
previous = ""
for chunk in stream:
delta = chunk[len(previous) :]
if delta:
sys.stdout.write(delta)
sys.stdout.flush()
previous = chunk
sys.stdout.write("\n")
sys.stdout.flush()
# A running Studio server keeps the model warm between runs, which is
# exactly what a one-shot command wants.
load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
if chat_backend is None:
chat_backend = load_chat_backend(model, **load_opts)
try:
stream = chat_backend.stream(
[{"role": "user", "content": prompt}],
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
enable_thinking = think,
)
typer.echo("Assistant:")
stream_to_stdout(stream, show_thinking = think)
finally:
chat_backend.close()

96
unsloth_cli/commands/studio.py
View file

@ -597,6 +597,11 @@ def studio_default(
f"defaults to {_PARALLEL_DEFAULT_RUN}."
),
),
cloudflare: bool = typer.Option(
True,
"--cloudflare/--no-cloudflare",
help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).",
),
):
"""Launch the Unsloth Studio server."""
# Runs before every subcommand (run/setup/update/...).
@ -614,6 +619,16 @@ def studio_default(
err = True,
)
raise typer.Exit(2)
# Same for --no-cloudflare: it would not reach the subcommand.
if not cloudflare:
typer.echo(
f"Error: --no-cloudflare on `unsloth studio` applies to the "
f"plain-server path only. For `unsloth studio "
f"{ctx.invoked_subcommand}`, put it after the subcommand: "
f"`unsloth studio {ctx.invoked_subcommand} --no-cloudflare ...`",
err = True,
)
raise typer.Exit(2)
return
# Use the studio venv if it exists and we aren't already in it.
@ -648,6 +663,8 @@ def studio_default(
args.append("--silent")
if api_only:
args.append("--api-only")
# Forward the explicit polarity (matches run.py's BooleanOptionalAction).
args.append("--cloudflare" if cloudflare else "--no-cloudflare")
# On Windows os.execvp keeps the parent alive, so Ctrl+C
# would orphan the child; use Popen+wait instead.
if sys.platform == "win32":
@ -689,6 +706,7 @@ def studio_default(
silent = silent,
api_only = api_only,
llama_parallel_slots = parallel,
cloudflare = cloudflare,
)
if frontend is not None:
run_kwargs["frontend_path"] = frontend
@ -861,6 +879,11 @@ def run(
f"{_PARALLEL_DEFAULT_RUN} (pre-PR hardcoded value)."
),
),
cloudflare: bool = typer.Option(
True,
"--cloudflare/--no-cloudflare",
help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).",
),
):
"""Start Studio, load a model, print an API key -- one-liner server.
@ -980,6 +1003,8 @@ def run(
# Typer claims --parallel outside ctx.args; without this the
# child reverts to its default and silently drops the value.
args.extend(["--parallel", str(parallel)])
# Forward the explicit polarity (same rationale as --load-in-4bit above).
args.append("--cloudflare" if cloudflare else "--no-cloudflare")
# llama-server pass-through extras → child ctx.args → load payload.
if extra_llama_args:
args.extend(extra_llama_args)
@ -997,7 +1022,13 @@ def run(
# ── 2. Start server (always suppress built-in banner) ─────────────
from studio.backend.run import run_server, _resolve_external_ip
run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = parallel)
run_kwargs = dict(
host = host,
port = port,
silent = True,
llama_parallel_slots = parallel,
cloudflare = cloudflare,
)
if frontend is not None:
run_kwargs["frontend_path"] = frontend
app = run_server(**run_kwargs)
@ -1011,32 +1042,41 @@ def run(
set_tool_policy(enable_tools)
# 3. Wait for server health.
if not silent:
typer.echo("Starting Unsloth Studio...")
if not _wait_for_server(actual_port):
typer.echo("Error: server did not become healthy within 30 seconds.", err = True)
raise typer.Exit(1)
# Steps 3-5 can abort (health timeout, model-load error, or Ctrl+C during the
# slow load); tear the server and its children (llama-server, cloudflared) down
# on any abort so they never orphan.
from studio.backend.run import _graceful_shutdown, _server
# 4. Create API key in-process.
api_key = _create_api_key_inprocess(api_key_name)
# 5. Load model via HTTP.
if not silent:
typer.echo(f"Loading model: {model}...")
try:
result = _load_model_via_http(
port = actual_port,
api_key = api_key,
model = model,
gguf_variant = gguf_variant,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
llama_extra_args = extra_llama_args,
)
except RuntimeError as exc:
typer.echo(f"Error: {exc}", err = True)
raise typer.Exit(1)
# 3. Wait for server health.
if not silent:
typer.echo("Starting Unsloth Studio...")
if not _wait_for_server(actual_port):
typer.echo("Error: server did not become healthy within 30 seconds.", err = True)
raise typer.Exit(1)
# 4. Create API key in-process.
api_key = _create_api_key_inprocess(api_key_name)
# 5. Load model via HTTP.
if not silent:
typer.echo(f"Loading model: {model}...")
try:
result = _load_model_via_http(
port = actual_port,
api_key = api_key,
model = model,
gguf_variant = gguf_variant,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
llama_extra_args = extra_llama_args,
)
except RuntimeError as exc:
typer.echo(f"Error: {exc}", err = True)
raise typer.Exit(1)
except BaseException:
_graceful_shutdown(_server)
raise
loaded_model = result.get("model", model)
display_variant = f" ({gguf_variant})" if gguf_variant else ""
@ -1045,6 +1085,8 @@ def run(
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
base_url = f"http://{display_host}:{actual_port}"
sdk_base_url = f"{base_url}/v1"
# run_server started the tunnel during the silent run above (0.0.0.0 only).
_cf_url = getattr(app.state, "cloudflare_url", None)
# Orange so the tool-policy notice stands out; printed under
# --silent / --yes too so the policy is never invisible.
@ -1074,6 +1116,8 @@ def run(
typer.echo("")
typer.echo("=" * 56)
typer.echo(f" Unsloth Studio running at {base_url}")
if _cf_url:
typer.echo(f" Secure link access via Cloudflare: {_cf_url}")
typer.echo(f" Model loaded: {loaded_model}{display_variant}")
typer.echo(f" API Key: {api_key}")
typer.echo("")
@ -1107,6 +1151,8 @@ def run(
else:
# Silent still prints URL + API key + tool-status policy.
typer.echo(f"URL: {base_url}")
if _cf_url:
typer.echo(f"Secure link access via Cloudflare: {_cf_url}")
typer.echo(f"API Key: {api_key}")
typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True)

401
unsloth_cli/tests/test_inference_chat.py Normal file
View file

@ -0,0 +1,401 @@
# 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 `unsloth chat` / `unsloth inference` CLI — fakes only, no model loads."""
from __future__ import annotations
import inspect
import sys
import types
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
sys.path.insert(0, str(_REPO_ROOT))
import typer
from rich.console import Console
from typer.testing import CliRunner
import unsloth_cli.commands.chat as chatmod
from unsloth_cli._inference import (
ChatBackend,
HttpChatBackend,
collect_stream,
render_columns,
visible_text,
)
class _FakeConfig:
is_gguf = False
is_lora = True
display_name = "fake-model"
base_model = "fake/base"
path = None
def _chat_app():
cli = typer.Typer()
cli.command()(chatmod.chat)
return cli
def test_visible_text_passthrough_when_shown():
text = "<think>reasoning</think>answer"
assert visible_text(text, show_thinking = True) == text
def test_visible_text_strips_closed_think_block():
text = "<think>step 1\nstep 2</think>The answer is 42."
assert visible_text(text, show_thinking = False) == "The answer is 42."
def test_visible_text_holds_unclosed_think():
# An open <think> is held back so partial reasoning never leaks mid-stream.
assert visible_text("<think>still thinking", show_thinking = False) == ""
assert visible_text("done.<think>more thinking", show_thinking = False) == "done."
def test_visible_text_holds_partial_think_prefix():
# Streams are cumulative, so the opening tag can arrive as "<", "<thi",
# then "<think>". Hold possible tag prefixes until they are disambiguated.
assert visible_text("<", show_thinking = False) == ""
assert visible_text("<thi", show_thinking = False) == ""
assert visible_text("done.<thi", show_thinking = False) == "done."
assert visible_text("2 < 3", show_thinking = False) == "2 < 3"
def _option(command_fn, name):
return inspect.signature(command_fn).parameters[name].default
def test_inference_think_defaults_off():
from unsloth_cli.commands.inference import inference
opt = _option(inference, "think")
assert getattr(opt, "default", None) is False
# typer stores a flag/--no-flag pair as one combined decl.
assert "--think/--no-think" in (getattr(opt, "param_decls", None) or [])
def test_chat_command_is_registered_with_options():
params = inspect.signature(chatmod.chat).parameters
assert "model" in params
think = _option(chatmod.chat, "think")
assert "--think/--no-think" in (getattr(think, "param_decls", None) or [])
compare = _option(chatmod.chat, "compare")
assert "--compare/--no-compare" in (getattr(compare, "param_decls", None) or [])
verbose = _option(chatmod.chat, "verbose")
assert {"--verbose", "-v"} <= set(getattr(verbose, "param_decls", None) or [])
class _FakeBackend:
def __init__(self):
self.calls = []
def generate_chat_response(self, **kwargs):
self.calls.append(("plain", None, kwargs))
return iter(["hi"])
def generate_with_adapter_control(self, *, use_adapter, **kwargs):
self.calls.append(("adapter", use_adapter, kwargs))
return iter(["hi"])
_STREAM_KWARGS = dict(
system_prompt = "",
temperature = 0.7,
top_p = 0.9,
top_k = 40,
max_new_tokens = 8,
repetition_penalty = 1.1,
enable_thinking = False,
)
def test_chatbackend_routes_compare_to_adapter_control():
fake = _FakeBackend()
backend = ChatBackend("unsloth", fake)
list(backend.stream([{"role": "user", "content": "x"}], use_adapter = False, **_STREAM_KWARGS))
list(backend.stream([{"role": "user", "content": "x"}], use_adapter = True, **_STREAM_KWARGS))
assert [(path, flag) for path, flag, _ in fake.calls] == [
("adapter", False),
("adapter", True),
]
def test_chatbackend_normal_path_skips_adapter_control():
fake = _FakeBackend()
backend = ChatBackend("unsloth", fake)
list(backend.stream([{"role": "user", "content": "x"}], **_STREAM_KWARGS))
assert fake.calls[0][0] == "plain"
def test_collect_stream_returns_last_cumulative_think_stripped():
stream = iter(["<think>r</think>hel", "<think>r</think>hello"])
assert collect_stream(stream, show_thinking = False) == "hello"
def test_render_columns_emits_both_answers_with_separator(capsys):
render_columns("base", "alpha", "tuned", "beta")
out = capsys.readouterr().out
assert "base" in out and "tuned" in out
assert "alpha" in out and "beta" in out
assert "│" in out
def test_you_prompt_matches_readline_backend(monkeypatch):
gnu = types.ModuleType("readline")
gnu.__doc__ = "Importing this module enables command line editing using GNU readline."
monkeypatch.setitem(sys.modules, "readline", gnu)
prompt = chatmod._you_prompt(colors = True)
assert "You: " in prompt and "\001" in prompt
libedit = types.ModuleType("readline")
libedit.__doc__ = "Importing this module enables command line editing using libedit readline."
monkeypatch.setitem(sys.modules, "readline", libedit)
assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m"
assert chatmod._you_prompt(colors = False) == "\nYou: "
# Windows: no readline module at all; the console's own line editing
# handles backspace, so plain ANSI color (no markers) is safe.
monkeypatch.setitem(sys.modules, "readline", None)
assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m"
assert chatmod._you_prompt(colors = False) == "\nYou: "
def test_chat_registered_on_app():
from unsloth_cli import app
# cmd.name is None until typer resolves it from the callback name.
names = {(cmd.name or cmd.callback.__name__) for cmd in app.registered_commands}
assert "chat" in names
def test_chat_exits_cleanly_on_slash_exit(monkeypatch):
closed = []
class _FakeChatBackend:
def stream(self, *a, **k):
return iter(["hello"])
def close(self):
closed.append(True)
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
runner = CliRunner()
for args in (["fake-model"], ["fake-model", "--compare"]):
closed.clear()
result = runner.invoke(_chat_app(), args, input = "hi\n/exit\n")
assert result.exit_code == 0, result.output
assert closed == [True]
assert "Bye." in result.output
# The prompt must go through input() (readline-safe), not a print.
assert "You: " in result.output
assert "You: You:" not in result.output
def test_pick_trained_model_lists_and_selects(monkeypatch):
fake_models = types.ModuleType("utils.models")
fake_models.scan_trained_models = lambda: [
("run-new", "outputs/run-new", "lora"),
("run-old", "outputs/run-old", "merged"),
]
monkeypatch.setitem(sys.modules, "utils.models", fake_models)
monkeypatch.setattr("builtins.input", lambda prompt = "": "2")
assert chatmod._pick_trained_model(Console()) == "outputs/run-old"
monkeypatch.setattr("builtins.input", lambda prompt = "": "")
assert chatmod._pick_trained_model(Console()) == "outputs/run-new"
def test_chat_no_arg_chats_with_picked_trained_model(monkeypatch):
class _FakeChatBackend:
def stream(self, *a, **k):
return iter(["hello"])
def close(self):
pass
resolved = []
monkeypatch.setattr(chatmod, "_pick_trained_model", lambda console: "outputs/run-42")
monkeypatch.setattr(
chatmod,
"resolve_model_config",
lambda model, **k: (resolved.append(model), _FakeConfig())[1],
)
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
result = CliRunner().invoke(_chat_app(), [], input = "/exit\n")
assert result.exit_code == 0, result.output
assert resolved == ["outputs/run-42"]
def test_find_studio_server_none_when_not_running(monkeypatch):
import urllib.request
from unsloth_cli import _inference
def refuse(*a, **k):
raise OSError("connection refused")
monkeypatch.setattr(urllib.request, "urlopen", refuse)
assert _inference.find_studio_server() is None
class _FakeSSEResponse:
def __init__(self, lines):
self._lines = lines
def __iter__(self):
return iter(self._lines)
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def test_http_backend_streams_cumulative_text(monkeypatch):
backend = HttpChatBackend("http://localhost:8888", "token")
response = _FakeSSEResponse(
[
b'data: {"choices":[{"delta":{"content":"He"}}]}\n',
b"\n",
b'data: {"choices":[{"delta":{"content":"llo"}}]}\n',
b"data: [DONE]\n",
]
)
monkeypatch.setattr(backend, "_request", lambda *a, **k: response)
out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS))
assert out == ["He", "Hello"]
def test_http_backend_merges_emoji_split_across_deltas(monkeypatch):
backend = HttpChatBackend("http://localhost:8888", "token")
response = _FakeSSEResponse(
[
b'data: {"choices":[{"delta":{"content":"hi "}}]}\n',
b'data: {"choices":[{"delta":{"content":"\\ud83d"}}]}\n',
b'data: {"choices":[{"delta":{"content":"\\ude0a"}}]}\n',
b"data: [DONE]\n",
]
)
monkeypatch.setattr(backend, "_request", lambda *a, **k: response)
out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS))
# The lone high surrogate is held back, then merged with its other half.
assert out == ["hi ", "hi ", "hi 😊"]
def test_chat_prefers_running_studio_server(monkeypatch):
closed = []
class _FakeHttpBackend:
def stream(self, *a, **k):
return iter(["hello"])
def close(self):
closed.append("http")
local_loads = []
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend())
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: local_loads.append(1))
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n")
assert result.exit_code == 0, result.output
assert local_loads == []
assert "stays warm" in result.output
assert closed == ["http"]
def test_chat_server_mode_compare_loads_base_locally(monkeypatch):
streamed, closed, base_loads = [], [], []
class _FakeHttpBackend:
def stream(self, *a, **k):
streamed.append("tuned")
return iter(["tuned-answer"])
def close(self):
closed.append("http")
class _FakeBaseBackend:
def stream(self, *a, **k):
streamed.append("base")
return iter(["base-answer"])
def close(self):
closed.append("base")
def fake_local_load(model, **kwargs):
base_loads.append((model, kwargs.get("fresh_backend", False)))
return _FakeBaseBackend()
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend())
monkeypatch.setattr(chatmod, "load_chat_backend", fake_local_load)
result = CliRunner().invoke(_chat_app(), ["tuned-run"], input = "/compare\nhi\n/exit\n")
assert result.exit_code == 0, result.output
assert "(compare on)" in result.output
# Only the base model loaded locally, on its own private backend.
assert base_loads == [("fake/base", True)]
assert streamed == ["base", "tuned"]
assert set(closed) == {"http", "base"}
def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch):
loads, streamed, closed = [], [], []
class _FakeLocalBackend:
def __init__(self, role):
self.role = role
def stream(self, *a, **k):
streamed.append((self.role, k.get("use_adapter")))
return iter([f"{self.role}-answer"])
def close(self):
closed.append(self.role)
def fake_load(model, **kwargs):
fresh = kwargs.get("fresh_backend", False)
loads.append((model, fresh))
return _FakeLocalBackend("base" if fresh else "tuned")
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(chatmod, "load_chat_backend", fake_load)
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: True)
monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
result = CliRunner().invoke(_chat_app(), ["tuned-run", "--compare"], input = "hi\n/exit\n")
assert result.exit_code == 0, result.output
assert loads == [("tuned-run", False), ("fake/base", True)]
# Both models answered the turn, via plain generation (no adapter toggle).
assert ("base", None) in streamed and ("tuned", None) in streamed
assert set(closed) == {"tuned", "base"}

292
unsloth_cli/tests/test_studio_cloudflare_flag.py Normal file
View file

@ -0,0 +1,292 @@
# 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/--no-cloudflare` Studio flag.
Pins the typer Option (default on) on both `unsloth studio` and
`unsloth studio run`, and that the chosen polarity reaches the re-exec'd
child and run_server. Modeled on test_studio_run_parallel_flag.py.
"""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
from typer.testing import CliRunner
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
sys.path.insert(0, str(_REPO_ROOT))
def _studio():
from unsloth_cli.commands import studio as _studio_mod
return _studio_mod
_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"]
# ── option registration ──────────────────────────────────────────────
def test_run_exposes_cloudflare_option_default_on():
import inspect
sig = inspect.signature(_studio().run)
assert "cloudflare" in sig.parameters
opt = sig.parameters["cloudflare"].default
decls = set(getattr(opt, "param_decls", []) or [])
assert "--cloudflare/--no-cloudflare" in decls
assert getattr(opt, "default", None) is True
def test_studio_default_exposes_cloudflare_option_default_on():
import inspect
sig = inspect.signature(_studio().studio_default)
assert "cloudflare" in sig.parameters
opt = sig.parameters["cloudflare"].default
assert getattr(opt, "default", None) is True
# ── re-exec forwarding: `unsloth studio run` ─────────────────────────
class _ExecCaptured(SystemExit):
def __init__(self, argv):
super().__init__(0)
self.argv = list(argv)
def _install_run_reexec_capture(monkeypatch, *, platform = "linux"):
studio_mod = _studio()
captured = []
monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
fake_venv = Path("/fake/studio/venv/unsloth_studio")
monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
fake_bin = fake_venv / "bin" / "unsloth"
real_is_file = Path.is_file
monkeypatch.setattr(
Path,
"is_file",
lambda self: True if str(self) == str(fake_bin) else real_is_file(self),
)
from unsloth_cli import _tool_policy as _tp_mod
monkeypatch.setattr(
_tp_mod,
"resolve_tool_policy",
lambda host, flag, yes, silent: False if flag is None else bool(flag),
)
monkeypatch.setattr(sys, "platform", platform)
def fake_execvp(file, argv):
captured.append(list(argv))
raise _ExecCaptured(argv)
monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
return captured
def _invoke_run(monkeypatch, args):
import typer as _typer
studio_mod = _studio()
captured = _install_run_reexec_capture(monkeypatch)
app = _typer.Typer()
app.command(
context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
)(studio_mod.run)
CliRunner().invoke(app, args, catch_exceptions = True)
return captured
@pytest.mark.parametrize(
"user_flag,expected,unexpected",
[
(None, "--cloudflare", "--no-cloudflare"), # default on
("--cloudflare", "--cloudflare", "--no-cloudflare"),
("--no-cloudflare", "--no-cloudflare", "--cloudflare"),
],
)
def test_run_reexec_forwards_cloudflare_polarity(monkeypatch, user_flag, expected, unexpected):
extras = [user_flag] if user_flag else []
captured = _invoke_run(monkeypatch, _BASE + extras)
assert len(captured) == 1, captured
argv = captured[0]
assert expected in argv, f"expected {expected} in child argv; got {argv}"
assert unexpected not in argv, f"unexpected {unexpected} in child argv; got {argv}"
# ── re-exec forwarding: plain `unsloth studio` ───────────────────────
def _invoke_studio_default(
monkeypatch,
args,
*,
platform = "linux",
):
import typer as _typer
studio_mod = _studio()
captured = []
monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
monkeypatch.setattr(studio_mod, "_ensure_studio_env_exported", lambda: None)
fake_venv = Path("/fake/studio/venv/unsloth_studio")
monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
monkeypatch.setattr(studio_mod, "_find_run_py", lambda: Path("/fake/studio/run.py"))
monkeypatch.setattr(studio_mod, "_find_frontend_dist", lambda: None)
monkeypatch.setattr(sys, "platform", platform)
def fake_execvp(file, argv):
captured.append(list(argv))
raise _ExecCaptured(argv)
monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
app = _typer.Typer()
app.command()(studio_mod.studio_default)
CliRunner().invoke(app, args, catch_exceptions = True)
return captured
@pytest.mark.parametrize(
"user_flag,expected,unexpected",
[
(None, "--cloudflare", "--no-cloudflare"),
("--no-cloudflare", "--no-cloudflare", "--cloudflare"),
],
)
def test_studio_default_reexec_forwards_cloudflare(monkeypatch, user_flag, expected, unexpected):
extras = [user_flag] if user_flag else []
captured = _invoke_studio_default(monkeypatch, ["-H", "0.0.0.0"] + extras)
assert len(captured) == 1, captured
argv = captured[0]
assert expected in argv, f"expected {expected}; got {argv}"
assert unexpected not in argv, f"unexpected {unexpected}; got {argv}"
# ── in-venv path forwards cloudflare into run_server ─────────────────
class _RunServerCaptured(SystemExit):
def __init__(self, kwargs):
super().__init__(0)
self.kwargs = dict(kwargs)
@pytest.mark.parametrize("user_flag,expected", [(None, True), ("--no-cloudflare", False)])
def test_run_in_venv_passes_cloudflare_to_run_server(monkeypatch, user_flag, expected):
import types
studio_mod = _studio()
fake_venv = Path("/fake/studio/venv/unsloth_studio")
monkeypatch.setattr(sys, "prefix", str(fake_venv))
monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent)
from unsloth_cli import _tool_policy as _tp_mod
monkeypatch.setattr(
_tp_mod,
"resolve_tool_policy",
lambda host, flag, yes, silent: False if flag is None else bool(flag),
)
captured: dict = {}
def fake_run_server(**kwargs):
captured.update(kwargs)
raise _RunServerCaptured(kwargs)
fake_backend_run = sys.modules.setdefault(
"studio.backend.run", types.ModuleType("studio.backend.run")
)
fake_backend_run.run_server = fake_run_server
fake_backend_run._resolve_external_ip = lambda: "127.0.0.1"
import typer as _typer
app = _typer.Typer()
app.command(
context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
)(studio_mod.run)
extras = [user_flag] if user_flag else []
CliRunner().invoke(app, _BASE + extras, catch_exceptions = True)
assert captured.get("cloudflare") is expected, captured
# ── parent-level --no-cloudflare with a subcommand is rejected ───────
def test_studio_default_rejects_no_cloudflare_with_subcommand(monkeypatch):
# `unsloth studio --no-cloudflare run ...` would not reach the subcommand,
# so it must error (mirrors --parallel) rather than silently still tunnel.
import typer as _typer
studio_mod = _studio()
app = _typer.Typer()
app.add_typer(studio_mod.studio_app, name = "studio")
result = CliRunner().invoke(app, ["studio", "--no-cloudflare", "run", "--model", "X"])
assert result.exit_code == 2, result.output
combined = (result.output or "") + (getattr(result, "stderr", "") or "")
assert "--no-cloudflare" in combined, combined
# ── run() tears the server + tunnel down if startup aborts ───────────
def test_run_in_venv_shuts_down_on_startup_abort(monkeypatch):
import types
studio_mod = _studio()
fake_venv = Path("/fake/studio/venv/unsloth_studio")
monkeypatch.setattr(sys, "prefix", str(fake_venv))
monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent)
from unsloth_cli import _tool_policy as _tp_mod
monkeypatch.setattr(
_tp_mod,
"resolve_tool_policy",
lambda host, flag, yes, silent: False if flag is None else bool(flag),
)
class _App:
class state:
server_port = 8888
shutdown_calls = []
backend = sys.modules.setdefault("studio.backend.run", types.ModuleType("studio.backend.run"))
backend.run_server = lambda **k: _App()
backend._resolve_external_ip = lambda: "1.2.3.4"
backend._server = object()
backend._shutdown_event = None
backend._graceful_shutdown = lambda server: shutdown_calls.append(server)
# set_tool_policy is imported as `from state.tool_policy import set_tool_policy`.
state_mod = sys.modules.setdefault("state", types.ModuleType("state"))
tp_mod = sys.modules.setdefault("state.tool_policy", types.ModuleType("state.tool_policy"))
tp_mod.set_tool_policy = lambda *a, **k: None
state_mod.tool_policy = tp_mod
# Force the health check to fail so startup aborts after run_server().
monkeypatch.setattr(studio_mod, "_wait_for_server", lambda *a, **k: False)
import typer as _typer
app = _typer.Typer()
app.command(
context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
)(studio_mod.run)
result = CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0"], catch_exceptions = True)
assert result.exit_code == 1, result.output
assert len(shutdown_calls) == 1, "startup abort must call _graceful_shutdown"