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

Merge remote-tracking branch 'origin/main' into ig_merge

Browse source 
# Conflicts:
#	scripts/scan_packages_baseline.json
This commit is contained in:
Daniel Han 2026-07-10 06:28:04 +00:00
commit b9ebfe089b
46 changed files with 8229 additions and 2289 deletions
Download patch file Download diff file Expand all files Collapse all files

148
.github/scripts/agent-guides-drive.sh vendored
View file

@ -527,6 +527,154 @@ case "$MODE" in
echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
;;
# ── resume: does a launched agent's session survive exit and resume? ────
# Unlike the other modes, this drives the real LAUNCH path (`unsloth start
# <agent> ...`, the interactive default), not the --no-launch recipe. That
# path relocates each agent's home to a throwaway temp dir wiped on exit, so
# a session cannot be resumed -- unless --persist routes it to the stable
# Unsloth agents dir instead. We run one headless turn per pass and check
# whether the turn left a session in a persistent store (deterministic, no
# reliance on the model recalling anything), for a baseline pass and a
# --persist pass, and assert the expected split for this agent.
resume)
CODEWORD="PLATYPUS7"
T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK."
T2="What codeword did I ask you to remember? Reply with just that word."
WORK="$WORKDIR_BASE/${AGENT}-resume"
# STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to.
# Read it from a --no-launch probe (which also writes the agent's config
# there). codex/pi relocate their whole home/HOME here; opencode/claude keep
# their session data in a fixed user dir, so STABLE_HOME stays empty for them.
parse_connect
case "$AGENT" in
codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;;
pi) STABLE_HOME="$(raw_env HOME)" ;;
*) STABLE_HOME="" ;;
esac
# The persistent stores a session would land in if it were NOT wiped. We
# count files here before/after each turn; a positive delta means the
# session persisted (is resumable), zero means it went to a wiped temp dir.
resume_tracked_dirs() {
case "$AGENT" in
codex) printf '%s\n' "$HOME/.codex" ;;
opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;;
claude) printf '%s\n' "$HOME/.claude" ;;
pi) printf '%s\n' "$HOME/.pi" ;;
*) : ;;
esac
[ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME"
}
count_session_files() {
local total=0 d n
while IFS= read -r d; do
[ -n "$d" ] && [ -d "$d" ] || continue
n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n))
done < <(resume_tracked_dirs)
echo "$total"
}
# The headless first-turn subcommand per agent (mirrors file-edit's map),
# forwarded verbatim through the launch path as passthrough args.
set_t1_cmd() {
case "$AGENT" in
claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;;
codex) T1_CMD=(exec "$T1") ;;
opencode) T1_CMD=(run "$T1") ;;
pi) T1_CMD=(-p "$T1") ;;
*) guide_fail "resume mode does not cover agent '$AGENT'" ;;
esac
}
# Run one headless turn through the launch path. $1=outfile, $2="" or
# "--persist", rest = the agent subcommand. --yolo auto-approves so no tool
# prompt can hang; --api-key attaches to the already-served CI model.
launch_turn() {
local out="$1" rflag="$2"; shift 2
local flag=(); [ -n "$rflag" ] && flag=("$rflag")
run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \
--api-key "$UNSLOTH_API_KEY" "$@"
local rc=$?
redact "$out"
return "$rc"
}
# One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED
# from the session-store delta. Runs in the main shell (not a command
# substitution) so a hang's guide_fail actually fails the job and the
# progress lines reach the CI log. $1 = "" (baseline) or "--persist".
RESULT=""
run_pass() {
local rflag="$1" label="baseline"
[ -n "$rflag" ] && label="resume"
rm -rf "$WORK"; mkdir -p "$WORK"
set_t1_cmd
local out="$LOGS_DIR/${AGENT}-resume-${label}.txt"
local before after rc
before="$(count_session_files)"
pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK"
launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$?
popd >/dev/null || true
after="$(count_session_files)"
echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})"
# The turn must succeed for the delta to mean anything: an agent that writes a
# session file then errors would otherwise be misread as PERSISTED. Mirror the
# file-edit mode and fail the pass on a non-zero launch (the flagship codex recall
# below stays WARN-only, driven by its own launch_turn calls).
[ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \
guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; }
if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi
}
run_pass ""; BASELINE="$RESULT"
# Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix.
# opencode/claude persist either way, so the baseline already proves it and a
# second full CPU turn only risks a timeout; skip it for them.
case "$AGENT" in
codex|pi) run_pass "--persist"; RESUME="$RESULT" ;;
*) RESUME="n/a (persists either way)" ;;
esac
# Expected: codex/pi relocate their whole home to the temp dir, so a plain
# launch is WIPED and only --persist PERSISTS. opencode/claude keep their
# session data in a fixed user dir, so the baseline already PERSISTS.
case "$AGENT" in
codex|pi) EXPECT_BASELINE="WIPED" ;;
opencode|claude) EXPECT_BASELINE="PERSISTED" ;;
esac
echo "──────────────────────────────────────────────"
echo "[$AGENT] RESUME EXPERIMENT"
echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})"
echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}"
echo "──────────────────────────────────────────────"
[ "$BASELINE" = "$EXPECT_BASELINE" ] \
|| guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}"
case "$AGENT" in
codex|pi)
[ "$RESUME" = "PERSISTED" ] \
|| guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;;
esac
# Flagship behavioral proof (codex only, WARN-only): after a --persist plant,
# resume the session and check the model actually recalls the codeword. A
# miss is not a failure (the CI model is small); the mechanism gate above is
# the real assertion.
if [ "$AGENT" = "codex" ]; then
rm -rf "$WORK"; mkdir -p "$WORK"
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true
if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then
echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}"
else
echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed"
fi
fi
echo "[$AGENT] resume OK"
;;
*)
echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2
exit 2

170
.github/workflows/local-agent-guides-ci.yml vendored
View file

@ -471,6 +471,176 @@ jobs:
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job: resume
# Does a conversation started with `unsloth start <agent>` survive exit
# and resume? This drives the REAL launch path (not the --no-launch
# recipe the other jobs use). A plain launch relocates the agent home to
# a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the
# session to the stable Unsloth agents dir so it persists. opencode/claude
# keep their session data in a fixed user dir, so they persist either way.
# Dispatch-only: it is an end-to-end experiment, not a PR gate.
# ═════════════════════════════════════════════════════════════════════
resume:
name: resume (${{ matrix.agent }})
if: github.event_name == 'workflow_dispatch'
runs-on: ubuntu-latest
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
# codex/pi relocate their whole home (resume broken without --persist);
# opencode/claude keep session data in a fixed dir (resume already works).
# One agent from each class proves the split end to end; openclaw/hermes
# share codex's relocation mechanism and are covered by the unit tests.
agent: [codex, opencode, claude, pi]
env:
GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF
GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf
STUDIO_PORT: '18904'
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- name: Linux deps for llama.cpp prebuilt
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libcurl4-openssl-dev libssl-dev jq
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- name: Restore GGUF model file
id: cache-gguf
uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
continue-on-error: true
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Download GGUF if cache miss
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
- name: Save GGUF model file
if: always() && steps.download-gguf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
unsloth studio reset-password
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
--extra "--seed $UNSLOTH_SEED --temp 0" \
--health-timeout 900
- name: Preflight the agent's API dialect (class-a isolation)
env:
AGENT: ${{ matrix.agent }}
run: |
set -uo pipefail
B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY"
preflight_fail() {
echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**.";
exit 1
}
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \
-H "Authorization: Bearer $K") || true
[ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code"
case "$AGENT" in
claude)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code"
;;
codex)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true
[ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code"
;;
*)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code"
;;
esac
echo "preflight OK for $AGENT"
- name: Install agent CLI (class-b isolation)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-install.sh "$AGENT"
- name: Resume experiment (launch path)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT"
- name: Collect server logs (debug)
if: always()
run: |
mkdir -p logs/studio-logs
cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
done
fi
- name: Stop Studio
if: always()
run: |
if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then
kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true
fi
sleep 2
ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true
- name: Upload logs
if: always()
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: resume-${{ matrix.agent }}-log
path: |
logs/
agent-workdir/
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job 3: prompt-cache
# (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0

6
.github/workflows/security-audit.yml vendored
View file

@ -2,8 +2,8 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Multi-language supply-chain audit. Triggers:
# - PRs touching any dependency manifest (Python / npm / Cargo) or
# this workflow file,
# - PRs touching any dependency manifest (Python / npm / Cargo), a
# scanner or its allowlist baseline, or this workflow file,
# - push to main / pip,
# - nightly @ 04:13 UTC so newly-published advisories surface even
# when no PR opens,
@ -57,7 +57,9 @@ on:
- 'studio/src-tauri/Cargo.lock'
- 'pyproject.toml'
- 'scripts/scan_packages.py'
- 'scripts/scan_packages_baseline.json'
- 'scripts/scan_npm_packages.py'
- 'scripts/scan_npm_packages_baseline.json'
- '.github/workflows/security-audit.yml'
push:
branches: [main, pip]

2
README.md
View file

@ -212,7 +212,7 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i
```bash
unsloth studio --secure -p 8888
```
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network.
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind.
```bash
unsloth studio -H 0.0.0.0 -p 8888
```

24
install.ps1
View file

@ -469,6 +469,17 @@ function Install-UnslothStudio {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when the command pins an index, clear every uv index env var so
# it wins, then restore in finally. Other installs keep the user's mirror.
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@ -488,6 +499,7 @@ function Install-UnslothStudio {
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
}
}
@ -2200,7 +2212,7 @@ exit 0
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
# still completes; Studio setup retries ROCm afterwards.
@ -2209,7 +2221,7 @@ exit 0
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2223,7 +2235,7 @@ exit 0
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2306,7 +2318,7 @@ exit 0
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
# is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
if (-not $SkipTorch) {
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
@ -2322,7 +2334,7 @@ exit 0
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
@ -2331,7 +2343,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)

28
install.sh
View file

@ -159,6 +159,12 @@ run_maybe_quiet() {
run_install_cmd() {
_label="$1"
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
esac
if _is_verbose; then
"$@" && return 0
_rc=$?
@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() {
esac
}
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX /
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
# resolves (torch + every transitive dep) via --default-index -- the same URLs the
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
# Repair when flavor is wrong AND the index is plain --default-index reinstallable
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi

3104
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

110
studio/backend/core/inference/_vulkan_probe.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. See /studio/LICENSE.AGPL-3.0
"""Standalone free-VRAM probe for the bundled ggml Vulkan backend.
Run in a short-lived subprocess (``python _vulkan_probe.py <bindir>``) so the
Vulkan instance never lives in the long-running backend process. Loads the
bundled ggml Vulkan backend from ``<bindir>`` and prints one
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` line per device to stdout.
Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
Uses only the standard library so it stays runnable as a bare script.
"""
import ctypes
import os
import sys
# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ...
_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
def _igpu_flags(base, lib, count: int) -> list[bool]:
"""Per-device integrated-GPU flags via ggml's backend registry.
The Vulkan reg enumerates devices in the same order as
``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
i``), so reg index == device ordinal. Returns all-False on any failure so
the reader never over-caps a discrete card.
"""
flags = [False] * count
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t
base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p]
base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p
base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
base.ggml_backend_dev_type.restype = ctypes.c_int
base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p]
reg = lib.ggml_backend_vk_reg()
if not reg:
return flags
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
except Exception:
# Best-effort: any failure degrades to "discrete" so the memory
# readings still get through instead of crashing the probe.
pass
return flags
def main() -> int:
if len(sys.argv) < 2:
return 0
bindir = sys.argv[1]
# Hold add_dll_directory's handle for the rest of main() (the documented
# idiom) so bindir stays on the search path while the sibling ggml DLLs
# resolve below.
_dll_dir = None
if sys.platform == "win32":
base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll"
try:
_dll_dir = os.add_dll_directory(bindir)
except Exception:
pass
else:
base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
# RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr
# falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode).
_rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0)
try:
base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global)
lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global)
except OSError as e:
print(f"ggml-vulkan load failed: {e}", file = sys.stderr)
return 1
lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int
lib.ggml_backend_vk_get_device_count.argtypes = []
lib.ggml_backend_vk_get_device_memory.restype = None
lib.ggml_backend_vk_get_device_memory.argtypes = [
ctypes.c_int,
ctypes.POINTER(ctypes.c_size_t),
ctypes.POINTER(ctypes.c_size_t),
]
count = lib.ggml_backend_vk_get_device_count()
igpu = _igpu_flags(base, lib, count)
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
sys.stdout.write("\n".join(rows))
return 0
if __name__ == "__main__":
raise SystemExit(main())

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

@ -100,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError):
Subclasses RuntimeError so existing handlers still catch it."""
class _LlamaStreamCancelled(Exception):
"""Internal signal for an expected client/request cancellation."""
# Shared so the from_identifier preflight and the load-time raise stay in sync.
LLAMA_SERVER_NOT_FOUND_DETAIL = (
"This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
@ -1436,6 +1440,50 @@ def _backfill_usage_from_timings(usage, timings):
return out
def _vulkan_lib_filename() -> str:
return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit
# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared
# system RAM, so hold back the same margin rather than inventing a larger one.
_IGPU_HOST_RESERVE_MIB = 1024
def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int:
"""Reserve host headroom on an integrated (shared-memory) Vulkan GPU.
An iGPU's reported free "VRAM" is really free system RAM, so sizing
context/offload against all of it would push the host into swap or the OOM
killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from
ggml's device type, so a discrete card is never touched; only ever reduces.
"""
if not is_igpu:
return free_mib
return max(0, free_mib - _IGPU_HOST_RESERVE_MIB)
def _llama_lib_dir(binary: str) -> Path:
# The installer exposes llama-server as a top-level entrypoint into build/bin/,
# where the ggml backend libs live, so callers looking for sibling libs (Vulkan
# detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a
# symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to
# a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it
# cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's
# exec target too, so a wrapper-based install still finds build/bin.
resolved = Path(binary).resolve()
try:
with open(resolved, "rb") as _f:
_head = _f.read(256)
if _head.startswith(b"#!"):
_m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore"))
if _m:
return (resolved.parent / _m.group(1)).resolve().parent
except OSError:
pass
return resolved.parent
def _is_external_link(path: Path) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
@ -2278,6 +2326,30 @@ class LlamaCppBackend:
return total
@staticmethod
def _is_vulkan_backend(binary: Optional[str] = None) -> bool:
"""True if the installed llama.cpp build is Vulkan-only.
The official prebuilts are single-backend, so the Vulkan ggml lib next
to llama-server identifies a Vulkan build. Keeps the free-memory probe
and GPU pin in ggml's Vulkan device-index space. For a custom
multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer
to that backend (torch-usable, better-understood probe/pin).
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
return False
lib_dir = _llama_lib_dir(binary)
if not (lib_dir / _vulkan_lib_filename()).is_file():
return False
for _backend in ("cuda", "hip"):
sibling = (
f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so"
)
if (lib_dir / sibling).is_file():
return False
return True
@staticmethod
def _resolve_visible_physical_ids() -> Optional[list[int]]:
"""Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on
@ -2440,11 +2512,42 @@ class LlamaCppBackend:
return True
@staticmethod
def _get_gpu_free_memory() -> list[tuple[int, int]]:
def _visible_devices_mask(env_name: str) -> Optional[set[int]]:
"""Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset.
``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask
("") yields an empty set (all devices hidden), distinct from an unset
var (None, no mask). Used by the nvidia-smi probe.
"""
raw = os.environ.get(env_name)
if raw is None:
return None
try:
return set(int(x.strip()) for x in raw.split(",") if x.strip())
except ValueError:
return None
@staticmethod
def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]:
"""``--device Vulkan<i>,...`` to pin a Vulkan launch to selected GPUs.
The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory
reports and the registry names ``Vulkan<i>``). Pin by that name, NOT via
GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw
vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices
and deduplicating ICDs), so a compact ordinal there could select a
different physical device or the CPU rasterizer.
"""
if not gpu_indices:
return []
return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)]
@staticmethod
def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]:
"""Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
index; empty if no supported GPU is reachable. Thin wrapper over
``_get_gpu_memory`` for callers that only need free VRAM."""
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)]
@staticmethod
def _apple_metal_memory_budget_bytes() -> int:
@ -2475,7 +2578,7 @@ class LlamaCppBackend:
return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION)
@staticmethod
def _get_gpu_memory() -> list[tuple[int, int, int]]:
def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
"""Query free AND total memory per GPU.
Order:
@ -2487,9 +2590,18 @@ class LlamaCppBackend:
probe returned [] on AMD) and NVIDIA hosts missing
``nvidia-smi`` from PATH.
On a Vulkan build the ggml Vulkan probe is authoritative, so the indices
are ggml's compact Vulkan ordinals (the space the pin selects via
``--device Vulkan<i>``). It reports ``total`` for discrete cards and 0
for an iGPU (shared RAM) so the fit falls back to free*frac there.
Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm.
Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no
supported GPU is reachable. ``total`` lets the fit reserve absolute headroom.
supported GPU is reachable.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if LlamaCppBackend._is_vulkan_backend(binary):
return LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
# ── NVIDIA via nvidia-smi ────────────────────────────────────
try:
result = subprocess.run(
@ -2505,16 +2617,7 @@ class LlamaCppBackend:
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode == 0:
allowed: Optional[set[int]] = None
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd is not None:
try:
# `if x.strip()` filters trailing-comma masks ("0,1,").
# Empty mask (CVD="") yields an empty set -> all GPUs
# filtered out, per codebase convention.
allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
except ValueError:
pass
allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES")
gpus: list[tuple[int, int, int]] = []
for line in result.stdout.strip().splitlines():
parts = [p.strip() for p in line.split(",")]
@ -2579,6 +2682,91 @@ class LlamaCppBackend:
logger.debug(f"torch GPU probe failed: {e}")
return []
@staticmethod
def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
"""Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance
in this process) and returns (device_index, free_mib, total_mib) sorted
by index. The index is ggml's compact Vulkan ordinal -- the one the
registry names ``Vulkan<index>`` and load_model pins with ``--device``,
NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set
``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the
list already reflects it. iGPUs leave a host-RAM margin (see
``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
their real total through. [] when no Vulkan build or device is reachable.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
return []
binary_dir = _llama_lib_dir(binary)
if not (binary_dir / _vulkan_lib_filename()).is_file():
return []
env = child_env_without_native_path_secret()
# Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so
# the probe enumerates the same device list the launch will, named
# Vulkan0..N in the compact order reported here and pinned by that name
# via --device -- probe, mask, and pin stay in one index space. Do NOT
# filter the mask in Python: ggml parses the env var in raw
# vkEnumeratePhysicalDevices space while this probe reports the compact
# post-filter ordinal, so a Python filter would compare mismatched spaces.
if sys.platform != "win32":
# Let the loader resolve sibling ggml libs next to the binary.
existing_ld = env.get("LD_LIBRARY_PATH", "")
env["LD_LIBRARY_PATH"] = (
f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir)
)
probe_script = Path(__file__).with_name("_vulkan_probe.py")
try:
result = subprocess.run(
[sys.executable, str(probe_script), str(binary_dir)],
capture_output = True,
text = True,
timeout = 15,
env = env,
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
logger.debug(
f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}"
)
return []
except Exception as e:
logger.debug(f"vulkan GPU probe failed: {e}")
return []
gpus: list[tuple[int, int, int]] = []
for line in result.stdout.strip().splitlines():
parts = line.split("\t")
if len(parts) != 4:
continue
try:
idx = int(parts[0])
free_mib = int(parts[1]) // (1024 * 1024)
is_igpu = parts[2] == "1"
# iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
# fit stays on free*frac (the host reserve below is its
# headroom); a discrete card passes its real total through.
total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024)
except ValueError:
continue
capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
if capped < free_mib:
logger.info(
f"Vulkan device VK{idx} is an integrated GPU sharing system "
f"RAM; reserving {free_mib - capped}MiB host headroom "
f"({free_mib}->{capped}MiB usable)"
)
gpus.append((idx, capped, total_mib))
gpus.sort(key = lambda g: g[0])
if gpus:
logger.info(
"Vulkan GPU memory detected: "
+ ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus)
)
return gpus
@staticmethod
def _available_system_memory_mib() -> Optional[int]:
"""Available system RAM in MiB (psutil, then /proc/meminfo), or None if
@ -2807,7 +2995,8 @@ class LlamaCppBackend:
def _llama_server_env_for_binary(binary: str) -> dict[str, str]:
"""Build a subprocess env that lets llama-server resolve native libs."""
env = child_env_without_native_path_secret()
binary_dir = str(Path(binary).parent)
# _llama_lib_dir resolves the llama-server symlink to the real build/bin.
binary_dir = str(_llama_lib_dir(binary))
if sys.platform == "win32":
# Ordering: see _build_windows_path_dirs. #5106.
@ -4488,6 +4677,29 @@ class LlamaCppBackend:
cancel_event = cancel_event,
)
def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
"""A drafter already in this repo's local HF cache, reused offline when a
fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all
cached snapshots; else an existing ``MTP/`` copy (any precision -- the
target verifies every drafted token). None if none is cached."""
try:
from utils.models.model_config import _iter_hf_cache_snapshots
roots: list[Path] = []
subdirs: list[Path] = []
for snap in _iter_hf_cache_snapshots(hf_repo): # newest first
for f in sorted(_gguf_snapshot_files(snap)):
if _is_companion_gguf_path(f) and "mmproj" not in f.lower():
(roots if "/" not in f else subdirs).append(snap / f)
# Keep snapshot order (newest first), root before any MTP/ copy, so a
# newer main GGUF pairs with the newest cached drafter, not a stale one.
for cand in roots + subdirs:
if cand.is_file():
return str(cand)
except Exception as e:
logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e)
return None
def _download_mtp(
self,
*,
@ -4504,11 +4716,25 @@ class LlamaCppBackend:
are intentionally skipped. Returns the local path, or None.
"""
# Offline, reuse any drafter already on disk (a fresh copy can't be
# fetched). Online, _download_companion_gguf/hf_hub_download reuse the
# current cached file and refetch a changed one, so skip the probe here
# rather than pair new weights with a stale draft.
if _hf_env_offline():
cached = self._cached_repo_mtp_drafter(hf_repo)
if cached:
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
return cached
def _pick_mtp(candidates: list[str]) -> Optional[str]:
# Root-level only: MTP/ subdir copies now share the mtp- prefix but
# are explicit-selection, not auto-fetch (they'd sort ahead of root).
mtp_files = sorted(
f
for f in candidates
if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-")
if f.lower().endswith(".gguf")
and "/" not in f
and Path(f).name.lower().startswith("mtp-")
)
return mtp_files[0] if mtp_files else None
@ -5210,6 +5436,7 @@ class LlamaCppBackend:
# Resolve llama-server now but defer a not-found error: a block-diffusion
# GGUF uses the diffusion runner, and its arch is only known after the header.
binary = self._find_llama_server_binary()
is_vulkan_backend = self._is_vulkan_backend(binary)
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
# mtp_draft_path arrives set for local Gemma loads (detected
@ -5449,7 +5676,8 @@ class LlamaCppBackend:
model_size = gguf_size + mmproj_size
# 2-tuple gpus for existing logic + a total map for the absolute
# per-GPU headroom (correct when the GPU is already partly used).
_gpu_mem = self._get_gpu_memory()
# Pass binary so a Vulkan build probes ggml's Vulkan ordinals.
_gpu_mem = self._get_gpu_memory(binary)
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
@ -6222,7 +6450,12 @@ class LlamaCppBackend:
# cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
# oversize load the OS would otherwise kill mid-flight. Base model
# only: an optional MTP drafter is dropped by the MTP-drop fallback.
if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices):
# CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals.
if (
model_size is not None
and not is_vulkan_backend
and self._amd_apu_wants_unified_memory(gpu_indices)
):
_ram_msg = self._apu_ram_shortfall_message(
model_size, self._available_system_memory_mib()
)
@ -6485,6 +6718,12 @@ class LlamaCppBackend:
", ".join(unsupported_cache_flags),
)
# Vulkan pins via --device (a cmd arg, unlike the env-based
# CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's
# last-wins parsing lets a user --device override Studio's pick.
if is_vulkan_backend and gpu_indices is not None:
cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices)
# User pass-through args go last so llama.cpp's last-wins parsing
# lets the user override Studio's auto-set flags. Already
# validated by the route via validate_extra_args().
@ -6536,23 +6775,25 @@ class LlamaCppBackend:
env.setdefault("OMP_NUM_THREADS", "2")
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
# shared system RAM. setdefault so a user value wins.
if self._amd_apu_wants_unified_memory(gpu_indices):
# shared system RAM. setdefault so a user value wins. Not on Vulkan
# (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids.
if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices):
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")
# DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU).
# See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1.
if self._apply_datacenter_env(env, gpu_indices):
if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices):
multi_gpu = self._effective_gpu_count(gpu_indices) > 1
logger.info(
f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
)
# Pin to selected GPU(s). On ROCm, narrowing only
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
# set, so set HIP_VISIBLE_DEVICES too.
if gpu_indices is not None:
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so
# set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device
# (above), not here.
if gpu_indices is not None and not is_vulkan_backend:
pinned = ",".join(str(i) for i in gpu_indices)
env["CUDA_VISIBLE_DEVICES"] = pinned
try:
@ -8379,7 +8620,7 @@ class LlamaCppBackend:
):
"""Open one streaming POST and let cancel interrupt prefill or reads."""
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
raise _LlamaStreamCancelled
_cancel_closed = threading.Event()
_response_ref: list = [None]
@ -8424,13 +8665,13 @@ class LlamaCppBackend:
) as response:
_response_ref[0] = response
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
raise _LlamaStreamCancelled
yield response
return
except (httpx.RequestError, RuntimeError):
# Response was closed by the cancel watcher
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
raise _LlamaStreamCancelled
raise
finally:
_cancel_closed.set()
@ -8633,6 +8874,8 @@ class LlamaCppBackend:
"finish_reason": _metadata_finish_reason,
}
except _LlamaStreamCancelled:
return
except httpx.ConnectError as e:
# Server already down. If this was an MTP+tensor crash, recover by
# reloading without MTP (scheduled in the background) and fail this
@ -9757,6 +10000,8 @@ class LlamaCppBackend:
break
continue
except _LlamaStreamCancelled:
return
except httpx.ConnectError:
# Mark unresolved provisional cards as failed before raising.
for _pid, _pname in provisional_started_tool_calls.items():
@ -9939,6 +10184,8 @@ class LlamaCppBackend:
if _meta is not None:
yield _meta
except _LlamaStreamCancelled:
return
except httpx.ConnectError:
raise RuntimeError("Lost connection to llama-server")
except Exception as e:

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

@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
"""The separate MTP drafter to fetch with every variant: the repo-root
``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
(Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
prefix, first in sort order) so download and load resolve the same file;
the higher-precision ``MTP/`` subdir copies are for explicit selection and
are not auto-fetched. None for repos with the head baked into the main
GGUF (Qwen)."""
(Gemma 4). Same pick as the loader's drafter resolution (root-level
``mtp-`` prefix, first in sort order) so download and load resolve the same
file; the higher-precision ``MTP/`` subdir copies are for explicit
selection and are not auto-fetched. None for repos with the head baked into
the main GGUF (Qwen)."""
# Root-level only: the MTP/ subdir copies now share the mtp- prefix too.
candidates = sorted(
(
s
for s in siblings
if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-")
),
key = lambda s: getattr(s, "rfilename"),
)

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

@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel):
tool_use_id: str
content: Union[str, list] = ""
@field_validator("content", mode = "before")
@classmethod
def _coerce_null_content(cls, v):
# Some clients send null content for an empty tool result; the str|list
# union would 400 on it, so treat null as "".
return "" if v is None else v
# Block types the converter translates explicitly. Anything else (thinking /
# redacted_thinking, a provider block a resumed session replays, or a future type)
# is accepted as an unknown block and dropped by the converter, rather than 400-ing
# the whole request on strict validation.
_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"})
class AnthropicUnknownBlock(BaseModel):
type: str
model_config = {"extra": "allow"}
@field_validator("type")
@classmethod
def _only_unknown_types(cls, v):
# Known types parse as their typed models above (so a malformed known block
# still fails cleanly); this fallback only catches the rest.
if v in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError("known block type handled by its typed model")
return v
AnthropicContentBlock = Union[
AnthropicTextBlock,
AnthropicImageBlock,
AnthropicToolUseBlock,
AnthropicToolResultBlock,
AnthropicUnknownBlock,
]
@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel):
role: Literal["user", "assistant"]
content: Union[str, list[AnthropicContentBlock]]
@model_validator(mode = "before")
@classmethod
def _normalize_content(cls, data):
# Role-aware leniency that never silently drops real user input:
# - assistant: a resumed tool-only turn's null content -> "" (str|list would
# 400 on null; "" keeps the converter's `for block in content` safe).
# Unknown blocks (thinking / future types) validate via
# AnthropicUnknownBlock and are dropped by the converter.
# - user: keep strict. Null user content stays None so str|list rejects it
# (400) rather than forwarding an empty prompt; and reject block types the
# converter cannot translate, since it silently skips unknown user blocks
# -- a user turn made only of them would validate yet send no content
# (silent data loss).
if not isinstance(data, dict):
return data
content = data.get("content")
if data.get("role") == "assistant":
# Coerce only an explicit null (resumed tool-only turn). A missing
# content key stays malformed so the required-field check still 400s.
if "content" in data and content is None:
return {**data, "content": ""}
return data
if isinstance(content, list):
for block in content:
btype = (
block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
)
# Guard the value: a non-string type is unsupported too, and a
# membership test on an unhashable value would raise TypeError
# (escaping as a 500 instead of a clean 400).
if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError(f"unsupported content block type {btype!r} in a user message")
return data
class AnthropicTool(BaseModel):
# Client tools have input_schema; server tools may only have type/name.

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

File diff suppressed because it is too large Load diff

184
studio/backend/tests/test_anthropic_messages.py
View file

@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "plain"
def test_resumed_session_thinking_and_null_content_do_not_400():
# A resumed session replays assistant turns with `thinking` (and sometimes null)
# content. Those must be accepted (thinking dropped by the converter), not 400ed.
from pydantic import ValidationError
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "secret reasoning", "signature": "s"},
{"type": "text", "text": "the answer"},
{"type": "tool_use", "id": "t1", "name": "f", "input": {}},
],
},
{"role": "assistant", "content": None}, # tool-only turn serialized as null
],
)
# Known blocks still parse as their typed models; only the unknown one is loose.
assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock"
assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock"
assert req.messages[2].content == "" # null coerced
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content"))
assert assistant["content"] == "the answer"
assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded
# A malformed KNOWN block still fails cleanly instead of being swallowed.
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}],
)
def test_user_null_content_rejected():
# The null->"" leniency is assistant-only; a null user content must be rejected
# at the boundary, not coerced into an empty prompt and forwarded to the model.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": None}],
)
def test_user_unknown_block_rejected_not_silently_dropped():
# The converter skips user blocks it cannot translate, so a user turn whose only
# block is unknown would validate yet forward no content. Reject at the boundary
# to avoid that silent data loss (the assistant fallback is unaffected).
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "document", "source": {}}]},
],
)
def test_user_translatable_blocks_still_accepted():
# text / image / tool_result are translatable, so a real user message built from
# them must still pass; the unknown-block guard only trips on other types.
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": "AA"},
},
{"type": "tool_result", "tool_use_id": "t1", "content": "ok"},
],
}
],
)
assert [type(b).__name__ for b in req.messages[0].content] == [
"AnthropicTextBlock",
"AnthropicImageBlock",
"AnthropicToolResultBlock",
]
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai)
def test_user_malformed_known_block_still_rejected():
# The guard only allow-lists a user block's *type*; the union still validates its
# shape, so a known-but-malformed block (tool_result without tool_use_id) fails.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
],
)
def test_user_content_block_non_string_type_rejected_cleanly():
# A user block whose `type` is a non-string (unhashable list / dict, or a stray
# int) must fail as a clean validation error, not raise TypeError from the
# frozenset membership test and escape as a 500.
from pydantic import ValidationError
for bad_type in ([], {}, 5):
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": [{"type": bad_type}]}],
)
def test_assistant_missing_content_key_still_rejected():
# The null -> "" leniency is only for an EXPLICIT null. An assistant message that
# omits content entirely stays malformed and must fail required-field validation.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant"}],
)
# An explicit null is still accepted and coerced (regression guard).
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": None},
],
)
assert req.messages[1].content == ""
def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch):
# user -> assistant(null) -> user is now accepted: the null assistant turn coerces
# to "" and is dropped. The route must then coalesce the two remaining user turns
# so a strict GGUF chat template does not 400 on non-alternating roles.
backend = _mock_backend(monkeypatch, context_length = 2048)
class _Req:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/messages")
method = "POST"
async def is_disconnected(self):
return False
payload = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "first question"},
{"role": "assistant", "content": None},
{"role": "user", "content": "please continue"},
],
)
response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t"))
assert response.status_code == 200
[(_path, kwargs)] = backend.calls
user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"]
assert len(user_turns) == 1 # the two user turns were merged, not left adjacent
merged = user_turns[0]["content"]
if isinstance(merged, list):
merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
assert "first question" in merged and "please continue" in merged

2
studio/backend/tests/test_inference_dispatcher_resilience.py
View file

@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
continue
calls.append(node)
assert len(calls) == 4
assert len(calls) == 5
for call in calls:
assert any(
kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False

170
studio/backend/tests/test_install_resolve_prebuilt.py
View file

@ -55,6 +55,27 @@ def _host(**kw):
return ilp.HostInfo(**base)
def test_force_cpu_clears_all_gpu_attributes_including_intel():
# --cpu-fallback is the "select the CPU prebuilt even when a GPU is present"
# escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or
# the planner still prepends the Vulkan asset on an Intel-GPU host.
host = _host(
is_linux = True,
is_x86_64 = True,
has_usable_nvidia = True,
has_physical_nvidia = True,
has_rocm = True,
rocm_gfx_target = "gfx1100",
has_intel_gpu = True,
)
forced = ilp._apply_host_overrides(host, force_cpu = True)
assert forced.has_usable_nvidia is False
assert forced.has_physical_nvidia is False
assert forced.has_rocm is False
assert forced.rocm_gfx_target is None
assert forced.has_intel_gpu is False
def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
pre26 = _host(
system = "Darwin",
@ -313,3 +334,152 @@ def test_sm103_host_drops_cuda128_windows_build():
)
kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129])
assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name]
def _upstream_release(tag, asset_names):
return {
"tag_name": tag,
"assets": [
{"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names
],
}
def test_direct_upstream_arm64_intel_prefers_vulkan():
# Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU
# second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset).
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-arm64" in kinds
assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only():
# A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical
# True, usable False) + an Intel iGPU must NOT get the Vulkan archive even
# when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES
# and could grab the reserved card. It falls through to the CPU asset.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
def test_direct_upstream_arm64_without_intel_is_cpu_only():
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64")
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-arm64"]
def test_direct_upstream_x86_intel_prefers_vulkan():
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-cpu" in kinds
def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
# The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
# libs so a valid Vulkan install is not re-flagged unhealthy every check.
choice = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz",
url = "https://example/x",
source_label = "upstream",
install_kind = "linux-vulkan",
)
groups = ilp.runtime_payload_health_groups(choice)
assert ["libggml-cpu*.so*"] in groups
assert ["libggml-cpu-*.so*"] not in groups
def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
assert repo == UPSTREAM
assert tag == ""
assert routed.has_intel_gpu is True
def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
# A pin set WITH an explicit upstream repo is already on upstream -> kept.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
_routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
assert repo == UPSTREAM
assert tag == "b9596"
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed is host
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
# Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
# An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_non_intel_unchanged():
host = _host(is_linux = True, is_x86_64 = True)
routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
assert routed is host
def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
# The --resolve-prebuilt probe must agree with the install path: an
# auto-detected Intel host resolves against upstream (Vulkan), not the fork.
monkeypatch.setattr(
ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == UPSTREAM
assert out["repo"] == UPSTREAM

81
studio/backend/tests/test_llama_cpp_stream_cancel.py Normal file
View file

@ -0,0 +1,81 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import contextlib
import os
import sys
import threading
import httpx
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled
def _backend_stub() -> LlamaCppBackend:
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
backend._port = 48848
backend._effective_context_length = 4096
backend._supports_reasoning = False
backend._reasoning_always_on = False
backend._reasoning_style = "enable_thinking"
backend._supports_preserve_thinking = False
return backend
def test_stream_cancel_uses_internal_exception_not_generator_exit():
class FakeResponse:
status_code = 200
def close(self):
pass
class FakeStream:
def __enter__(self):
return FakeResponse()
def __exit__(self, *_args):
return False
class FakeClient:
def stream(self, *_args, **_kwargs):
return FakeStream()
cancel_event = threading.Event()
with pytest.raises(Exception) as exc_info:
with LlamaCppBackend._stream_with_retry(
FakeClient(),
"http://llama.test/v1/chat/completions",
{},
cancel_event,
):
cancel_event.set()
raise httpx.ReadError("client closed")
assert exc_info.type is _LlamaStreamCancelled
assert not issubclass(exc_info.type, GeneratorExit)
def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch):
backend = _backend_stub()
@contextlib.contextmanager
def fake_open_stream(*_args, **_kwargs):
raise _LlamaStreamCancelled
monkeypatch.setattr(backend, "_open_stream", fake_open_stream)
chunks = list(
backend.generate_chat_completion(
[{"role": "user", "content": "hi"}],
cancel_event = threading.Event(),
)
)
assert chunks == []

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

@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
# A Vulkan install (marker asset carries 'vulkan') must re-assert
# UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to
# CUDA/ROCm and silently replaces the Vulkan build.
install_dir = tmp_path / "llama.cpp"
binary = _write_install(
install_dir,
"b9493",
repo = "ggml-org/llama.cpp",
asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz",
)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
def _on_start(cmd):
_write_install(
install_dir,
"b9518",
repo = "ggml-org/llama.cpp",
asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
)
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
lines = ["installed\n"],
on_start = _on_start,
captured_kwargs = popen_kwargs,
)
assert upd.start_update()["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")

193
studio/backend/tests/test_llama_cpp_vulkan_probe.py Normal file
View file

@ -0,0 +1,193 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Vulkan free-VRAM reader regression tests on a synthetic probe output.
Covers the post-probe handling in
``LlamaCppBackend._get_gpu_free_memory_vulkan``:
* integrated GPUs (probe reports is_igpu=1) leave a flat per-device host
margin matching llama.cpp's --fit-target, so context auto-sizing can't
over-commit shared RAM, and report total 0 (shared RAM is not a budget),
* discrete GPUs (is_igpu=0) keep their free untouched and pass their real
total through so the fit can reserve absolute headroom,
* an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged
(ggml applies it), not stripped or filtered in Python -- the probe reports
ggml's compact ordinal, which load_model pins with ``--device Vulkan<i>``.
The ggml Vulkan library is never loaded: subprocess.run is mocked to emit
the tab-separated lines the real ``_vulkan_probe.py`` would print.
"""
from __future__ import annotations
import subprocess
import sys
import types as _types
from pathlib import Path
from unittest import mock
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import importlib as _importlib # noqa: E402
def _maybe_stub(name: str, builder):
try:
_importlib.import_module(name)
except ImportError:
sys.modules[name] = builder()
def _build_loggers_stub():
m = _types.ModuleType("loggers")
m.get_logger = lambda name: __import__("logging").getLogger(name)
return m
_maybe_stub("loggers", _build_loggers_stub)
_maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
from core.inference import llama_cpp as _llama_mod # noqa: E402
from core.inference.llama_cpp import ( # noqa: E402
LlamaCppBackend,
_llama_lib_dir,
_vulkan_lib_filename,
)
MIB = 1024 * 1024
GIB = 1024 * MIB
def _make_vulkan_install(tmp_path: Path) -> str:
"""A binary whose sibling dir holds the Vulkan ggml lib, so the
reader's ``is_vulkan_backend`` sibling-file check passes."""
bindir = tmp_path / "build" / "bin"
bindir.mkdir(parents = True)
binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server")
binary.write_bytes(b"stub")
(bindir / _vulkan_lib_filename()).write_bytes(b"stub")
return str(binary)
def _mock_probe(rows: list[str], captured_env: dict | None = None):
"""Patch subprocess.run so the _vulkan_probe.py call returns ``rows``
(already tab-formatted), recording the env it was launched with."""
real_run = subprocess.run
def fake_run(cmd, *args, **kwargs):
if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd):
if captured_env is not None:
captured_env.clear()
captured_env.update(kwargs.get("env") or {})
return subprocess.CompletedProcess(
args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = ""
)
return real_run(cmd, *args, **kwargs)
return mock.patch("subprocess.run", side_effect = fake_run)
def _row(
idx: int,
free_bytes: int,
is_igpu: int,
total_bytes: int = 0,
) -> str:
return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
def test_integrated_gpu_leaves_host_margin(tmp_path):
binary = _make_vulkan_install(tmp_path)
# iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target).
# total stays 0: shared system RAM is not a VRAM budget for the fit.
rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus
def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path):
binary = _make_vulkan_install(tmp_path)
# 6 GiB free on a partially occupied 24 GiB card: free is untouched and the
# real total flows through so the fit reserves absolute headroom (CUDA/ROCm
# parity) instead of the looser free*frac budget.
rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus
def test_large_discrete_gpu_is_untouched(tmp_path):
binary = _make_vulkan_install(tmp_path)
# A 48 GiB discrete card stays untouched regardless of size; only the
# iGPU flag triggers the host margin, never a VRAM/RAM ratio.
rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus
def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch):
# The mask is NOT stripped or filtered in Python: ggml parses it in raw
# physical-device space while this probe reports the compact post-filter
# ordinal, so mixing spaces would be wrong. It is passed through unchanged
# so ggml applies it to the same device list the launch will enumerate.
binary = _make_vulkan_install(tmp_path)
monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1")
captured: dict = {}
rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows, captured_env = captured):
LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured
def test_vulkan_pin_args_uses_device_names_not_env_mask():
# Pin by compact device name via --device (the space the probe reports and
# the registry names), never by writing a compact ordinal into the raw
# GGML_VK_VISIBLE_DEVICES index space.
assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"]
assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"]
assert LlamaCppBackend._vulkan_pin_args(None) == []
assert LlamaCppBackend._vulkan_pin_args([]) == []
def test_vulkan_only_build_is_detected(tmp_path):
binary = _make_vulkan_install(tmp_path)
assert LlamaCppBackend._is_vulkan_backend(binary) is True
def test_multi_backend_build_is_not_vulkan_only(tmp_path):
# A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be
# treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan
# device; defer to the CUDA/HIP path instead.
binary = _make_vulkan_install(tmp_path)
cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so"
(_llama_lib_dir(binary) / cuda).write_bytes(b"stub")
assert LlamaCppBackend._is_vulkan_backend(binary) is False
@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX")
def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path):
# create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root
# when it cannot symlink; _find_llama_server_binary returns that root entrypoint,
# so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else
# _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently
# never engage on a valid Vulkan install.
import os
binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib
bindir = Path(binary).parent
wrapper = tmp_path / "llama-server"
wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n')
os.chmod(wrapper, 0o755)
assert _llama_lib_dir(str(wrapper)) == bindir
assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))

84
studio/backend/tests/test_llama_route_timeouts.py
View file

@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel():
asyncio.run(_run(False))
asyncio.run(_run(True))
def test_stream_stall_timeout_callable_re_resolved_each_read():
# The OpenAI passthrough passes a callable so the stall bound can switch to
# the short post-terminal grace mid-stream; it must be re-resolved per read,
# not captured once at generator start.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
values = iter([100.0, 2.0])
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 3:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 1,
post_first_item_read_timeout_s = lambda: next(values, 5.0),
):
seen.append(response.request.extensions["timeout"].get("read"))
assert len(seen) == 3
# The callable is resolved right after the first item (arming the
# post-first window) and again before each later read, consuming
# successive values.
assert seen[0] == 100.0
assert 1.0 <= seen[1] <= 2.0
assert 4.0 <= seen[2] <= 5.0
asyncio.run(_run())
def test_stream_stall_timeout_disabled_clears_read_timeout():
# UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so
# the callable returns None. Once a chunk has arrived the leftover
# first-token read timeout must be cleared, else a long post-first-chunk gap
# trips a stale deadline the operator asked to turn off.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 2:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 5,
post_first_item_read_timeout_s = lambda: None,
):
seen.append(response.request.extensions["timeout"].get("read"))
# The first-token path armed a finite read timeout; after the first chunk
# with the guard disabled, it is cleared to None on every subsequent read.
assert seen == [None, None], seen
asyncio.run(_run())

183
studio/backend/tests/test_mtp_drafter_companion.py
View file

@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path:
from hub.utils.download_manifest import ExpectedFile
from hub.utils.gguf import is_mtp_drafter_path
from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files
from hub.utils.gguf_plan import (
build_gguf_variant_plans,
plan_from_expected_files,
preferred_mtp_sibling,
)
from utils.models.model_config import (
_is_mtp_drafter,
detect_gguf_model,
@ -37,6 +41,8 @@ from utils.models.model_config import (
DRAFTER_CASES = [
("mtp-gemma-4-12b-it.gguf", True),
("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True),
# New-scheme MTP/ copies carry the mtp- basename prefix too.
("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True),
("foo/MTP/bar.gguf", True),
("gemma-4-12b-it-Q8_0.gguf", False),
# Baked-in Qwen MTP repos: the head is inside the main GGUF, the file
@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path):
assert detect_gguf_model(str(copy)) is None
# Selecting the MTP dir itself must not surface the copies as models.
assert detect_gguf_model(str(sub)) is None
# ── Root drafter wins over new-scheme MTP/ copies ────────────────────
# The MTP/ copies were renamed to share the mtp- basename prefix (e.g.
# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the
# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes
# lowercase, so the subdir path would otherwise win).
NEW_SCHEME_SIBLINGS = [
_sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"),
_sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"),
_sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"),
_sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"),
_sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"),
_sib("mmproj-F16.gguf", 500, "mmproj"),
]
def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies():
picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS)
assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf"
def test_variant_plans_new_scheme_uses_root_drafter():
plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS)
assert set(plans) == {"q4_k_m", "q8_0"}
for plan in plans.values():
assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames
assert not any("MTP/" in name for name in plan.target_filenames)
assert "drafter" in plan.companion_hashes
# Download size = main + mmproj + root drafter (not the 200-byte BF16 copy).
assert plans["q4_k_m"].download_size_bytes == 4_600
def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch):
# _pick_mtp is nested; capture it via the companion-download seam.
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe
captured = {}
def _fake_companion(
*,
hf_repo,
hf_token,
pick,
label,
cancel_event = None,
):
captured["pick"] = pick
return None
b = LlamaCppBackend()
b._download_companion_gguf = _fake_companion
b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
repo_files = [
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
"MTP/mtp-gemma-4-E4B-it-Q4_0.gguf",
"MTP/mtp-gemma-4-E4B-it-Q8_0.gguf",
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
"mmproj-F16.gguf",
"mtp-gemma-4-E4B-it.gguf",
]
assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf"
# ── Reuse an on-disk drafter offline; fetch fresh online ─────────────
def _seed_snapshot(tmp_path, names):
snap = tmp_path / "snap"
for rel in names:
f = snap / rel
f.parent.mkdir(parents = True, exist_ok = True)
f.write_bytes(b"x")
return snap
def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch):
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snap = _seed_snapshot(
tmp_path,
[
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
"mtp-gemma-4-E4B-it.gguf",
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
"mmproj-F16.gguf",
],
)
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch):
# Pre-fix build may have fetched only the MTP/ copy; reuse it offline.
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snap = _seed_snapshot(
tmp_path,
[
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
],
)
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf"
def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch):
# A newer partial snapshot holds only the MTP/ copy; an older one has the
# root. Must still return the small root, not the large subdir copy.
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"])
snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"])
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch):
# Two snapshots both hold a root drafter; newest-first order must win so a
# fresh main GGUF is not paired with a stale drafter revision.
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"])
oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"])
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).parent.parent.name == "newest"
def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
# Online, do not reuse a cached copy: go to the download path so a changed
# drafter is refetched (hf_hub_download checks the current revision).
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"])
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
reached = {}
def _fake_companion(
*,
hf_repo,
hf_token,
pick,
label,
cancel_event = None,
):
reached["hit"] = True
return None
b = LlamaCppBackend()
b._download_companion_gguf = _fake_companion
assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None
assert reached.get("hit") is True

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

File diff suppressed because it is too large Load diff

1
studio/backend/tests/test_passthrough_healing.py
View file

@ -515,6 +515,7 @@ class ScriptedClient:
_url,
json = None,
timeout = None,
headers = None,
):
self.posts.append(json)
return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])

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

@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
logger.info("llama update: installing", cmd = " ".join(cmd))
# Stream progress lines into job["progress"].
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm
# box would otherwise re-route and silently replace the Vulkan build.
# Re-assert it via the same env flag setup uses (mirrors
# _rocm_install_args).
if asset and "vulkan" in asset.lower():
env["UNSLOTH_FORCE_VULKAN"] = "1"
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,

409
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -2,6 +2,10 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getAuthToken } from "@/features/auth";
import {
loadRememberedLoadSettings,
rememberedLoadSettingsKey,
} from "@/components/assistant-ui/model-selector/remembered-load-settings";
import { projectHasSources } from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
import { parseParamCountB } from "@/lib/model-size";
@ -63,11 +67,17 @@ import {
listStoredChatThreads,
updateStoredChatThread,
} from "../utils/chat-history-storage";
import {
readLastLocalModelLoad,
recordLastLocalModelLoad,
type LastLocalModelKind,
} from "../utils/last-local-model-load";
import { getImageInputUnavailableReason } from "../utils/image-input-support";
import {
hasClosedThinkTag,
parseAssistantContent,
} from "../utils/parse-assistant-content";
import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
import {
generateAudio,
listCachedGguf,
@ -1309,6 +1319,30 @@ const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi;
const GGUF_KNOWN_QUANT_RE =
/(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i;
type AutoLoadCandidate = {
id: string;
kind: LastLocalModelKind;
ggufVariant: string | null;
maxSeqLength: number;
successLabel: string;
};
function autoLoadCandidateKey(
kind: LastLocalModelKind,
id: string,
ggufVariant?: string | null,
): string {
return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`;
}
function findCachedRepo<T extends { repo_id: string }>(
repos: T[],
id: string,
): T | undefined {
const normalized = id.toLowerCase();
return repos.find((repo) => repo.repo_id.toLowerCase() === normalized);
}
function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean {
const normalized = filename.replace(/\\/g, "/").toLowerCase();
const separatorIndex = normalized.lastIndexOf("/");
@ -1357,14 +1391,18 @@ async function autoLoadSmallestModel(): Promise<{
const hfToken = store.hfToken || null;
const trustRemoteCode = store.params.trustRemoteCode ?? false;
const specSettings = resolveSpeculativeSettingsForLoad();
const lastLoaded = readLastLocalModelLoad();
const toastId = toast("Loading a model…", {
description: "Auto-selecting the smallest downloaded model.",
description: lastLoaded
? "Loading last used model."
: "Auto-selecting the smallest downloaded model.",
duration: 5000,
closeButton: true,
});
let blockedByTrustRemoteCode = false;
let hadNonTrustFailure = false;
let loadAttempts = 0;
const skippedAutoLoadCandidates = new Set<string>();
async function canAutoLoad(payload: {
model_path: string;
@ -1389,12 +1427,224 @@ async function autoLoadSmallestModel(): Promise<{
}
return true;
}
async function loadAutoLoadCandidate(
candidate: AutoLoadCandidate,
): Promise<boolean> {
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
return false;
}
const currentStore = useChatRuntimeStore.getState();
const remembered = loadRememberedLoadSettings(
rememberedLoadSettingsKey({
id: candidate.id,
ggufVariant: candidate.ggufVariant,
}),
);
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId: candidate.id,
ggufVariant: candidate.ggufVariant,
isGguf: candidate.kind === "gguf",
customContextLength: remembered?.contextLength ?? null,
ggufContextLength: null,
currentCheckpoint: currentStore.params.checkpoint,
activeGgufVariant: currentStore.activeGgufVariant,
maxSeqLength: candidate.maxSeqLength,
presetSource: currentStore.activePresetSource,
});
const effectiveSpeculativeType =
remembered?.speculativeType ?? specSettings.speculativeType;
const effectiveSpecDraftNMax =
remembered?.specDraftNMax ?? specSettings.specDraftNMax;
if (
!(await canAutoLoad({
model_path: candidate.id,
max_seq_length: effectiveMaxSeqLength,
is_lora: false,
gguf_variant: candidate.ggufVariant,
}))
) {
skippedAutoLoadCandidates.add(
autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant),
);
return false;
}
loadAttempts += 1;
const loadResp = await loadModel({
model_path: candidate.id,
hf_token: hfToken,
max_seq_length: effectiveMaxSeqLength,
load_in_4bit: true,
is_lora: false,
gguf_variant: candidate.ggufVariant,
trust_remote_code: trustRemoteCode,
cache_type_kv: remembered?.kvCacheDtype ?? null,
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
tensor_parallel: remembered?.tensorParallel ?? false,
});
saveSpeculativeType(effectiveSpeculativeType);
useChatRuntimeStore
.getState()
.setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
const store = useChatRuntimeStore.getState();
store.setModelRequiresTrustRemoteCode(
loadResp.requires_trust_remote_code ?? false,
);
store.setParams({
...store.params,
maxTokens:
candidate.kind === "gguf"
? loadResp.context_length ?? 131072
: effectiveMaxSeqLength,
});
const autoModel: ChatModelSummary = {
id: candidate.id,
name: loadResp.display_name ?? candidate.id,
isVision: loadResp.is_vision ?? false,
isLora: loadResp.is_lora ?? false,
isGguf: loadResp.is_gguf ?? candidate.kind === "gguf",
isAudio: loadResp.is_audio ?? false,
audioType: loadResp.audio_type ?? null,
hasAudioInput: loadResp.has_audio_input ?? false,
};
if (!store.models.some((m) => m.id === candidate.id)) {
store.setModels([...store.models, autoModel]);
}
if (candidate.kind === "gguf") {
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
ggufMaxContextLength:
loadResp.max_context_length ?? loadResp.context_length ?? 131072,
ggufNativeContextLength: loadResp.native_context_length ?? null,
supportsReasoning: loadResp.supports_reasoning ?? false,
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
reasoningEnabled: loadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(loadResp),
supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
supportsTools: loadResp.supports_tools ?? false,
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
kvCacheDtype: loadResp.cache_type_kv ?? null,
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
...resolveLoadedSpeculativeSettings(loadResp),
});
} else {
useChatRuntimeStore.setState({
supportsReasoning: loadResp.supports_reasoning ?? false,
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
reasoningEnabled: loadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(loadResp),
supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
supportsTools: loadResp.supports_tools ?? false,
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
kvCacheDtype: loadResp.cache_type_kv ?? null,
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
...resolveLoadedSpeculativeSettings(loadResp),
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
});
}
if (!(loadResp.is_lora ?? false)) {
recordLastLocalModelLoad({
id: candidate.id,
kind: candidate.kind,
ggufVariant: candidate.ggufVariant,
});
}
toast.success(candidate.successLabel, { id: toastId });
return true;
}
try {
const [ggufRepos, modelRepos] = await Promise.all([
listCachedGguf().catch(() => []),
listCachedModels().catch(() => []),
]);
if (lastLoaded) {
if (lastLoaded.kind === "gguf") {
const repo = findCachedRepo(ggufRepos, lastLoaded.id);
if (repo && lastLoaded.ggufVariant) {
try {
const variants = await listGgufVariants(repo.repo_id);
const variant = variants.variants.find(
(entry) =>
entry.downloaded &&
entry.quant?.toLowerCase() ===
lastLoaded.ggufVariant?.toLowerCase() &&
isAutoLoadableGgufVariant(entry),
);
if (variant) {
toast("Loading last used model…", {
id: toastId,
description: `${repo.repo_id} (${variant.quant})`,
duration: 5000,
});
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "gguf",
ggufVariant: variant.quant,
maxSeqLength: 0,
successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
}
} catch {
hadNonTrustFailure = true;
skippedAutoLoadCandidates.add(
autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant),
);
}
}
} else {
const repo = findCachedRepo(modelRepos, lastLoaded.id);
if (repo) {
try {
toast("Loading last used model…", {
id: toastId,
description: repo.repo_id,
duration: 5000,
});
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "model",
ggufVariant: null,
maxSeqLength: store.params.maxSeqLength,
successLabel: `Loaded ${repo.repo_id}`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
} catch {
hadNonTrustFailure = true;
skippedAutoLoadCandidates.add(
autoLoadCandidateKey("model", repo.repo_id),
);
}
}
}
toast("Loading a model…", {
id: toastId,
description: "Auto-selecting the smallest downloaded model.",
duration: 5000,
});
}
// GGUF first: smallest-total-size repo, then its smallest variant.
if (ggufRepos.length > 0) {
const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes);
@ -1408,82 +1658,23 @@ async function autoLoadSmallestModel(): Promise<{
if (downloaded.length > 0) {
const variant = downloaded[0];
if (
!(await canAutoLoad({
model_path: repo.repo_id,
max_seq_length: 0,
is_lora: false,
gguf_variant: variant.quant,
}))
skippedAutoLoadCandidates.has(
autoLoadCandidateKey("gguf", repo.repo_id, variant.quant),
)
) {
continue;
}
loadAttempts += 1;
const loadResp = await loadModel({
model_path: repo.repo_id,
hf_token: hfToken,
max_seq_length: 0,
load_in_4bit: true,
is_lora: false,
gguf_variant: variant.quant,
trust_remote_code: trustRemoteCode,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
});
saveSpeculativeType(specSettings.speculativeType);
useChatRuntimeStore
.getState()
.setCheckpoint(repo.repo_id, variant.quant);
const store = useChatRuntimeStore.getState();
store.setModelRequiresTrustRemoteCode(
loadResp.requires_trust_remote_code ?? false,
);
store.setParams({
...store.params,
maxTokens: loadResp.context_length ?? 131072,
});
// Add to store so the selector shows the name.
const autoModel: ChatModelSummary = {
id: repo.repo_id,
name: loadResp.display_name ?? repo.repo_id,
isVision: loadResp.is_vision ?? false,
isLora: loadResp.is_lora ?? false,
isGguf: loadResp.is_gguf ?? false,
isAudio: loadResp.is_audio ?? false,
audioType: loadResp.audio_type ?? null,
hasAudioInput: loadResp.has_audio_input ?? false,
};
const existingModels = store.models;
if (!existingModels.some((m) => m.id === repo.repo_id)) {
store.setModels([...existingModels, autoModel]);
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "gguf",
ggufVariant: variant.quant,
maxSeqLength: 0,
successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
ggufMaxContextLength:
loadResp.max_context_length ??
loadResp.context_length ??
131072,
supportsReasoning: loadResp.supports_reasoning ?? false,
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
reasoningEnabled: loadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(loadResp),
supportsPreserveThinking:
loadResp.supports_preserve_thinking ?? false,
supportsTools: loadResp.supports_tools ?? false,
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
kvCacheDtype: loadResp.cache_type_kv ?? null,
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
loadedIsMultimodal: isMultimodalResponse(loadResp),
...resolveLoadedSpeculativeSettings(loadResp),
});
toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, {
id: toastId,
});
return { loaded: true, blockedByTrustRemoteCode: false };
}
} catch {
hadNonTrustFailure = true;
@ -1501,64 +1692,23 @@ async function autoLoadSmallestModel(): Promise<{
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
try {
if (
!(await canAutoLoad({
model_path: repo.repo_id,
max_seq_length: 4096,
is_lora: false,
gguf_variant: null,
}))
skippedAutoLoadCandidates.has(
autoLoadCandidateKey("model", repo.repo_id),
)
) {
continue;
}
loadAttempts += 1;
const sfLoadResp = await loadModel({
model_path: repo.repo_id,
hf_token: hfToken,
max_seq_length: 4096,
load_in_4bit: true,
is_lora: false,
gguf_variant: null,
trust_remote_code: trustRemoteCode,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
});
saveSpeculativeType(specSettings.speculativeType);
useChatRuntimeStore.getState().setCheckpoint(repo.repo_id);
const store = useChatRuntimeStore.getState();
store.setModelRequiresTrustRemoteCode(
sfLoadResp.requires_trust_remote_code ?? false,
);
store.setParams({ ...store.params, maxTokens: 4096 });
useChatRuntimeStore.setState({
supportsReasoning: sfLoadResp.supports_reasoning ?? false,
reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false,
reasoningEnabled: sfLoadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(sfLoadResp),
supportsPreserveThinking:
sfLoadResp.supports_preserve_thinking ?? false,
supportsTools: sfLoadResp.supports_tools ?? false,
// Parity with the GGUF branch above.
...resolveToolsEnabledOnLoad(sfLoadResp.supports_tools ?? false),
defaultChatTemplate: sfLoadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
...resolveLoadedSpeculativeSettings(sfLoadResp),
});
const sfModel: ChatModelSummary = {
id: repo.repo_id,
name: sfLoadResp.display_name ?? repo.repo_id,
isVision: sfLoadResp.is_vision ?? false,
isLora: sfLoadResp.is_lora ?? false,
isGguf: sfLoadResp.is_gguf ?? false,
};
if (!store.models.some((m) => m.id === repo.repo_id)) {
store.setModels([...store.models, sfModel]);
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "model",
ggufVariant: null,
maxSeqLength: 4096,
successLabel: `Loaded ${repo.repo_id}`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
useChatRuntimeStore.setState({
loadedIsMultimodal: isMultimodalResponse(sfLoadResp),
});
toast.success(`Loaded ${repo.repo_id}`, { id: toastId });
return { loaded: true, blockedByTrustRemoteCode: false };
} catch {
hadNonTrustFailure = true;
continue;
@ -1650,6 +1800,11 @@ async function autoLoadSmallestModel(): Promise<{
loadedIsMultimodal: isMultimodalResponse(loadResp),
...resolveLoadedSpeculativeSettings(loadResp),
});
recordLastLocalModelLoad({
id: "unsloth/Qwen3.5-4B-MTP-GGUF",
kind: "gguf",
ggufVariant: "UD-Q4_K_XL",
});
toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId });
return { loaded: true, blockedByTrustRemoteCode: false };
} catch {

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

@ -44,6 +44,7 @@ import {
mergeBackendRecommendedInference,
resolveLoadMaxSeqLength,
} from "../presets/preset-policy";
import { recordLastLocalModelLoad } from "../utils/last-local-model-load";
import {
isMultimodalResponse,
} from "../types/api";
@ -818,6 +819,23 @@ export function useChatModelRuntime() {
}
}
await refresh({ signal: abortCtrl.signal });
if (
!isLora &&
!(loadResponse.is_lora ?? false) &&
!nativePathToken &&
!isLocalModelPath(modelId) &&
!isExternalModelId(modelId)
) {
if (loadResponse.is_gguf || isGguf || ggufVariant) {
recordLastLocalModelLoad({
id: modelId,
kind: "gguf",
ggufVariant: ggufVariant ?? null,
});
} else {
recordLastLocalModelLoad({ id: modelId, kind: "model" });
}
}
// A successful load owns the shared (pick-unscoped) settings fields,
// so any surviving stage is stale: the just-loaded pick itself, or a
// pick queued for a different model mid-load whose knobs this load

86
studio/frontend/src/features/chat/utils/last-local-model-load.ts Normal file
View file

@ -0,0 +1,86 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export type LastLocalModelKind = "gguf" | "model";
export type LastLocalModelLoad = {
id: string;
kind: LastLocalModelKind;
ggufVariant: string | null;
loadedAt: number;
};
const STORAGE_KEY = "unsloth.last-local-model-load.v1";
function storage(): Storage | null {
try {
return typeof localStorage === "undefined" ? null : localStorage;
} catch {
return null;
}
}
function isLastLocalModelKind(value: unknown): value is LastLocalModelKind {
return value === "gguf" || value === "model";
}
export function readLastLocalModelLoad(): LastLocalModelLoad | null {
try {
const raw = storage()?.getItem(STORAGE_KEY);
if (!raw) {
return null;
}
const parsed = JSON.parse(raw) as Partial<LastLocalModelLoad>;
if (
typeof parsed.id !== "string" ||
!parsed.id.trim() ||
!isLastLocalModelKind(parsed.kind) ||
typeof parsed.loadedAt !== "number"
) {
return null;
}
if (
parsed.kind === "gguf" &&
(typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim())
) {
return null;
}
return {
id: parsed.id,
kind: parsed.kind,
ggufVariant:
typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null,
loadedAt: parsed.loadedAt,
};
} catch {
return null;
}
}
export function recordLastLocalModelLoad(input: {
id: string;
kind: LastLocalModelKind;
ggufVariant?: string | null;
}): void {
const id = input.id.trim();
if (!id) {
return;
}
const ggufVariant = input.ggufVariant?.trim() || null;
if (input.kind === "gguf" && !ggufVariant) {
return;
}
try {
storage()?.setItem(
STORAGE_KEY,
JSON.stringify({
id,
kind: input.kind,
ggufVariant: input.kind === "gguf" ? ggufVariant : null,
loadedAt: Date.now(),
} satisfies LastLocalModelLoad),
);
} catch {
// Ignore disabled storage / quota errors; auto-load falls back to size order.
}
}

399
studio/install_llama_prebuilt.py
View file

@ -10,6 +10,7 @@ import argparse
import atexit
import errno
import fnmatch
import glob
import hashlib
import json
import os
@ -265,6 +266,7 @@ class HostInfo:
has_physical_nvidia: bool
has_usable_nvidia: bool
has_rocm: bool = False
has_intel_gpu: bool = False
rocm_gfx_target: str | None = None
# (major, minor) from platform.mac_ver(); None off macOS or if unparseable.
# Skips a macos prebuilt whose minimum-OS exceeds this host.
@ -1284,162 +1286,6 @@ def synthetic_checksums_for_release(
)
def parse_direct_linux_release_bundle(
repo: str, release: dict[str, Any]
) -> PublishedReleaseBundle | None:
release_tag = release.get("tag_name")
if not isinstance(release_tag, str) or not release_tag:
return None
assets = release_asset_map(release)
artifacts: list[PublishedLlamaArtifact] = []
inferred_labels: list[str] = []
linux_asset_re = re.compile(
r"^app-(?P<label>.+)-(?P<target>linux-x64(?:-cpu)?|linux-x64-cuda\d+-(?:older|newer|portable))\.tar\.gz$"
)
for asset_name in sorted(assets):
match = linux_asset_re.fullmatch(asset_name)
if not match:
continue
inferred_labels.append(match.group("label"))
target = match.group("target")
if target in {"linux-x64", "linux-x64-cpu"}:
artifacts.append(
PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = "linux-cpu",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
rank = 1000,
)
)
continue
bundle_profile = target.removeprefix("linux-x64-")
profile = _resolve_linux_bundle_profile(bundle_profile)
if profile is None:
continue
artifacts.append(
PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = "linux-cuda",
runtime_line = str(profile["runtime_line"]),
coverage_class = str(profile["coverage_class"]),
supported_sms = [str(value) for value in profile["supported_sms"]],
min_sm = int(profile["min_sm"]),
max_sm = int(profile["max_sm"]),
bundle_profile = bundle_profile,
rank = int(profile["rank"]),
)
)
if not artifacts:
return None
upstream_tag = (
release_tag
if is_release_tag_like(release_tag)
else inferred_labels[0]
if len(set(inferred_labels)) == 1 and inferred_labels
else release_tag
)
selection_log = [
f"published_release: repo={repo}",
f"published_release: tag={release_tag}",
f"published_release: upstream_tag={upstream_tag}",
"published_release: direct_asset_scan=linux",
]
return PublishedReleaseBundle(
repo = repo,
release_tag = release_tag,
upstream_tag = upstream_tag,
assets = assets,
manifest_asset_name = DEFAULT_PUBLISHED_MANIFEST_ASSET,
artifacts = artifacts,
selection_log = selection_log,
)
def direct_linux_release_plan(
release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
) -> InstallReleasePlan | None:
bundle = parse_direct_linux_release_bundle(repo, release)
if bundle is None:
return None
if not direct_release_matches_request(
release_tag = bundle.release_tag,
llama_tag = bundle.upstream_tag,
requested_tag = requested_tag,
):
return None
attempts: list[AssetChoice] = []
if host.has_usable_nvidia:
# Prefer the cudart major Studio loads at runtime (torch's bundled
# libcudart), not the newest on disk. Otherwise a stray cuda13
# runtime outranks the torch cuda12 the binary links against.
torch_preference = detect_torch_cuda_runtime_preference(host)
selection = linux_cuda_choice_from_release(
host,
bundle,
preferred_runtime_line = torch_preference.runtime_line,
selection_preamble = torch_preference.selection_log,
)
if selection is not None:
attempts.extend(selection.attempts)
elif not host.has_rocm:
# A ROCm-only host gets no CPU asset: leaving attempts empty lets the
# raise below trigger a HIP source build instead of shipping a CPU
# binary on a GPU host (this ggml-org path has no per-gfx ROCm asset).
cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
if cpu_choice is not None:
attempts.append(cpu_choice)
# NVIDIA hosts whose CUDA selection produced nothing fall through to the
# raise below (mirroring the ROCm policy above): the caller then walks
# back to an older release that still ships a usable CUDA line instead of
# silently installing a CPU binary on a GPU host. Today's walk-back only
# works because partial releases ship no CPU bundle; this keeps it working
# if a future partial release does.
if not attempts:
raise PrebuiltFallback("no compatible Linux prebuilt asset was found")
approved_checksums = synthetic_checksums_for_release(
repo,
bundle.release_tag,
bundle.upstream_tag,
)
resolved_upstream_tag = bundle.upstream_tag
if DEFAULT_PUBLISHED_SHA256_ASSET in bundle.assets and not is_release_tag_like(
bundle.upstream_tag
):
approved_checksums = load_approved_release_checksums(repo, bundle.release_tag)
# Require exact source provenance for branch/pull/commit releases.
# Mirrors validated_checksums_for_bundle so incomplete metadata fails
# closed instead of degrading to the legacy branch-as-tag source
# hydration path this PR eliminates.
if (
not approved_checksums.source_commit
or exact_source_archive_hash(approved_checksums) is None
or source_clone_url_from_checksums(approved_checksums) is None
):
raise PrebuiltFallback(
f"approved checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET} for "
f"{repo}@{bundle.release_tag} did not contain exact source provenance"
)
attempts = apply_approved_hashes(attempts, approved_checksums)
return InstallReleasePlan(
requested_tag = requested_tag,
llama_tag = resolved_upstream_tag,
release_tag = bundle.release_tag,
attempts = attempts,
approved_checksums = approved_checksums,
)
def direct_upstream_release_plan(
release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
) -> InstallReleasePlan | None:
@ -1482,6 +1328,24 @@ def direct_upstream_release_plan(
install_kind = "windows-hip",
)
)
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. Gate
# on no PHYSICAL NVIDIA (not just no usable one): a host that hid NVIDIA
# via CUDA_VISIBLE_DEVICES must not reach Vulkan, which ignores that mask
# and could enumerate the reserved card. Falls through to CPU below.
elif host.has_intel_gpu and not host.has_physical_nvidia:
vulkan_asset = f"llama-{release_tag}-bin-win-vulkan-x64.zip"
vulkan_url = assets.get(vulkan_asset)
if vulkan_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = vulkan_asset,
url = vulkan_url,
source_label = "upstream",
install_kind = "windows-vulkan",
)
)
cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip"
cpu_url = assets.get(cpu_asset)
if cpu_url:
@ -1545,6 +1409,23 @@ def direct_upstream_release_plan(
# ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm
# asset, so they fall through to the empty-attempts raise (HIP source
# build) rather than silently getting a CPU binary on a GPU host.
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. The
# elif already excludes usable NVIDIA and ROCm; also require no PHYSICAL
# NVIDIA so a CUDA-hidden card isn't reached through Vulkan (CPU below).
if host.has_intel_gpu and not host.has_physical_nvidia:
vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-x64.tar.gz"
vulkan_url = assets.get(vulkan_asset)
if vulkan_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = vulkan_asset,
url = vulkan_url,
source_label = "upstream",
install_kind = "linux-vulkan",
)
)
asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz"
asset_url = assets.get(asset_name)
if asset_url:
@ -1564,6 +1445,23 @@ def direct_upstream_release_plan(
# selector returned 0 attempts and the installer fell back to a
# source build on every Linux ARM64 host (DGX Spark, Ampere
# Altra, GitHub-hosted ubuntu-24.04-arm runners, etc.).
# Intel (or other non-NVIDIA/non-AMD) GPU: prefer the Vulkan prebuilt,
# mirroring the x86_64 branch. Upstream ships bin-ubuntu-vulkan-arm64.
# No physical NVIDIA: don't reach a CUDA-hidden card through Vulkan.
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-arm64.tar.gz"
vulkan_url = assets.get(vulkan_asset)
if vulkan_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = vulkan_asset,
url = vulkan_url,
source_label = "upstream",
install_kind = "linux-vulkan",
)
)
asset_name = f"llama-{release_tag}-bin-ubuntu-arm64.tar.gz"
asset_url = assets.get(asset_name)
if asset_url:
@ -3075,6 +2973,40 @@ def detect_host() -> HostInfo:
# Note: amdhip64.dll presence alone is NOT treated as GPU evidence
# since the HIP SDK can be installed without an AMD GPU.
# Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs
# vendor id (0x8086); Windows queries the WMI video controller list. Only
# probed with no usable NVIDIA and no ROCm (matching the Vulkan branches),
# keeping the probe (notably the Windows powershell call) off that path.
has_intel_gpu = False
if not has_usable_nvidia and not has_rocm:
if is_linux:
for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"):
try:
with open(_vendor_file) as _vf:
if _vf.read().strip().lower() == "0x8086":
has_intel_gpu = True
break
except OSError:
continue
elif is_windows:
_ps = shutil.which("powershell") or shutil.which("pwsh")
if _ps:
try:
_result = run_capture(
[
_ps,
"-NoProfile",
"-Command",
"Get-CimInstance Win32_VideoController | "
"Select-Object -ExpandProperty Name",
],
timeout = 15,
)
if _result.returncode == 0 and "intel" in _result.stdout.lower():
has_intel_gpu = True
except Exception:
pass
return HostInfo(
system = system,
machine = machine,
@ -3090,6 +3022,7 @@ def detect_host() -> HostInfo:
has_physical_nvidia = has_physical_nvidia,
has_usable_nvidia = has_usable_nvidia,
has_rocm = has_rocm,
has_intel_gpu = has_intel_gpu,
rocm_gfx_target = rocm_gfx_target,
macos_version = macos_version,
)
@ -3126,6 +3059,7 @@ def _apply_host_overrides(
has_physical_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
has_intel_gpu = False,
)
gfx = _normalize_forwarded_gfx(override_rocm_gfx)
if gfx:
@ -3866,6 +3800,23 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
"falling back to source build with HIP support"
)
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. No
# physical NVIDIA (not just no usable one): a CUDA-hidden card must not
# be reached through Vulkan, which ignores CUDA_VISIBLE_DEVICES.
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz"
if vulkan_name in upstream_assets:
log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}")
return AssetChoice(
repo = UPSTREAM_REPO,
tag = llama_tag,
name = vulkan_name,
url = upstream_assets[vulkan_name],
source_label = "upstream",
install_kind = "linux-vulkan",
)
log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU")
upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz"
if upstream_name not in upstream_assets:
raise PrebuiltFallback("upstream Linux CPU asset was not found")
@ -3908,6 +3859,24 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
)
log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU")
# Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No
# physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan.
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
vulkan_name = f"llama-{llama_tag}-bin-win-vulkan-x64.zip"
if vulkan_name in upstream_assets:
log(
f"Intel GPU detected on Windows -- using upstream Vulkan prebuilt {vulkan_name}"
)
return AssetChoice(
repo = UPSTREAM_REPO,
tag = llama_tag,
name = vulkan_name,
url = upstream_assets[vulkan_name],
source_label = "upstream",
install_kind = "windows-vulkan",
)
log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU")
upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip"
if upstream_name not in upstream_assets:
raise PrebuiltFallback("upstream Windows CPU asset was not found")
@ -4503,6 +4472,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
"linux-arm64-cuda",
"linux-rocm",
"linux-arm64",
"linux-vulkan",
}:
return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"]
if choice.install_kind in {"macos-arm64", "macos-x64"}:
@ -4516,6 +4486,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
"windows-cpu",
"windows-cuda",
"windows-hip",
"windows-vulkan",
"windows-rocm",
"windows-arm64",
}:
@ -5731,8 +5702,10 @@ def validate_server(
"linux-cuda",
"linux-arm64-cuda",
"linux-rocm",
"linux-vulkan",
"windows-cuda",
"windows-hip",
"windows-vulkan",
"windows-rocm",
"macos-arm64",
}
@ -6354,6 +6327,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libmtmd.so*"],
["libggml-hip.so*"],
]
if choice.install_kind == "linux-vulkan":
return [
["libllama-common.so*"],
["libllama.so*"],
["libggml.so*"],
["libggml-base.so*"],
# Match the sibling globs (linux-cuda/-rocm): x64 bundles ship
# arch-suffixed libggml-cpu-<variant>.so, arm64 may ship a bare
# libggml-cpu.so; the '-' form missed the latter and re-flagged
# the install unhealthy on every check.
["libggml-cpu*.so*"],
["libmtmd.so*"],
["libggml-vulkan.so*"],
]
if choice.install_kind in {"windows-cpu", "windows-arm64"}:
return [["llama.dll"]]
if choice.install_kind == "windows-cuda":
@ -6373,6 +6360,8 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
return groups
if choice.install_kind in {"windows-hip", "windows-rocm"}:
return [["llama.dll"], ["*hip*.dll"]]
if choice.install_kind == "windows-vulkan":
return [["llama.dll"], ["ggml-vulkan.dll"]]
return []
@ -6654,6 +6643,89 @@ def validate_prebuilt_attempts(
raise PrebuiltFallback("no prebuilt bundle passed validation")
def force_vulkan_requested() -> bool:
"""Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp
prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can
run the Vulkan build for inference). Scoped to the llama.cpp backend; the
torch/training stack installs separately and still sees the real GPU.
"""
return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in (
"1",
"true",
"yes",
)
def _vulkan_only_host(host: HostInfo) -> HostInfo:
"""Rewrite ``host`` so the asset selectors take their Vulkan branch.
That branch fires on ``has_intel_gpu and not nvidia and not rocm``, so clear
the CUDA/ROCm flags and raise the integrated-GPU flag. The synthetic flag
never leaves install planning -- it only routes the llama.cpp prebuilt
choice, not the torch/training stack.
"""
return dataclasses_replace(
host,
has_usable_nvidia = False,
has_physical_nvidia = False,
has_rocm = False,
has_intel_gpu = True,
)
def _route_to_vulkan_prebuilt(
host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool
) -> tuple[HostInfo, str, str]:
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
from UPSTREAM_REPO. Two triggers route here, both suppressed under
--cpu-fallback (the explicit "give me CPU" last resort wins):
* UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
Applied by BOTH the install path and the --resolve-prebuilt probe so the
"is a prebuilt available" answer matches what actually gets installed.
Returns the (possibly rewritten) host, repo, and release tag.
"""
forced = force_vulkan_requested()
# Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed
# NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps
# has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores
# CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the
# reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides.
auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
if force_cpu or not (forced or auto_intel):
return host, published_repo, published_release_tag
if host.is_macos:
if forced:
log(
"UNSLOTH_FORCE_VULKAN is set but ignored on macOS "
"(Metal is used; there is no Vulkan prebuilt)"
)
return host, published_repo, published_release_tag
if forced:
log(
"UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan "
"llama.cpp prebuilt instead of the detected GPU backend"
)
# Forcing may override a detected NVIDIA/ROCm host, so normalize it to
# Vulkan-only; an auto-detected Intel host already is.
host = _vulkan_only_host(host)
else:
log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt")
# Swapping the fork for upstream invalidates a fork release pin: the two use
# different tag namespaces (fork b9596-mix-<sha> vs upstream b9596), so a
# pinned fork tag would make the upstream resolver query a nonexistent
# release and fall back to source. Drop it and let the upstream resolver
# pick by the requested llama tag. A pin already on an explicit upstream repo
# (repo unchanged here) is preserved.
if published_repo != UPSTREAM_REPO:
published_release_tag = ""
return host, UPSTREAM_REPO, published_release_tag
def diffusion_visual_server_backfill_needed(
install_dir: Path, host: HostInfo, choice: AssetChoice
) -> bool:
@ -6696,6 +6768,9 @@ def install_prebuilt(
override_rocm_gfx = override_rocm_gfx,
force_cpu = force_cpu,
)
host, published_repo, published_release_tag = _route_to_vulkan_prebuilt(
host, published_repo, published_release_tag, force_cpu = force_cpu
)
choice: AssetChoice | None = None
try:
with install_lock(install_lock_path(install_dir)):
@ -6708,7 +6783,9 @@ def install_prebuilt(
f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install"
)
# Single resolver: every fork host selects from the release manifest;
# an explicit ggml-org override selects by asset filename instead.
# an explicit ggml-org override selects by asset filename instead. A
# forced-Vulkan host already has published_repo pointed at
# UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch.
requested_tag, release_plans = resolve_simple_install_release_plans(
llama_tag,
host,
@ -6994,10 +7071,14 @@ def main() -> int:
override_rocm_gfx = args.rocm_gfx,
force_cpu = args.cpu_fallback,
)
repo = args.published_repo
# Same Vulkan routing the install path applies, so the probe's answer
# matches what would install (an Intel/forced-Vulkan host -> upstream).
host, repo, release_tag = _route_to_vulkan_prebuilt(
host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback
)
try:
_requested, plans = resolve_simple_install_release_plans(
args.resolve_prebuilt, host, repo, args.published_release_tag or ""
args.resolve_prebuilt, host, repo, release_tag
)
choice = plans[0].attempts[0] if plans and plans[0].attempts else None
if choice is None:

13
studio/setup.ps1
View file

@ -2621,7 +2621,18 @@ function Fast-Install {
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
if ($UseUv) {
$VenvPy = (Get-Command python).Source
$result = & uv pip install --python $VenvPy @Args_ 2>&1
# An explicit --index-url must win. Inherited uv index env vars otherwise
# override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
# them only for index-pinned installs; mirrors still apply elsewhere.
$saved = @{}
if (@($Args_) -contains '--index-url') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
}
try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
if ($LASTEXITCODE -eq 0) { return }
}
& python -m pip install @Args_ 2>&1

122
tests/python/test_fast_sentence_transformer_embedding_parity.py Normal file
View file

@ -0,0 +1,122 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text
like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a
"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building
via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings;
`_create_transformer_module` uses `Transformer.load(...)` instead.
Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST
is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity,
opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected).
"""
from __future__ import annotations
import inspect
import os
import pytest
def test_transformer_load_signature_supports_unsloth_kwargs():
"""Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs
the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back
to Transformer(...) there, so mirror that gate and skip."""
models = pytest.importorskip("sentence_transformers.models")
load = getattr(models.Transformer, "load", None)
assert callable(load), (
"sentence_transformers Transformer.load is missing; the #6881 fix in "
"unsloth.models.sentence_transformer._create_transformer_module depends on it."
)
params = inspect.signature(load).parameters
accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
# Mirror _create_transformer_module's hub_capable gate.
hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
if not hub_capable:
pytest.skip(
"legacy Transformer.load(input_path); production path falls back to Transformer(...)"
)
unsupported = [
k
for k in ("token", "cache_folder", "revision", "trust_remote_code")
if not (accepts_var_kw or k in params)
]
assert not unsupported, (
f"installed sentence_transformers Transformer.load no longer accepts {unsupported} "
f"and has no **kwargs; update _create_transformer_module (#6881) before it silently "
f"falls back to Transformer(...)."
)
def _probe_texts():
return [
"roasted chickpeas in 20 kg bags",
"The capital of France is Paris.",
"A fast brown fox jumps over the lazy dog.",
"recette de tarte aux pommes traditionnelle",
]
def test_fast_sentence_transformer_matches_stock_st():
"""End-to-end: FastSentenceTransformer embeddings and tokenization must match a
stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and
GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners."""
model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL")
if not model_id:
pytest.skip(
"set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model "
"(HF id or local path) to run the #6881 parity test"
)
torch = pytest.importorskip("torch")
if not torch.cuda.is_available():
pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
np = pytest.importorskip("numpy")
pytest.importorskip("sentence_transformers")
from sentence_transformers import SentenceTransformer
device = "cuda"
# Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native
# embedders such as EmbeddingGemma (Gemma3), which would mask real parity.
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
texts = _probe_texts()
max_seq_length = 256
# Control FIRST, before importing unsloth, so its global import patches never
# touch the stock reference (mirrors the issue's "restart runtime" repro).
ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
ctrl.max_seq_length = max_seq_length
ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
ctrl_emb = np.asarray(
ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
)
import unsloth # noqa: F401
from unsloth import FastSentenceTransformer
fast = FastSentenceTransformer.from_pretrained(
model_id,
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = False,
load_in_16bit = True,
)
fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist()
fast_emb = np.asarray(
fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
)
# Identical tokenization = no chat-template wrapping slipped in (the #6881 defect).
assert fast_ids == ctrl_ids, (
f"tokenization diverged (chat-template wrapping regressed?):\n"
f" stock: {ctrl_ids}\n fast: {fast_ids}"
)
cos = (ctrl_emb * fast_emb).sum(1) / (
np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1)
)
assert float(cos.min()) > 0.99, (
f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 "
f"(per-text {[round(float(c), 5) for c in cos]})"
)

51
tests/python/test_tokenizers_and_torch_constraint.py
View file

@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
_INSTALL_SH = _REPO_ROOT / "install.sh"
_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged:
assert '"torch>=2.4,<2.11.0"' in self._ps1
class TestInstallPs1UvDefaultIndex:
"""Installer-managed torch indexes must override inherited uv defaults."""
_ps1 = _read(_INSTALL_PS1)
def test_torch_installs_use_default_index(self):
assert "--default-index $TorchIndexUrl" in self._ps1
assert "--default-index $ROCmIndexUrl" in self._ps1
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert "--index-url $TorchIndexUrl" not in self._ps1
assert "--index-url $ROCmIndexUrl" not in self._ps1
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# Extra-index vars outrank --default-index, so pinned installs must clear them.
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
assert 'Remove-Item "Env:$n"' in self._ps1
class TestSetupPs1FastInstallIndex:
"""setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
_ps1 = _read(_SETUP_PS1)
def test_fast_install_clears_all_uv_index_env_vars(self):
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
# Must truly remove the vars (child sees no value), not set them empty.
assert 'Remove-Item "Env:$n"' in self._ps1
class TestInstallShUvDefaultIndex:
"""Linux/Mac installer torch indexes must override inherited uv defaults."""
_sh = _read(_INSTALL_SH)
def test_torch_installs_use_default_index(self):
assert '--default-index "$TORCH_INDEX_URL"' in self._sh
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# --default-index installs run with all uv index env vars unset via `env -u`.
assert (
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
)
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
class TestTorchConstraintShell:
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""

2
tests/sh/test_mac_intel_compat.sh
View file

@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then
else
echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
TORCH_EOF

62
tests/studio/install/test_selection_logic.py
View file

@ -2684,68 +2684,6 @@ class TestBlackwellCuda124Exclusion:
assert kept == [cpu]
# N.1c3. direct_linux_release_plan -- no silent CPU on NVIDIA hosts
class TestDirectLinuxNvidiaCpuGate:
"""A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
def _bundle_cpu_only(self):
return make_release(
[
make_artifact(
"llama-b8508-bin-ubuntu-x64.tar.gz",
install_kind = "linux-cpu",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
),
]
)
def _patch(self, monkeypatch):
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"parse_direct_linux_release_bundle",
lambda repo, release: self._bundle_cpu_only(),
)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"detect_torch_cuda_runtime_preference",
lambda host: CudaRuntimePreference(runtime_line = None, selection_log = []),
)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"detected_linux_runtime_lines",
lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
)
def test_nvidia_host_without_cuda_line_raises_for_walkback(self, monkeypatch):
self._patch(monkeypatch)
host = make_host(driver_cuda_version = (13, 1), compute_caps = ["100"])
with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt"):
INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
)
def test_cpu_host_still_gets_cpu_bundle(self, monkeypatch):
self._patch(monkeypatch)
host = make_host(
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
has_physical_nvidia = False,
has_usable_nvidia = False,
)
plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
)
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
class TestLinuxPublishedAttemptsNvidiaCpuGate:
"""Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""

96
tests/studio/playwright_chat_ui.py
View file

@ -1296,16 +1296,92 @@ with sync_playwright() as p:
# still abort or interrupt this navigation, so the field wait below is the
# final confirmation that we reached /login.
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
try:
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
except Exception as exc:
if not any(t in str(exc) for t in _tolerated_nav):
raise
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
pw_field = page.locator("#password")
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW2)
page.locator('button[type="submit"]').click()
# A slow CI runner can make this re-login navigation time out even with the
# server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors
# the change-password retry above). wait_for_health is a diagnostic pre-gate.
wait_for_health(BASE, timeout = 30.0, info = info)
relogin_err: Exception | None = None
for _relogin_attempt in range(3):
try:
try:
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
except Exception as exc:
if not any(t in str(exc) for t in _tolerated_nav):
raise
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
pw_field = page.locator("#password")
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW2)
# Wait on the login POST so a transient 4xx/5xx is caught and retried
# here, not swallowed until the out-of-loop composer wait.
status, _ = click_and_wait_for_response(
page,
url_substr = "/api/auth/login",
method = "POST",
do_click = lambda: page.locator('button[type="submit"]').click(),
timeout_ms = 30_000,
info = lambda m: print(f"[ui] {m}", flush = True),
)
if status is not None and status >= 400:
raise AssertionError(
f"login POST returned {status}; see console_errors={console_errors[:1]!r}"
)
relogin_err = None
break
except Exception as e:
relogin_err = e
try:
cur_url = page.url
except Exception:
cur_url = "<page closed>"
print(
f"[ui] re-login attempt {_relogin_attempt + 1} failed: "
f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; "
f"page_errors={len(page_errors)} console_errors={len(console_errors)}",
flush = True,
)
if console_errors:
print(
f"[ui] first console.error: {console_errors[0][:200]!r}",
flush = True,
)
if page_errors:
print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True)
try:
shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail")
except Exception:
pass
if _relogin_attempt < 2:
# ERR_NO_BUFFER_SPACE needs the OS to recover socket
# buffers; back off 5s then 15s before retrying.
if "ERR_NO_BUFFER_SPACE" in str(e):
backoff_s = 5 if _relogin_attempt == 0 else 15
print(
f"[ui] ENOBUFS detected; sleeping {backoff_s}s "
f"before retry to let OS recover socket buffers...",
flush = True,
)
time.sleep(backoff_s)
# Replace the page if it died; otherwise next iteration's
# page.goto() handles the reload.
old_page = page
page = recover_or_replace_page(
page,
ctx,
default_timeout_ms = 60_000,
info = lambda m: print(f"[ui] recovery: {m}", flush = True),
)
# A freshly created replacement page loses the pageerror/console
# listeners; re-attach so error tracking survives recovery.
if page is not old_page:
page.on("pageerror", lambda e: page_errors.append(str(e)))
page.on("console", _on_console)
if relogin_err is not None:
raise relogin_err
# Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the
# retry: the loop breaks right after submit, so we never re-goto /login once login
# has set tokens -- that would hit the guest guard, redirect to /chat, and make a
# merely-slow composer look like a broken login.
composer = page.locator('textarea[aria-label="Message input"]')
composer.wait_for(state = "visible", timeout = 60_000)
shoot("18-relogin-with-NEW2")

63
tests/test_fast_gemv_dispatch.py Normal file
View file

@ -0,0 +1,63 @@
# 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.
"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is
already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the
bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes.
"""
from types import SimpleNamespace
import pytest
import torch
# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up
# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection
# errors instead of producing a skip. Any other import error still surfaces as a failure.
pytest.importorskip("bitsandbytes")
import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers)
from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES
_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None
def _proj(weight, weight_scale = None):
proj = SimpleNamespace(weight = weight, bias = None, merged = False)
if weight_scale is not None:
proj.weight_scale = weight_scale
return proj
def test_bf16_weight_scale_not_used_as_quant_state():
"""A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None."""
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2))
W, W_quant = get_lora_parameters_bias(proj)[:2]
assert W_quant is None
def test_fp8_weight_keeps_scale():
"""An actual fp8 weight still resolves its weight_scale as the quant state."""
if _FP8 is None:
pytest.skip("no float8 dtype in this torch build")
scale = torch.rand(2, 2)
proj = _proj(torch.randn(4, 4).to(_FP8), scale)
W, W_quant = get_lora_parameters_bias(proj)[:2]
assert W_quant is scale
def test_plain_bf16_has_no_quant_state():
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16))
W, W_quant = get_lora_parameters_bias(proj)[:2]
assert W_quant is None

358
tests/test_fp8_restore_dropped_scale.py Normal file
View file

@ -0,0 +1,358 @@
# 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.
"""Restoring dropped block-fp8 `weight_scale_inv` tensors on load (#6200).
Some block-scale fp8 checkpoints leave a Linear (e.g. `mlp.gate_proj`) unconverted, so its raw
quantized values land in a plain bf16 weight and its `weight_scale_inv` is dropped, producing a
garbage un-scaled weight. `_restore_dropped_fp8_scales` dequantizes such orphaned weights in place
using the scale from the checkpoint. Runs offline on CPU with synthetic checkpoints.
"""
import json
import os
import tempfile
from types import SimpleNamespace
import torch
from torch import nn
from safetensors.torch import save_file
# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
import unsloth
from unsloth.models.loader_utils import _restore_dropped_fp8_scales, _FP8_DTYPES
_SHARD = "model-00001-of-00001.safetensors"
_FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None
def _write_checkpoint(
path,
tensors,
filename = _SHARD,
include_index = True,
):
save_file(tensors, os.path.join(path, filename))
if include_index:
weight_map = {name: filename for name in tensors}
with open(os.path.join(path, "model.safetensors.index.json"), "w") as f:
json.dump({"weight_map": weight_map}, f)
def _fp8_config(block = (2, 2)):
return SimpleNamespace(
quantization_config = {
"quant_method": "fp8",
"weight_block_size": list(block),
}
)
def _fp8_anchor():
"""A module carrying a real fp8 weight, so the model looks like a genuine fp8 load."""
m = nn.Linear(2, 2, bias = False)
m.weight = nn.Parameter(torch.randn(2, 2).to(_FP8), requires_grad = False)
return m
def _bf16_linear(out_f, in_f, raw):
m = nn.Linear(in_f, out_f, bias = False).to(torch.bfloat16)
with torch.no_grad():
m.weight.copy_(raw)
return m
def _expand(scale, block, shape):
bs0, bs1 = block
expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
return expanded[: shape[0], : shape[1]]
def test_restore_dequantizes_orphaned_scale():
"""A plain bf16 weight whose scale was dropped is dequantized in place."""
if _FP8 is None:
return
torch.manual_seed(0)
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(
d,
{
"layer.weight": raw.to(torch.float32),
"layer.weight_scale_inv": scale,
},
)
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_skips_already_fp8_weight():
"""A correctly converted fp8 weight is skipped, never double-scaled."""
if _FP8 is None:
return
weight = torch.randn(4, 4).to(_FP8)
before = weight.clone()
model = nn.Module()
model.config = _fp8_config((2, 2))
model.layer = nn.Linear(4, 4, bias = False)
model.layer.weight = nn.Parameter(weight, requires_grad = False)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": torch.rand(2, 2)})
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 0 and skipped == 1
assert torch.equal(model.layer.weight.data.float(), before.float())
def test_skips_offloaded_meta_weight():
"""A disk-offloaded layer (weight on the meta device) is skipped without error or restore."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = nn.Linear(4, 4, bias = False)
# Simulate an offloaded weight living on the meta device.
model.layer.weight = nn.Parameter(
torch.empty(4, 4, dtype = torch.bfloat16, device = "meta"), requires_grad = False
)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(
d,
{
"layer.weight": raw.to(torch.float32),
"layer.weight_scale_inv": scale,
},
)
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 0
assert model.layer.weight.device.type == "meta"
def test_noop_when_fully_dequantized():
"""If the model has no fp8 weights at all (e.g. load_in_16bit dequantize), do not rescale."""
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.layer = _bf16_linear(4, 4, raw) # no fp8 anchor -> looks dequantized
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert (restored, skipped) == (0, 0)
assert torch.equal(model.layer.weight.data, raw) # untouched
def test_non_block_divisible_shape():
"""Block scale is expanded then sliced to a non-divisible weight shape."""
if _FP8 is None:
return
raw = torch.randn(3, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(3, 4, raw) # weight shape [3, 4]
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (3, 4))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_transposed_scale_layout():
"""A scale stored in the transposed block grid is transposed before use."""
if _FP8 is None:
return
raw = torch.randn(4, 2, dtype = torch.bfloat16) # weight [4, 2] -> grid (2, 1)
scale_correct = torch.rand(2, 1, dtype = torch.float32) + 0.1
scale_stored = scale_correct.t().contiguous() # stored transposed as (1, 2)
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 2, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale_stored})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale_correct, (2, 2), (4, 2))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_single_file_checkpoint_without_index():
"""Unsharded model.safetensors (no index) is still scanned for dropped scales."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(
d, {"layer.weight_scale_inv": scale}, filename = "model.safetensors", include_index = False
)
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_scalar_block_size_config():
"""A scalar weight_block_size (not a list) is handled without error."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = SimpleNamespace(
quantization_config = {"quant_method": "fp8", "weight_block_size": 2}
)
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
def test_text_only_prefix_mapping():
"""Checkpoint keys with a language_model prefix match the stripped text-only module names."""
if _FP8 is None:
return
raw = torch.randn(2, 2, dtype = torch.bfloat16)
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.model = nn.Module()
model.model.gate_proj = _bf16_linear(2, 2, raw) # module lacks the language_model prefix
with tempfile.TemporaryDirectory() as d:
# checkpoint key carries the language_model wrapper the text-only load stripped
_write_checkpoint(d, {"model.language_model.gate_proj.weight_scale_inv": scale})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
assert torch.equal(model.model.gate_proj.weight.data, expected)
def test_skips_variant_load():
"""A variant load (variant="fp8") is skipped to avoid applying default-checkpoint scales."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
result = _restore_dropped_fp8_scales(model, d, local_files_only = True, variant = "fp8")
assert result == (0, 0)
assert torch.equal(model.layer.weight.data, raw) # untouched
def test_vlm_language_model_model_alias():
"""A checkpoint key language_model.model.* matches a model.language_model.* module."""
if _FP8 is None:
return
raw = torch.randn(2, 2, dtype = torch.bfloat16)
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.model = nn.Module()
model.model.language_model = nn.Module()
model.model.language_model.gate_proj = _bf16_linear(
2, 2, raw
) # -> model.language_model.gate_proj
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"language_model.model.gate_proj.weight_scale_inv": scale})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
assert torch.equal(model.model.language_model.gate_proj.weight.data, expected)
def test_noop_without_scale_keys():
if _FP8 is None:
return
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight": torch.randn(4, 4)})
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
def test_noop_without_index_or_single_file():
if _FP8 is None:
return
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
with tempfile.TemporaryDirectory() as d:
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
def test_noop_when_not_block_fp8():
"""A non-fp8 (or non-block) quantization config is ignored."""
scale = torch.rand(2, 2)
model = nn.Module()
model.config = SimpleNamespace(quantization_config = {"quant_method": "compressed-tensors"})
model.layer = nn.Linear(4, 4, bias = False)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)

33
tests/utils/test_rope_scaling_drift.py
View file

@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu():
), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
def test_extended_rope_scaling_keeps_llama3_and_carries_theta():
# Long-context extension keeps native llama3, but falls back to linear for every other
# type (the patched attention constructor only rebuilds linear/llama3/longrope), and the
# linear dict carries rope_theta so transformers v5 does not fall back to base 10000.
from types import SimpleNamespace
from unsloth.models.llama import _extended_rope_scaling
# llama3 model: keep native scaling, do not synthesize linear.
scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0)
assert (
scaling is None and native == "llama3"
), "must keep native llama3 scaling instead of overwriting it with linear."
# yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native.
yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0)
scaling, _ = _extended_rope_scaling(yarn, 2.0)
assert scaling == {
"type": "linear",
"factor": 2.0,
"rope_theta": 500000.0,
}, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}."
# plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta.
v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0})
scaling, _ = _extended_rope_scaling(v5, 2.0)
assert scaling == {
"type": "linear",
"factor": 2.0,
"rope_theta": 1000000.0,
}, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000."
def test_extended_rotary_reads_config_factor():
# LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8
# (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405).

38
tests/version_compat/test_sentence_transformers_pinned_symbols.py
View file

@ -19,6 +19,8 @@ ST_TAGS = [
"v5.2.3",
"v5.3.0",
"v5.4.1",
"v5.5.1",
"v5.6.0",
"master",
]
@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str):
)
# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881).
@pytest.mark.parametrize("tag", ST_TAGS)
def test_st_transformer_load_accepts_unsloth_kwargs(tag: str):
"""unsloth builds saved ST models via Transformer.load(...) so the saved
modality_config is honored (#6881). If .load stops accepting the hub kwargs it
passes (and has no **kwargs), update the fix before it silently regresses. Not
locating .load is a SKIP (may be inherited); the live test guards the install."""
candidates = [
"sentence_transformers/models/Transformer.py",
"sentence_transformers/models/transformer.py",
"sentence_transformers/base/modules/transformer.py",
"sentence_transformers/base/modules/module.py",
]
for p in candidates:
src = fetch_text("UKPLab/sentence-transformers", tag, p)
if src is None or not has_def(src, "load", "func"):
continue
m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S)
if m is None:
continue
sig = m.group(1)
accepts_var_kw = "**" in sig
missing = [
kw
for kw in ("token", "cache_folder", "revision", "trust_remote_code")
if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig))
]
assert not missing, (
f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no "
f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module "
f"(#6881) before it silently falls back to Transformer(...)."
)
return
pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)")
# sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls.
@pytest.mark.parametrize("tag", ST_TAGS)
def test_st_util_helpers(tag: str):

9
unsloth/_gpu_init.py
View file

@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules]
# Fix some issues before importing other packages
from .import_fixes import (
fix_message_factory_issue,
fix_torch_check_is_size,
check_fbgemm_gpu_version,
disable_broken_causal_conv1d,
disable_broken_vllm,
@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection()
disable_broken_causal_conv1d()
disable_broken_vllm()
fix_message_factory_issue()
fix_torch_check_is_size()
check_fbgemm_gpu_version()
torchvision_compatibility_check()
fix_diffusers_warnings()
@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection
del disable_broken_causal_conv1d
del disable_broken_vllm
del fix_message_factory_issue
del fix_torch_check_is_size
del check_fbgemm_gpu_version
del torchvision_compatibility_check
del fix_diffusers_warnings
@ -173,6 +176,7 @@ from .import_fixes import (
fix_vllm_guided_decoding_params,
fix_vllm_pdl_blackwell,
fix_triton_compiled_kernel_missing_attrs,
fix_dynamo_config_thread_visibility,
patch_trunc_normal_precision_issue,
ignore_logger_messages,
patch_ipykernel_hf_xet,
@ -203,6 +207,10 @@ fix_vllm_guided_decoding_params()
fix_trl_vllm_ascend()
fix_vllm_pdl_blackwell()
fix_triton_compiled_kernel_missing_attrs()
# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary
# patches raise the dynamo recompile limits, so those settings reach the
# autograd worker threads on torch >= 2.12.
fix_dynamo_config_thread_visibility()
patch_trunc_normal_precision_issue()
ignore_logger_messages()
patch_ipykernel_hf_xet()
@ -233,6 +241,7 @@ del fix_vllm_guided_decoding_params
del fix_trl_vllm_ascend
del fix_vllm_pdl_blackwell
del fix_triton_compiled_kernel_missing_attrs
del fix_dynamo_config_thread_visibility
del patch_trunc_normal_precision_issue
del ignore_logger_messages
del patch_ipykernel_hf_xet

160
unsloth/import_fixes.py
View file

@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING:
# Deprecation warnings from torchao
warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated")
warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated")
# torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant)
warnings.filterwarnings(
"ignore", message = r"_check_is_size will be removed", category = FutureWarning
)
# TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752)
warnings.filterwarnings(
@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING:
)
def fix_torch_check_is_size():
"""Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit
dequant calls it). The FutureWarning is silenced in suppress_cuda_printf."""
try:
import torch
if hasattr(torch, "_check_is_size"):
return
def _check_is_size(
i,
message = None,
*,
max = None,
):
torch._check(i >= 0, message)
if max is not None:
torch._check(i <= max, message)
torch._check_is_size = _check_is_size
except Exception:
return
# Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'
# MUST do this at the start primarily due to tensorflow causing issues
def fix_message_factory_issue():
@ -1064,6 +1092,135 @@ def fix_triton_compiled_kernel_missing_attrs():
)
def fix_dynamo_config_thread_visibility():
"""torch 2.12 made torch._dynamo/_inductor config overrides thread-local
(ContextVars), so `config.recompile_limit = 1024` set on the main thread is
invisible to the autograd worker threads that run backward. Gradient
checkpointing recompiles fullgraph gpt-oss kernels there against the default
limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config
assignments into the process-global entry default (torch <= 2.11 semantics).
config.patch(...) and config.load_config(...) also assign via __setattr__ but
are thread-local by design, so skip mirroring while inside one (tracked per
thread). No-op below torch 2.12 and on any torch without this internal layout.
"""
try:
import torch
if Version(torch.__version__) < Version("2.12.0"):
return
import torch._dynamo.config as _dynamo_config
from torch.utils._config_module import ConfigModule
from contextvars import ContextVar
except Exception:
return
try:
probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None)
if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar):
# Overrides are not context-local on this torch; nothing to fix.
return
original_setattr = ConfigModule.__setattr__
if getattr(original_setattr, "__unsloth_patched__", False):
return
except Exception:
return
mirrored_modules = ("torch._dynamo.config", "torch._inductor.config")
# config.patch(...) and config.load_config(...) also assign via __setattr__, but
# their writes are thread-local by design; a per-thread depth counter marks them
# so they are not mirrored into the process-global default.
import threading
_scoped_depth = threading.local()
def _in_scoped_write():
return getattr(_scoped_depth, "n", 0) > 0
def _bump(delta):
_scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta
original_patch = ConfigModule.patch
if not getattr(original_patch, "__unsloth_patched__", False):
@functools.wraps(original_patch)
def _patched_patch(self, *args, **kwargs):
ctx = original_patch(self, *args, **kwargs)
try:
cls = type(ctx) # patch() builds a fresh ConfigPatch class each call
if not getattr(cls, "__unsloth_patch_wrapped__", False):
_enter0, _exit0 = cls.__enter__, cls.__exit__
def _enter(s, _e = _enter0):
_bump(1)
try:
return _e(s)
finally:
_bump(-1)
def _exit(
s,
*a,
_x = _exit0,
):
_bump(1)
try:
return _x(s, *a)
finally:
_bump(-1)
cls.__enter__, cls.__exit__ = _enter, _exit
cls.__unsloth_patch_wrapped__ = True
except Exception:
pass
return ctx
_patched_patch.__unsloth_patched__ = True
ConfigModule.patch = _patched_patch
# load_config restores a saved config by calling setattr per key (thread-local).
original_load_config = getattr(ConfigModule, "load_config", None)
if callable(original_load_config) and not getattr(
original_load_config, "__unsloth_patched__", False
):
@functools.wraps(original_load_config)
def _patched_load_config(self, *args, **kwargs):
_bump(1)
try:
return original_load_config(self, *args, **kwargs)
finally:
_bump(-1)
_patched_load_config.__unsloth_patched__ = True
ConfigModule.load_config = _patched_load_config
@functools.wraps(original_setattr)
def _patched_setattr(self, name, value):
original_setattr(self, name, value)
if _in_scoped_write():
return # transient patch / load_config write: keep it thread-local
# Aliases (cache_size_limit -> recompile_limit) re-enter with the real name.
if self.__dict__.get("__name__", None) in mirrored_modules:
try:
entry = self.__dict__["_config"].get(name, None)
if entry is not None and entry.alias is None:
entry.default = value
except Exception:
pass
_patched_setattr.__unsloth_patched__ = True
ConfigModule.__setattr__ = _patched_setattr
# No replay of existing overrides: unsloth installs this before it sets any
# dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying
# would also bake a still-active config.patch override into the global default.
logger.info(
"Unsloth: Patched torch config modules so dynamo/inductor settings "
"(e.g. recompile_limit) apply across threads on torch >= 2.12."
)
def patch_trunc_normal_precision_issue():
"""
Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
@ -1323,8 +1480,7 @@ def fix_vllm_pdl_blackwell():
if patched:
logger.info(
f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - "
f"patched: {', '.join(patched)}"
f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}"
)
else:
# Just set the env var - vLLM might be an older version without supports_pdl

27
unsloth/kernels/utils.py
View file

@ -282,6 +282,21 @@ def QUANT_STATE(W):
return getattr(W, "quant_state", None)
# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a
# quant state when the weight itself is still fp8. compressed-tensors layers expose an
# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around;
# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv /
# fast_dequantize path, which then reads a missing `absmax` and crashes.
_FP8_WEIGHT_DTYPES = tuple(
dtype
for dtype in (
getattr(torch, "float8_e4m3fn", None),
getattr(torch, "float8_e5m2", None),
)
if dtype is not None
)
def get_lora_parameters(proj):
"""Return (weight, weight quant_state, lora A, lora B, lora scale).
With QAT enabled, also fake-quantizes the base layer and lora weights.
@ -298,9 +313,11 @@ def get_lora_parameters(proj):
if weight_fake_quantizer is not None:
W = weight_fake_quantizer(W)
# Get quant state for 4bit or FP8
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
# must not carry a scale as its quant state or fast_gemv will crash on it.
W_quant = getattr(W, "quant_state", None)
if W_quant is None:
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj):
) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
# Get quant state for 4bit or FP8
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
# must not carry a scale as its quant state or fast_gemv will crash on it.
W_quant = getattr(W, "quant_state", None)
if W_quant is None:
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)

91
unsloth/models/llama.py
View file

@ -28,7 +28,11 @@ from ._utils import (
is_bfloat16_supported,
get_quant_type,
)
from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
from .loader_utils import (
_exclude_rope_inv_freq_from_ddp,
_get_fp8_mode_and_check_settings,
_restore_dropped_fp8_scales,
)
from ..utils.packing import (
get_packed_info_from_kwargs,
mask_packed_sequence_boundaries,
@ -1651,6 +1655,26 @@ def _rope_scaling_as_dict(rope_scaling):
return {}
def _extended_rope_scaling(config, factor):
"""RoPE scaling to extend a model past its native window. Keeps native llama3 as-is
(linear extension is far worse for long context); everything else gets linear. Returns
(scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so
transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000.
Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/
longrope and its longrope branch needs a top-level original_max_position_embeddings."""
existing = _rope_scaling_as_dict(
getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {}
)
existing_type = existing.get("rope_type") or existing.get("type")
if existing_type == "llama3":
return None, existing_type
return {
"type": "linear",
"factor": factor,
"rope_theta": _get_rope_theta(config),
}, existing_type
def _llama3_inv_freq_from_config(
config,
rope_scaling,
@ -2518,34 +2542,33 @@ class FastLlamaModel:
max_seq_length = model_max_seq_length
if (rope_scaling is None) and (max_seq_length > model_max_seq_length):
rope_scaling = max_seq_length / model_max_seq_length
factor = max_seq_length / model_max_seq_length
if fast_inference:
raise NotImplementedError(
"Unsloth: Fast inference does not yet work with RoPE Scaling."
)
logger.warning_once(
f"Unsloth: {model_name} can only handle sequence lengths of at most "
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
f"{round(rope_scaling, 3)}, it can be magically be extended to "
f"{max_seq_length}!"
)
# Warn RoPE scaling isn't allowed
if not has_rope_scaling:
raise RuntimeError(
f"However, {model_name} doesn't support RoPE Scaling!\n"
"Please file a feature request at https://github.com/unslothai/unsloth."
linear_scaling, native_type = _extended_rope_scaling(model_config, factor)
if linear_scaling is not None:
logger.warning_once(
f"Unsloth: {model_name} can only handle sequence lengths of at most "
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
f"{round(factor, 3)}, it can be magically be extended to "
f"{max_seq_length}!"
)
if not has_rope_scaling:
raise RuntimeError(
f"However, {model_name} doesn't support RoPE Scaling!\n"
"Please file a feature request at https://github.com/unslothai/unsloth."
)
kwargs["rope_scaling"] = linear_scaling
else:
# Native llama3 scaling already handles long context; just widen the window.
logger.warning_once(
f"Unsloth: extending {model_name} to {max_seq_length} using its native "
f"{native_type} RoPE scaling."
)
rope_scaling = {
"type": "linear",
"factor": rope_scaling,
}
# Add to kwargs
kwargs["rope_scaling"] = rope_scaling
from .loader_utils import (
check_and_disable_bitsandbytes_loading,
@ -2659,6 +2682,18 @@ class FastLlamaModel:
offload_embedding = False,
fast_inference = fast_inference,
)
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
_restore_dropped_fp8_scales(
model,
model_name,
local_files_only = kwargs.get("local_files_only", False),
token = token,
# Weights load from the default branch (revision not forwarded), so read scales from there too.
revision = None,
subfolder = kwargs.get("subfolder"),
cache_dir = kwargs.get("cache_dir"),
variant = kwargs.get("variant"),
)
elif not fast_inference:
if user_config is not None:
# Transformers 5.x @strict model init rejects extra kwargs next
@ -2697,6 +2732,18 @@ class FastLlamaModel:
offload_embedding = False,
fast_inference = False,
)
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
_restore_dropped_fp8_scales(
model,
model_name,
local_files_only = kwargs.get("local_files_only", False),
token = token,
# Weights load from the default branch (revision not forwarded), so read scales from there too.
revision = None,
subfolder = kwargs.get("subfolder"),
cache_dir = kwargs.get("cache_dir"),
variant = kwargs.get("variant"),
)
model.fast_generate = make_fast_generate_wrapper(model.generate)
model.fast_generate_batches = None
else:

281
unsloth/models/loader_utils.py
View file

@ -367,6 +367,287 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
pass
_FP8_DTYPES = tuple(
dtype
for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
if dtype is not None
)
def _fp8_block_size_from_config(model):
"""Return the [block_out, block_in] block size of an fp8 checkpoint, or None if not block-fp8."""
config = getattr(model, "config", None)
quant = getattr(config, "quantization_config", None)
if quant is None:
return None
if hasattr(quant, "to_dict"):
quant = quant.to_dict()
if not isinstance(quant, dict):
return None
if quant.get("quant_method") != "fp8":
return None
block = quant.get("weight_block_size")
if not block:
return None
if isinstance(block, (int, float)):
block = [block, block]
elif isinstance(block, (list, tuple)):
if len(block) == 1:
block = [block[0], block[0]]
elif len(block) < 2:
return None
else:
return None
return [int(block[0]), int(block[1])]
def _load_fp8_weight_map(
model_name,
local_files_only,
token,
revision = None,
subfolder = None,
cache_dir = None,
):
"""Return the checkpoint's tensor->file map, using the same snapshot the load used.
Prefers the sharded `model.safetensors.index.json`; falls back to a single `model.safetensors`
(every tensor maps to that one file) so unsharded checkpoints are covered too.
"""
def _local_path(filename):
return (
os.path.join(model_name, subfolder, filename)
if subfolder
else os.path.join(model_name, filename)
)
def _remote_path(filename):
from huggingface_hub import hf_hub_download
return hf_hub_download(
model_name,
filename,
revision = revision,
subfolder = subfolder,
cache_dir = cache_dir,
local_files_only = local_files_only,
token = token,
)
index_file = "model.safetensors.index.json"
single_file = "model.safetensors"
is_local = os.path.isdir(model_name)
# Sharded checkpoint.
if is_local and os.path.exists(_local_path(index_file)):
index_path = _local_path(index_file)
elif not is_local:
try:
index_path = _remote_path(index_file)
except Exception:
index_path = None
else:
index_path = None
if index_path is not None:
import json
with open(index_path, "r") as f:
return json.load(f).get("weight_map", None)
# Unsharded single file: map every tensor to it.
try:
if is_local and os.path.exists(_local_path(single_file)):
single_path = _local_path(single_file)
elif not is_local:
single_path = _remote_path(single_file)
else:
return None
from safetensors import safe_open
with safe_open(single_path, framework = "pt") as f:
return {key: single_file for key in f.keys()}
except Exception:
return None
def _resolve_fp8_shard(
model_name,
shard,
local_files_only,
token,
revision = None,
subfolder = None,
cache_dir = None,
):
"""Resolve a checkpoint shard filename to a local path (repo id or local dir)."""
if os.path.isdir(model_name):
return (
os.path.join(model_name, subfolder, shard)
if subfolder
else os.path.join(model_name, shard)
)
from huggingface_hub import hf_hub_download
return hf_hub_download(
model_name,
shard,
revision = revision,
subfolder = subfolder,
cache_dir = cache_dir,
local_files_only = local_files_only,
token = token,
)
def _match_fp8_module(module_by_name, base):
"""Resolve a checkpoint module name to a live module, allowing for VLM key remappings.
VLM loads can name the text tower differently from the checkpoint keys: `text_only=True`
strips the `language_model.` wrapper (so `model.language_model.layers.*` -> `model.layers.*`),
and full VLM loads may expose `model.language_model.*` while the checkpoint stores
`language_model.model.*`. Try the raw key first, then a few safe remappings.
"""
if base in module_by_name:
return module_by_name[base]
candidates = []
if "language_model." in base:
candidates.append(base.replace("language_model.", "", 1)) # text-only: drop wrapper
if "language_model.model." in base:
candidates.append(base.replace("language_model.model.", "model.language_model.", 1))
if base.startswith("language_model."):
candidates.append("model." + base) # add model. prefix
for candidate in candidates:
if candidate in module_by_name:
return module_by_name[candidate]
return None
def _restore_dropped_fp8_scales(
model,
model_name,
*,
local_files_only = False,
token = None,
revision = None,
subfolder = None,
cache_dir = None,
variant = None,
):
"""Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
On some block-scale fp8 checkpoints (e.g. Qwen3.6-27B-FP8, issue #6200) transformers fails to
convert a Linear (such as `mlp.gate_proj`) to an fp8 module, loading the raw quantized values
into a plain bf16 weight and discarding its `weight_scale_inv` as an unexpected key. The weight
is then used un-scaled, producing a garbage model. For every checkpoint scale whose live weight
is not fp8, dequantize the orphaned weight in place. Modules that were converted correctly keep
an fp8 weight and are skipped, so a healthy checkpoint is a no-op. Returns (restored, skipped).
"""
try:
block = _fp8_block_size_from_config(model)
if block is None or not _FP8_DTYPES:
return (0, 0)
# A variant load reads variant-named files; skip to avoid applying default scales to them.
if variant:
return (0, 0)
# No fp8 params means the checkpoint was dequantized on purpose (e.g. load_in_16bit);
# re-applying a scale would corrupt those already-correct 16bit weights, so do nothing.
if not any(p.dtype in _FP8_DTYPES for p in model.parameters()):
return (0, 0)
weight_map = _load_fp8_weight_map(
model_name, local_files_only, token, revision, subfolder, cache_dir
)
if not weight_map:
return (0, 0)
scale_keys = {k: v for k, v in weight_map.items() if k.endswith(".weight_scale_inv")}
if not scale_keys:
return (0, 0)
module_by_name = dict(model.named_modules())
bs0, bs1 = block
restored = 0
skipped = 0
failed = 0
offloaded = 0
shard_cache = {}
for scale_key, shard in scale_keys.items():
base = scale_key[: -len(".weight_scale_inv")]
module = _match_fp8_module(module_by_name, base)
if module is None:
continue
weight = getattr(module, "weight", None)
if not isinstance(weight, torch.Tensor) or weight.ndim != 2:
continue
if weight.device.type == "meta":
# Disk-offloaded layer: weight lives on meta until forward, so it cannot be
# scaled in place here. Count and warn rather than silently leave it unscaled.
offloaded += 1
continue
if weight.dtype in _FP8_DTYPES:
# Correctly converted fp8 module: the fp8 path already handles the scale.
skipped += 1
continue
# Errors after this point are per-tensor: warn and continue, never abort or hide them.
try:
if shard not in shard_cache:
from safetensors import safe_open
shard_path = _resolve_fp8_shard(
model_name,
shard,
local_files_only,
token,
revision,
subfolder,
cache_dir,
)
shard_cache[shard] = safe_open(shard_path, framework = "pt")
scale = shard_cache[shard].get_tensor(scale_key).to(torch.float32)
out_features, in_features = weight.shape
out_blocks = (out_features + bs0 - 1) // bs0
in_blocks = (in_features + bs1 - 1) // bs1
if tuple(scale.shape) == (out_blocks, in_blocks):
pass
elif tuple(scale.shape) == (in_blocks, out_blocks) and out_blocks != in_blocks:
# Transposed block layout: same handling as the fp8 forward path.
scale = scale.t().contiguous()
else:
# Shape does not match the block grid: skip rather than apply a wrong scale.
continue
scale = scale.to(weight.device)
with torch.no_grad():
if out_features % bs0 == 0 and in_features % bs1 == 0:
# Memory-frugal path: multiply block views in place against the broadcast
# fp32 scale, avoiding a full expanded scale and fp32 copy that could OOM.
# The in-place multiply promotes to fp32, matching the fallback exactly.
module.weight.data.view(out_blocks, bs0, in_blocks, bs1).mul_(
scale[:, None, :, None]
)
else:
scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(
bs1, dim = 1
)[:out_features, :in_features]
module.weight.data = (weight.to(torch.float32) * scale_expanded).to(
weight.dtype
)
restored += 1
except Exception:
failed += 1
continue
if restored > 0:
print(f"Unsloth: Restored {restored} dropped FP8 weight_scale_inv tensor(s) on load")
if failed > 0:
print(f"Unsloth: {failed} dropped FP8 weight_scale_inv tensor(s) could not be restored")
if offloaded > 0:
print(
f"Unsloth: {offloaded} dropped FP8 weight_scale_inv tensor(s) skipped because the "
"layer is disk-offloaded; load without disk offload so the scales can be restored"
)
return (restored, skipped)
except Exception:
return (0, 0)
def check_and_disable_bitsandbytes_loading(
model_config,
load_in_4bit = True,

65
unsloth/models/sentence_transformer.py
View file

@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel):
return None
@staticmethod
def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code):
def _create_transformer_module(
model_name,
model,
tokenizer,
max_seq_length,
trust_remote_code,
token = None,
cache_dir = None,
revision = None,
module_subfolder = "",
):
"""Helper to create and configure a Transformer module."""
from sentence_transformers.models import Transformer
@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel):
elif "tokenizer_args" in transformer_init_params:
transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy()
transformer_module = Transformer(model_name, **transformer_kwargs)
# Build via Transformer.load so the saved modality_config is honored: plain
# Transformer(...) makes ST 5.x infer a "message" modality for chat-template
# models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings
# (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or
# **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality
# bug, so fall back to the constructor.
transformer_module = None
transformer_load = getattr(Transformer, "load", None)
has_modules_json = (
FastSentenceTransformer._module_path(
model_name, token, cache_dir = cache_dir, revision = revision
)
is not None
)
if callable(transformer_load) and has_modules_json:
load_params = inspect.signature(transformer_load).parameters
accepts_var_kw = any(
p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values()
)
hub_capable = accepts_var_kw or any(
key in load_params for key in ("token", "cache_folder", "revision")
)
if hub_capable:
load_kwargs = {
"token": token,
"cache_folder": cache_dir,
"revision": revision,
"trust_remote_code": trust_remote_code,
**transformer_kwargs,
}
# Resolve config/tokenizer from the module's saved subfolder
# (modules.json "path"), like stock ST; "" (root) is a no-op.
if module_subfolder:
load_kwargs["subfolder"] = module_subfolder
if not accepts_var_kw:
load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
transformer_module = Transformer.load(model_name, **load_kwargs)
if transformer_module is None:
transformer_module = Transformer(model_name, **transformer_kwargs)
finally:
# Restore original Auto* loading immediately
AutoModel.from_pretrained = original_model_from_pretrained
@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel):
tokenizer,
max_seq_length,
trust_remote_code,
token,
cache_dir,
revision,
module_subfolder = module_config.get("path") or "",
)
modules[name] = transformer_module
else:
@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel):
)
transformer_module = FastSentenceTransformer._create_transformer_module(
model_name, model, tokenizer, max_seq_length, trust_remote_code
model_name,
model,
tokenizer,
max_seq_length,
trust_remote_code,
token,
cache_dir,
revision,
)
modules["0"] = transformer_module

17
unsloth/models/vision.py
View file

@ -41,7 +41,11 @@ from ._utils import (
set_task_config_attr,
)
from ._utils import *
from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
from .loader_utils import (
_exclude_rope_inv_freq_from_ddp,
_get_fp8_mode_and_check_settings,
_restore_dropped_fp8_scales,
)
from ..save import patch_saving_functions
from ..models.loader_utils import is_distributed
from unsloth_zoo.gradient_checkpointing import (
@ -1192,6 +1196,17 @@ class FastBaseModel:
offload_embedding = offload_embedding,
fast_inference = fast_inference,
)
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
_restore_dropped_fp8_scales(
model,
model_name,
local_files_only = local_files_only,
token = token,
revision = kwargs.get("revision"),
subfolder = kwargs.get("subfolder"),
cache_dir = kwargs.get("cache_dir"),
variant = kwargs.get("variant"),
)
if hasattr(model, "generate"):
model.fast_generate = make_fast_generate_wrapper(model.generate)
model.fast_generate_batches = error_out_no_vllm

72
unsloth_cli/commands/start.py
View file

@ -133,6 +133,21 @@ _YOLO_OPTION = typer.Option(
"flag/config. Any of the three spellings works for any agent."
),
)
_PERSIST_OPTION = typer.Option(
False,
"--persist/--no-persist",
help = (
"Keep this agent's Unsloth-managed session dir so you can resume it later. "
"codex/openclaw/hermes/pi have their whole home relocated into an Unsloth dir "
"that is a throwaway temp dir (wiped on exit) by default; with --persist it "
"lives under the Unsloth agents dir and survives, so their own resume can reopen "
"it. claude and opencode keep sessions in your own stores (~/.claude, "
"~/.local/share/opencode), so they already resume regardless. To reopen a "
"session, pass the agent's own resume command through, e.g. "
"`unsloth start codex --persist resume` or `claude --resume <id>`; those flow to "
"the agent unchanged."
),
)
# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no
# such flag (config only) and are handled in their config writers, so they are absent.
@ -989,6 +1004,12 @@ def _refresh_windows_path() -> None:
os.environ["PATH"] = os.pathsep.join(entries)
def _install_source(install_hint: str) -> Optional[str]:
"""The first http(s) URL an install hint fetches, or None (e.g. an npm install)."""
match = re.search(r"https?://[^\s'\")]+", install_hint)
return match.group(0) if match else None
def _install_agent(name: str, install_hint: str) -> Optional[str]:
# Missing agent under --launch: offer to run its documented install command, then
# re-resolve it on PATH. Consent-based (we never auto-run a remote install script
@ -997,7 +1018,18 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
if not sys.stdin.isatty():
return None
typer.echo(f"`{name}` is not installed.")
if not typer.confirm(f"Install it now with `{install_hint}`?", default = False):
# Make the supply-chain risk explicit before the prompt: these are the vendors'
# own installers (curl | bash, irm | iex, npm), run with the user's privileges,
# and nothing checks a signature or hash on the fetched content. Naming the source
# turns a blind "yes" into informed consent.
source = _install_source(install_hint)
warning = (
f"This will download and RUN a script from {source} with your privileges"
if source
else f"This will RUN `{install_hint}` with your privileges"
)
typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True)
if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False):
return None
# Run each hint through the shell it is written for: PowerShell (irm | iex, or npm)
# on Windows, /bin/sh (curl | bash, or npm) everywhere else.
@ -1116,15 +1148,20 @@ def _agents_config_root() -> Path:
@contextlib.contextmanager
def _session_config(agent: str, launch: bool):
def _session_config(
agent: str,
launch: bool,
persist: bool = False,
):
"""Yield a private directory for an agent's session config (never the user's own).
launch: an ephemeral temp dir removed after the agent process exits, so nothing
persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later
on this machine), reused across runs. Either way the user's real ~/.<agent>
config is left untouched.
launch (default): an ephemeral temp dir removed after the agent process exits, so
nothing persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run
later on this machine), reused across runs. persist (from --persist): use that same
stable dir even for a launch, so the agent's session survives the exit and can be
resumed next time. Either way the user's real ~/.<agent> config is left untouched.
"""
if launch:
if launch and not persist:
path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-"))
try:
yield path
@ -1436,6 +1473,7 @@ def claude(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point Claude Code at the running Studio server and start it."""
base, key, entry = _connect(
@ -1480,6 +1518,9 @@ def claude(
# --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions.
# IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a
# sandbox is detected, and we don't want to falsely claim one on the user's host.
# claude keeps its history in ~/.claude/projects, which --settings/env never
# relocate, so a session already survives exit; resume it with `claude --continue`
# or `--resume <id>` passed through.
command = [
"claude",
"--model",
@ -1516,6 +1557,7 @@ def codex(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point OpenAI Codex at the running Studio server and start it."""
base, key, entry = _connect(
@ -1541,7 +1583,7 @@ def codex(
*_yolo_command_flags("codex", yolo),
*ctx.args,
]
with _session_config("codex", launch) as home:
with _session_config("codex", launch, persist = persist) as home:
write_codex_config(base, entry, home)
env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)}
_run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex")
@ -1559,6 +1601,7 @@ def openclaw(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point OpenClaw at the running Studio server and start it."""
base, key, entry = _connect(
@ -1584,7 +1627,7 @@ def openclaw(
if os.name == "nt"
else "curl -fsSL https://openclaw.ai/install.sh | bash"
)
with _session_config("openclaw", launch) as cfg:
with _session_config("openclaw", launch, persist = persist) as cfg:
config_path = cfg / "openclaw.json"
# key lives in the config, not the env; --yolo writes the exec policy here too.
write_openclaw_config(base, key, entry, config_path, yolo = yolo)
@ -1605,6 +1648,7 @@ def opencode(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point OpenCode at the running Studio server and start it."""
base, key, entry = _connect(
@ -1628,7 +1672,9 @@ def opencode(
command = ["opencode", "--model", opencode_model]
else:
command = ["opencode"]
with _session_config("opencode", launch) as cfg:
# opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume
# already survives exit; reopen the last one by passing `opencode --continue` through.
with _session_config("opencode", launch, persist = persist) as cfg:
config_path = cfg / "opencode.json"
# OPENCODE_CONFIG is an overlay (loaded between the user's global and project
# configs), so this adds the Unsloth provider/model for the session without
@ -1680,6 +1726,7 @@ def hermes(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point Hermes (Nous Research) at the running Studio server and start it."""
base, key, entry = _connect(
@ -1691,7 +1738,7 @@ def hermes(
)
command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args]
install_hint = _hermes_install_hint()
with _session_config("hermes", launch) as home:
with _session_config("hermes", launch, persist = persist) as home:
# HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state)
# like CODEX_HOME, so the user's ~/.hermes is left untouched for the session.
write_hermes_config(base, entry, home / "config.yaml")
@ -1711,6 +1758,7 @@ def pi(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point Pi (coding agent) at the running Studio server and start it."""
base, key, entry = _connect(
@ -1735,7 +1783,7 @@ def pi(
# --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs
# no install scripts), so accepting the prompt skips dependency lifecycle scripts.
install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent"
with _session_config("pi", launch) as home:
with _session_config("pi", launch, persist = persist) as home:
# Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers
# it over $HOME/.pi/agent), so pin it at the session dir: an inherited
# PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real

168
unsloth_cli/tests/test_start.py
View file

@ -128,6 +128,32 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch):
assert ran == [["powershell", "-NoProfile", "-Command", install_hint]]
def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys):
# Before the confirm, a remote installer must name the URL it fetches so the
# user consents to a specific source rather than blindly accepting.
monkeypatch.setattr(start.os, "name", "nt")
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) # decline: nothing runs
hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup"
assert start._install_agent("hermes", hint) is None
err = capsys.readouterr().err
assert "https://hermes-agent.nousresearch.com/install.ps1" in err
assert "download and RUN" in err
assert "signature or hash" in err
def test_install_agent_warns_for_package_installer(monkeypatch, capsys):
# An npm-style installer has no URL to fetch, but still runs with the user's
# privileges, so the warning names the command instead.
monkeypatch.setattr(start.os, "name", "posix")
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False)
assert start._install_agent("codex", "npm install -g @openai/codex") is None
err = capsys.readouterr().err
assert "npm install -g @openai/codex" in err
assert "with your privileges" in err
def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch):
monkeypatch.setattr(start.os, "name", "nt")
@ -2522,3 +2548,145 @@ def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path
with start._session_config("codex", launch = False) as home2:
assert home2 == home
assert (home2 / "sessions" / "live.sqlite").read_text() == "state"
# ── --persist: persist the agent session so it can be resumed ────────────────
def test_session_config_persist_uses_stable_dir_and_survives(monkeypatch, tmp_path):
# --persist routes a launch to the stable Unsloth agents dir (the one --no-launch
# already uses) instead of a throwaway temp dir, and never wipes it on exit.
monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents")
with start._session_config("codex", launch = True, persist = True) as home:
assert home == tmp_path / "agents" / "codex"
(home / "marker").write_text("kept")
assert home.exists()
assert (home / "marker").read_text() == "kept"
def test_session_config_default_launch_is_ephemeral():
# Default launch (no --persist) still uses a throwaway temp dir wiped on exit.
with start._session_config("codex", launch = True) as home:
assert home.exists()
assert "unsloth-codex-" in home.name
assert not home.exists()
# The temp-dir agents: --persist points each one's home/state env at the stable dir;
# without it, at an ephemeral temp path. opencode is handled separately (only its
# config overlay is relocated; its session data was never in the temp dir).
_RESUME_ENV_VAR = {
"codex": "CODEX_HOME",
"openclaw": "OPENCLAW_STATE_DIR",
"hermes": "HERMES_HOME",
"pi": "HOME",
}
def _capture_launch(monkeypatch, argv):
captured = {}
def run(
command,
env = None,
**kwargs,
):
captured["command"] = command
captured["env"] = env
return SimpleNamespace(returncode = 0)
monkeypatch.setattr(start.subprocess, "run", run)
result = CliRunner().invoke(start.start_app, argv)
assert result.exit_code == 0, result.output
return captured
@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
def test_resume_persists_agent_home_to_stable_dir(agent, fake_studio, tmp_path, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
captured = _capture_launch(monkeypatch, [agent, "--persist"])
stable = tmp_path / "agents" / agent
assert captured["env"][_RESUME_ENV_VAR[agent]] == str(stable)
# The stable dir survives the agent exit, so the session can be resumed.
assert stable.exists()
@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
captured = _capture_launch(monkeypatch, [agent])
home = captured["env"][_RESUME_ENV_VAR[agent]]
assert f"unsloth-{agent}-" in home
assert str(tmp_path / "agents") not in home
def test_resume_opencode_config_in_stable_dir(fake_studio, tmp_path, monkeypatch):
# opencode's session data lives in ~/.local/share/opencode (never relocated), so
# resume already survives exit; --persist also stabilizes its config overlay dir.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
stable = tmp_path / "agents" / "opencode"
assert captured["env"]["OPENCODE_CONFIG"] == str(stable / "opencode.json")
assert stable.exists()
def test_persist_bare_codex_launch_has_no_resume_token(fake_studio, monkeypatch):
# A bare `--persist` only persists the session dir; it must NOT auto-append a native
# resume token, or the very first launch (no session yet) would send codex down its
# no-session error path. The user resumes explicitly: `unsloth start codex --persist resume`.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
captured = _capture_launch(monkeypatch, ["codex", "--persist"])
assert "resume" not in captured["command"]
# command[0] is the resolved executable path; assert the argv after it.
assert captured["command"][1:] == ["--oss", "--profile", start._CODEX_PROFILE]
def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
assert "--continue" not in captured["command"]
assert captured["command"][1:] == ["--model", f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"]
def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
monkeypatch.setattr(start, "_claude_flags", lambda: [])
captured = _capture_launch(monkeypatch, ["claude", "--persist"])
assert "--continue" not in captured["command"]
assert captured["command"][1:] == ["--model", MODEL["id"]]
def test_resume_with_passthrough_does_not_auto_append(fake_studio, monkeypatch):
# When the caller drives their own subcommand, --persist only persists the dir; it
# must not inject a resume token that would collide with the user's command.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
captured = _capture_launch(monkeypatch, ["codex", "--persist", "exec", "hello"])
assert "resume" not in captured["command"]
assert captured["command"][-2:] == ["exec", "hello"]
def test_default_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
captured = _capture_launch(monkeypatch, ["codex"])
assert "resume" not in captured["command"]
def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch):
# openclaw/hermes persist their session dir but have no non-interactive resume
# selector, so --persist must not append a token; their own picker resumes.
for agent in ("openclaw", "hermes"):
monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}")
captured = _capture_launch(monkeypatch, [agent, "--persist"])
assert "resume" not in captured["command"]
assert "--continue" not in captured["command"]
def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
# The persistence flag is --persist, NOT --resume, so an agent's own
# `--resume <id>` (e.g. `unsloth start claude --resume <guid>`) still flows
# through to the agent verbatim and is not swallowed as a Studio option.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
monkeypatch.setattr(start, "_claude_flags", lambda: [])
captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"])
assert captured["command"][-2:] == ["--resume", "some-session-guid"]
# Studio never auto-appends its own resume token when the user drives resume.
assert captured["command"].count("--resume") == 1
assert "--continue" not in captured["command"]