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

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-07-23 03:05:55 -07:00
commit ef4894c559
454 changed files with 49860 additions and 10491 deletions
Download patch file Download diff file Expand all files Collapse all files

33
.github/workflows/cross-platform-parity-ci.yml vendored
View file

@ -1,18 +1,16 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Runs installer parity and autostart opt-out tests on Windows and macOS.
# Runs installer parity and autostart opt-out tests across all three platforms.
#
# Why: that test is the guard that install.sh and install.ps1 stay in
# sync, but today it only runs on ubuntu-latest (auto-discovered by
# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both
# installer scripts, and on Windows Path.read_text() defaults to the
# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already
# contains a U+274C) raises UnicodeDecodeError there even though Linux and
# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
# #6166; this job keeps that from silently regressing by exercising the
# test on the platforms it claims parity for. Pure pytest, no GPU,
# sub-second, so the matrix is cheap.
# Why: the parity test guards that install.sh and install.ps1 stay in sync.
# It originally ran only on ubuntu-latest through studio-backend-ci.yml.
# On Windows, Path.read_text() defaults to the cp1252 locale encoding, so a
# non-cp1252 byte in install.sh raises UnicodeDecodeError even though Linux
# and macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
# #6166; this matrix keeps that from silently regressing. Pure pytest, no GPU,
# sub-second, so the matrix is cheap. Linux also runs the POSIX rollback test
# under dash, matching the supported curl-to-sh installer path.
name: Cross-platform parity
@ -23,6 +21,8 @@ on:
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- 'tests/sh/test_install_rollback_lifecycle.sh'
- 'tests/studio/test_install_rollback_lifecycle.ps1'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
branches: [main]
@ -31,6 +31,8 @@ on:
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- 'tests/sh/test_install_rollback_lifecycle.sh'
- 'tests/studio/test_install_rollback_lifecycle.ps1'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
@ -47,7 +49,7 @@ jobs:
strategy:
fail-fast: false
matrix:
os: [windows-latest, macos-latest]
os: [ubuntu-latest, windows-latest, macos-latest]
runs-on: ${{ matrix.os }}
timeout-minutes: 10
steps:
@ -67,3 +69,10 @@ jobs:
tests/python/test_cross_platform_parity.py
tests/test_installer_skip_autostart.py
-q
- name: PowerShell rollback lifecycle tests
if: runner.os == 'Windows'
shell: pwsh
run: pwsh -NoProfile -File tests/studio/test_install_rollback_lifecycle.ps1
- name: POSIX rollback lifecycle tests
if: runner.os == 'Linux'
run: sh tests/sh/test_install_rollback_lifecycle.sh

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

@ -231,12 +231,69 @@ jobs:
mkdir -p logs/playwright_extra
python tests/studio/playwright_extra_ui.py
- name: UI font size scaling regression (Playwright)
env:
BASE_URL: http://127.0.0.1:18894
STUDIO_PW: ${{ env.STUDIO_EXTRA_NEW_PW }}
PW_ART_DIR: logs/playwright_fontscale
run: |
mkdir -p logs/playwright_fontscale
python tests/studio/playwright_ui_font_scale.py
- name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
# Model-picker per-model-config regression (PR #7207 re-land of #6647).
# Fourth Unsloth on its own port; loads the tiny GGUF and drives the
# picker's run-settings surface: Context Length persists across a reload,
# Reset clears the stored override (never pins it), and the infra models
# (RAG embedder + llama.cpp probe) stay hidden from the picker.
- name: Reset auth + boot Unsloth for model-config tests (port 18898)
run: |
unsloth studio reset-password
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \
> logs/studio_modelcfg.log 2>&1 &
echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV"
- name: Wait for /api/health on 18898
run: |
for i in $(seq 1 180); do
if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then
jq -e '.status == "healthy"' /tmp/health4.json && break
fi
sleep 1
done
jq -e '.status == "healthy"' /tmp/health4.json
- name: Pass bootstrap pw for model-config test
run: |
NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
echo "::add-mask::$NEW"
echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive model-picker per-model-config with Playwright
env:
BASE_URL: http://127.0.0.1:18898
STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }}
PW_ART_DIR: logs/playwright_modelcfg
STUDIO_UI_STRICT: '1'
GGUF_REPO: ${{ env.GGUF_REPO }}
GGUF_VARIANT: ${{ env.GGUF_VARIANT }}
STUDIO_MODEL_HINT: gemma-3-270m
run: |
mkdir -p logs/playwright_modelcfg
python tests/studio/playwright_model_config.py
- name: Stop fourth Unsloth
if: always()
run: |
kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true
sleep 2
# IME + multilingual paste regression (issue #5318 / PR #5327).
# Third Unsloth on its own port so a hang here cannot poison the
# earlier UI tests. No GGUF -- the bug surface is the composer.
@ -297,12 +354,15 @@ jobs:
path: |
logs/studio.log
logs/studio_extra.log
logs/studio_modelcfg.log
logs/studio_ime.log
logs/install.log
logs/server-logs/
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/playwright_fontscale
logs/playwright_modelcfg
logs/playwright_ime
logs/studio-permissions-*.log
retention-days: 7

7
README.md
View file

@ -86,6 +86,13 @@ Replace `claude` with any supported agent:
| OpenCode | `unsloth start opencode` |
| Pi Coding Agent | `unsloth start pi` |
Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local
subagent:
```bash
unsloth start claude --as-subagent --model unsloth/model-GGUF:quant
```
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.

123
install.ps1
View file

@ -1416,13 +1416,82 @@ exit 0
$suffix++
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix"
}
Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
$script:StudioVenvRollbackDir = $candidate
$script:StudioVenvRollbackTarget = $ExistingDir
$script:StudioVenvRollbackActive = $true
# Publish the rollback state before the atomic rename so interruption
# cannot land after Move-Item but before cleanup knows where the old venv went.
try {
Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
} catch {
# A collision or ordinary rename failure leaves the original in place.
# Keep state active only when the rename happened before interruption.
if (Test-Path -LiteralPath $ExistingDir) {
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
}
throw
}
substep "previous environment preserved for rollback"
}
function Remove-StudioVenvTreeWithRetry {
param(
[Parameter(Mandatory = $true)][string]$Path,
[Parameter(Mandatory = $true)][string]$Label
)
$lastError = $null
for ($attempt = 1; $attempt -le 3; $attempt++) {
try {
Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
} catch {
$lastError = $_.Exception.Message
}
if (-not (Test-Path -LiteralPath $Path)) { return $true }
if ($attempt -lt 3) { Start-Sleep -Milliseconds (250 * $attempt) }
}
Write-Host "[WARN] Could not remove $Label at $Path" -ForegroundColor Yellow
if ($lastError) { Write-Host " $lastError" -ForegroundColor Yellow }
return $false
}
function Test-StudioVenvRollbackMustBePreserved {
param([Parameter(Mandatory = $true)][System.IO.FileSystemInfo]$Rollback)
# Preserve anything outside the installer's timestamp.PID[.suffix] format.
if ($Rollback.Name -notmatch '^unsloth_studio\.rollback\.[0-9]{14}\.([0-9]+)(?:\.[0-9]+)?$') {
return $true
}
$ownerPid = 0
if (-not [int]::TryParse($Matches[1], [ref]$ownerPid)) { return $true }
if ($ownerPid -eq $PID) { return $true }
return $null -ne (Get-Process -Id $ownerPid -ErrorAction SilentlyContinue)
}
function Remove-StaleStudioVenvRollbacks {
try {
$rollbacks = @(
Get-ChildItem -LiteralPath $StudioHome -Directory -Force -ErrorAction Stop |
Where-Object { $_.Name -like 'unsloth_studio.rollback.*' }
)
} catch {
Write-Host "[WARN] Could not inspect stale environment rollbacks in $StudioHome" -ForegroundColor Yellow
Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow
return
}
foreach ($rollback in $rollbacks) {
if (($rollback.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -ne 0) {
Write-Host "[WARN] Refusing to remove rollback reparse point $($rollback.FullName)" -ForegroundColor Yellow
continue
}
# A concurrent installer may have moved its live venv aside. The PID
# in the generated name keeps this run from deleting its rescue copy.
if (Test-StudioVenvRollbackMustBePreserved -Rollback $rollback) { continue }
if (Remove-StudioVenvTreeWithRetry -Path $rollback.FullName -Label "stale environment rollback") {
substep "removed stale environment rollback $($rollback.Name)"
}
}
}
function Restore-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
@ -1434,7 +1503,9 @@ exit 0
substep "restoring previous environment after failed install..." "Yellow"
try {
if (Test-Path -LiteralPath $target) {
Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue
if (-not (Remove-StudioVenvTreeWithRetry -Path $target -Label "incomplete environment")) {
throw "Could not remove incomplete environment at $target"
}
}
Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop
substep "restored previous environment"
@ -1449,13 +1520,17 @@ exit 0
function Complete-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
if ($backup -and (Test-Path -LiteralPath $backup)) {
Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue
}
# The replacement is committed. Disable restoration before deleting the
# backup so interruption cannot restore a partially deleted environment.
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
if ($backup -and (Test-Path -LiteralPath $backup)) {
Remove-StudioVenvTreeWithRetry -Path $backup -Label "environment rollback" | Out-Null
}
}
$studioVenvReplacementCommitted = $false
try {
if (Test-Path -LiteralPath $VenvPython) {
# why: matching guard to the .venv branch below -- in env-mode
# $StudioHome is a user-chosen workspace, so refuse to nuke an
@ -1844,7 +1919,7 @@ exit 0
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
@{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
@{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
@ -2030,16 +2105,18 @@ exit 0
# _strip_index_url_credentials (install.sh / py / setup.ps1).
function Remove-IndexUrlCredentials {
param([string]$Url)
$sep = $Url.IndexOf('://')
# Ordinal, not culture-aware: on non-English locales (e.g. th-TH) linguistic
# IndexOf treats "://" as ignorable, mis-locates it, and crashes Substring (issue #7279).
$sep = $Url.IndexOf('://', [System.StringComparison]::Ordinal)
if ($sep -lt 0) { return $Url }
$scheme = $Url.Substring(0, $sep)
$rest = $Url.Substring($sep + 3)
# Drop query / fragment (may hold auth tokens).
$q = $rest.IndexOfAny([char[]]('?', '#'))
if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
$slash = $rest.IndexOf('/')
$slash = $rest.IndexOf('/', [System.StringComparison]::Ordinal)
$authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
$at = $authority.LastIndexOf('@')
$at = $authority.LastIndexOf('@', [System.StringComparison]::Ordinal)
$host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
return "${scheme}://${host_}"
@ -2128,6 +2205,10 @@ exit 0
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
"gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all"
"gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all"
"gfx1030" = "gfx103X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in
@ -2266,7 +2347,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2280,7 +2361,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2354,7 +2435,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2366,7 +2447,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2394,7 +2475,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@ -2420,6 +2501,13 @@ exit 0
}
}
$installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim()
if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) {
step $PackageName "$installedPackageVersion installed"
} else {
substep "[WARN] installed $PackageName version could not be determined" "Yellow"
}
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
@ -2675,6 +2763,13 @@ exit 0
}
Refresh-SessionPath # sync current session with registry
Complete-StudioVenvRollback
$studioVenvReplacementCommitted = $true
Remove-StaleStudioVenvRollbacks
} finally {
if (-not $studioVenvReplacementCommitted) {
Restore-StudioVenvRollback
}
}
# Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy
# User PATH entry (Machine > User > current $env:Path) would win.

488
install.sh
View file

@ -475,14 +475,20 @@ _start_studio_venv_replacement() {
_stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time")
_candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$"
_suffix=0
while [ -e "$_candidate" ]; do
while [ -e "$_candidate" ] || [ -L "$_candidate" ]; do
_suffix=$((_suffix + 1))
_candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix"
done
mv "$_existing_dir" "$_candidate"
_VENV_ROLLBACK_DIR="$_candidate"
_VENV_ROLLBACK_TARGET="$_existing_dir"
_VENV_ROLLBACK_ACTIVE=true
# Publish the rollback state before the atomic rename so a signal cannot
# land after mv but before the exit handlers know where the old venv went.
if ! mv "$_existing_dir" "$_candidate"; then
_VENV_ROLLBACK_ACTIVE=false
_VENV_ROLLBACK_DIR=""
return 1
fi
substep "previous environment preserved for rollback"
}
@ -503,13 +509,68 @@ _restore_studio_venv_replacement() {
fi
}
_commit_studio_venv_replacement() {
[ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0
if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then
rm -rf "$_VENV_ROLLBACK_DIR" || true
_studio_venv_rollback_must_be_preserved() {
_rollback_name=${1##*/}
_rollback_metadata=${_rollback_name#unsloth_studio.rollback.}
_rollback_stamp=${_rollback_metadata%%.*}
_rollback_process=${_rollback_metadata#*.}
# Preserve anything outside the installer's timestamp.PID[.suffix] format.
[ "$_rollback_process" != "$_rollback_metadata" ] || return 0
case "$_rollback_stamp" in
time) ;;
''|*[!0-9]*) return 0 ;;
*) [ "${#_rollback_stamp}" -eq 14 ] || return 0 ;;
esac
_rollback_pid=${_rollback_process%%.*}
case "$_rollback_pid" in
''|*[!0-9]*) return 0 ;;
esac
_rollback_suffix=${_rollback_process#*.}
if [ "$_rollback_suffix" != "$_rollback_process" ]; then
case "$_rollback_suffix" in ''|*[!0-9]*) return 0 ;; esac
fi
_VENV_ROLLBACK_ACTIVE=false
_VENV_ROLLBACK_DIR=""
kill -0 "$_rollback_pid" 2>/dev/null
}
_prune_stale_studio_venv_rollbacks() {
for _stale_rollback in "$STUDIO_HOME"/unsloth_studio.rollback.*; do
[ -d "$_stale_rollback" ] || continue
if [ -L "$_stale_rollback" ]; then
echo "⚠️ Refusing to remove rollback symlink $_stale_rollback" >&2
continue
fi
# A concurrent installer may have moved its live venv aside. The PID in
# the generated name keeps this successful run from deleting its rescue copy.
_studio_venv_rollback_must_be_preserved "$_stale_rollback" && continue
if rm -rf "$_stale_rollback"; then
substep "removed stale environment rollback ${_stale_rollback##*/}"
else
echo "⚠️ Could not remove stale environment rollback $_stale_rollback" >&2
fi
done
}
_commit_studio_venv_replacement() {
if [ "$_VENV_ROLLBACK_ACTIVE" = true ]; then
_rollback_to_remove="$_VENV_ROLLBACK_DIR"
# The new environment is already committed. Clear the restore state
# before deletion so an interrupt cannot replace it with a half-deleted backup.
_VENV_ROLLBACK_ACTIVE=false
_VENV_ROLLBACK_DIR=""
if [ -n "$_rollback_to_remove" ] && [ -d "$_rollback_to_remove" ]; then
if ! rm -rf "$_rollback_to_remove"; then
echo "⚠️ Could not remove environment rollback $_rollback_to_remove" >&2
fi
fi
fi
# Only prune older orphaned copies after the replacement has succeeded, so
# an interrupted install never discards the last known-good environment.
_prune_stale_studio_venv_rollbacks
}
_cleanup_install_temporaries() {
[ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
[ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true
}
_on_install_exit() {
@ -517,15 +578,28 @@ _on_install_exit() {
if [ "$_status" -ne 0 ]; then
_restore_studio_venv_replacement
fi
[ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
[ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true
_cleanup_install_temporaries
exit "$_status"
}
_on_install_signal() {
_signal_status="$1"
# EXIT is disabled to avoid a second cleanup pass. Ignore further termination
# signals until the old environment is back in place.
trap - EXIT
trap '' HUP INT TERM
_restore_studio_venv_replacement
_cleanup_install_temporaries
exit "$_signal_status"
}
# Empty so an inherited value never reaches the trap's rm; only temp paths this
# script creates below (spaced-path dir, torch-trio overrides) are removed.
_UV_OVERRIDE_TMPDIR=""
_UNSLOTH_TORCH_OVERRIDES=""
trap _on_install_exit EXIT
trap '_on_install_signal 129' HUP
trap '_on_install_signal 130' INT
trap '_on_install_signal 143' TERM
# ── Helper: download a URL to a file (supports curl and wget) ──
download() {
@ -1636,7 +1710,7 @@ _maybe_reroute_strixhalo_to_2404() {
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
if _has_usable_nvidia_gpu; then return 0; fi
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
@ -2115,18 +2189,148 @@ _has_amd_rocm_gpu() {
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
return 0
elif [ -e /dev/kfd ] && \
awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \
gpu && amd { found=1 } END{ exit !found }' \
awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
# vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver
# 560+) can register KFD topology nodes with non-zero gpu_id but
# vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting
# NVIDIA-only hosts to the ROCm install path.
# vendor_id 4098 = 0x1002 (AMD) marks a GPU node: the KFD CPU node
# reports vendor_id 0, so any 4098 node is an AMD GPU. NVIDIA's open
# kernel module (driver 560+) registers KFD nodes as vendor_id 4318
# (0x10DE), so this never false-positives on NVIDIA-only hosts.
# The prior check also required a gpu_id line, but gpu_id is a SIBLING
# sysfs file, not a line in properties -- it never matched, so the
# fallback silently missed every ROCm-less AMD host (issue: fresh
# Arch/CachyOS boxes reporting "no GPU detected").
return 0
fi
return 1
}
# Returns 0 if an AMD display GPU is on the PCI bus even when ROCm can't use it
# (e.g. a Strix Halo iGPU with no /dev/kfd). Only sharpens the "no GPU detected"
# hint. vendor 0x1002 = AMD/ATI; class 0x03* = display controller.
_amd_gpu_present_via_pci() {
[ -d /sys/bus/pci/devices ] || return 1
for _pci_vendor in /sys/bus/pci/devices/*/vendor; do
[ -r "$_pci_vendor" ] || continue
read -r _v < "$_pci_vendor" 2>/dev/null || continue
[ "$_v" = "0x1002" ] || continue
_cls="${_pci_vendor%vendor}class"
[ -r "$_cls" ] || continue
read -r _c < "$_cls" 2>/dev/null || continue
case "$_c" in 0x03*) return 0 ;; esac
done
return 1
}
# Map a gfx arch to the AMD pip index family (mirrors install.ps1 $archFamilyMap).
_amd_arch_index_family_for_gfx() {
case "$1" in
gfx1201|gfx1200) echo gfx120X-all ;;
gfx1151) echo gfx1151 ;;
gfx1150) echo gfx1150 ;;
gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;;
gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;;
gfx90a) echo gfx90a ;;
gfx908) echo gfx908 ;;
*) return 1 ;;
esac
}
# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable).
_infer_amd_gfx_arch_from_gpu_name() {
case "$1" in
*"9070 XT"*|*9080*) echo gfx1201 ;;
*9070*|*9060*) echo gfx1200 ;;
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;;
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;;
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;;
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;;
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;;
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;;
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;;
*"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) echo gfx1034 ;;
*) return 1 ;;
esac
}
# Best-effort gfx inference when ROCm tools can't see the GPU (unslothai#7301).
# Mirrors install.ps1 arch resolution on Windows ($HasROCm false, $ROCmGfxArch set).
_infer_linux_amd_gfx_arch() {
if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
printf '%s\n' "$(printf '%s' "$UNSLOTH_ROCM_GFX_ARCH" | tr '[:upper:]' '[:lower:]')"
return 0
fi
# On WSL /proc/cpuinfo and lspci still report the host APU, but without the
# ROCDXG bridge (librocdxg over /dev/dxg) the AMD wheels can't reach the GPU;
# keep the CPU fallback there unless that runtime is present (the explicit
# override above still wins). Mirrors install_python_stack.py.
_gpu_evidence=""
if [ -e /dev/dxg ] || grep -qi microsoft /proc/version 2>/dev/null; then
for _d in /opt/rocm/lib /opt/rocm/lib64 /opt/rocm-*/lib /opt/rocm-*/lib64; do
{ [ -e "$_d/librocdxg.so" ] || [ -e "$_d/librocdxg.so.1" ]; } && _rocdxg=1 && break
done
[ -n "${_rocdxg:-}" ] || return 1
# WSL enumerates no PCI display device; /dev/dxg + librocdxg IS the
# GPU evidence there.
_gpu_evidence=1
elif _amd_gpu_present_via_pci; then
_gpu_evidence=1
fi
# /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received
# no AMD GPU, so the CPU-model text alone is not GPU evidence: require an
# AMD display device (PCI vendor 0x1002, class 0x03*) before trusting it.
# The lspci fallback below needs no gate; an AMD display line IS evidence.
if [ -n "$_gpu_evidence" ] && grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null; then
echo gfx1151
return 0
fi
if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
echo gfx1150
return 0
fi
if command -v lspci >/dev/null 2>&1; then
# A non-AMD controller can enumerate first (Intel/ASPEED before an AMD
# dGPU), so scan every display-class line and take the first AMD one
# that maps. The vendor guard is case-SENSITIVE (a -i "ATI" would match
# "CorporATIon" on every Intel/NVIDIA line); whole-line matching also
# survives the 0000: PCI domain prefix. Mirrors install_python_stack.py.
_amd_disp=$(lspci -nn 2>/dev/null | grep -E 'VGA compatible controller|3D controller|Display controller' | grep -E 'AMD|ATI' || true)
while IFS= read -r _ln; do
[ -n "$_ln" ] || continue
if _gfx=$(_infer_amd_gfx_arch_from_gpu_name "$_ln"); then
echo "$_gfx"
return 0
fi
done <<EOF
$_amd_disp
EOF
fi
return 1
}
# Reads the AMD gfx arch for wheel-index decisions: a user-set
# UNSLOTH_ROCM_GFX_ARCH is authoritative (lowercased), else rocminfo, then
# amd-smi. rocminfo/amd-smi honor ROCR/HIP_VISIBLE_DEVICES, so a container mask
# (e.g. ROCR_VISIBLE_DEVICES=-1) would hide a GPU that the env-independent KFD
# detection still sees -- the tool probes run with the masks cleared. Prints the
# gfx token(s) or nothing when unreadable, and always returns 0 (a failing probe
# as the last command would trip set -e in callers' assignments). Shared by
# get_torch_index_url's gfx gate and the runtime-less reroute gate so the two
# can never disagree on what "readable" means.
_probe_amd_gfx_arch() {
_ensure_rocm_probe_env
_pg=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]')
if [ -z "$_pg" ] && command -v rocminfo >/dev/null 2>&1; then
_pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_pg" ] && command -v amd-smi >/dev/null 2>&1; then
_pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
if [ -z "$_pg" ]; then
_pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
fi
printf '%s\n' "$_pg"
}
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@ -2180,6 +2384,29 @@ get_torch_index_url() {
if ! _has_amd_rocm_gpu; then
echo "$_base/cpu"; return
fi
# A generic rocm index is only safe when the gfx arch is readable: the
# Strix reroute (gfx1150/1151 -> arch-specific index) learns gfx from
# rocminfo/amd-smi, so if those are missing OR do not enumerate the GPU, an
# unknown-arch box might be Strix and would get the broken _grouped_mm
# wheels. Probe via the shared helper (override first, then rocminfo/amd-smi
# with visibility masks cleared); if the arch is unreadable, never guess a
# rocm index. A KFD-only host whose arch is still inferable from hardware
# IDs (PCI/cpuinfo/lspci) returns the cpu index and lets the runtime-less
# reroute below upgrade it to AMD per-arch wheels -- the reroute gate uses
# this same probe, so the handoff can't misfire. Only when inference fails
# too is CPU final, with the actionable warning.
_amd_gfx_probe=$(_probe_amd_gfx_arch)
if [ -z "$_amd_gfx_probe" ]; then
if _amd_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null) && \
[ -n "$_amd_inferred_gfx" ] && \
_amd_arch_index_family_for_gfx "$_amd_inferred_gfx" >/dev/null 2>&1; then
echo "[WARN] AMD GPU detected but rocminfo/amd-smi can't read its gfx arch -- inferring $_amd_inferred_gfx from hardware IDs." >&2
echo "$_base/cpu"; return
fi
echo "[WARN] AMD GPU detected but its gfx arch can't be read (rocminfo/amd-smi missing or not enumerating the GPU) -- installing CPU-only PyTorch." >&2
echo "[WARN] For GPU PyTorch, install or repair rocminfo/amd-smi (e.g. sudo pacman -S rocm-hip-sdk) and re-run this installer." >&2
echo "$_base/cpu"; return
fi
# AMD GPU confirmed -- detect ROCm version
_rocm_tag=""
_rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \
@ -2196,7 +2423,11 @@ get_torch_index_url() {
{ command -v rpm >/dev/null 2>&1 && \
ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \
[ -n "$ver" ] && \
printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null
printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null || _rocm_tag=""
# ^ || guard: when EVERY version source is missing (e.g. rocminfo present
# but rocm-core not installed, so dpkg-query/rpm exit 1), the whole ||
# chain fails and set -e would kill the installer BEFORE the actionable
# no-version WARN below -- exactly the fresh-install case it exists for.
# Validate _rocm_tag: must match "rocmX.Y" with major >= 1
case "$_rocm_tag" in
rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1)
@ -2232,12 +2463,27 @@ get_torch_index_url() {
esac
return
fi
# AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be
# read from any source (amd-smi, /opt/rocm/.info/version, hipconfig,
# dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch.
echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2
echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2
echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
# AMD GPU confirmed (rocminfo/amd-smi or the KFD topology fallback) but
# no ROCm/HIP install was found to read the version from (amd-smi,
# /opt/rocm/.info/version, hipconfig, dpkg, rpm). This is the common
# fresh-install case: the GPU is real, but with no ROCm userspace the
# correct PyTorch build can't be selected. Warn with an actionable fix
# rather than silently installing CPU PyTorch.
# A user-set UNSLOTH_ROCM_GFX_ARCH seeded the probe above, so rocminfo/
# amd-smi may still be unable to see the GPU; when the named arch maps to
# a wheel family, the runtime-less reroute (gated on the override) will
# install the AMD per-arch wheels -- a CPU-only warning here would be
# false for that path. Defer like the inferable-arch branch does.
if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] && \
_amd_arch_index_family_for_gfx "$_amd_gfx_probe" >/dev/null 2>&1; then
echo "[WARN] AMD GPU detected with no readable ROCm version, but UNSLOTH_ROCM_GFX_ARCH=$_amd_gfx_probe is set -- routing to AMD per-arch wheels." >&2
echo "$_base/cpu"; return
fi
echo "[WARN] AMD GPU detected, but no ROCm/HIP install was found to select the matching GPU PyTorch build -- falling back to CPU-only PyTorch." >&2
echo "[WARN] Install the ROCm/HIP SDK, then re-run this installer:" >&2
echo "[WARN] Arch / CachyOS : sudo pacman -S rocm-hip-sdk" >&2
echo "[WARN] other distros : https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "[WARN] Minimum required for version detection: amd-smi, hipconfig, /opt/rocm/.info/version, or the rocm-core package." >&2
echo "$_base/cpu"; return
fi
# Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P).
@ -2649,7 +2895,7 @@ _maybe_bootstrap_rocm_wsl() {
[ -e /dev/dxg ] || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
@ -2735,6 +2981,72 @@ fi
TORCH_INDEX_URL=$(get_torch_index_url)
# Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo
# in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's
# per-arch wheels like install.ps1 does on Windows (unslothai#7301).
# Gated on the runtime probes NOT naming a gfx: either no AMD GPU is detected at
# all (_has_amd_rocm_gpu false), or the GPU is visible only through the
# env-independent KFD topology while rocminfo/amd-smi can't read its arch
# (KFD-only host, unslothai#7314 -- before the KFD detection fix these hosts
# reached this reroute via the false branch, so the empty-probe condition
# preserves that routing). A */cpu index chosen WITH a readable gfx
# (unsupported/unreadable ROCm version, after its own warning) is a deliberate
# fallback -- rerouting it would contradict that decision, and stays excluded
# because the shared probe returns its gfx. An explicit UNSLOTH_ROCM_GFX_ARCH
# override stays authoritative either way.
if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
! _has_usable_nvidia_gpu && \
{ [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu || \
[ -z "$(_probe_amd_gfx_arch)" ]; } && \
case "$(uname -s)" in Linux) true ;; *) false ;; esac && \
case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then
# ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other
# arches, so an inferred/overridden gfx must not reroute arm64 to AMD wheels.
case "$TORCH_INDEX_URL" in
*/cpu)
_linux_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null || true)
if [ -n "$_linux_inferred_gfx" ]; then
_amd_family=$(_amd_arch_index_family_for_gfx "$_linux_inferred_gfx") || _amd_family=""
if [ -n "$_amd_family" ]; then
_amd_mirror="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}"
while [ "${_amd_mirror%/}" != "$_amd_mirror" ]; do
_amd_mirror="${_amd_mirror%/}"
done
TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/"
# Hand the inferred arch to setup.sh (llama.cpp): it re-probes
# ROCm on its own, and on these runtime-less hosts its probes
# find nothing, so without this it classifies the box as
# non-ROCm and installs the CPU prebuilt while torch just got
# AMD per-arch wheels. setup.sh and install_llama_prebuilt.py
# both honor UNSLOTH_ROCM_GFX_ARCH, so exporting it is the
# whole handoff (a user-set override re-exports unchanged).
export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"
case "$_linux_inferred_gfx" in
gfx1201|gfx1200|gfx1151|gfx1150)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
;;
esac
echo "" >&2
# KFD-only hosts reach this reroute with /dev/kfd present
# (that's what detected them), so don't claim it's missing.
if _has_amd_rocm_gpu; then
echo " [WARN] AMD GPU visible via the kernel driver (KFD) but rocminfo/amd-smi can't read its gfx arch; using $_linux_inferred_gfx." >&2
else
echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2
fi
echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2
echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2
echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2
echo " [WARN] Tip: set UNSLOTH_ROCM_GFX_ARCH=$_linux_inferred_gfx to skip inference next time." >&2
echo "" >&2
fi
fi
;;
esac
fi
# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that
# downstream scripts (setup.sh -> install_python_stack.py) know what was
# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts.
@ -2818,29 +3130,64 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ───────
# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug
# that causes a segfault in torch._grouped_mm (moe_utils.py line 167).
# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when
# _amd_gpu_radeon=true the installer silently lands on the broken combo.
# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2.
case "$TORCH_INDEX_URL" in
*/rocm7.1|*/rocm7.1.*)
# 0 when a rocmX.Y index leaf ($1, the final path segment) is older than floor
# $2.$3 (int compare, so rocm7.2 < rocm7.13). Non-rocm leaves (gfx*, cu*, cpu) and
# non-numeric versions return 1. Leaf-based (like $_torch_index_leaf) so a mirror
# base holding its own rocm token compares the family leaf, not the base path.
_rocm_leaf_below() {
case "$1" in rocm[0-9]*.[0-9]*) : ;; *) return 1 ;; esac
_rb=${1#rocm}; _maj=${_rb%%.*}; _min=${_rb#*.}; _min=${_min%%.*}
case "$_maj$_min" in *[!0-9]*) return 1 ;; esac
if [ "$_maj" -lt "$2" ]; then return 0; fi
if [ "$_maj" -eq "$2" ] && [ "$_min" -lt "$3" ]; then return 0; fi
return 1
}
# ── Strix Halo / Strix Point: route to the AMD arch-specific index ───────────
# gfx1151/gfx1150 need torch 2.11+rocm7.13 from repo.amd.com/rocm/whl/gfx<arch>/,
# which carries AMD's real fixes (the rocm7.1 _grouped_mm segfault, moe_utils.py:167,
# and later Strix kernel bugs). Every generic pytorch.org index below rocm7.13 lacks
# them (and the Radeon repo can be offline, unslothai#7264), so reroute a detected
# Strix GPU whenever the picked index is older than the arch build -- covers today's
# rocm6.0-7.2 and any future 7.x < 7.13; rocm7.13+ already has the fixes, so leave it.
case "$_torch_index_leaf" in
rocm[0-9]*)
# Collect every gfx token in rocminfo / amd-smi enumeration order
# (skip duplicates), then index by HIP_VISIBLE_DEVICES /
# ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box
# where the user selected the dGPU does NOT get rerouted to the
# Strix per-gfx index.
_gfx_all=""
if command -v rocminfo >/dev/null 2>&1; then
_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
# || true on each probe: no gfx match makes grep exit 1, which under
# set -euo pipefail would abort the installer before the next fallback
# runs (now that the case matches every rocm* index, not just rocm7.1).
# A user-supplied UNSLOTH_ROCM_GFX_ARCH overrides probing (mirrors setup.sh
# and the display block), so a Strix override still reaches the arch index.
_gfx_all=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]')
if [ -z "$_gfx_all" ] && command -v rocminfo >/dev/null 2>&1; then
_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
# PowerShell paths also probe `amd-smi static --asic`; mirror it
# so a host with hipinfo-less amd-smi reports the gfx target.
if [ -z "$_gfx_all" ]; then
_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
fi
# get_torch_index_url reads the arch with ROCR/HIP masks cleared, so a
# mask hiding every agent (e.g. ROCR_VISIBLE_DEVICES=-1) still lands
# here on a generic rocm index; re-probe unmasked or a masked-out Strix
# box keeps the broken generic wheels. Partial masks never get here
# (they enumerate at least one agent above) and keep their selection.
# ${VAR+x} (not :-): a SET-but-empty mask also hides every agent and
# must trigger the re-probe too.
if [ -z "$_gfx_all" ] && [ -n "${ROCR_VISIBLE_DEVICES+x}${HIP_VISIBLE_DEVICES+x}" ]; then
if command -v rocminfo >/dev/null 2>&1; then
_gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
_gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
[ -z "$_gfx_all" ] && \
_gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
fi
_runtime_gfx=""
@ -2865,13 +3212,14 @@ case "$TORCH_INDEX_URL" in
case "$_runtime_gfx" in
gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
esac
if [ -n "$_strix_gfx" ]; then
# Skip rocm7.13+ generic indexes: they already ship the fixes, so the
# arch build (rocm7.13) would be a downgrade rather than a rescue.
if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then
echo "" >&2
echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2
echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2
echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2
echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2
echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo " [WARN] $_strix_gfx (Strix) detected -- routing to the AMD arch-specific index" >&2
echo " [WARN] torch 2.11+rocm7.13 has AMD's real gfx1150/gfx1151 fixes (the ROCm 7.1" >&2
echo " [WARN] _grouped_mm segfault, moe_utils.py:167, and later Strix kernel bugs)," >&2
echo " [WARN] and is more reliable than the rocm7.2 index or an offline Radeon repo." >&2
echo "" >&2
# AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's
# actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred
@ -2960,7 +3308,7 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
case "$_gpu_disp_mkt" in
*"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
@ -2995,6 +3343,17 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
# Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only.
step "gpu" "Apple Silicon (Metal, unified memory)"
elif _has_amd_rocm_gpu; then
if [ "$_torch_index_pinned" = true ]; then
# An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY pin skipped all probing;
# do not claim ROCm is unusable when a CPU/other index was requested.
step "gpu" "AMD GPU (torch index pinned: $_torch_index_leaf)" "$C_WARN"
else
# AMD GPU visible to the kernel but the torch index stayed CPU: no usable
# ROCm userspace to pick a wheel. "none" would repeat the false diagnosis
# this installer used to give.
step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)" "$C_WARN"
fi
else
step "gpu" "none (CPU-only)" "$C_WARN"
fi
@ -3003,8 +3362,17 @@ fi
case "$TORCH_INDEX_URL" in
*/cpu)
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
if [ "$OS" = "wsl" ]; then
if [ "$_torch_index_pinned" = true ]; then
# An explicit CPU pin is a request, not a detection failure:
# skip the SDK guidance (ROCm may be perfectly healthy here).
substep "CPU-only PyTorch (index pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY)."
elif _has_amd_rocm_gpu; then
substep "AMD GPU detected, but no usable ROCm/HIP install -- installing CPU-only PyTorch." "$C_WARN"
substep "Install the ROCm/HIP SDK and re-run this installer for GPU PyTorch." "$C_WARN"
else
substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
fi
if [ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]; then
# WSL + no GPU detected (detection above found nothing). Common
# cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet --
# /dev/dxg present (graphics) but no ROCm runtime.
@ -3031,6 +3399,13 @@ case "$TORCH_INDEX_URL" in
substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself."
else
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
# Only when ROCm truly can't see the GPU: a detected-but-too-old
# ROCm (rocminfo works, wheels need 6.0+) has its own guidance.
if ! _has_amd_rocm_gpu && _amd_gpu_present_via_pci; then
substep "An AMD GPU is on the PCI bus but ROCm cannot see it (no /dev/kfd," "$C_WARN"
substep " rocminfo, or amd-smi). Install the ROCm kernel stack so /dev/kfd exists;"
substep " Strix Halo (gfx1151/gfx1150) needs a recent kernel (6.11+) and ROCm 7.x."
fi
fi
substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):"
substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
@ -3096,7 +3471,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
"unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -3113,7 +3488,7 @@ if [ "$_MIGRATED" = true ]; then
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-}
"unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-}
[ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
_UNSLOTH_TORCH_OVERRIDES=""
fi
@ -3337,7 +3712,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
"unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -3356,7 +3731,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
--upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -3384,7 +3759,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -3396,6 +3771,15 @@ else
fi
fi
_installed_package_version=$("$_VENV_PY" -c \
'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' \
"$PACKAGE_NAME" 2>/dev/null || true)
if [ -n "$_installed_package_version" ]; then
step "$PACKAGE_NAME" "$_installed_package_version installed"
else
substep "[WARN] installed $PACKAGE_NAME version could not be determined" "$C_WARN"
fi
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on

71
scripts/build_whisper_cpp.sh Executable file
View file

@ -0,0 +1,71 @@
#!/bin/sh
# Build whisper.cpp's whisper-server for Studio's GGUF dictation engine.
#
# Installs into the managed Studio home so the backend's binary discovery
# (core/inference/stt_ggml_sidecar.py::find_whisper_server_binary) picks it up:
# <UNSLOTH_STUDIO_HOME>/whisper.cpp/build/bin/whisper-server (custom home)
# ~/.unsloth/whisper.cpp/build/bin/whisper-server (default)
#
# Usage:
# ./scripts/build_whisper_cpp.sh # build the pinned tag
# WHISPER_CPP_TAG=v1.9.0 ./scripts/build_whisper_cpp.sh
#
# Requires: git, cmake, a C/C++ toolchain (the same prerequisites as a
# llama.cpp source build). GPU backends are auto-detected by whisper.cpp's
# CMake (Metal on macOS; set GGML_CUDA=1 to force a CUDA build on Linux).
set -eu
WHISPER_CPP_SOURCE="${WHISPER_CPP_SOURCE:-https://github.com/ggml-org/whisper.cpp}"
WHISPER_CPP_TAG="${WHISPER_CPP_TAG:-v1.9.1}"
STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}"
CUSTOM_STUDIO_HOME=false
if [ -n "$STUDIO_HOME" ]; then
CUSTOM_STUDIO_HOME=true
INSTALL_DIR="$STUDIO_HOME/whisper.cpp"
else
INSTALL_DIR="$HOME/.unsloth/whisper.cpp"
fi
command -v git >/dev/null 2>&1 || { echo "ERROR: git is required" >&2; exit 1; }
command -v cmake >/dev/null 2>&1 || { echo "ERROR: cmake is required" >&2; exit 1; }
# Same policy as studio/setup.sh's _assert_studio_owned_or_absent: never delete
# a directory under a custom Studio home unless Studio itself created it (the
# marker file below). Protects a user-managed whisper.cpp/src from rm -rf.
STUDIO_OWNED_MARKER=".unsloth-studio-owned"
if [ "$CUSTOM_STUDIO_HOME" = true ] && [ -e "$INSTALL_DIR" ] && \
[ ! -f "$INSTALL_DIR/$STUDIO_OWNED_MARKER" ]; then
echo "ERROR: $INSTALL_DIR already exists and is not marked as an Unsloth-owned whisper.cpp build tree." >&2
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2
exit 1
fi
echo "==> Building whisper.cpp ($WHISPER_CPP_TAG) into $INSTALL_DIR"
mkdir -p "$INSTALL_DIR"
: > "$INSTALL_DIR/$STUDIO_OWNED_MARKER"
if [ ! -d "$INSTALL_DIR/src/.git" ]; then
rm -rf "$INSTALL_DIR/src"
git clone --depth 1 --branch "$WHISPER_CPP_TAG" "$WHISPER_CPP_SOURCE" "$INSTALL_DIR/src"
else
git -C "$INSTALL_DIR/src" fetch --depth 1 origin "$WHISPER_CPP_TAG"
git -C "$INSTALL_DIR/src" checkout FETCH_HEAD
fi
CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF"
if [ "${GGML_CUDA:-0}" = "1" ]; then
CMAKE_FLAGS="$CMAKE_FLAGS -DGGML_CUDA=ON"
fi
# shellcheck disable=SC2086
cmake -S "$INSTALL_DIR/src" -B "$INSTALL_DIR/src/build" $CMAKE_FLAGS
NCPU="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)"
cmake --build "$INSTALL_DIR/src/build" --config Release --target whisper-server -j"$NCPU"
mkdir -p "$INSTALL_DIR/build/bin"
cp "$INSTALL_DIR/src/build/bin/whisper-server" "$INSTALL_DIR/build/bin/whisper-server"
echo "==> Installed $INSTALL_DIR/build/bin/whisper-server"
"$INSTALL_DIR/build/bin/whisper-server" --help >/dev/null 2>&1 && echo "==> Binary runs OK"

72
scripts/scan_packages_baseline.json
View file

@ -1545,6 +1545,78 @@
"severity": "HIGH",
"evidence": "Obfusc: L836: code = compile(module, \"<werkzeug routing>\", \"exec\")\nExec: L736: exec(code, globs, locs)",
"evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d"
},
{
"package": "unsloth-zoo",
"file": "tests/test_mlx_save_export_regressions.py",
"check": "Writes to /tmp and executes (staged dropper)",
"severity": "CRITICAL",
"evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:ab5c587f9ec31a0cc10ee55698ab133a417148d9d3f371bbc81b1e13fa119c13",
"evidence_hash": "93a11159147aad94f353ec4d2e0b8486b256abef88cd96d741813222cd32b138"
},
{
"package": "unsloth-zoo",
"file": "tests/test_vision_collator_audio.py",
"check": "Writes to /tmp and executes (staged dropper)",
"severity": "CRITICAL",
"evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:2efe23ffbe2b91b8403aec9b700736919b59e5ca770f8e1f5501651b44b7d398",
"evidence_hash": "d416b79dd17b24214f3f7653ac01354507d7bf0fc464dee30a4a4b8998f063ba"
},
{
"package": "openai",
"file": "openai/_base_client.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6",
"evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66"
},
{
"package": "openai",
"file": "openai/auth/_workload.py",
"check": "Accesses cloud metadata/IMDS AND makes network calls",
"severity": "CRITICAL",
"evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()",
"evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0"
},
{
"package": "openai",
"file": "openai/resources/beta/responses/responses.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd",
"evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f"
},
{
"package": "openai",
"file": "openai/resources/realtime/realtime.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5",
"evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650"
},
{
"package": "openai",
"file": "openai/resources/responses/responses.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac",
"evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f"
},
{
"package": "unsloth-zoo",
"file": "tests/test_gemma4_forced_float32_ple_dtype.py",
"check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
"severity": "HIGH",
"evidence": "Obfusc: L277: compile(rewritten + _GEMMA4_PLE_CAST_HELPER, \"<gemma4-ple-generated>\", \"exec\") | L440: compile(on, \"<gemma4-ple-append>\", \"exec\") | L468: compile(generated, \"<gemma4-ple-crosspath>\", \"exec\")\nExec: L19: exec(_GEMMA4_PLE_CAST_HELPER, namespace)",
"evidence_hash": "a85e24d8e7c431563cbd83b70f91a3b971abde0f37083d68e70984147960cc70"
},
{
"package": "unsloth-zoo",
"file": "tests/test_vision_collator_audio.py",
"check": "Writes to /tmp and executes (staged dropper)",
"severity": "CRITICAL",
"evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:022f81dd21acfc6a35a058de96132834c218404a9e37b3d09a7768a8c8f6c728",
"evidence_hash": "2d1e75446af120d9133a42aa8af426a839d3434d9dc109cc1d6c1b22ca1ddb75"
}
]
}

4
studio/backend/auth/terminal_prompt.py
View file

@ -236,6 +236,10 @@ def prompt_for_password_change(
out.write(f"Password must be at least {min_length} characters; try again.\n")
out.flush()
continue
if any(ch.isspace() for ch in new_password):
out.write("Password cannot contain spaces; try again.\n")
out.flush()
continue
if is_current_password(new_password):
out.write(
"New password must differ from the current bootstrap password; try again.\n"

88
studio/backend/cloudflare_tunnel.py
View file

@ -20,6 +20,7 @@ import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Optional, Tuple
@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
# A registered edge connection does not mean the hostname resolves yet, so the
# URL is fetched once before it is advertised.
_PUBLIC_PROBE_PATH = "/api/health"
_PUBLIC_PROBE_MARKER = "Unsloth UI Backend"
# One deadline for DNS propagation + the health probe, bounding the startup stall.
_PUBLIC_PROBE_TIMEOUT = 45.0
_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0
_PUBLIC_PROBE_RETRY_DELAY = 1.0
# Wait for the hostname via DoH first: an early OS lookup negative-caches the
# NXDOMAIN for up to 30 min.
_DNS_POLL_DELAY = 2.0
# Retry transient DoH failures, but give up fast when DoH is blocked outright.
_DNS_MAX_DOH_ERRORS = 3
_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A"
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]:
return None
def _wait_for_dns(host: str, deadline: float) -> None:
import json
import urllib.request
errors = 0
while True:
answered = False
try:
req = urllib.request.Request(
_DOH_URL.format(host = host),
headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"},
)
with urllib.request.urlopen(req, timeout = 5) as response:
answered = bool(json.loads(response.read(65536)).get("Answer"))
errors = 0
except Exception:
errors += 1
if errors >= _DNS_MAX_DOH_ERRORS:
return
if answered:
return
remaining = deadline - time.monotonic()
if remaining <= 0:
return
time.sleep(min(_DNS_POLL_DELAY, remaining))
def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool:
import json
import urllib.request
from urllib.parse import urlsplit
deadline = time.monotonic() + timeout
host = urlsplit(url).hostname
if host:
_wait_for_dns(host, deadline)
probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}"
while True:
try:
req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response:
body = response.read(4096)
if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER:
return True
except Exception:
pass
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining))
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:<port>. Best-effort throughout.
@ -322,11 +392,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
Waits for cloudflared to both mint the URL and register an edge connection
before returning, so the caller never advertises a URL that yields Cloudflare
error 1033 (HTTP 530). If a URL is minted but no connection registers within
the window (e.g. quic is blocked on this network), retries once forcing the
http2 protocol. On any failure the tunnel is stopped and None is returned.
Waits for cloudflared to both mint the URL and register an edge connection,
then fetches /api/health over the public URL, so the caller never advertises
a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host.
If a URL is minted but no connection registers within the window (e.g. quic
is blocked on this network), retries once forcing the http2 protocol. On any
failure the tunnel is stopped and None is returned.
"""
global _active_tunnel, _shutdown_requested
binary = ensure_cloudflared()
@ -349,9 +420,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
registered = False
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
registered = url is not None
if url and not verify_public_url(url):
url = None
except Exception:
url = None
if url:
@ -371,6 +446,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
# http2 will not help, so do not burn another window on it.
if not saw_url:
return None
# probe failure after registering is DNS propagation; http2 would not help
if registered:
return None
return None

11
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -27,7 +27,6 @@ from .constants import (
)
from .parse import apply_update, coerce_event, parse_log_message
from .types import Job
from .worker import run_job_process
from loggers import get_logger
logger = get_logger(__name__)
@ -169,12 +168,18 @@ class JobManager:
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
with native_path_secret_removed_for_child_start():
cache_env = get_hf_cache_paths().child_env({})
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
mp_q = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_job_process,),
args = ("core.data_recipe.jobs.worker", "run_job_process", cache_env),
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon = True,
)

10
studio/backend/core/export/orchestrator.py
View file

@ -230,16 +230,20 @@ class ExportOrchestrator:
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
from .worker import run_export_process
cache_env = get_hf_cache_paths().child_env({})
with native_path_secret_removed_for_child_start():
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_export_process,),
args = ("core.export.worker", "run_export_process", cache_env),
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,

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

@ -76,8 +76,14 @@ class AudioCodecManager:
if self._snac_model is not None:
return
from snac import SNAC
from utils.hf_cache_settings import active_hf_hub_cache
self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
# Route weights to the selected cache; this can run in the main process.
self._snac_model = (
SNAC.from_pretrained("hubertsiuzdak/snac_24khz", cache_dir = active_hf_hub_cache())
.to(device)
.eval()
)
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(

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

@ -247,6 +247,59 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
return out
def _bundled_hip_present(binary_dir: str) -> bool:
"""True when a prebuilt bundle ships its own HIP backend library."""
if not binary_dir:
return False
try:
# Glob the version suffix (libggml-hip.so, .so.0, .so.0.11.1) the same
# way the installer's runtime health check matches libggml-hip.so*.
return any(Path(str(binary_dir)).glob("libggml-hip.so*"))
except OSError:
return False
def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]":
"""System ROCm lib dir(s) to prepend before a prebuilt's bundled HIP, on native Linux.
The bundled bare-metal HIP runtime can mismatch the host amdkfd driver and crash
in hsa_init(); prepending the whole system ROCm lib dir loads a driver-matched,
version-consistent stack (libhsa-runtime64 / libamdhip64 / librocblas) ahead of it.
The whole dir is deliberate: mixing the bundle's rocBLAS with a different-version
system HIP/ROCR risks missing symbols. UNSLOTH_LLAMA_NO_SYSTEM_ROCM=1 keeps the pure
bundle (for a host whose system ROCm lacks this arch); no-op on WSL / non-Linux.
"""
if os.environ.get("UNSLOTH_LLAMA_NO_SYSTEM_ROCM") == "1":
return []
if sys.platform != "linux" or os.path.exists("/dev/dxg"):
return []
if not os.path.exists("/dev/kfd"):
return []
if not _bundled_hip_present(binary_dir):
return []
# Env-configured ROCm root first; /opt/rocm only as a fallback so a stale
# /opt/rocm doesn't shadow the driver-matching install these vars point at.
candidates = []
for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
val = os.environ.get(var)
if val:
candidates.append(val)
candidates.append("/opt/rocm")
out: "list[str]" = []
seen: "set[str]" = set()
for base in candidates:
for lib_sub in ("lib", "lib64"):
d = os.path.join(base, lib_sub)
if d in seen:
continue
seen.add(d)
if os.path.exists(os.path.join(d, "libhsa-runtime64.so")) or os.path.exists(
os.path.join(d, "libhsa-runtime64.so.1")
):
out.append(d)
return out
# Plan-without-action re-prompt state now lives in tool_call_parser (imported above).
# Default max_tokens to the effective context when known. The floor is high
@ -526,7 +579,14 @@ def _swa_entry_from_layer_types(lt) -> Optional[object]:
def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
try:
from huggingface_hub import hf_hub_download
cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
from utils.hf_cache_settings import active_hf_hub_cache
cfg_path = hf_hub_download(
repo_id,
"config.json",
repo_type = "model",
cache_dir = active_hf_hub_cache(),
)
with open(cfg_path) as f:
cfg = json.load(f)
except Exception:
@ -928,6 +988,7 @@ def _cached_hf_snapshot_file(
filename: str,
*,
expected_size: Optional[int] = None,
cache_dir: Optional[str] = None,
) -> Optional[str]:
"""Return a cached snapshot file even when HF's current-ref probe misses it."""
if not filename:
@ -936,8 +997,22 @@ def _cached_hf_snapshot_file(
if not parts or any(part in (".", "..") for part in parts):
return None
try:
from utils.models.model_config import _iter_hf_cache_snapshots
for snap in _iter_hf_cache_snapshots(repo_id):
if cache_dir is None:
from utils.models.model_config import _iter_hf_cache_snapshots
snapshots = _iter_hf_cache_snapshots(repo_id)
else:
from hub.utils.hf_cache_state import iter_active_repo_cache_dirs
snapshots = (
snapshot
for repo_dir in iter_active_repo_cache_dirs(
"model",
repo_id,
root = Path(cache_dir),
)
for snapshot in (repo_dir / "snapshots").glob("*")
if snapshot.is_dir()
)
for snap in snapshots:
candidate = snap.joinpath(*parts)
if not candidate.is_file():
continue
@ -1179,6 +1254,16 @@ def _snapshot_dir_of(path: str) -> Optional[Path]:
return None
def _hub_cache_dir_for_snapshot_path(path: Optional[str]) -> Optional[str]:
"""Return the HF Hub cache root that owns a snapshot-contained path."""
if not path:
return None
snapshot = _snapshot_dir_of(path)
if snapshot is None or snapshot.parent.name != "snapshots":
return None
return str(snapshot.parent.parent.parent)
def _companion_snapshot_sibling(
near_path: str, pick: Callable[[list[str]], Optional[str]]
) -> Optional[str]:
@ -2859,12 +2944,25 @@ class LlamaCppBackend:
on the ordinal->physical mapping."""
try:
import torch
is_rocm = getattr(torch.version, "hip", None) is not None
# Same ROCm detection as _emit_child_gpu_visibility: AMD SDK wheels
# leave version.hip unset but encode "rocm" in __version__. The two
# must agree, else an inherited ROCR mask reads back as "no mask",
# ordinal 0 is labelled physical 0, and the child's new ROCR pin
# re-exposes the GPU the inherited mask was hiding.
is_rocm = (
getattr(torch.version, "hip", None) is not None
or "rocm" in getattr(torch, "__version__", "").lower()
)
except Exception:
is_rocm = False
if is_rocm:
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
# ROCR_VISIBLE_DEVICES is a Linux ROCr variable; Windows HIP has no
# ROCr layer, so a stray ROCR var there does not mask the runtime and
# must not be read as the ordinal->physical mapping (mirrors the
# Windows gate in _emit_child_gpu_visibility).
rocr_v = None if sys.platform == "win32" else os.environ.get("ROCR_VISIBLE_DEVICES")
cvd = (
hip_v
if hip_v is not None
@ -2882,20 +2980,52 @@ class LlamaCppBackend:
return None
@staticmethod
def _emit_child_gpu_visibility(env: dict, pinned: str) -> None:
"""Write the child's GPU visibility mask (CUDA, plus the HIP mirror on
ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child
seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP
mask at different layers, so the same indices apply twice -- ROCR reduces
and re-indexes from 0, then a non-zero HIP pin points out of range, HIP
enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone
narrows correctly; clear any inherited ROCR mask so it can't double up."""
def _emit_child_gpu_visibility(
env: dict,
pinned: str,
*,
prefer_rocr: bool = False,
) -> None:
"""Write the child's GPU visibility mask: CUDA, plus a ROCm mirror on AMD
(masking only CUDA_VISIBLE_DEVICES leaves an AMD child seeing every GPU).
Default: HIP_VISIBLE_DEVICES, clearing any inherited ROCR mask so the two
can't stack (ROCR re-indexes from 0, then a non-zero HIP pin points out of
range, HIP sees 0 devices, and llama.cpp falls back to CPU).
prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask
filters only AFTER the HSA runtime enumerates every agent, and that
enumeration segfaults at startup on a GPU the build has no kernels for
(e.g. a gfx1103 iGPU under a gfx110X prebuilt), before llama-server logs a
line. ROCR drops the device at the driver layer, consuming physical ids.
The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps
the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a
Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin
would be dead there while the cleared HIP mask stops selecting."""
env["CUDA_VISIBLE_DEVICES"] = pinned
try:
import torch as _torch
if getattr(_torch.version, "hip", None) is not None:
env["HIP_VISIBLE_DEVICES"] = pinned
env.pop("ROCR_VISIBLE_DEVICES", None)
# torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may
# leave it unset but encode "rocm" in __version__ (mirrors detect_hardware).
if (
getattr(_torch.version, "hip", None) is not None
or "rocm" in getattr(_torch, "__version__", "").lower()
):
if prefer_rocr and pinned != "-1" and sys.platform != "win32":
env["ROCR_VISIBLE_DEVICES"] = pinned
env.pop("HIP_VISIBLE_DEVICES", None)
# ROCR re-indexes the visible agents from 0, and with HIP
# cleared HIP honours CUDA_VISIBLE_DEVICES -- so it must carry
# the post-ROCR ordinals (0..N-1), not the physical ids, else a
# non-zero pick points out of range and HIP sees 0 devices (the
# same stacking the default path avoids by clearing ROCR).
env["CUDA_VISIBLE_DEVICES"] = ",".join(
str(i) for i in range(len(pinned.split(",")))
)
else:
env["HIP_VISIBLE_DEVICES"] = pinned
env.pop("ROCR_VISIBLE_DEVICES", None)
except Exception as e:
logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
@ -2930,7 +3060,21 @@ class LlamaCppBackend:
logger.debug("Could not read reported GPU order for split pin: %s", e)
if order is None:
order = sorted(inherited)
LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order))
# Re-emit at the layer that produced the mapping. A parent masked only
# via ROCR_VISIBLE_DEVICES hides agents at the driver layer, and the
# default HIP re-emission clears that mask -- HSA then enumerates every
# agent again and can segfault at startup on an unsupported GPU the
# parent was hiding (the crash prefer_rocr exists to avoid). Linux-only,
# mirroring _resolve_visible_physical_ids: on Windows a stray ROCR var
# is dead and was not the mapping's source.
prefer_rocr = (
sys.platform != "win32"
and env.get("HIP_VISIBLE_DEVICES") is None
and env.get("ROCR_VISIBLE_DEVICES") is not None
)
LlamaCppBackend._emit_child_gpu_visibility(
env, ",".join(str(i) for i in order), prefer_rocr = prefer_rocr
)
@staticmethod
def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
@ -3592,6 +3736,9 @@ class LlamaCppBackend:
lib_dirs.extend(_wsl_system_rocm_lib_dirs())
if lib_dirs:
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
# Native Linux AMD: system ROCm libs before the bundle's HIP runtime,
# which can be incompatible with the host amdkfd driver.
lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))
lib_dirs.append(binary_dir)
_arch = platform.machine() # x86_64, aarch64, etc.
@ -4955,6 +5102,9 @@ class LlamaCppBackend:
touching the shared one; defaults to the shared event.
"""
cancel_event = cancel_event if cancel_event is not None else self._cancel_event
from utils.hf_cache_settings import get_hf_cache_paths
download_cache_dir = str(get_hf_cache_paths().hub_cache)
try:
import huggingface_hub # noqa: F401 -- presence check only
except ImportError:
@ -5050,7 +5200,11 @@ class LlamaCppBackend:
if not p.size:
continue
try:
cached_path = try_to_load_from_cache(hf_repo, p.path)
cached_path = try_to_load_from_cache(
hf_repo,
p.path,
cache_dir = download_cache_dir,
)
except Exception:
cached_path = None
if (
@ -5061,6 +5215,7 @@ class LlamaCppBackend:
hf_repo,
p.path,
expected_size = p.size,
cache_dir = download_cache_dir,
)
if isinstance(cached_path, str) and os.path.exists(cached_path):
try:
@ -5074,12 +5229,8 @@ class LlamaCppBackend:
total_download_bytes = max(0, total_bytes - already_cached_bytes)
if total_download_bytes > 0:
cache_dir = os.environ.get(
"HF_HUB_CACHE",
str(Path.home() / ".cache" / "huggingface" / "hub"),
)
Path(cache_dir).mkdir(parents = True, exist_ok = True)
free_bytes = shutil.disk_usage(cache_dir).free
Path(download_cache_dir).mkdir(parents = True, exist_ok = True)
free_bytes = shutil.disk_usage(download_cache_dir).free
total_gb = total_download_bytes / (1024**3)
free_gb = free_bytes / (1024**3)
@ -5097,7 +5248,7 @@ class LlamaCppBackend:
# surface the disk shortfall for the requested variant.
raise RuntimeError(
f"Not enough disk space to download {gguf_filename}. "
f"Only {free_gb:.1f} GB free in {cache_dir}"
f"Only {free_gb:.1f} GB free in {download_cache_dir}"
)
smaller = self._find_smallest_fitting_variant(
hf_repo,
@ -5128,7 +5279,7 @@ class LlamaCppBackend:
else:
raise RuntimeError(
f"Not enough disk space to download any variant. "
f"Only {free_gb:.1f} GB free in {cache_dir}"
f"Only {free_gb:.1f} GB free in {download_cache_dir}"
)
except RuntimeError:
raise
@ -5151,6 +5302,7 @@ class LlamaCppBackend:
cancel_event = cancel_event,
on_status = lambda m: logger.info(m),
force_download = force,
cache_dir = download_cache_dir,
)
for shard in gguf_extra_shards:
if cancel_event.is_set():
@ -5162,6 +5314,7 @@ class LlamaCppBackend:
hf_token,
cancel_event = cancel_event,
force_download = force,
cache_dir = download_cache_dir,
)
except Exception as e:
if isinstance(e, RuntimeError) and "Cancelled" in str(e):
@ -5207,6 +5360,12 @@ class LlamaCppBackend:
logger.info("Reusing cached %s: %s", label, cached)
return cached
from utils.hf_cache_settings import get_hf_cache_paths
companion_cache_dir = _hub_cache_dir_for_snapshot_path(near_path) or str(
get_hf_cache_paths().hub_cache
)
if _hub_download_in_flight(hf_repo):
logger.info("Skipping %s download while a hub download is active", label)
return None
@ -5241,7 +5400,7 @@ class LlamaCppBackend:
if target is None:
try:
from utils.models.model_config import _iter_hf_cache_snapshots
for snap in _iter_hf_cache_snapshots(hf_repo):
for snap in _iter_hf_cache_snapshots(hf_repo, companion_cache_dir):
rel_files = _gguf_snapshot_files(snap)
target = pick(rel_files)
if target is not None:
@ -5259,7 +5418,11 @@ class LlamaCppBackend:
# hf_hub_download with hf_repo would miss the canonical file and silently
# drop the companion. _cached_hf_snapshot_file scans every case variant.
if _hf_env_offline():
cached = _cached_hf_snapshot_file(hf_repo, target)
cached = _cached_hf_snapshot_file(
hf_repo,
target,
cache_dir = companion_cache_dir,
)
if cached:
logger.info("Resolved %s from local HF cache: %s", label, cached)
return cached
@ -5272,6 +5435,7 @@ class LlamaCppBackend:
target,
hf_token,
cancel_event = cancel_event,
cache_dir = companion_cache_dir,
)
except Exception as e:
logger.warning(f"Could not download {label}: {e}")
@ -5302,7 +5466,12 @@ class LlamaCppBackend:
near_path = near_path,
)
def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
def _cached_repo_mtp_drafter(
self,
hf_repo: str,
*,
cache_dir: Optional[str] = None,
) -> 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
@ -5312,7 +5481,12 @@ class LlamaCppBackend:
roots: list[Path] = []
subdirs: list[Path] = []
for snap in _iter_hf_cache_snapshots(hf_repo): # newest first
snapshots = (
_iter_hf_cache_snapshots(hf_repo)
if cache_dir is None
else _iter_hf_cache_snapshots(hf_repo, cache_dir)
)
for snap in snapshots: # 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)
@ -5365,7 +5539,10 @@ class LlamaCppBackend:
# 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)
cached = self._cached_repo_mtp_drafter(
hf_repo,
cache_dir = _hub_cache_dir_for_snapshot_path(near_path),
)
if cached:
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
return cached
@ -7684,7 +7861,12 @@ class LlamaCppBackend:
# default FASTEST_FIRST order (#5025).
if gpu_ids:
env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices))
# Mask on AMD at the ROCr/HSA layer: HIP-only masking still
# enumerates every agent first, which segfaults on a deselected
# unsupported GPU (e.g. gfx1103 iGPU under a gfx110X prebuilt).
self._emit_child_gpu_visibility(
env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True
)
elif manual_tensor_split_emitted and not is_vulkan_backend:
# A manual per-GPU ratio across ALL GPUs (no explicit pick, so
# no CUDA_VISIBLE_DEVICES mask above): the UI built the
@ -8046,6 +8228,20 @@ class LlamaCppBackend:
# an OS-killed text-only retry still gets the OOM message.
_retry_rc = self._process.poll() if self._process is not None else None
self._kill_process()
# If the text-only retry ALSO hard-crashed (a signal, not
# OOM/timeout), the vision projector was never the cause:
# llama-server is faulting during GPU/driver init. Say so
# -- with the ROCm fix -- instead of blaming the mmproj.
if self._is_signal_crash(_retry_rc):
raise RuntimeError(
"llama-server crashed at startup on both the vision "
"and text-only attempts -- a GPU driver/runtime "
"initialization crash, not a model or vision-projector "
"problem. This often means an unsupported secondary "
"GPU; on AMD/ROCm, hide it with ROCR_VISIBLE_DEVICES "
"(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first "
"GPU) before launching Unsloth Studio."
)
raise RuntimeError(
"Vision projector incompatible with this llama.cpp "
"build, and the text-only retry also failed: "

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

@ -146,6 +146,7 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
_is_hidden_model,
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
from utils.hf_cache_settings import known_hf_hub_caches
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
@ -174,7 +175,12 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
except Exception as exc:
logger.debug("auto-switch: ./models scan failed: %s", exc)
try:
for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()):
for hf_dir in (
*known_hf_hub_caches(),
_resolve_hf_cache_dir(),
legacy_hf_cache_dir(),
hf_default_cache_dir(),
):
found += _scan_hf_once(hf_dir)
except Exception as exc:
logger.debug("auto-switch: HF cache scan failed: %s", exc)

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

@ -181,19 +181,27 @@ def _vlm_messages_have_tool_history(messages):
)
def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
def _build_generation_stats(
prompt_n,
prompt_tps,
gen_n,
gen_tps,
cached_n = 0,
):
"""Map mlx stream stats onto the usage/timings shape llama-server emits."""
prompt_n = int(prompt_n or 0)
gen_n = int(gen_n or 0)
cached_n = int(cached_n or 0)
prompt_tps = float(prompt_tps or 0.0)
gen_tps = float(gen_tps or 0.0)
prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0
predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0
total_prompt_n = prompt_n + cached_n
return {
"usage": {
"prompt_tokens": prompt_n,
"prompt_tokens": total_prompt_n,
"completion_tokens": gen_n,
"total_tokens": prompt_n + gen_n,
"total_tokens": total_prompt_n + gen_n,
},
"timings": {
"prompt_n": prompt_n,
@ -204,11 +212,123 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"predicted_ms": predicted_ms,
"predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0,
"predicted_per_second": gen_tps,
"cache_n": 0,
"cache_n": cached_n,
},
}
PROMPT_CACHE_ENTRIES = 6
PROMPT_CACHE_MEMORY_FRACTION = 0.15
PROMPT_CACHE_FALLBACK_BYTES = 2 * 1024**3
def _mlx_prompt_cache_api():
try:
from mlx_lm.models.cache import (
LRUPromptCache,
can_trim_prompt_cache,
make_prompt_cache,
trim_prompt_cache,
)
except ImportError:
return None
return LRUPromptCache, make_prompt_cache, can_trim_prompt_cache, trim_prompt_cache
def _prompt_cache_max_bytes(recommended_gb = None):
override = os.environ.get("UNSLOTH_MLX_PROMPT_CACHE_BYTES")
if override:
try:
return max(int(override), 0)
except ValueError:
logger.warning("Ignoring non-integer UNSLOTH_MLX_PROMPT_CACHE_BYTES=%r", override)
if recommended_gb:
return int(recommended_gb * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
return PROMPT_CACHE_FALLBACK_BYTES
def _flatten_kv_entries(cache):
for entry in cache:
nested = getattr(entry, "caches", None)
if nested is None:
yield entry
else:
yield from _flatten_kv_entries(nested)
def _kv_prefix_coverage(cache):
covered = None
for entry in _flatten_kv_entries(cache):
offset = getattr(entry, "offset", None)
if offset is None:
return None
if getattr(entry, "start_position", 0):
return None
window = getattr(entry, "max_size", None)
if window is not None and offset > window:
return None
if covered is None:
covered = offset
elif covered != offset:
return None
return covered
class _MLXPromptCacheHistory:
def __init__(self, max_entries, max_bytes):
api = _mlx_prompt_cache_api()
if api is None:
raise RuntimeError("mlx-lm is too old for LRUPromptCache")
lru_cls, make, can_trim, trim = api
self._make_prompt_cache = make
self._can_trim = can_trim
self._trim = trim
self._max_bytes = max_bytes
self._lru = lru_cls(max_size = max_entries, max_bytes = max_bytes)
def fetch(self, model, key, tokens):
cache, rest = self._lru.fetch_nearest_cache(key, list(tokens))
if cache is not None:
if rest:
return cache, list(rest)
if self._can_trim(cache) and self._trim(cache, 1) == 1:
return cache, list(tokens[-1:])
if len(tokens) > 1:
head = list(tokens[:-1])
cache, rest = self._lru.fetch_nearest_cache(key, head)
if cache is not None:
covered = len(head) - len(rest)
return cache, list(tokens[covered:])
return self._make_prompt_cache(model), list(tokens)
def insert(self, key, tokens, cache):
# An over-budget entry evicts itself and every other conversation.
nbytes = sum(getattr(entry, "nbytes", 0) for entry in cache)
if nbytes > self._max_bytes:
logger.debug(
"MLX prompt cache: skipping %.2f GB entry over the %.2f GB budget",
nbytes / 1e9,
self._max_bytes / 1e9,
)
return
covered = _kv_prefix_coverage(cache)
if covered is None:
logger.debug("MLX prompt cache: skipping cache with unverifiable prefix coverage")
return
tokens = list(tokens)
if covered > len(tokens):
logger.debug(
"MLX prompt cache: cache covers %d tokens but only %d were tracked",
covered,
len(tokens),
)
return
tokens = tokens[:covered]
if not tokens:
return
self._lru.insert_cache(key, tokens, cache)
def _mlx_distributed_rank_size(group = None):
"""Return ``(rank, world_size)`` for an optional MLX distributed group."""
if group is None:
@ -313,6 +433,55 @@ class MLXInferenceBackend:
# Recorded for unload to release pinned memory back to the OS.
self._memory_limits_applied = {}
self._prompt_cache_history = None
self._prompt_cache_unavailable = False
def _prompt_cache(self):
if self._prompt_cache_history is not None or self._prompt_cache_unavailable:
return self._prompt_cache_history
max_bytes = _prompt_cache_max_bytes(self._memory_limits_applied.get("recommended_gb"))
if max_bytes <= 0:
self._prompt_cache_unavailable = True
logger.info("MLX prompt cache disabled by budget")
return None
try:
self._prompt_cache_history = _MLXPromptCacheHistory(
PROMPT_CACHE_ENTRIES,
max_bytes,
)
except Exception as exc:
self._prompt_cache_unavailable = True
logger.info("MLX prompt cache unavailable (%s); prefilling every request", exc)
return None
logger.info(
"MLX prompt cache: %d entries, %.2f GB budget",
PROMPT_CACHE_ENTRIES,
max_bytes / 1e9,
)
return self._prompt_cache_history
def _clear_prompt_cache(self):
self._prompt_cache_history = None
self._prompt_cache_unavailable = False
def _prepare_prompt_cache(self, prompt, adapter_state):
history = self._prompt_cache()
if history is None:
return prompt, None, None, None, 0
try:
tokenizer = self._tokenizer
bos = getattr(tokenizer, "bos_token", None)
add_special_tokens = bos is None or not prompt.startswith(bos)
tokens = list(tokenizer.encode(prompt, add_special_tokens = add_special_tokens))
if not tokens:
return prompt, None, None, None, 0
key = f"{self.active_model_name}|{adapter_state!r}"
cache, rest = history.fetch(self._model, key, tokens)
except Exception as exc:
logger.debug("MLX prompt cache lookup failed: %s", exc)
return prompt, None, None, None, 0
return rest, cache, key, tokens, len(tokens) - len(rest)
def _configure_memory_limits(self):
"""Apply Metal memory caps before loading a model.
@ -535,6 +704,7 @@ class MLXInferenceBackend:
self._distributed_world_size = 1
if self.active_model_name == model_name:
self.active_model_name = None
self._clear_prompt_cache()
gc.collect()
mx.clear_cache()
@ -731,24 +901,34 @@ class MLXInferenceBackend:
# <think> prefix on every native-protocol snapshot just as the normal
# decoding path does below.
normalized_output = think_prefix
logger.info(
"Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s",
len(prompt),
max_new_tokens,
type(self._model).__name__,
type(self._tokenizer).__name__,
)
with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
(
gen_prompt,
prompt_cache,
cache_key,
prompt_tokens,
cached_n,
) = self._prepare_prompt_cache(prompt, _adapter_state)
logger.info(
"Generating: prompt_len=%d, cached=%d, max_tokens=%d, model=%s, tokenizer=%s",
len(prompt),
cached_n,
max_new_tokens,
type(self._model).__name__,
type(self._tokenizer).__name__,
)
final_response = None
try:
# Enter request-scoped model state before yielding any response.
if think_prefix:
yield think_prefix
gen_kwargs = dict(
prompt = prompt,
prompt = gen_prompt,
max_tokens = max_new_tokens,
sampler = sampler,
)
if prompt_cache is not None:
gen_kwargs["prompt_cache"] = prompt_cache
if logits_processors is not None:
gen_kwargs["logits_processors"] = logits_processors
for response in stream_generate(
@ -757,6 +937,7 @@ class MLXInferenceBackend:
**gen_kwargs,
):
final_response = response
token_ids.append(response.token)
if preserve_native_channels:
piece = getattr(response, "text", None) or ""
delta = normalizer.feed(piece)
@ -764,7 +945,6 @@ class MLXInferenceBackend:
normalized_output += delta
yield normalized_output
else:
token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
@ -773,6 +953,13 @@ class MLXInferenceBackend:
if cancel_event and cancel_event.is_set():
break
if prompt_cache is not None and prompt_tokens is not None:
history = self._prompt_cache_history
if history is not None:
try:
history.insert(cache_key, prompt_tokens + token_ids, prompt_cache)
except Exception as exc:
logger.debug("MLX prompt cache insert failed: %s", exc)
except Exception as e:
import traceback
logger.error("stream_generate failed:\n%s", traceback.format_exc())
@ -785,6 +972,7 @@ class MLXInferenceBackend:
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
cached_n,
)
if normalizer is not None:
cancelled = cancel_event is not None and cancel_event.is_set()

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

@ -54,9 +54,8 @@ class GenStreamError(str):
"""A stream chunk carrying a real backend/generation error, not model text.
Subclasses str so existing display/logging consumers are unaffected, while
callers that must abort a distributed run on error (raise_on_streamed_error)
can distinguish a real error from model output whose visible text starts with
"Error:" by checking isinstance(chunk, GenStreamError).
callers can distinguish a real error from model output whose visible text
starts with "Error:" by checking isinstance(chunk, GenStreamError).
"""
__slots__ = ("public",)
@ -218,10 +217,14 @@ class InferenceOrchestrator:
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
from .worker import run_inference_process
cache_env = get_hf_cache_paths().child_env({})
with native_path_secret_removed_for_child_start():
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._cancel_event = _CTX.Event()
@ -229,7 +232,7 @@ class InferenceOrchestrator:
self._proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_inference_process,),
args = ("core.inference.worker", "run_inference_process", cache_env),
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,

876
studio/backend/core/inference/stt_ggml_sidecar.py Normal file
View file

@ -0,0 +1,876 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""whisper.cpp (GGML/GGUF) speech-to-text sidecar for Studio dictation.
Runs the same curated Whisper checkpoints as the Transformers sidecar
(stt_sidecar.py) through whisper.cpp's `whisper-server`, ~2.5x faster at
identical quality on Apple Silicon and CPU because its Metal/CPU kernels run
the weights in f16 where PyTorch MPS requires fp32.
Owns a single `whisper-server` subprocess bound to 127.0.0.1 on an ephemeral
port; the model loads on demand, stays warm between dictations, and unloads
after the same keep-alive as the Transformers sidecar. Curated GGML checkpoints
are single files from `unslothai/whisper-*-GGUF`, downloaded directly rather
than through the Model Hub (whose variant planner only handles `.gguf` chat
layouts).
Binary discovery mirrors `_find_llama_server_binary`: env override, then managed
Studio home, then PATH. With no binary the engine is unavailable and dictation
falls back to the Transformers sidecar; `scripts/build_whisper_cpp.sh` installs
the binary.
"""
from __future__ import annotations
import io
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import threading
import time
import urllib.request
import uuid
import wave
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, Optional
from loggers import get_logger
from core.inference.stt_sidecar import (
STT_KEEP_ALIVE_SECONDS,
SttAudioDecodeError,
SttLanguageError,
SttLoadCancelledError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
_decode_audio_bounded,
_known_whisper_languages,
_TARGET_SAMPLE_RATE,
_training_active,
normalize_whisper_language,
)
from utils.prebuilt.child_env import isolate_home, scrub_env, wsl_system_rocm_lib_dirs
from utils.prebuilt.runtime_libs import dedupe_existing_dirs
from utils.prebuilt.whisper_layout import lookup_marker
from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid
logger = get_logger(__name__)
# Curated GGML checkpoints, one repo per model. Keys match the Transformers
# sidecar's ids so the frontend reuses one picker; values are the single file
# inside each repo.
GGML_STT_REPOS: dict[str, str] = {
"tiny": "unslothai/whisper-tiny-GGUF",
"base": "unslothai/whisper-base-GGUF",
"small": "unslothai/whisper-small-GGUF",
"large-v3-turbo": "unslothai/whisper-large-v3-turbo-GGUF",
"large-v3": "unslothai/whisper-large-v3-GGUF",
}
GGML_STT_MODELS: dict[str, str] = {
"tiny": "whisper-tiny.bin",
"base": "whisper-base.bin",
"small": "whisper-small.bin",
"large-v3-turbo": "whisper-large-v3-turbo.bin",
"large-v3": "whisper-large-v3.bin",
}
DEFAULT_GGML_STT_MODEL = "small"
_SERVER_START_TIMEOUT_SECONDS = 120.0
_TRANSCRIBE_TIMEOUT_SECONDS = 600.0
class SttEngineUnavailableError(SttUnavailableError):
"""whisper-server is not installed; the GGUF dictation engine is off."""
def resolve_ggml_model_id(model: Optional[str]) -> str:
"""Validate a curated GGML model id. Custom repos are not supported here."""
if model is None or not str(model).strip():
return DEFAULT_GGML_STT_MODEL
normalized = str(model).strip()
if normalized in GGML_STT_MODELS:
return normalized
raise SttModelIdError(
f"STT model '{model}' is not a curated GGUF dictation model. "
f"Choose one of: {', '.join(GGML_STT_MODELS)}."
)
def _managed_whisper_cpp_dir() -> Path:
"""`<STUDIO_HOME>/whisper.cpp` in custom mode, else `~/.unsloth/whisper.cpp`.
Mirrors `managed_node_dir` / `_find_llama_server_binary` so managed runtimes
share one parent directory.
"""
legacy = Path.home() / ".unsloth" / "whisper.cpp"
try:
from utils.paths.storage_roots import studio_root
resolved = studio_root()
legacy_studio = Path.home() / ".unsloth" / "studio"
try:
is_legacy = resolved.resolve() == legacy_studio.resolve()
except (OSError, ValueError):
is_legacy = resolved == legacy_studio
return legacy if is_legacy else (resolved / "whisper.cpp")
except (ImportError, OSError, ValueError):
override = (
os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") or ""
).strip()
if override:
try:
return Path(override).expanduser().resolve() / "whisper.cpp"
except (OSError, ValueError):
return Path(override).expanduser() / "whisper.cpp"
return legacy
def find_whisper_server_binary() -> Optional[str]:
"""Locate the whisper-server binary.
Search order:
1. WHISPER_SERVER_PATH environment variable (direct path to binary)
2. UNSLOTH_WHISPER_CPP_PATH env var (custom whisper.cpp install dir)
3. managed dir: <STUDIO_HOME or ~/.unsloth>/whisper.cpp/{,build/bin/}whisper-server
4. whisper-server on PATH
"""
binary_name = "whisper-server.exe" if sys.platform == "win32" else "whisper-server"
def _layout_candidates(d: Path) -> list[Path]:
cands = [d / binary_name, d / "build" / "bin" / binary_name]
if sys.platform == "win32":
cands.append(d / "build" / "bin" / "Release" / binary_name)
return cands
env_path = os.environ.get("WHISPER_SERVER_PATH")
if env_path:
p = Path(env_path)
if _is_runnable(p):
return str(p)
custom_dir = os.environ.get("UNSLOTH_WHISPER_CPP_PATH")
if custom_dir:
for p in _layout_candidates(Path(custom_dir)):
if _is_runnable(p):
return str(p)
for p in _layout_candidates(_managed_whisper_cpp_dir()):
if _is_runnable(p):
return str(p)
return shutil.which(binary_name)
def _is_runnable(p: Path) -> bool:
"""A real whisper-server is an executable file. On Windows os.access(X_OK) is
effectively an existence check; on Unix it rejects a non-executable stub so a
half-written or wrong-mode file isn't mistaken for the server."""
return p.is_file() and (sys.platform == "win32" or os.access(p, os.X_OK))
def _whisper_install_marker(binary: str) -> Optional[dict]:
"""The prebuilt install marker above ``binary``, or None (source/custom builds)."""
return lookup_marker(binary).marker
def slim_runtime_intact(binary: str) -> bool:
"""True unless the marker says slim and the linked ggml runtime is missing
beside the server. New markers record the exact wired filenames
(linked_libraries), all of which must be present; legacy markers without the
field fall back to the per-OS core ggml name globs. A broken slim install
reads as engine-unavailable (reinstall via `unsloth studio update`), never a
crash at load."""
lookup = lookup_marker(binary)
marker = lookup.marker
if lookup.invalid or marker is None:
return not lookup.slim_collision
if not marker or marker.get("install_kind") != "slim":
return True
if lookup.authoritative:
valid = marker.get("component") == "whisper.cpp"
valid = valid and isinstance(marker.get("schema_version"), int)
valid = valid and all(
isinstance(marker.get(key), str) and marker[key]
for key in ("release_tag", "backend", "paired_llama_tag")
)
valid = valid and isinstance(marker.get("linked_libraries"), list)
valid = valid and bool(marker.get("linked_libraries"))
valid = valid and all(
isinstance(name, str) and name and Path(name).name == name
for name in marker["linked_libraries"]
)
if not valid:
return False
bin_dir = Path(binary).parent
linked = marker.get("linked_libraries")
if isinstance(linked, list) and linked and all(isinstance(name, str) for name in linked):
intact = all((bin_dir / name).is_file() for name in linked)
else:
if sys.platform == "win32":
required = ("ggml.dll", "ggml-base.dll")
elif sys.platform == "darwin":
required = ("libggml*.dylib", "libggml-base*.dylib")
else:
required = ("libggml.so*", "libggml-base.so*")
intact = all(any(p.is_file() for p in bin_dir.glob(pattern)) for pattern in required)
runtime_dirs = marker.get("linked_runtime_directories")
if intact and isinstance(runtime_dirs, list) and runtime_dirs:
intact = all(
isinstance(name, str)
and name
and (bin_dir / name).is_dir()
and any(path.is_file() for path in (bin_dir / name).rglob("*"))
for name in runtime_dirs
)
if intact and marker.get("backend") == "rocm":
expected_runtime_dirs = set() if sys.platform == "win32" else {"hipblaslt", "rocblas"}
intact = (
marker.get("runtime_wiring_version") == 2
and isinstance(runtime_dirs, list)
and set(runtime_dirs) == expected_runtime_dirs
)
if not intact:
logger.warning(
"slim whisper install is missing its linked ggml runtime at "
f"{bin_dir}; run `unsloth studio update` to reinstall it"
)
return intact
def is_available() -> bool:
binary = find_whisper_server_binary()
if binary is None:
return False
if not slim_runtime_intact(binary):
return False
try:
import av # noqa: F401
except Exception:
# No PyAV means every transcription 501s on decode.
return False
return True
def ensure_engine_available() -> str:
binary = find_whisper_server_binary()
if binary is None:
raise SttEngineUnavailableError(
"The local transcription runtime is not installed. Run "
"`unsloth studio update` to install it."
)
if not slim_runtime_intact(binary):
raise SttEngineUnavailableError(
"The local transcription runtime is missing its paired ggml "
"libraries. Run `unsloth studio update` to reinstall it."
)
return binary
# ---------------------------------------------------------------------------
# whisper-server child-process environment
# ---------------------------------------------------------------------------
# Build the whisper-server env: prepend the binary dir (co-located libs win, and
# a backstop where the loader ignores the rpath) and scrub secret-bearing vars the
# binary never needs. On WSL2 ROCm the system HIP libs go first, since a bundle's
# bare-metal HIP cannot drive /dev/dxg. A CUDA bundle ships libggml-cuda.so but not
# libcudart/libcublas (paired with the user's PyTorch), so add the
# CUDA-from-PyTorch runtime dirs the selection gated on, else the backend cannot
# resolve a runtime that lives only in wheels. Mirrors llama's binary_env(); the
# scrub/WSL/dedupe helpers live in utils.prebuilt.
# Module-level aliases keep the historical patch points for tests and callers.
_wsl_system_rocm_lib_dirs = wsl_system_rocm_lib_dirs
_dedupe_existing_dirs = dedupe_existing_dirs
def _whisper_server_child_env(binary: str) -> dict[str, str]:
"""Env for the whisper-server subprocess: secrets scrubbed, home/profile vars
repointed at a managed scratch dir (a downloaded binary must not see the real
home's token caches), co-located libs on the loader path, WSL system HIP first
on WSL2 ROCm."""
env = scrub_env(os.environ)
isolate_home(env, str(_managed_whisper_cpp_dir() / ".child_home"))
bin_dir = str(Path(binary).parent)
# A CUDA bundle needs the CUDA-from-PyTorch wheel dirs so libcudart/libcublas
# resolve at launch when they live only in site-packages/nvidia/*/lib. Placed
# after bin_dir so co-located libs still win; empty for other bundles.
cuda_runtime_dirs: list[str] = []
bundle_dir = Path(bin_dir)
has_cuda_module = any(
path.is_file()
for pattern in ("libggml-cuda.so*", "ggml-cuda*.dll")
for path in bundle_dir.glob(pattern)
)
if has_cuda_module:
try:
from utils.prebuilt.runtime_libs import python_runtime_dirs
cuda_runtime_dirs = python_runtime_dirs()
except Exception:
cuda_runtime_dirs = []
if sys.platform == "win32":
var, lead = "PATH", [bin_dir, *cuda_runtime_dirs]
elif sys.platform == "darwin":
var, lead = "DYLD_LIBRARY_PATH", [bin_dir]
else:
var, lead = "LD_LIBRARY_PATH", [bin_dir, *cuda_runtime_dirs]
wsl_rocm = _wsl_system_rocm_lib_dirs()
if wsl_rocm:
lead = [*wsl_rocm, bin_dir, *cuda_runtime_dirs]
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
existing = [p for p in env.get(var, "").split(os.pathsep) if p]
env[var] = os.pathsep.join(_dedupe_existing_dirs([*lead, *existing]))
return env
# ---------------------------------------------------------------------------
# Model file download (single files; deliberately outside the Model Hub flow)
# ---------------------------------------------------------------------------
def _cached_model_path(model_id: str) -> Optional[str]:
"""Path of a fully downloaded GGML file in the shared HF cache, else None."""
from huggingface_hub import hf_hub_download
try:
return hf_hub_download(
repo_id = GGML_STT_REPOS[model_id],
filename = GGML_STT_MODELS[model_id],
local_files_only = True,
)
except Exception:
return None
class _GgmlDownloadState:
"""Tracks one background hf_hub_download of a curated GGML file."""
def __init__(self) -> None:
self._lock = threading.Lock()
self._thread: Optional[threading.Thread] = None
self._model_id: Optional[str] = None
self._error: Optional[str] = None
self._total_bytes: Optional[int] = None
self._etag: Optional[str] = None
def status(self) -> dict:
with self._lock:
downloading = self._thread is not None and self._thread.is_alive()
return {
"downloading": downloading,
"model": self._model_id if downloading else None,
"error": self._error,
"bytes_total": self._total_bytes if downloading else None,
"bytes_done": self._incomplete_bytes() if downloading else None,
}
def _incomplete_bytes(self) -> Optional[int]:
"""Best-effort progress: size of the in-flight blob in the HF cache.
hf_hub_download writes ``blobs/<etag>.incomplete``; prefer this file's
etag, else the largest in-flight blob.
"""
try:
from huggingface_hub.constants import HF_HUB_CACHE
# Caller may hold the non-reentrant self._lock; bare reads are safe.
model_id = self._model_id
if not model_id:
return None
repo_dir = (
Path(HF_HUB_CACHE)
/ f"models--{GGML_STT_REPOS[model_id].replace('/', '--')}"
/ "blobs"
)
if not repo_dir.is_dir():
return None
etag = self._etag
if etag:
target = repo_dir / f"{etag}.incomplete"
if target.is_file():
return target.stat().st_size
sizes = [p.stat().st_size for p in repo_dir.glob("*.incomplete") if p.is_file()]
return max(sizes) if sizes else None
except Exception:
return None
def start(
self,
model_id: str,
hf_token: Optional[str] = None,
) -> None:
model_id = resolve_ggml_model_id(model_id)
with self._lock:
if self._thread is not None and self._thread.is_alive():
if self._model_id == model_id:
return
raise SttModelIdError(
f"Another GGUF dictation model ('{self._model_id}') is still "
"downloading; wait for it to finish."
)
self._model_id = model_id
self._error = None
self._total_bytes = None
self._etag = None
thread = threading.Thread(target = self._run, args = (model_id, hf_token), daemon = True)
self._thread = thread
thread.start()
def _run(self, model_id: str, hf_token: Optional[str]) -> None:
repo_id = GGML_STT_REPOS[model_id]
filename = GGML_STT_MODELS[model_id]
try:
from huggingface_hub import (
get_hf_file_metadata,
hf_hub_download,
hf_hub_url,
)
try:
# One HEAD request for the total and etag.
meta = get_hf_file_metadata(hf_hub_url(repo_id, filename), token = hf_token or None)
with self._lock:
self._total_bytes = meta.size
self._etag = meta.etag
except Exception:
pass
hf_hub_download(
repo_id = repo_id,
filename = filename,
token = hf_token or None,
)
except Exception as exc:
logger.warning("GGUF STT download failed for %s: %s", model_id, exc)
with self._lock:
self._error = f"Download failed for '{model_id}'."
_download_state = _GgmlDownloadState()
def start_model_download(model: Optional[str], hf_token: Optional[str] = None) -> None:
_download_state.start(resolve_ggml_model_id(model), hf_token)
def download_status() -> dict:
return _download_state.status()
# ---------------------------------------------------------------------------
# WAV packaging
# ---------------------------------------------------------------------------
def _pcm_to_wav_bytes(decoded_audio) -> bytes:
"""Wrap decoded float32 mono 16 kHz PCM into an in-memory 16-bit WAV."""
import numpy as np
clipped = np.clip(decoded_audio, -1.0, 1.0)
pcm16 = (clipped * 32767.0).astype("<i2")
buf = io.BytesIO()
with wave.open(buf, "wb") as w:
w.setnchannels(1)
w.setsampwidth(2)
w.setframerate(_TARGET_SAMPLE_RATE)
w.writeframes(pcm16.tobytes())
return buf.getvalue()
# ---------------------------------------------------------------------------
# Sidecar
# ---------------------------------------------------------------------------
class GgmlSttSidecar:
"""Owns one whisper-server subprocess and proxies dictation to it."""
def __init__(self, keep_alive_seconds: float = STT_KEEP_ALIVE_SECONDS) -> None:
self._lock = threading.RLock()
self._process: Optional[subprocess.Popen] = None
self._port: Optional[int] = None
self._model_id: Optional[str] = None
self._idle_timer: Optional[threading.Timer] = None
self._idle_generation = 0
self._keep_alive_seconds = keep_alive_seconds
# Set while whisper-server starts so training admission can account for
# the accelerator memory it is about to bind. Read without the lock.
self._loading = False
# A still-starting whisper-server is cancellable so training can preempt
# it before it binds accelerator memory. Assigned inside self._lock but
# acted on without it: cancel_pending_load() runs while load() holds the
# lock, so the event is the source of truth and terminating the process
# is a best-effort fast path.
self._load_cancel_event: Optional[threading.Event] = None
self._starting_process: Optional[subprocess.Popen] = None
# Set before the updater waits for _lock, then kept set while it owns
# the lock and atomically replaces the managed install tree. New loads
# fail fast instead of starting a process from files being swapped.
self._update_in_progress = False
@property
def loaded_model(self) -> Optional[str]:
# Lock-free status read (like stt_sidecar.py): transcribe() holds
# self._lock for the whole inference call (up to
# _TRANSCRIBE_TIMEOUT_SECONDS), and status polls plus training admission
# must not block behind it. _process_alive() snapshots self._process
# before poll(), which subprocess guards with _waitpid_lock, so a
# concurrent unload is safe.
return self._model_id if self._process_alive() else None
@property
def device(self) -> Optional[str]:
return "whisper.cpp" if self._process_alive() else None
def is_loading(self) -> bool:
# True only while whisper-server is starting (seconds to bind its GPU
# backend); load() sets and clears the flag around that window.
return self._loading
@property
def keep_alive_seconds(self) -> float:
return self._keep_alive_seconds
def _process_alive(self) -> bool:
# Snapshot self._process once: a concurrent unload() nulls it under the
# lock, so lock-free readers would otherwise re-read None between the
# truthiness check and .poll().
process = self._process
return process is not None and process.poll() is None
# -- idle unload ------------------------------------------------------
def _cancel_idle_unload_locked(self) -> None:
self._idle_generation += 1
if self._idle_timer is not None:
self._idle_timer.cancel()
self._idle_timer = None
def _schedule_idle_unload_locked(self) -> None:
self._cancel_idle_unload_locked()
if not self._process_alive():
return
generation = self._idle_generation
timer = threading.Timer(self._keep_alive_seconds, self._idle_unload, args = (generation,))
timer.daemon = True
self._idle_timer = timer
timer.start()
def _idle_unload(self, generation: int) -> None:
with self._lock:
if generation != self._idle_generation:
return
logger.info("Unloading idle GGUF STT model %s", self._model_id)
self._release_locked()
# -- process lifecycle -------------------------------------------------
def _release_locked(self) -> None:
self._cancel_idle_unload_locked()
process = self._process
self._process = None
self._port = None
self._model_id = None
if process is not None and process.poll() is None:
process.terminate()
try:
process.wait(timeout = 10)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout = 10)
if process is not None:
forget_pid(process.pid)
def unload(self) -> None:
with self._lock:
self._release_locked()
def _raise_if_update_in_progress(self) -> None:
if self._update_in_progress:
raise SttEngineUnavailableError(
"The local transcription runtime is being updated. Try dictation again shortly."
)
@contextmanager
def update_maintenance(self) -> Iterator[bool]:
"""Block new loads while the managed whisper.cpp tree is replaced.
The flag is published before waiting for an existing transcription to
release ``_lock``. Holding that lock across the yielded installer phase
prevents Windows from relocking the executable and prevents every host
from starting a process against a partially swapped tree. The yielded
value records whether a warm model had to be unloaded.
"""
self._update_in_progress = True
try:
with self._lock:
model_was_active = self._process_alive()
self._release_locked()
yield model_was_active
finally:
self._update_in_progress = False
def cancel_pending_load(self) -> bool:
# Preempt a starting whisper-server so training does not launch while it
# binds accelerator memory. load() holds self._lock for the whole startup,
# so act without the lock: signal abort and terminate the starting
# process. _wait_for_server observes the event and raises, then load()
# reaps the process and releases the lock.
if not self._loading:
return False
event = self._load_cancel_event
if event is None:
return False
event.set()
process = self._starting_process
if process is not None and process.poll() is None:
try:
process.terminate()
except Exception:
pass
return True
def wait_for_load_to_settle(self) -> None:
# load() holds self._lock across startup and cancel cleanup, so acquiring
# it blocks until a cancelled server is killed, reaped, and its
# accelerator memory released.
with self._lock:
pass
@staticmethod
def _reserve_free_port() -> tuple[socket.socket, int]:
"""Bind an ephemeral port and keep the socket held.
The caller closes the reservation immediately before spawning
whisper-server, shrinking the window in which another local process
could bind the port. SO_REUSEADDR lets the child rebind right after.
"""
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
s.bind(("127.0.0.1", 0))
return s, s.getsockname()[1]
def _ensure_model_downloaded(self, model_id: str) -> str:
path = _cached_model_path(model_id)
if path is None:
raise SttModelNotDownloadedError(
f"STT model '{model_id}' (GGUF) is not downloaded. "
"Download it in Settings, then Voice, before loading it."
)
return path
def load(self, model: Optional[str] = None) -> None:
"""Start (or switch) whisper-server for the requested curated model."""
self._raise_if_update_in_progress()
model_id = resolve_ggml_model_id(model)
with self._lock:
self._raise_if_update_in_progress()
binary = ensure_engine_available()
if self._process_alive() and self._model_id == model_id:
self._schedule_idle_unload_locked()
return
model_path = self._ensure_model_downloaded(model_id)
self._release_locked()
reservation, port = self._reserve_free_port()
command = [binary, "-m", model_path, "--host", "127.0.0.1", "--port", str(port)]
marker = _whisper_install_marker(binary)
if _training_active():
# Keep whisper.cpp off the accelerator during training (like the
# Transformers sidecar's CPU choice) so a mid-training dictation
# cannot reclaim the VRAM training just freed.
command.append("--no-gpu")
elif marker is not None and marker.get("backend") == "cpu":
# A deliberate CPU install must stay CPU: the slim wiring links
# every llama ggml backend (including CUDA/ROCm), so without
# this flag a cpu-selected install would still grab the GPU.
command.append("--no-gpu")
logger.info(
"Starting whisper-server for STT model %s on 127.0.0.1:%s",
model_id,
port,
)
cancel_event = threading.Event()
self._load_cancel_event = cancel_event
self._loading = True
try:
# Release the reservation as late as possible: whisper-server
# binds the port moments after this close.
reservation.close()
process = subprocess.Popen(
command,
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
stdin = subprocess.DEVNULL,
# Co-located GPU libs on the loader path (WSL system HIP first),
# secrets scrubbed from the downloaded binary's env.
env = _whisper_server_child_env(binary),
# Die with Studio (Linux PDEATHSIG, Windows job) so a crash
# never orphans a server holding the model.
**child_popen_kwargs(),
)
self._starting_process = process
adopt_pid(process.pid) # terminate_all backstop for graceful exits
try:
self._wait_for_server(process, port, cancel_event)
except Exception:
if process.poll() is None:
process.kill()
process.wait(timeout = 10)
forget_pid(process.pid)
raise
self._process = process
self._port = port
self._model_id = model_id
self._schedule_idle_unload_locked()
finally:
reservation.close() # no-op when already released before spawn
self._loading = False
self._load_cancel_event = None
self._starting_process = None
@staticmethod
def _wait_for_server(
process: subprocess.Popen,
port: int,
cancel_event: Optional[threading.Event] = None,
) -> None:
deadline = time.monotonic() + _SERVER_START_TIMEOUT_SECONDS
while time.monotonic() < deadline:
if cancel_event is not None and cancel_event.is_set():
raise SttLoadCancelledError(
"GGUF STT model loading was cancelled so training could start."
)
if process.poll() is not None:
raise SttEngineUnavailableError(
"The local transcription runtime exited before becoming "
"ready; the model file may be corrupt or unsupported."
)
# Require a whisper-server-specific response twice, with the managed
# child alive around each probe. An arbitrary local process that won
# the bind race would otherwise be mistaken for the sidecar and
# receive the user's microphone audio.
if GgmlSttSidecar._probe_is_whisper_server(process, port) and (
GgmlSttSidecar._probe_is_whisper_server(process, port)
):
return
time.sleep(0.2)
raise SttEngineUnavailableError("The local transcription runtime did not start in time.")
@staticmethod
def _probe_is_whisper_server(process: subprocess.Popen, port: int) -> bool:
"""One readiness probe: our child is alive and the responder looks like
whisper.cpp's server (its index page and errors identify whisper)."""
if process.poll() is not None:
return False
try:
req = urllib.request.Request(f"http://127.0.0.1:{port}/", method = "GET")
with urllib.request.urlopen(req, timeout = 2) as response:
body = response.read(65536)
except Exception:
return False
if process.poll() is not None:
return False
return b"whisper" in body.lower()
# -- transcription ------------------------------------------------------
def transcribe(
self,
audio: bytes,
model: Optional[str] = None,
language: Optional[str] = None,
fast: bool = False,
) -> dict:
"""Transcribe encoded audio bytes via whisper-server.
Accepts any container PyAV can decode (same validation and caps as the
Transformers sidecar). Returns {text, language, duration, model}.
"""
self._raise_if_update_in_progress()
ensure_engine_available()
model_id = resolve_ggml_model_id(model)
lang = normalize_whisper_language(language)
known_languages = _known_whisper_languages()
if lang is not None and known_languages is not None and lang not in known_languages:
raise SttLanguageError(
f"Language '{language}' is not supported by STT model '{model_id}'."
)
# Reject a missing model before decoding so a long clip does not burn CPU
# only to 409 (matches the Transformers sidecar's preflight).
self._ensure_model_downloaded(model_id)
decoded_audio = _decode_audio_bounded(audio)
wav_bytes = _pcm_to_wav_bytes(decoded_audio)
with self._lock:
try:
self.load(model_id)
text = self._post_inference(wav_bytes, lang, fast)
finally:
self._schedule_idle_unload_locked()
duration = (len(decoded_audio) / _TARGET_SAMPLE_RATE) if len(decoded_audio) else None
return {
"text": text,
"language": lang,
"duration": duration,
"model": model_id,
}
def _post_inference(self, wav_bytes: bytes, lang: Optional[str], fast: bool) -> str:
boundary = uuid.uuid4().hex
fields = {
"temperature": "0.0",
"response_format": "json",
# Match the Transformers sidecar: 5-way beam search, greedy for fast.
"beam_size": "1" if fast else "5",
"language": lang or "auto",
}
parts: list[bytes] = []
for name, value in fields.items():
parts.append(
(
f"--{boundary}\r\nContent-Disposition: form-data; "
f'name="{name}"\r\n\r\n{value}\r\n'
).encode()
)
parts.append(
(
f"--{boundary}\r\nContent-Disposition: form-data; "
'name="file"; filename="dictation.wav"\r\n'
"Content-Type: audio/wav\r\n\r\n"
).encode()
+ wav_bytes
+ b"\r\n"
)
parts.append(f"--{boundary}--\r\n".encode())
body = b"".join(parts)
req = urllib.request.Request(
f"http://127.0.0.1:{self._port}/inference",
data = body,
headers = {"Content-Type": f"multipart/form-data; boundary={boundary}"},
)
try:
with urllib.request.urlopen(req, timeout = _TRANSCRIBE_TIMEOUT_SECONDS) as resp:
payload = json.load(resp)
except SttAudioDecodeError:
raise
except Exception as exc:
raise SttEngineUnavailableError(
"The local transcription runtime did not answer the request."
) from exc
text = payload.get("text")
if not isinstance(text, str):
raise SttAudioDecodeError("Could not decode the audio.")
# whisper.cpp joins segments with newlines; dictation wants one line.
return " ".join(part.strip() for part in text.splitlines() if part.strip()).strip()
_sidecar: Optional[GgmlSttSidecar] = None
def get_ggml_stt_sidecar() -> GgmlSttSidecar:
global _sidecar
if _sidecar is None:
_sidecar = GgmlSttSidecar()
return _sidecar

1130
studio/backend/core/inference/stt_sidecar.py Normal file
View file

File diff suppressed because it is too large Load diff

9
studio/backend/core/rag/embed_llama_server.py
View file

@ -188,7 +188,14 @@ class LlamaServerBackend:
match = [f for f in files if variant in f.lower()] or files
filename = sorted(match, key = len)[0]
logger.info("resolving GGUF embedder %s/%s", repo, filename)
self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token)
from utils.hf_cache_settings import active_hf_hub_cache
self._model_path = hf_hub_download(
repo_id = repo,
filename = filename,
token = token,
cache_dir = active_hf_hub_cache(),
)
self._model_repo = desired
self._dim = None
return self._model_path

80
studio/backend/core/rag/embeddings.py
View file

@ -22,6 +22,7 @@ from typing import Callable
from utils.hardware.hardware import DeviceType, get_device
from utils.transformers_dtype import dtype_kwargs
from utils.utils import hf_env_offline
from . import config
@ -103,9 +104,15 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
else:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
from utils.hf_cache_settings import active_hf_hub_cache
try:
local = hf_hub_download(name, "modules.json", token = token or None)
local = hf_hub_download(
name,
"modules.json",
token = token or None,
cache_dir = active_hf_hub_cache(),
)
except EntryNotFoundError:
return ()
data = json.loads(open(local).read())
@ -119,30 +126,55 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
return ()
def _guard_model_security(name: str) -> None:
def _guard_model_security(name: str, local_only: bool = False) -> None:
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
/settings gate (a name can also arrive via env/default); local paths and unreachable
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
``local_only`` (offline) inspects the local cache; subdir probes are skipped (they'd hit the
network and hang, and the offline gate walks the whole snapshot anyway).
"""
try:
from utils.security import evaluate_file_security, security_load_subdirs
token = _ambient_hf_token()
# Union the audio-model load roots with the ST module dirs so a flagged pickle
# directly under a Transformer module dir (0_Transformer/) blocks instead of
# passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
)
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
if local_only:
load_subdirs = ()
else:
# Union audio-model load roots with ST module dirs so a flagged pickle under a
# Transformer module dir blocks instead of passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys(
(*security_load_subdirs(name, token), *_st_module_subdirs(name, token))
)
)
blocked = evaluate_file_security(
name, hf_token = token, load_subdirs = load_subdirs, local_only_load = local_only
).blocked
except Exception:
return
if blocked:
raise UnsafeEmbeddingModelError(
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
"scan; refusing to load. Set a different RAG embedding model."
reason = (
"has cached pickle weights that cannot be security-scanned offline and no "
"safetensors alternative"
if local_only
else "is flagged as unsafe by Hugging Face's security scan"
)
raise UnsafeEmbeddingModelError(
f"Embedding model {name!r} {reason}; refusing to load. "
"Set a different RAG embedding model."
)
def _st_accepts_local_files_only(st_cls) -> bool:
"""Whether this SentenceTransformer version accepts local_files_only; passing it to an
older constructor raises, so gate on the signature."""
try:
import inspect
return "local_files_only" in inspect.signature(st_cls.__init__).parameters
except Exception:
return False
def _get(model_name: str | None = None):
@ -150,15 +182,35 @@ def _get(model_name: str | None = None):
for a ~1.5x speedup at negligible accuracy loss."""
global _model, _name
name = model_name or config.effective_embedding_model()
# Capture offline state once so the gate and the load agree (no window where the gate is
# skipped as offline but the constructor then reaches the network).
local_only = hf_env_offline()
with _lock:
if _model is None or _name != name:
_install_torchao_stub_once()
from sentence_transformers import SentenceTransformer
from utils.hf_cache_settings import active_hf_hub_cache
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
_guard_model_security(name, local_only)
st_kwargs = dict(
device = device,
cache_folder = active_hf_hub_cache(),
model_kwargs = dtype_kwargs("float16"),
)
load_target = name
if local_only:
from utils.utils import hf_cache_snapshot_dir
snapshot = hf_cache_snapshot_dir(name)
if snapshot is not None:
# Load from the local snapshot dir: a local path never touches the Hub, so
# this is offline-safe on ANY sentence-transformers version (even ones
# predating local_files_only).
load_target = str(snapshot)
elif _st_accepts_local_files_only(SentenceTransformer):
st_kwargs["local_files_only"] = True
_model = SentenceTransformer(load_target, **st_kwargs)
_name = name
return _model

173
studio/backend/core/training/resume.py
View file

@ -4,6 +4,8 @@
"""Helpers for validating resumable training outputs."""
import json
import pickletools
import zipfile
from pathlib import Path
from typing import Optional
@ -33,21 +35,158 @@ def _checkpoint_step(path: Path) -> int:
return -1
def get_resume_checkpoint_path(path_value: str) -> Optional[str]:
_MODEL_FILES = (
"adapter_model.safetensors",
"adapter_model.bin",
"model.safetensors",
"pytorch_model.bin",
)
_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json")
def _valid_state_file(path: Path, require_tensor: bool = True) -> bool:
try:
if not path.is_file() or path.stat().st_size == 0:
return False
if path.suffix == ".safetensors":
try:
from safetensors import SafetensorError, safe_open
except ImportError:
return False
try:
with safe_open(str(path), framework = "np") as state:
return bool(state.keys())
except SafetensorError:
return False
if path.suffix in {".bin", ".pt"}:
with zipfile.ZipFile(path) as state:
infos = state.infolist()
names = [info.filename for info in infos]
data_name = next(
(name for name in names if name == "data.pkl" or name.endswith("/data.pkl")),
None,
)
if data_name is None:
return False
data_prefix = data_name.removesuffix("data.pkl") + "data/"
operations = list(pickletools.genops(state.read(data_name)))
if not operations or operations[-1][0].name != "STOP":
return False
if not require_tensor:
return True
# Require a non-empty tensor record; a zero-byte one fails torch.load.
return any(
info.filename.startswith(data_prefix)
and not info.is_dir()
and info.file_size > 0
for info in infos
)
# Unrecognized state-file formats are not usable resume state.
return False
except (OSError, ValueError, zipfile.BadZipFile):
return False
def _checkpoint_state(path: Path) -> Optional[int]:
try:
state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8"))
step = state.get("global_step") if isinstance(state, dict) else None
except (OSError, UnicodeDecodeError, json.JSONDecodeError):
return None
if isinstance(step, bool) or not isinstance(step, int) or step < 0:
return None
directory_step = _checkpoint_step(path)
return step if directory_step < 0 or step == directory_step else None
_INDEX_SHARD_SUFFIX = {
"model.safetensors.index.json": ".safetensors",
"pytorch_model.bin.index.json": ".bin",
}
def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool:
# Shard must be a relative, in-format path contained in the checkpoint dir.
if not isinstance(shard, str) or not shard:
return False
if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix:
return False
try:
root = checkpoint.resolve(strict = True)
candidate = (checkpoint / shard).resolve(strict = True)
candidate.relative_to(root)
except (OSError, ValueError):
return False
return _valid_state_file(candidate)
def _has_model_state(path: Path) -> bool:
if any(_valid_state_file(path / name) for name in _MODEL_FILES):
return True
for name in _MODEL_INDEXES:
try:
index = json.loads((path / name).read_text(encoding = "utf-8"))
shards = set(index["weight_map"].values())
except (
AttributeError,
OSError,
KeyError,
TypeError,
UnicodeDecodeError,
json.JSONDecodeError,
):
continue
expected_suffix = _INDEX_SHARD_SUFFIX[name]
if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards):
return True
return False
def is_resume_checkpoint_valid(
path: Path,
expected_step: Optional[int] = None,
backend: Optional[str] = None,
) -> bool:
step = _checkpoint_state(path) if path.is_dir() else None
step_valid = step is not None and (expected_step is None or step == expected_step)
if backend == "mlx":
valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file(
path / "optimizer_state.safetensors"
)
else:
valid_bundle = (
_has_model_state(path)
# optimizer/scheduler state can be validly tensor-free (e.g. SGD without
# momentum); _has_model_state still requires real model tensors.
and _valid_state_file(path / "optimizer.pt", require_tensor = False)
and _valid_state_file(path / "scheduler.pt", require_tensor = False)
)
if backend is None and not valid_bundle:
valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file(
path / "optimizer_state.safetensors"
)
return step_valid and valid_bundle
def get_resume_checkpoint_path(
path_value: str, expected_step: Optional[int] = None
) -> Optional[str]:
path = resolve_output_dir(path_value)
if not _is_under_outputs(path) or not path.is_dir():
return None
if (path / "trainer_state.json").is_file():
if is_resume_checkpoint_valid(path, expected_step):
return str(path)
checkpoints = [
child
for child in path.glob("checkpoint-*")
if child.is_dir() and (child / "trainer_state.json").is_file()
]
if not checkpoints:
return None
return str(max(checkpoints, key = _checkpoint_step))
checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True)
return next(
(
str(checkpoint)
for checkpoint in checkpoints
if _checkpoint_step(checkpoint) >= 0
and is_resume_checkpoint_valid(checkpoint, expected_step)
),
None,
)
def normalize_resume_output_dir(path_value: str) -> str:
@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool:
def can_resume_run(run: dict) -> bool:
if run.get("resumed_later"):
return False
# Set when a stop-and-save failed to write a current-step checkpoint.
if run.get("resume_blocked"):
return False
if _uses_s3_dataset(run):
return False
status = run.get("status")
if status == "error":
# A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability.
return has_resume_state(run.get("output_dir"))
final_step = run.get("final_step")
total_steps = run.get("total_steps")
has_remaining_steps = (
@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool:
or total_steps <= 0
or final_step < total_steps
)
return (
run.get("status") == "stopped"
and has_remaining_steps
and has_resume_state(run.get("output_dir"))
)
return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir"))

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

@ -754,6 +754,9 @@ class TrainingBackend:
def __init__(self):
# Subprocess state
self._proc: Optional[mp.Process] = None
# True from the sidecar-swap handshake until the worker is recorded, so
# installs and STT loads treat the startup window as active.
self._spawn_in_progress: bool = False
self._event_queue: Any = None
self._stop_queue: Any = None
self._pump_thread: Optional[threading.Thread] = None
@ -761,6 +764,7 @@ class TrainingBackend:
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
self._lock = threading.Lock()
self._run_intent_lock = threading.RLock()
# Stop watchdog: after a stop is requested, escalates to force_terminate()
# if the worker does not exit on its own within a bounded time. The watched
@ -773,6 +777,7 @@ class TrainingBackend:
self._progress = TrainingProgress()
self._should_stop = False
self._cancel_requested = False # True only for stop(save=False)
self._cancel_cleanup_output_dir: Optional[str] = None
# Throttled training-status logging to the server log (not one line/step).
self._last_progress_log_ts: float = 0.0
@ -792,6 +797,8 @@ class TrainingBackend:
# Job metadata
self.current_job_id: Optional[str] = None
self._output_dir: Optional[str] = None
self._resume_source_run_id: Optional[str] = None
self._terminal_finalize_payload: Optional[dict] = None
# DB persistence
self._metric_buffer: list[dict] = []
@ -819,6 +826,7 @@ class TrainingBackend:
job_id: str,
*,
before_spawn = None,
resume_source_run_id: Optional[str] = None,
**kwargs,
) -> bool:
"""Spawn a subprocess to run the full training pipeline.
@ -924,16 +932,21 @@ class TrainingBackend:
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
cache_env = get_hf_cache_paths().child_env({})
try:
with native_path_secret_removed_for_child_start():
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
args = ("core.training.worker", "run_training_process", cache_env),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
@ -956,6 +969,7 @@ class TrainingBackend:
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._cancel_cleanup_output_dir = None
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
@ -972,7 +986,10 @@ class TrainingBackend:
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._output_dir = config.get("output_dir") if resume_source_run_id else None
self._progress.output_dir = self._output_dir
self._resume_source_run_id = resume_source_run_id
self._terminal_finalize_payload = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
@ -982,6 +999,7 @@ class TrainingBackend:
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._last_hf_cache_env = cache_env
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
@ -990,6 +1008,17 @@ class TrainingBackend:
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
if resume_source_run_id and not self._db_run_created:
if proc.is_alive():
proc.terminate()
proc.join(timeout = 5.0)
if proc.is_alive():
proc.kill()
proc.join(timeout = 2.0)
self._progress.is_training = False
self._progress.error = "Resume checkpoint is no longer available."
self._spawn_in_progress = False
return False
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
@ -1011,28 +1040,75 @@ class TrainingBackend:
def stop_training(self, save: bool = True) -> bool:
"""Send stop signal to the training subprocess."""
self._should_stop = True
if not save:
self._cancel_requested = True
with self._lock:
if self._stop_queue is not None:
try:
self._stop_queue.put({"type": "stop", "save": save})
except (OSError, ValueError):
pass
# Update progress immediately for responsive UI.
self._progress.status_message = (
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
)
# Guarantee the run finalizes even if the worker wedges after saving.
self._start_stop_watchdog(cancel = not save)
with self._run_intent_lock:
with self._lock:
run_id = self.current_job_id
if not save and run_id:
persist_error: Optional[Exception] = None
for attempt in range(_DB_FINALIZE_RETRIES):
try:
from storage.studio_db import mark_run_cancel_requested
self._ensure_db_run_created()
with self._lock:
terminal_payload = self._terminal_finalize_payload
if (
terminal_payload
and terminal_payload.get("expected_job_id") == run_id
):
return False
if not mark_run_cancel_requested(run_id):
if self._db_run_created:
return False
raise RuntimeError(
"Training run disappeared before cancellation persisted"
)
if self.current_job_id != run_id:
return False
self._should_stop = self._cancel_requested = True
self._cancel_cleanup_output_dir = self._output_dir
self._output_dir = self._progress.output_dir = None
persist_error = None
break
except Exception as exc:
persist_error = exc
if attempt + 1 < _DB_FINALIZE_RETRIES:
time.sleep(_DB_FINALIZE_RETRY_S)
if persist_error is not None:
raise RuntimeError("Failed to persist Stop-without-Save") from persist_error
with self._lock:
if self.current_job_id != run_id:
return False
if save or not run_id:
self._should_stop = True
if not save and not run_id:
self._cancel_requested = True
self._cancel_cleanup_output_dir = self._output_dir
self._output_dir = self._progress.output_dir = None
if self._stop_queue is not None:
try:
self._stop_queue.put({"type": "stop", "save": save})
except (OSError, ValueError):
pass
self._progress.status_message = (
"Stopping training and saving checkpoint..."
if save
else "Cancelling training..."
)
self._start_stop_watchdog(cancel = not save, expected_job_id = run_id)
return True
def _start_stop_watchdog(self, cancel: bool) -> None:
def _start_stop_watchdog(
self,
cancel: bool,
expected_job_id: Optional[str] = None,
) -> None:
"""Start a daemon that force-terminates the worker if a requested stop does not
exit on its own. No-op if no worker is alive or a live watchdog already watches
this proc (a stale watchdog on an old proc never blocks a new run's watcher)."""
with self._lock:
if expected_job_id is not None and self.current_job_id != expected_job_id:
return
proc = self._proc
if proc is None or not proc.is_alive():
return
@ -1113,8 +1189,9 @@ class TrainingBackend:
watched_job_id: Optional[str] = None,
) -> None:
"""Finalize parent state after a force-terminate so the UI leaves "Stopping..."
even if the worker is wedged in driver teardown; preserves output_dir so a saved
checkpoint is kept. No-ops if a new run already replaced the watched worker, so a
even if the worker is wedged in driver teardown; preserves output_dir on a save so
the checkpoint is kept, and clears it on a cancel (Stop without saving must not
offer resume/export). No-ops if a new run already replaced the watched worker, so a
stale watchdog never marks a fresh run stopped or drops its handle.
Supersession is checked on both the watched proc and job id: start_training sets
@ -1134,7 +1211,18 @@ class TrainingBackend:
return # a new run is already starting up; leave its state alone
run_id = self.current_job_id # == watched_job_id
self._progress.is_training = False
self._progress.status_message = "Training stopped."
terminal_payload = self._terminal_finalize_kwargs()
status = terminal_payload["status"]
error_message = terminal_payload.get("error_message")
output_dir = terminal_payload["output_dir"]
clear_output_dir = terminal_payload["clear_output_dir"]
resume_blocked = bool(terminal_payload.get("resume_blocked"))
with self._lock:
if self.current_job_id != run_id:
return
self._progress.status_message = error_message or "Training stopped."
if error_message:
self._progress.error = error_message
# Create the row if a start-time create failed (no-op otherwise; skips when the pump
# is mid-create, in which case its create-then-finalize records the run instead).
self._ensure_db_run_created()
@ -1148,7 +1236,8 @@ class TrainingBackend:
batch: list = []
final_step = final_loss = duration = None
loss_history: list = []
output_dir = self._output_dir
if clear_output_dir:
self._output_dir = self._progress.output_dir = None
if claim:
self._run_finalized = True # claim this run's finalize
batch = list(self._metric_buffer)
@ -1161,7 +1250,17 @@ class TrainingBackend:
loss_history = list(self.loss_history)
if claim:
self._finish_stopped_run(
run_id, output_dir, batch, final_step, final_loss, duration, loss_history
run_id,
output_dir,
batch,
final_step,
final_loss,
duration,
loss_history,
status = status,
error_message = error_message,
clear_output_dir = clear_output_dir,
resume_blocked = resume_blocked,
)
with self._lock:
if target_proc is None or self._proc is target_proc:
@ -1176,6 +1275,10 @@ class TrainingBackend:
final_loss: Optional[float],
duration: Optional[float],
loss_history: list,
status: str = "stopped",
error_message: Optional[str] = None,
clear_output_dir: bool = False,
resume_blocked: bool = False,
) -> None:
"""Record a force-stopped run finished by its captured id, from state snapshotted
under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE,
@ -1194,14 +1297,16 @@ class TrainingBackend:
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
status = "stopped",
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = None,
error_message = error_message,
clear_output_dir = clear_output_dir,
resume_blocked = resume_blocked,
)
return
except Exception:
@ -1231,7 +1336,7 @@ class TrainingBackend:
logger.info("Force-terminating training subprocess (pid=%s)", proc.pid)
proc.terminate()
cancelled = self._cancel_requested
output_dir = self._output_dir
output_dir = self._cancel_cleanup_output_dir or self._output_dir
if proc is not None:
proc.join(timeout = 5.0)
@ -1304,7 +1409,11 @@ class TrainingBackend:
self._last_full_config = config
logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall")
from .worker import run_training_process
cache_env = getattr(self, "_last_hf_cache_env", None)
if not cache_env:
from utils.hf_cache_settings import get_hf_cache_paths
cache_env = get_hf_cache_paths().child_env({})
from utils.hf_cache_settings import child_environment_for_spawn
# This run is active, so an install request 409s rather than proceeds: a reservation seen here
# is transient (an aborting install or short lazy repair). Wait it out instead of stranding the
@ -1336,12 +1445,15 @@ class TrainingBackend:
# crashed respawn cannot wedge is_training_active until restart.
try:
try:
with native_path_secret_removed_for_child_start():
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
args = ("core.training.worker", "run_training_process", cache_env),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
@ -1595,17 +1707,60 @@ class TrainingBackend:
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
terminal_payload = self._terminal_finalize_kwargs()
with self._lock:
if terminal_payload["clear_output_dir"]:
self._output_dir = self._progress.output_dir = None
if terminal_payload.get("error_message"):
self._progress.error = terminal_payload["error_message"]
self._progress.status_message = terminal_payload["error_message"]
self._finalize_run_in_db(**terminal_payload)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
self._pump_running = False
return
def _has_current_resume_checkpoint(self, output_dir, step) -> bool:
# A valid checkpoint at the current step means the stop-and-save landed on
# disk even if the worker died before confirming it.
if not output_dir or not isinstance(step, int) or step <= 0:
return False
from core.training.resume import get_resume_checkpoint_path
return get_resume_checkpoint_path(output_dir, expected_step = step) is not None
def _terminal_finalize_kwargs(self) -> dict:
with self._lock:
job_id = self.current_job_id
payload = self._terminal_finalize_payload
if payload and payload.get("expected_job_id") == job_id:
return dict(payload)
cancel, stopped = self._cancel_requested, self._should_stop
output_dir = None if cancel else self._output_dir
step = self._progress.step
existing_error = self._progress.error
status, error, blocked = (
("stopped", None, cancel)
if stopped
else (
"error",
existing_error or "Training process terminated unexpectedly",
False,
)
)
# Block only when no valid current-step checkpoint actually landed.
if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step):
status = "error"
error = "Stop and Save ended before a valid current-step checkpoint was written."
blocked = True
return {
"status": status,
"error_message": error,
"output_dir": output_dir,
"clear_output_dir": cancel,
"resume_blocked": blocked,
"expected_job_id": job_id,
}
def _handle_event(self, event: dict) -> None:
"""Apply a subprocess event to local state.
@ -1764,6 +1919,15 @@ class TrainingBackend:
elif etype == "eval_configured":
self.eval_enabled = True
elif etype == "output_dir":
event_output_dir = event.get("output_dir")
if self._cancel_requested:
self._cancel_cleanup_output_dir = event_output_dir
self._output_dir = self._progress.output_dir = None
else:
self._output_dir = event_output_dir
db_action = "persist_output_dir"
elif etype == "status":
self._progress.status_message = event.get("message", "")
self._progress.is_training = True
@ -1778,7 +1942,12 @@ class TrainingBackend:
self._complete_seen.set()
self._progress.is_training = False
self._progress.is_completed = not stopped
self._output_dir = event.get("output_dir")
event_output_dir = event.get("output_dir")
if self._cancel_requested:
self._cancel_cleanup_output_dir = event_output_dir
self._output_dir = None
else:
self._output_dir = event_output_dir
self._progress.output_dir = self._output_dir
self._progress.status_message = msg
if not self._db_run_created and self.current_job_id and self._db_config:
@ -1788,11 +1957,16 @@ class TrainingBackend:
db_action_kwargs = {
"status": "stopped" if stopped else "completed",
"output_dir": self._output_dir,
"clear_output_dir": self._cancel_requested,
"expected_job_id": self.current_job_id,
}
self._terminal_finalize_payload = dict(db_action_kwargs)
elif etype == "error":
self._progress.is_training = False
self._progress.error = event.get("error", "Unknown error")
if self._cancel_requested:
self._output_dir = self._progress.output_dir = None
logger.error("Training error: %s", event.get("error"))
stack = event.get("stack", "")
if stack:
@ -1801,29 +1975,36 @@ class TrainingBackend:
db_action = "create_and_finalize"
else:
db_action = "finalize"
stop_save_failed = (
self._should_stop
and not self._cancel_requested
and not self._has_current_resume_checkpoint(
self._output_dir, self._progress.step
)
)
db_action_kwargs = {
"status": "stopped" if self._should_stop else "error",
"status": "stopped"
if self._should_stop
and not stop_save_failed
and not event.get("keep_error_status")
else "error",
"error_message": event.get("error", "Unknown error"),
"output_dir": self._output_dir,
"clear_output_dir": self._cancel_requested,
"resume_blocked": stop_save_failed or bool(event.get("resume_blocked")),
"expected_job_id": self.current_job_id,
}
self._terminal_finalize_payload = dict(db_action_kwargs)
# --- DB I/O outside the lock ---
if db_action == "create_run":
try:
from storage.studio_db import create_run
create_run(
id = db_action_kwargs["job_id"],
model_name = db_action_kwargs["model_name"],
dataset_name = db_action_kwargs["dataset_name"],
config_json = db_action_kwargs["config_json"],
started_at = db_action_kwargs["started_at"],
total_steps = db_action_kwargs["total_steps"],
)
self._db_run_created = True
self._ensure_db_run_created()
if self._db_run_created:
if db_action_kwargs["total_steps"]:
self._db_total_steps_set = True
except Exception:
logger.warning("Failed to create DB run record", exc_info = True)
self._persist_output_dir()
elif db_action == "persist_output_dir":
self._persist_output_dir()
elif db_action == "create_and_finalize":
self._ensure_db_run_created()
self._finalize_run_in_db(**db_action_kwargs)
@ -1842,6 +2023,22 @@ class TrainingBackend:
if etype == "progress":
self._log_training_progress()
def _persist_output_dir(self) -> None:
with self._lock:
if (
not self._output_dir
or not self.current_job_id
or not self._db_run_created
or self._cancel_requested
):
return
run_id, output_dir = self.current_job_id, self._output_dir
try:
from storage.studio_db import update_run_output_dir
update_run_output_dir(run_id, output_dir)
except Exception:
logger.warning("Failed to persist output_dir", exc_info = True)
def _log_training_progress(self) -> None:
"""One throttled training-status line to the server log (the per-step stream
still goes to the UI via SSE): first step, then at most every 30s, plus the
@ -1875,6 +2072,7 @@ class TrainingBackend:
caller create at a time, and ``_db_run_created`` is published only after
``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a
not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running)."""
self._run_intent_lock.acquire()
with self._lock:
if (
self._db_run_created
@ -1882,6 +2080,7 @@ class TrainingBackend:
or not self.current_job_id
or not self._db_config
):
self._run_intent_lock.release()
return
self._db_create_in_progress = True # only one caller creates
job_id = self.current_job_id
@ -1898,6 +2097,12 @@ class TrainingBackend:
or _s3_dataset_name(db_config.get("s3_dataset"))
or "unknown"
)
with self._lock:
if self.current_job_id != job_id:
return
output_dir = self._output_dir
cancel_requested = self._cancel_requested
resumed_from_run_id = self._resume_source_run_id
create_run(
id = job_id,
model_name = db_config["model_name"],
@ -1905,6 +2110,9 @@ class TrainingBackend:
config_json = _json.dumps(db_config),
started_at = started_at,
total_steps = total_steps,
output_dir = output_dir,
cancel_requested = cancel_requested,
resumed_from_run_id = resumed_from_run_id,
)
created = True
except Exception:
@ -1919,12 +2127,15 @@ class TrainingBackend:
if created:
self._db_run_created = True # publish only after the insert commits
self._db_create_in_progress = False
self._run_intent_lock.release()
def _finalize_run_in_db(
self,
status: str,
error_message: Optional[str] = None,
output_dir: Optional[str] = None,
clear_output_dir: bool = False,
resume_blocked: bool = False,
expected_job_id: Optional[str] = None,
) -> None:
"""Flush remaining metrics and mark a run finished in the DB. Claims the finalize
@ -1947,26 +2158,33 @@ class TrainingBackend:
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
self._flush_metrics_to_db(run_id = run_id)
try:
from storage.studio_db import finish_run
from utils.downsample import downsample
for attempt in range(_DB_FINALIZE_RETRIES):
try:
from storage.studio_db import finish_run
from utils.downsample import downsample
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = error_message,
)
except Exception:
with self._lock:
self._run_finalized = False # unclaim so a later flush can retry
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
finish_run(
id = run_id,
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(downsample(loss_history, 50)),
output_dir = output_dir,
error_message = error_message,
clear_output_dir = clear_output_dir,
resume_blocked = resume_blocked,
)
return
except Exception:
if attempt + 1 < _DB_FINALIZE_RETRIES:
time.sleep(_DB_FINALIZE_RETRY_S)
continue
with self._lock:
if self.current_job_id == run_id:
self._run_finalized = False
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None:
"""Flush buffered metrics to the DB and update live progress. The target run id,

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

@ -90,6 +90,79 @@ _FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
# run_training_process() and isn't GC'd mid-run.
_WINDOWS_ROCM_GROUPED_MM_LIB = None
def _install_grouped_mm_cpu_fallback(torch_mod, logger, label):
"""Register a Python mm/bmm fallback for torch._grouped_mm and return the Library.
RDNA4 (gfx1200/gfx1201) ships a null HIP _grouped_mm kernel on ROCm <= 7.12
(fixed in 7.13; ROCm/TheRock #5284). JitDecomp dispatches _grouped_mm to the
null kernel and crashes; overriding the CUDA dispatch key bypasses it. Shared
by the Windows and Linux ROCm guards. Keep the returned Library referenced so
the registration outlives the caller.
"""
import warnings as _warnings
_gm_lib = torch_mod.library.Library("aten", "IMPL")
def _grouped_mm_safe_impl(
self,
mat2,
offs = None,
bias = None,
out_dtype = None,
):
"""Python mm/bmm fallback for _grouped_mm on gfx120X (null HIP kernel, ROCm <= 7.12)."""
_t = torch_mod
if offs is None:
# No offsets: 2-D -> mm, 3-D batched -> bmm (unconditional mm broke 3-D MoE).
if self.dim() == 3 and mat2.dim() == 3:
result = _t.bmm(self.contiguous(), mat2.contiguous())
elif self.dim() == 3 and mat2.dim() == 2:
result = _t.matmul(self.contiguous(), mat2.contiguous())
elif self.dim() == 2 and mat2.dim() == 3:
result = _t.matmul(self.contiguous(), mat2.contiguous())
else:
result = _t.mm(self.contiguous(), mat2.contiguous())
else:
# Grouped: offs[i] is the exclusive end-row of group i.
offs_list = offs.tolist()
pieces = []
prev = 0
for idx, end in enumerate(offs_list):
end = int(end)
a_part = self[prev:end].contiguous()
b_part = mat2[idx].contiguous() if mat2.dim() == 3 else mat2.contiguous()
pieces.append(_t.mm(a_part, b_part))
prev = end
# Include trailing rows not covered by offs.
if prev < self.shape[0]:
a_tail = self[prev:].contiguous()
b_tail = mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous()
pieces.append(_t.mm(a_tail, b_tail))
result = (
_t.cat(pieces, dim = 0)
if pieces
else _t.zeros(0, mat2.shape[-1], device = self.device, dtype = self.dtype)
)
if bias is not None:
result = result + bias
if out_dtype is not None:
result = result.to(out_dtype)
elif result.dtype != self.dtype:
result = result.to(self.dtype)
return result
with _warnings.catch_warnings():
_warnings.simplefilter("ignore")
_gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
logger.info(
"%s: patched _grouped_mm CUDA dispatch (null HIP kernel on gfx120X, "
"ROCm <= 7.12 -- bypassed with Python mm fallback)",
label,
)
return _gm_lib
# Subprocesses don't inherit os.add_dll_directory registrations. Replicate
# main.py's Windows ROCm DLL setup so the first `import torch` finds
# amdhip64.dll. Handles retained at module scope so they aren't GC'd.
@ -702,8 +775,9 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
3. Device-name substring match (last resort when all arch attrs absent;
AMD SDK / Radeon wheels may not populate them):
- gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
- gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
``Radeon 8050S`` (cut-down SKU)
- gfx1151 Strix Halo / Gorgon Halo: ``Radeon 8065S`` (Ryzen AI
Max+ 495), ``Radeon 8060S`` (Ryzen AI MAX+
395), ``Radeon 8050S`` (cut-down SKU)
"""
gcn_arch = ""
for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"):
@ -728,7 +802,11 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
# Arch attrs absent — fall back to device-name matching.
dev_lower = (getattr(props, "name", "") or "").lower()
is_unified = (
"890m" in dev_lower or "880m" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower
"890m" in dev_lower
or "880m" in dev_lower
or "8065s" in dev_lower
or "8060s" in dev_lower
or "8050s" in dev_lower
)
return gcn_arch, is_unified
@ -1840,8 +1918,15 @@ def _run_mlx_training(event_queue, stop_queue, config):
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import ensure_dir
output_dir = _resolve_mlx_output_dir(config, model_name)
# Resume must land in the original run dir even when config lacks output_dir.
resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint(
resume_from_checkpoint
)
output_dir = _resolve_mlx_output_dir(
{**config, "output_dir": resume_dir} if resume_dir else config, model_name
)
ensure_dir(Path(output_dir))
_emit_output_dir(event_queue, output_dir)
# ── 6. Create trainer ──
eval_steps_val = config.get("eval_steps", 0) or 0
@ -2067,6 +2152,17 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.add_eval_callback(_on_eval)
_opt_ref = [None]
_orig_build_optimizer = getattr(trainer, "_build_optimizer", None)
if callable(_orig_build_optimizer):
def _capture_optimizer(total_steps):
_opt_ref[0] = _orig_build_optimizer(total_steps)
return _opt_ref[0]
trainer._build_optimizer = _capture_optimizer
# ── 11. Run training ──
gc.collect()
mx.synchronize()
@ -2082,31 +2178,58 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.save_model = _save_model
# ── 12. Save and finalize ──
if trainer.stop_requested:
if not _stop_save[0]:
# Cancel (save=False): skip saving.
_send("complete", output_dir = None, status_message = "Training cancelled")
def _finish_tracking() -> None:
# Runs on every save/finalize exit so TB/W&B never leak on early return.
if tb_writer is not None:
try:
tb_writer.close()
except Exception:
pass
if wandb_run is not None:
try:
wandb_run.finish()
except Exception:
pass
def _stop_checkpoint_ok() -> bool:
if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir):
return True
_send(
"error",
error = (
"Failed to save a resumable checkpoint after stop. "
"Model files were saved, but this run cannot be resumed."
),
# A user stop finalizes as 'stopped'; keep this failure's error status so history explains it.
keep_error_status = True,
# Older checkpoints are stale; resuming would roll back past this stop.
resume_blocked = True,
)
return False
try:
if trainer.stop_requested:
if not _stop_save[0]:
# Cancel (save=False): skip saving.
_send("complete", output_dir = None, status_message = "Training cancelled")
else:
_send("status", status_message = "Saving stopped model...")
mx.synchronize()
trainer.save_model(output_dir)
# Stop-and-save promises a resumable checkpoint, not just model files.
if not _stop_checkpoint_ok():
return
_send("complete", output_dir = output_dir, status_message = "Training stopped")
else:
_send("status", status_message = "Saving stopped model...")
_send("status", status_message = "Saving model...")
mx.synchronize()
trainer.save_model(output_dir)
_send("complete", output_dir = output_dir, status_message = "Training stopped")
else:
_send("status", status_message = "Saving model...")
mx.synchronize()
trainer.save_model(output_dir)
_send("complete", output_dir = output_dir, status_message = "Training completed")
if tb_writer is not None:
try:
tb_writer.close()
except Exception:
pass
if wandb_run is not None:
try:
wandb_run.finish()
except Exception:
pass
# A save-stop can race the natural final save; it made the same promise.
if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok():
return
_send("complete", output_dir = output_dir, status_message = "Training completed")
finally:
_finish_tracking()
def _is_current_process_apple_silicon() -> bool:
@ -2644,80 +2767,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# so 7.13+ uses the real GPU kernel.
if not _hip_ver_at_least(7, 13):
try:
import warnings as _warnings
_gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
def _grouped_mm_safe_impl(
self,
mat2,
offs = None,
bias = None,
out_dtype = None,
):
"""Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
_t = _torch_for_rocm
if offs is None:
# No offsets: 2-D -> mm, 3-D batched -> bmm
# (unconditional mm broke 3-D MoE).
if self.dim() == 3 and mat2.dim() == 3:
result = _t.bmm(self.contiguous(), mat2.contiguous())
elif self.dim() == 3 and mat2.dim() == 2:
# Broadcast 2-D mat2 across the batch dim.
result = _t.matmul(self.contiguous(), mat2.contiguous())
elif self.dim() == 2 and mat2.dim() == 3:
# Broadcast 2-D self across batch via matmul.
result = _t.matmul(self.contiguous(), mat2.contiguous())
else:
result = _t.mm(self.contiguous(), mat2.contiguous())
else:
# Grouped: offs[i] is the exclusive end-row of group i.
offs_list = offs.tolist()
pieces = []
prev = 0
for idx, end in enumerate(offs_list):
end = int(end)
a_part = self[prev:end].contiguous()
if mat2.dim() == 3:
b_part = mat2[idx].contiguous()
else:
b_part = mat2.contiguous()
pieces.append(_t.mm(a_part, b_part))
prev = end
# Include trailing rows not covered by offs.
if prev < self.shape[0]:
a_tail = self[prev:].contiguous()
b_tail = (
mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous()
)
pieces.append(_t.mm(a_tail, b_tail))
result = (
_t.cat(pieces, dim = 0)
if pieces
else _t.zeros(
0,
mat2.shape[-1],
device = self.device,
dtype = self.dtype,
)
)
if bias is not None:
result = result + bias
if out_dtype is not None:
result = result.to(out_dtype)
elif result.dtype != self.dtype:
result = result.to(self.dtype)
return result
with _warnings.catch_warnings():
_warnings.simplefilter("ignore")
_gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
_WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
logger.info(
"Windows ROCm: patched _grouped_mm CUDA dispatch "
"(null HIP kernel on gfx1200, ROCm ≤ 7.12 — "
"bypassed with Python mm fallback)"
_WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(
_torch_for_rocm, logger, "Windows ROCm"
)
except Exception as _patch_exc:
logger.warning(
@ -2731,6 +2782,44 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
"skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
)
# ── 1f-linux. Linux ROCm RDNA4 _grouped_mm null kernel ──
# The win32 guard above misses Linux: RDNA4 (gfx1200/gfx1201) hits the same null
# HIP _grouped_mm kernel at ROCm <= 7.12 (fixed 7.13, ROCm/TheRock #5284). Gate on
# arch + HIP < 7.13 so NVIDIA/CUDA and non-RDNA4 AMD are untouched; no-op if fixed.
if sys.platform.startswith("linux") and _hw.IS_ROCM:
try:
_torch_lin = sys.modules.get("torch")
if _torch_lin is not None and _torch_lin.cuda.is_available():
# Prefer torch.version.hip, else rocmX.Y from torch.__version__ (AMD
# SDK / Radeon wheels leave version.hip unset). Unknown version on a
# gfx120X build -> assume affected unless it is a post-fix rocmsdk wheel.
_hip_str = str(getattr(getattr(_torch_lin, "version", None), "hip", "") or "")
_ver = getattr(_torch_lin, "__version__", "").lower()
_m = re.match(r"(\d+)\.(\d+)", _hip_str) or re.search(r"rocm(\d+)\.(\d+)", _ver)
if _m:
_hip_lt_713 = (int(_m.group(1)), int(_m.group(2))) < (7, 13)
else:
_hip_lt_713 = "rocmsdk" not in _ver
# Scan every visible GPU (device_map="balanced" can place layers on a
# later RDNA4 card, so device 0 is not enough). Match gfx120X by arch,
# or by RX 9000 / R9700 name when the wheel omits gcnArchName.
_rdna4 = False
for _i in range(_torch_lin.cuda.device_count()):
_props = _torch_lin.cuda.get_device_properties(_i)
_lin_arch, _ = _rocm_classify_unified_memory(_props)
_lin_name = (getattr(_props, "name", "") or "").lower()
if _lin_arch.lower() in ("gfx1200", "gfx1201") or (
not _lin_arch and re.search(r"rx\s*90[0-9]0|r9700", _lin_name)
):
_rdna4 = True
break
if _rdna4 and _hip_lt_713:
_WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(
_torch_lin, logger, "Linux ROCm gfx120X"
)
except Exception as _gm_lin_exc:
logger.warning("Linux ROCm gfx120X: could not patch _grouped_mm: %s", _gm_lin_exc)
# ── 1g. ROCm OOM guard ──
# On ROCm, exhausting VRAM can hang the HIP driver instead of raising.
# set_per_process_memory_fraction caps the allocator so PyTorch raises
@ -3177,6 +3266,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
_emit_output_dir(event_queue, output_dir)
tensorboard_dir = config.get("tensorboard_dir")
if config.get("enable_tensorboard", False):
@ -3296,6 +3386,61 @@ def _send_status(event_queue: Any, message: str) -> None:
)
def _emit_output_dir(event_queue: Any, output_dir: str) -> None:
try:
event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()})
except Exception:
pass
def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool:
if step <= 0:
return False
from core.training.resume import is_resume_checkpoint_valid
return is_resume_checkpoint_valid(
Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx"
)
def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool:
"""Write a full resume checkpoint for a stopped MLX run.
Returns True when a checkpoint for the current training step exists.
"""
step = int(getattr(trainer, "_global_step", 0) or 0)
# A periodic save or a resumed run may already cover the current step.
if _mlx_has_checkpoint_at_step(output_dir, step):
return True
if step <= 0 or optimizer is None:
return False
ckpt_dir = Path(output_dir) / f"checkpoint-{step}"
if ckpt_dir.is_symlink():
# Refuse a symlinked dir: it could redirect writes outside output_dir.
logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir)
return False
try:
ckpt_dir.mkdir(parents = True, exist_ok = True)
from unsloth_zoo.mlx.utils import (
save_optimizer_state,
save_trainable_adapters,
save_trainer_state,
)
save_trainable_adapters(trainer.model, str(ckpt_dir))
save_optimizer_state(optimizer, str(ckpt_dir))
save_trainer_state(
{
"global_step": step,
"train_loss_history": list(getattr(trainer, "_train_loss_history", [])),
},
str(ckpt_dir),
)
logger.info("Saved stop checkpoint to %s", ckpt_dir)
except Exception:
logger.exception("Failed to write stop checkpoint under %s", output_dir)
return _mlx_has_checkpoint_at_step(output_dir, step)
def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None:
"""Self-contained embedding model training pipeline.
@ -3660,6 +3805,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
_emit_output_dir(event_queue, output_dir)
num_epochs = config.get("num_epochs", 2)
batch_size = config.get("batch_size", 256)

6
studio/backend/hub/routes/datasets.py
View file

@ -61,9 +61,11 @@ async def list_cached_datasets(current_subject: str = Depends(get_current_subjec
@router.delete("/cached", response_model = DeleteCachedDatasetResponse)
async def delete_cached_dataset(
repo_id: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject)
repo_id: str = Body(..., embed = True),
cache_path: Optional[str] = Body(None, embed = True),
current_subject: str = Depends(get_current_subject),
):
return await cache_inventory.delete_cached_dataset_response(repo_id)
return await cache_inventory.delete_cached_dataset_response(repo_id, cache_path)
@router.get("/download-progress", response_model = DownloadProgressResponse)

14
studio/backend/hub/routes/inventory.py
View file

@ -28,6 +28,7 @@ from hub.schemas.inventory import (
CachedModelsResponse,
DeleteCachedModelResponse,
GgufVariantsResponse,
HiddenModelsResponse,
LocalModelListResponse,
ModelsFolderResponse,
RecommendedFoldersResponse,
@ -214,6 +215,16 @@ async def list_cached_models(
return await cache_inventory.list_cached_models_response(hf_token)
@router.get("/hidden-models", response_model = HiddenModelsResponse)
async def list_hidden_models(current_subject: str = Depends(get_current_subject)):
import asyncio
from routes.models import hidden_model_matchers
needles, exact_ids, exact_paths = await asyncio.to_thread(hidden_model_matchers)
return HiddenModelsResponse(needles = needles, exact_ids = exact_ids, exact_paths = exact_paths)
@router.delete(
"/delete-cached",
response_model = DeleteCachedModelResponse,
@ -222,7 +233,8 @@ async def list_cached_models(
async def delete_cached_model(
repo_id: str = Body(...),
variant: Optional[str] = Body(None),
cache_path: Optional[str] = Body(None),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await deletion.delete_cached_model_response(repo_id, variant, hf_token)
return await deletion.delete_cached_model_response(repo_id, variant, hf_token, cache_path)

11
studio/backend/hub/schemas/inventory.py
View file

@ -99,6 +99,10 @@ class LocalModelInfo(BaseModel):
None,
description = "HF repo id for cached models, e.g. org/model",
)
active_cache: Optional[bool] = Field(
None,
description = "Whether this HF entry belongs to the current download cache.",
)
base_model: Optional[str] = Field(
None,
description = "Base model from adapter_config.json when this is an adapter",
@ -160,6 +164,7 @@ class CachedRepoBase(BaseModel):
repo_id: str
size_bytes: int = 0
cache_path: Optional[str] = None
last_modified: Optional[float] = None
partial: bool = False
partial_transport: Optional[str] = None
inventory_id: Optional[str] = None
@ -189,6 +194,12 @@ class CachedModelsResponse(BaseModel):
cached: List[CachedModelRepo] = Field(default_factory = list)
class HiddenModelsResponse(BaseModel):
needles: List[str] = Field(default_factory = list)
exact_ids: List[str] = Field(default_factory = list)
exact_paths: List[str] = Field(default_factory = list)
class AddScanFolderRequest(BaseModel):
"""Request body for adding a custom scan folder."""

123
studio/backend/hub/services/datasets/cache_inventory.py
View file

@ -20,12 +20,11 @@ from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import (
purge_partial_repo,
purge_repo_cache_dirs,
resolve_delete_target_root,
resolve_destructive_case_matches,
)
from hub.utils.paths import (
hf_default_cache_dir,
is_valid_repo_id as _is_valid_repo_id,
legacy_hf_cache_dir,
resolve_cached_repo_id_case,
)
@ -43,38 +42,8 @@ def _collect_hf_cache_scans() -> tuple[list, set[str]]:
def _hf_hub_cache_roots() -> list[Path]:
roots: list[Path] = []
seen: set[str] = set()
def _add(path: Optional[Path]) -> None:
if path is None or not path.is_dir():
return
try:
resolved = str(path.resolve())
except OSError:
return
if resolved in seen:
return
seen.add(resolved)
roots.append(path)
try:
from huggingface_hub.constants import HF_HUB_CACHE
_add(Path(HF_HUB_CACHE))
except Exception:
pass
hf_hub_cache = os.environ.get("HF_HUB_CACHE")
if hf_hub_cache:
_add(Path(hf_hub_cache).expanduser())
hf_home = os.environ.get("HF_HOME")
if hf_home:
_add(Path(hf_home).expanduser() / "hub")
_add(legacy_hf_cache_dir())
_add(hf_default_cache_dir())
return roots
from hub.utils.hf_cache_state import hf_cache_roots
return hf_cache_roots()
def _repo_id_from_hub_dataset_dir(name: str) -> str | None:
@ -207,6 +176,21 @@ def _repo_id_from_datasets_cache_dir(name: str) -> str | None:
return repo_id if _is_valid_repo_id(repo_id) else None
def _is_processed_dataset_cache_path(repo_id: str, cache_path: str) -> bool:
"""True when *cache_path* is this repo's processed Arrow cache dir
(``<owner>___<repo>`` directly under an HF_DATASETS_CACHE root). Such rows
have no Hub ``datasets--`` layout, so they are deleted via the processed
path and must not be rejected as an invalid cache_path."""
try:
resolved = Path(cache_path).expanduser().resolve(strict = False)
except (OSError, RuntimeError, ValueError):
return False
if resolved.name.lower() != repo_id.replace("/", "___").lower():
return False
roots = {r.resolve(strict = False) for r in _hf_datasets_cache_roots()}
return resolved.parent.resolve(strict = False) in roots
def _processed_dataset_cache_size(path: Path) -> int:
total = 0
try:
@ -361,7 +345,7 @@ async def list_cached_datasets_response() -> dict:
) from exc
async def delete_cached_dataset_response(repo_id: str) -> dict:
async def delete_cached_dataset_response(repo_id: str, cache_path: Optional[str] = None) -> dict:
"""Remove a cached dataset repo from the HF cache."""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
@ -373,22 +357,40 @@ async def delete_cached_dataset_response(repo_id: str) -> dict:
detail = "Cancel the active download before deleting.",
)
try:
return await asyncio.to_thread(_delete_cached_dataset_blocking, repo_key)
return await asyncio.to_thread(_delete_cached_dataset_blocking, repo_key, cache_path)
finally:
downloads.registry.end_delete(repo_key)
hf_cache_scan.invalidate_hf_cache_scans()
def _delete_cached_dataset_blocking(repo_id: str) -> dict:
def _delete_cached_dataset_blocking(repo_id: str, cache_path: Optional[str] = None) -> dict:
scans, _seen_roots = _collect_hf_cache_scans()
candidate_entries = []
# Group this dataset's copies by owning cache root, then target exactly one
# cache so a delete never removes copies in other, previously selected caches.
owners: dict = {}
for hf_cache in scans:
for repo_info in hf_cache.repos:
if str(repo_info.repo_type) != "dataset":
continue
if repo_info.repo_id.lower() == repo_id.lower():
candidate_entries.append((hf_cache, repo_info))
if repo_info.repo_id.lower() != repo_id.lower():
continue
try:
owner = Path(repo_info.repo_path).parent.resolve(strict = False)
except (OSError, RuntimeError, ValueError):
continue
owners.setdefault(owner, []).append((hf_cache, repo_info))
target_root = resolve_delete_target_root("dataset", repo_id, cache_path, owners.keys())
# A processed-only dataset row sends its Arrow cache path (<owner>___<repo>
# under HF_DATASETS_CACHE), which is not a Hub datasets-- dir, so
# resolve_delete_target_root returns None. Accept it and fall through to the
# processed-cache delete rather than rejecting a legitimate row.
if target_root is None and not (
cache_path and _is_processed_dataset_cache_path(repo_id, cache_path)
):
raise HTTPException(status_code = 400, detail = "Invalid cache_path")
candidate_entries = owners.get(target_root, []) if target_root is not None else []
matched_repo_ids = resolve_destructive_repo_ids(
repo_id,
[str(repo_info.repo_id) for _hf_cache, repo_info in candidate_entries],
@ -414,7 +416,26 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict:
exc_info = True,
)
processed_deleted, processed_failures = _delete_processed_dataset_cache(repo_id)
# Restrict the processed Arrow-cache delete to the selected cache's datasets
# root so it never removes copies under other cache homes. A processed
# cache_path scopes to its own root; a Hub target scopes to the datasets root
# sharing its cache home; an unspecified cache_path stays global (legacy).
processed_roots: Optional[set[Path]]
if not cache_path:
processed_roots = None
elif _is_processed_dataset_cache_path(repo_id, cache_path):
processed_roots = {Path(cache_path).expanduser().resolve(strict = False).parent}
else:
home = target_root.parent if target_root is not None else None
processed_roots = {
root.resolve(strict = False)
for root in _hf_datasets_cache_roots()
if home is not None and root.resolve(strict = False).parent == home
}
processed_deleted, processed_failures = _delete_processed_dataset_cache(
repo_id, only_roots = processed_roots
)
failures.extend(processed_failures)
if failures:
raise HTTPException(
@ -427,15 +448,23 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict:
# ``scan_cache_dir()`` skips blob-only/corrupt repos the revision delete
# can't touch, yet the fallback scanner shows them; purge the whole dir.
cache_purged = purge_repo_cache_dirs("dataset", repo_id)
partial_purged = purge_partial_repo("dataset", repo_id)
state_purged = download_manifest.purge_all_state_for_repo("dataset", repo_id) > 0
# Only for a Hub cache target; a processed-only path has no Hub dir/state.
cache_purged = partial_purged = state_purged = False
if target_root is not None:
cache_purged = purge_repo_cache_dirs("dataset", repo_id, root = target_root)
partial_purged = purge_partial_repo("dataset", repo_id, root = target_root)
state_purged = (
download_manifest.purge_all_state_for_repo("dataset", repo_id, hub_cache = target_root)
> 0
)
if not (deleted or processed_deleted or cache_purged or partial_purged or state_purged):
raise HTTPException(status_code = 404, detail = "Dataset not found in cache")
return {"status": "deleted", "repo_id": repo_id}
def _delete_processed_dataset_cache(repo_id: str) -> tuple[bool, list[str]]:
def _delete_processed_dataset_cache(
repo_id: str, only_roots: Optional[set[Path]] = None
) -> tuple[bool, list[str]]:
import shutil
target = repo_id.replace("/", "___")
@ -443,6 +472,10 @@ def _delete_processed_dataset_cache(repo_id: str) -> tuple[bool, list[str]]:
deleted = False
failures: list[str] = []
for root in _hf_datasets_cache_roots():
# Scope to the selected cache's datasets root(s): a delete must not remove
# processed copies living under other, previously selected cache homes.
if only_roots is not None and root.resolve(strict = False) not in only_roots:
continue
try:
entries = [
entry

14
studio/backend/hub/services/datasets/downloads.py
View file

@ -159,12 +159,18 @@ async def download_dataset_response(
use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
transport = download_lifecycle.resolve_transport(use_xet)
from utils.hf_cache_settings import get_hf_cache_paths
cache_paths = get_hf_cache_paths()
cache_env = cache_paths.child_env({})
claimed, claim_state = _registry.claim(
key,
transport,
repo_type = "dataset",
repo_id = repo_id,
hub_cache = str(cache_paths.hub_cache),
xet_cache = str(cache_paths.xet_cache),
)
generation = _registry.current_generation(key)
if not claimed:
@ -176,7 +182,12 @@ async def download_dataset_response(
"accepted": _registry.adoptable(key),
"generation": generation,
}
download_manifest.clear_cancel_marker("dataset", repo_id, None)
download_manifest.clear_cancel_marker(
"dataset",
repo_id,
None,
hub_cache = cache_paths.hub_cache,
)
state = download_lifecycle.launch_worker(
_registry,
@ -185,6 +196,7 @@ async def download_dataset_response(
["--repo-id", repo_id, "--dataset"],
hf_token,
use_xet = use_xet,
cache_env = cache_env,
),
hf_token = hf_token,
label = repo_id,

34
studio/backend/hub/services/download_lifecycle.py
View file

@ -11,7 +11,7 @@ import sys
import time
import threading
from pathlib import Path
from typing import Callable, Optional
from typing import Callable, Mapping, Optional
from fastapi import HTTPException
@ -57,6 +57,7 @@ def spawn_worker(
*,
use_xet: bool,
protected_blob_hashes: Optional[frozenset[str]] = None,
cache_env: Optional[Mapping[str, str]] = None,
) -> subprocess.Popen:
"""Spawn the download worker.
@ -68,7 +69,11 @@ def spawn_worker(
"""
cwd = backend_dir()
mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
env = os.environ.copy()
from utils.hf_cache_settings import get_hf_cache_paths
env = get_hf_cache_paths().child_env()
if cache_env is not None:
env.update(cache_env)
if protected_blob_hashes:
env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes))
else:
@ -230,6 +235,7 @@ def finalize_worker_exit(
(stderr_data or b"").decode("utf-8", "replace").strip(),
hf_token = hf_token,
)
metadata = registry.get_job_metadata(key)
state = classify_exit(rc, cancel_requested = cancel_requested)
if state == "complete":
registry.set_job(key, "complete")
@ -252,13 +258,13 @@ def finalize_worker_exit(
repo_type,
repo_id,
download_registry.variant_from_key(key),
hub_cache = metadata.hub_cache if metadata is not None else None,
)
except Exception as exc:
logger.debug(f"clear_cancel_marker failed for {repo_id} (rc=0): {exc}")
elif state == "cancelled":
# Read metadata before the terminal set_job so a concurrent eviction
# can't drop it; the job key is the fallback variant label.
metadata = registry.get_job_metadata(key)
registry.set_job(key, "cancelled")
logger.info(f"{log_prefix} cancelled: {label} (rc={rc})")
download_registry.persist_cancel_marker(
@ -268,6 +274,7 @@ def finalize_worker_exit(
if metadata is not None and metadata.variant
else download_registry.variant_from_key(key),
cancel_marker_transport or transport,
hub_cache = metadata.hub_cache if metadata is not None else None,
logger = logger,
)
else:
@ -303,6 +310,7 @@ def _set_retry_failure_state(
metadata.transport
if metadata is not None and metadata.transport
else fallback_transport,
hub_cache = metadata.hub_cache if metadata is not None else None,
logger = logger,
)
return state
@ -371,6 +379,7 @@ def _try_http_retry(
repo_type,
repo_id,
progress_blob_hashes,
root = Path(original_metadata.hub_cache) if original_metadata.hub_cache else None,
)
if progress_blob_hashes
else 0
@ -403,6 +412,8 @@ def _try_http_retry(
generation = generation,
replace_active = True,
cancel_marker_transport = original_metadata.transport,
hub_cache = original_metadata.hub_cache,
xet_cache = original_metadata.xet_cache,
)
if claimed:
break
@ -446,11 +457,24 @@ def _try_http_retry(
label,
)
try:
cache_env = (
{
"HF_HUB_CACHE": original_metadata.hub_cache,
"HF_XET_CACHE": original_metadata.xet_cache,
}
if original_metadata.hub_cache and original_metadata.xet_cache
else None
)
spawn_kwargs = {
"use_xet": False,
"protected_blob_hashes": peer_hashes or None,
}
if cache_env is not None:
spawn_kwargs["cache_env"] = cache_env
proc = spawn_worker(
args,
hf_token,
use_xet = False,
protected_blob_hashes = peer_hashes or None,
**spawn_kwargs,
)
except Exception as exc:
scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token)

164
studio/backend/hub/services/models/cache_inventory.py
View file

@ -31,9 +31,9 @@ from hub.services.models.common import (
_is_checkpoint_weight_name,
_is_gguf_filename,
_is_main_gguf_filename,
_is_mmproj_filename,
_is_transformers_safetensors_weight_name,
_local_inventory_id,
_prefer_complete_larger,
_runtime_for_format,
)
@ -132,6 +132,39 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
def _blob_mtime(file_obj) -> float:
ts = getattr(file_obj, "blob_last_modified", None)
if isinstance(ts, (int, float)) and ts > 0:
return float(ts)
blob_path = getattr(file_obj, "blob_path", None)
if blob_path:
try:
return float(Path(blob_path).stat().st_mtime)
except OSError:
pass
return 0.0
def _repo_gguf_last_modified(repo_info) -> float:
latest = 0.0
for revision in repo_info.revisions:
for f in revision.files:
if _is_main_gguf_filename(f.file_name):
latest = max(latest, _blob_mtime(f))
return latest
def _repo_has_mmproj(repo_info) -> bool:
# An mmproj file only makes a repo vision-capable when it is an actual GGUF
# projector; a non-GGUF sidecar (e.g. mmproj_config.json) does not, and the
# runtime's projector detection is GGUF-only.
return any(
_is_gguf_filename(f.file_name) and _is_mmproj_filename(f.file_name)
for revision in repo_info.revisions
for f in revision.files
)
def _cached_repo_file_name(file_obj) -> str:
file_path = getattr(file_obj, "file_path", None)
if file_path:
@ -216,24 +249,46 @@ def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[
def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool:
if existing is None:
return True
return _prefer_complete_larger(
bool(candidate.get("partial")),
int(candidate.get("size_bytes") or 0),
bool(existing.get("partial")),
int(existing.get("size_bytes") or 0),
)
candidate_partial = bool(candidate.get("partial"))
existing_partial = bool(existing.get("partial"))
if candidate_partial != existing_partial:
return not candidate_partial
candidate_active = bool(candidate.get("active_cache"))
existing_active = bool(existing.get("active_cache"))
if candidate_active != existing_active:
return candidate_active
return int(candidate.get("size_bytes") or 0) > int(existing.get("size_bytes") or 0)
def _cache_inventory_fields(
repo_id: str,
model_format: ModelFormat,
*,
repo_path: Optional[Path] = None,
snapshot_path: Optional[Path] = None,
active_hub_cache: Optional[Path] = None,
partial: bool = False,
requires_variant: bool = False,
) -> dict:
load_id = repo_id
active_cache = True
if repo_path is not None:
try:
if active_hub_cache is None:
from utils.hf_cache_settings import get_hf_cache_paths
active_hub_cache = get_hf_cache_paths().hub_cache
active_root = active_hub_cache.resolve(strict = False)
cached_root = repo_path.parent.resolve(strict = False)
if cached_root != active_root:
active_cache = False
load_id = str(snapshot_path or repo_path.resolve(strict = False))
except (OSError, RuntimeError, ValueError):
active_cache = False
load_id = str(snapshot_path or repo_path)
return {
"inventory_id": _local_inventory_id("cache", model_format, repo_id),
"load_id": repo_id,
"load_id": load_id,
"active_cache": active_cache,
"model_format": model_format,
"runtime": _runtime_for_format(model_format),
"format_variant": None,
@ -260,6 +315,9 @@ def _is_hidden_infra_repo(*values: str | None) -> bool:
def _scan_cached_gguf() -> list[dict]:
"""Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread."""
cache_scans = all_hf_cache_scans()
from utils.hf_cache_settings import get_hf_cache_paths
active_hub_cache = get_hf_cache_paths().hub_cache
seen_lower: dict[str, dict] = {}
for hf_cache in cache_scans:
@ -271,7 +329,10 @@ def _scan_cached_gguf() -> list[dict]:
repo_path = Path(repo_info.repo_path)
snapshot_path = _cached_model_snapshot_path(repo_path)
total_size = _repo_gguf_size_bytes(repo_info)
has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id)
has_variant_state, variant_state_size = _gguf_variant_state_summary(
repo_id,
hub_cache = repo_path.parent,
)
is_hidden_infra = _is_hidden_infra_repo(
repo_id,
str(repo_path),
@ -291,6 +352,7 @@ def _scan_cached_gguf() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
last_modified = _repo_gguf_last_modified(repo_info)
row = {
"repo_id": repo_id,
"size_bytes": max(total_size, variant_state_size),
@ -300,19 +362,34 @@ def _scan_cached_gguf() -> list[dict]:
# per-variant detail lives on GgufVariantDetail.
"partial_transport": None,
}
last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0))
if last_modified > 0:
row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
"gguf",
repo_path = repo_path,
snapshot_path = snapshot_path,
active_hub_cache = active_hub_cache,
partial = bool(row["partial"]),
requires_variant = True,
)
)
if _repo_has_mmproj(repo_info):
row["capabilities"]["supports_vision"] = True
# Visible infra variants remain management-only.
if is_hidden_infra:
row["capabilities"]["can_chat"] = False
if _prefer_cache_row(row, existing):
if existing and existing["capabilities"].get("supports_vision"):
row["capabilities"]["supports_vision"] = True
seen_lower[key] = row
else:
if last_modified > existing.get("last_modified", 0.0):
existing["last_modified"] = last_modified
if row["capabilities"].get("supports_vision"):
existing["capabilities"]["supports_vision"] = True
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}")
@ -340,13 +417,14 @@ class _CachedNonGgufPayload(NamedTuple):
size_bytes: int
has_runnable_weights: bool
model_format: ModelFormat
last_modified: float
def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
all_weight_blobs: dict[str, int] = {}
adapter_blobs: dict[str, int] = {}
safetensors_blobs: dict[str, int] = {}
checkpoint_blobs: dict[str, int] = {}
all_weight_blobs: dict[str, tuple[int, float]] = {}
adapter_blobs: dict[str, tuple[int, float]] = {}
safetensors_blobs: dict[str, tuple[int, float]] = {}
checkpoint_blobs: dict[str, tuple[int, float]] = {}
has_config = False
has_adapter_config = False
has_adapter_weights = False
@ -354,12 +432,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
has_transformers_safetensors = False
has_checkpoint = False
def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None:
def _record_blob(
target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str
) -> None:
blob_path = getattr(file_obj, "blob_path", None)
size = int(file_obj.size_on_disk or 0)
key = str(blob_path) if blob_path else f"{rev_id}:{file_name}"
target[key] = size
all_weight_blobs[key] = size
value = (size, _blob_mtime(file_obj))
target[key] = value
all_weight_blobs[key] = value
for revision in repo_info.revisions:
rev_id = getattr(revision, "commit_hash", None) or str(id(revision))
@ -403,18 +484,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
or "unknown"
)
if model_format == "adapter":
size_bytes = sum(adapter_blobs.values())
selected_blobs = adapter_blobs
elif model_format == "safetensors":
size_bytes = sum(safetensors_blobs.values())
selected_blobs = safetensors_blobs
elif model_format == "checkpoint":
size_bytes = sum(checkpoint_blobs.values())
selected_blobs = checkpoint_blobs
else:
size_bytes = sum(all_weight_blobs.values())
selected_blobs = all_weight_blobs
return _CachedNonGgufPayload(
size_bytes = size_bytes,
size_bytes = sum(size for size, _mtime in selected_blobs.values()),
has_runnable_weights = model_format != "unknown",
model_format = model_format,
last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0),
)
@ -435,6 +517,19 @@ def _read_json_object(path: Path) -> dict:
return {}
def _is_whisper_model_config(config: object) -> bool:
if not isinstance(config, dict):
return False
model_type = config.get("model_type")
if isinstance(model_type, str) and model_type.strip().lower() == "whisper":
return True
architectures = config.get("architectures")
return isinstance(architectures, list) and any(
isinstance(name, str) and name == "WhisperForConditionalGeneration"
for name in architectures
)
def _read_model_card_frontmatter(path: Path) -> dict:
try:
text = path.read_text(encoding = "utf-8")
@ -465,6 +560,8 @@ def _cached_model_local_metadata(repo_path: Path) -> dict:
result: dict = {}
config = _read_json_object(snapshot / "config.json")
if _is_whisper_model_config(config):
result["_hidden_stt"] = True
quant_method = (
config.get("quantization_config", {}).get("quant_method")
if isinstance(config.get("quantization_config"), dict)
@ -491,11 +588,15 @@ def _cached_model_local_metadata(repo_path: Path) -> dict:
def _scan_cached_models() -> list[dict]:
"""Synchronous HF-cache disk walk for non-GGUF model repos; runs in a worker thread."""
cache_scans = all_hf_cache_scans()
from utils.hf_cache_settings import get_hf_cache_paths
active_hub_cache = get_hf_cache_paths().hub_cache
seen_lower: dict[str, dict] = {}
inspected = 0
skipped_gguf = 0
skipped_no_weights = 0
skipped_stt = 0
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
inspected += 1
@ -523,6 +624,10 @@ def _scan_cached_models() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
local_metadata = _cached_model_local_metadata(repo_path)
if local_metadata.pop("_hidden_stt", False):
skipped_stt += 1
continue
snapshot_partial = hf_cache_scan.is_snapshot_partial(
"model",
repo_id,
@ -542,27 +647,40 @@ def _scan_cached_models() -> list[dict]:
if snapshot_partial
else None
),
**_cached_model_local_metadata(repo_path),
**local_metadata,
}
last_modified = max(
payload.last_modified,
(existing or {}).get("last_modified", 0.0),
)
if last_modified > 0:
row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
payload.model_format,
repo_path = repo_path,
snapshot_path = snapshot_path,
active_hub_cache = active_hub_cache,
partial = bool(row["partial"]),
)
)
if _prefer_cache_row(row, existing):
seen_lower[key] = row
elif last_modified > existing.get("last_modified", 0.0):
existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
continue
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
logger.info(
"Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d returned=%d",
"Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d "
"skipped_stt=%d returned=%d",
inspected,
skipped_gguf,
skipped_no_weights,
skipped_stt,
len(cached),
)
return cached

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

@ -150,7 +150,9 @@ def _prefer_complete_larger(
return candidate_size_bytes > existing_size_bytes
def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]:
def _gguf_variant_state_summary(
repo_id: str, *, hub_cache: Optional[str | Path] = None
) -> tuple[bool, int]:
"""Whether GGUF variant-scoped state exists and its expected size; a cancelled/in-progress variant may have only manifests/markers/`.incomplete` blobs, which inventory needs to avoid a generic fallback row."""
from hub.utils import download_manifest
@ -159,10 +161,16 @@ def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]:
for variant, _path in download_manifest.iter_variant_manifests(
"model",
repo_id,
hub_cache = hub_cache,
):
key = variant.lower()
variant_keys.add(key)
manifest = download_manifest.read_manifest("model", repo_id, variant)
manifest = download_manifest.read_manifest(
"model",
repo_id,
variant,
hub_cache = hub_cache,
)
if manifest is None:
continue
size_by_variant[key] = max(
@ -172,6 +180,7 @@ def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]:
for variant, _path in download_manifest.iter_variant_markers(
"model",
repo_id,
hub_cache = hub_cache,
):
variant_keys.add(variant.lower())
return bool(variant_keys), sum(size_by_variant.values())
@ -432,8 +441,13 @@ def _local_model_info(
base_model_source: Optional[str] = None,
adapter_type: Optional[str] = None,
training_method: Optional[str] = None,
active_cache: Optional[bool] = None,
) -> LocalModelInfo:
load_id = model_id if source == "hf_cache" and model_id else str(load_path)
load_id = (
model_id
if source == "hf_cache" and model_id and active_cache is not False
else str(load_path)
)
semantic_id = model_id or str(load_path)
return LocalModelInfo(
id = load_id,
@ -445,6 +459,7 @@ def _local_model_info(
),
load_id = load_id,
model_id = model_id,
active_cache = active_cache if source == "hf_cache" else None,
display_name = display_name or (scan_path.stem if scan_path.is_file() else scan_path.name),
path = str(load_path),
size_bytes = max(0, int(size_bytes or 0)),
@ -476,6 +491,7 @@ def _classify_local_path(
model_id: Optional[str] = None,
updated_at: Optional[float] = None,
partial: bool = False,
active_cache: Optional[bool] = None,
) -> list[LocalModelInfo]:
load_path = load_path or scan_path
files = (
@ -512,6 +528,7 @@ def _classify_local_path(
requires_variant = scan_path.is_dir(),
format_variant = variant,
size_bytes = gguf_size_bytes,
active_cache = active_cache,
)
)
@ -574,6 +591,7 @@ def _classify_local_path(
),
adapter_type = adapter_type if model_format == "adapter" else None,
training_method = training_method if model_format == "adapter" else None,
active_cache = active_cache,
)
)
elif not rows:
@ -592,6 +610,7 @@ def _classify_local_path(
updated_at = updated_at,
partial = partial or trusted_hf_cache_repo,
size_bytes = size_bytes,
active_cache = active_cache,
)
)

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

@ -19,8 +19,10 @@ from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.gguf import extract_quant_label, extract_quant_token
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
iter_repo_cache_dirs,
purge_partial_repo,
purge_repo_cache_dirs,
resolve_delete_target_root,
)
from hub.utils.paths import (
is_valid_gguf_variant as _is_valid_gguf_variant,
@ -184,6 +186,7 @@ def _delete_gguf_variant_from_repos(
hf_token: Optional[str],
*,
sibling_active: bool = False,
root: Optional[Path] = None,
) -> dict:
failures: list[str] = []
removed_snapshots = 0
@ -265,6 +268,7 @@ def _delete_gguf_variant_from_repos(
hf_token,
extra_hashes = frozenset(completed_hashes),
companions = not sibling_active,
root = root,
)
if incomplete_result.unresolved:
raise HTTPException(
@ -276,7 +280,7 @@ def _delete_gguf_variant_from_repos(
),
)
state_purged = download_manifest.purge_state("model", repo_id, variant)
state_purged = download_manifest.purge_state("model", repo_id, variant, hub_cache = root)
# Reclaim the empty quant folder so it stops 404ing on delete.
removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant)
removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos)
@ -316,6 +320,8 @@ def reclaim_replaced_gguf_variant(
variant: str,
keep_main_hashes: frozenset[str],
hf_token: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> dict:
"""Prune stale main-GGUF files for a variant after a replacement verified.
@ -366,12 +372,22 @@ def reclaim_replaced_gguf_variant(
"reason": "scan_failed",
}
if hub_cache is None:
from utils.hf_cache_settings import get_hf_cache_paths
hub_cache = get_hf_cache_paths().hub_cache
try:
target_hub_cache = Path(hub_cache).expanduser().resolve(strict = False)
except (OSError, RuntimeError, ValueError):
target_hub_cache = Path(hub_cache).expanduser()
candidate_repos = [
repo_info
for hf_cache in cache_scans
for repo_info in hf_cache.repos
if str(getattr(repo_info, "repo_type", "")) == "model"
and str(getattr(repo_info, "repo_id", "")).lower() == repo_id.lower()
and getattr(repo_info, "repo_path", None)
and Path(repo_info.repo_path).parent.resolve(strict = False) == target_hub_cache
]
try:
matched_repo_ids = resolve_destructive_repo_ids(
@ -493,10 +509,24 @@ def reclaim_replaced_gguf_variant(
def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool:
"""True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``."""
"""Match a loaded repo ID or an on-disk path inside any copy of the repo."""
rid = repo_id.lower()
lid = loaded_id.lower()
return lid == rid or lid.startswith(f"{rid}/")
if lid == rid or lid.startswith(f"{rid}/"):
return True
try:
loaded_path = Path(loaded_id).expanduser().resolve(strict = False)
except (OSError, RuntimeError, ValueError):
return False
for repo_dir in iter_repo_cache_dirs("model", repo_id):
try:
resolved_repo = repo_dir.resolve(strict = False)
if loaded_path == resolved_repo or loaded_path.is_relative_to(resolved_repo):
return True
except (OSError, RuntimeError, ValueError):
continue
return False
def _loaded_repo_variant_blocks_delete(
@ -560,6 +590,7 @@ async def delete_cached_model_response(
repo_id: str,
variant: Optional[str] = None,
hf_token: Optional[str] = None,
cache_path: Optional[str] = None,
):
"""Delete a cached model repo (or a specific GGUF variant) from the HF cache.
@ -603,14 +634,19 @@ async def delete_cached_model_response(
)
raise HTTPException(status_code = 400, detail = detail)
try:
return await asyncio.to_thread(_delete_cached_model_blocking, repo_id, variant, hf_token)
return await asyncio.to_thread(
_delete_cached_model_blocking, repo_id, variant, hf_token, cache_path
)
finally:
downloads.registry.end_delete(repo_key, variant)
cache_inventory.invalidate_hf_cache_scans()
def _delete_cached_model_blocking(
repo_id: str, variant: Optional[str], hf_token: Optional[str]
repo_id: str,
variant: Optional[str],
hf_token: Optional[str],
cache_path: Optional[str] = None,
) -> dict:
try:
# If a sibling quant is downloading concurrently, restrict this delete to
@ -621,13 +657,26 @@ def _delete_cached_model_blocking(
cache_scans = cache_inventory.all_hf_cache_scans()
candidate_entries = []
# A repo can live in several remembered caches. Group its copies by the
# cache root that owns each, then target exactly one cache so a delete
# never removes copies in other, previously selected caches.
owners: dict = {}
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
if str(repo_info.repo_type) != "model":
continue
if repo_info.repo_id.lower() == repo_id.lower():
candidate_entries.append((hf_cache, repo_info))
if repo_info.repo_id.lower() != repo_id.lower():
continue
try:
owner = Path(repo_info.repo_path).parent.resolve(strict = False)
except (OSError, RuntimeError, ValueError):
continue
owners.setdefault(owner, []).append((hf_cache, repo_info))
target_root = resolve_delete_target_root("model", repo_id, cache_path, owners.keys())
if target_root is None:
raise HTTPException(status_code = 400, detail = "Invalid cache_path")
candidate_entries = owners.get(target_root, [])
matched_repo_ids = resolve_destructive_repo_ids(
repo_id,
@ -642,10 +691,15 @@ def _delete_cached_model_blocking(
if not target_entries:
if variant is None:
cache_purged = purge_repo_cache_dirs("model", repo_id) or purge_partial_repo(
"model", repo_id
cache_purged = purge_repo_cache_dirs(
"model", repo_id, root = target_root
) or purge_partial_repo("model", repo_id, root = target_root)
state_purged = (
download_manifest.purge_all_state_for_repo(
"model", repo_id, hub_cache = target_root
)
> 0
)
state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0
if cache_purged or state_purged:
return {"status": "deleted", "repo_id": repo_id}
if variant:
@ -654,6 +708,7 @@ def _delete_cached_model_blocking(
variant,
hf_token,
companions = not sibling_active,
root = target_root,
)
if incomplete_result.unresolved:
raise HTTPException(
@ -668,6 +723,7 @@ def _delete_cached_model_blocking(
"model",
repo_id,
variant,
hub_cache = target_root,
)
if incomplete_result.deleted > 0 or state_purged:
return {
@ -684,6 +740,7 @@ def _delete_cached_model_blocking(
[repo for _cache, repo in target_entries],
hf_token,
sibling_active = sibling_active,
root = target_root,
)
deleted_revisions = False
@ -702,9 +759,11 @@ def _delete_cached_model_blocking(
delete_strategy.execute()
deleted_revisions = True
cache_purged = purge_repo_cache_dirs("model", repo_id)
partial_purged = purge_partial_repo("model", repo_id)
state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0
cache_purged = purge_repo_cache_dirs("model", repo_id, root = target_root)
partial_purged = purge_partial_repo("model", repo_id, root = target_root)
state_purged = (
download_manifest.purge_all_state_for_repo("model", repo_id, hub_cache = target_root) > 0
)
if not (deleted_revisions or cache_purged or partial_purged or state_purged):
raise HTTPException(status_code = 404, detail = "No revisions found for model")

16
studio/backend/hub/services/models/downloads.py
View file

@ -90,6 +90,7 @@ def _spawn_download_worker(
hf_token: Optional[str],
use_xet: bool = True,
protected_blob_hashes: Optional[frozenset[str]] = None,
cache_env: Optional[dict[str, str]] = None,
) -> subprocess.Popen:
args = ["--repo-id", repo_id]
if variant:
@ -99,6 +100,7 @@ def _spawn_download_worker(
hf_token,
use_xet = use_xet,
protected_blob_hashes = protected_blob_hashes,
cache_env = cache_env,
)
@ -125,6 +127,10 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
key = _download_job_key(repo_id, variant)
use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
transport = download_lifecycle.resolve_transport(use_xet)
from utils.hf_cache_settings import get_hf_cache_paths
cache_paths = get_hf_cache_paths()
cache_env = cache_paths.child_env({})
variant_blob_hashes = frozenset()
variant_progress_blob_hashes = frozenset()
completed_baseline_bytes = 0
@ -175,6 +181,8 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
progress_blob_hashes = variant_progress_blob_hashes,
completed_baseline_bytes = completed_baseline_bytes,
admission_check = lambda: not _load_in_flight(repo_id),
hub_cache = str(cache_paths.hub_cache),
xet_cache = str(cache_paths.xet_cache),
)
generation = _registry.current_generation(key)
if not claimed:
@ -189,7 +197,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
"accepted": _registry.adoptable(key),
"generation": generation,
}
download_manifest.clear_cancel_marker("model", repo_id, variant)
download_manifest.clear_cancel_marker(
"model",
repo_id,
variant,
hub_cache = cache_paths.hub_cache,
)
# Blobs a concurrent same-repo variant is already writing (e.g. a shared
# mmproj). The worker must not purge these during cache preparation.
protected_blob_hashes = _registry.peer_blob_hashes(key) if variant else frozenset()
@ -204,6 +217,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
hf_token,
use_xet = use_xet,
protected_blob_hashes = protected_blob_hashes,
cache_env = cache_env,
),
hf_token = hf_token,
label = label,

11
studio/backend/hub/services/models/folder_browser.py
View file

@ -30,6 +30,7 @@ from hub.utils.paths import (
)
from utils.paths.external_media import (
linux_run_media_mount_roots,
macos_volume_roots,
windows_drive_roots,
)
from hub.services.models.common import _safe_is_dir
@ -187,7 +188,7 @@ def _build_browse_allowlist(
_add(Path.home())
if media_roots is None:
media_roots = linux_run_media_mount_roots()
media_roots = [*linux_run_media_mount_roots(), *macos_volume_roots()]
if drive_roots is None:
drive_roots = windows_drive_roots()
for p in media_roots:
@ -195,6 +196,12 @@ def _build_browse_allowlist(
for p in drive_roots:
_add(p)
_add(_resolve_hf_cache_dir())
try:
from utils.hf_cache_settings import known_hf_cache_homes
for cache_home in known_hf_cache_homes():
_add(cache_home)
except Exception: # noqa: BLE001 -- best-effort
pass
try:
_add(hf_default_cache_dir())
except Exception: # noqa: BLE001 -- best-effort
@ -431,7 +438,7 @@ def browse_folders_response(
# Probe removable-media and Windows drive roots once; the allowlist and
# chips reuse the result so a disconnected mapped drive isn't scanned twice.
media_roots = linux_run_media_mount_roots()
media_roots = [*linux_run_media_mount_roots(), *macos_volume_roots()]
drive_roots = windows_drive_roots()
# Build the allowlist once -- the sandbox check and suggestion chips share
# it so chips are always navigable.

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

@ -9,6 +9,7 @@ import asyncio
import threading
import time
from collections import OrderedDict
from pathlib import Path
from typing import NamedTuple, Optional
from fastapi import HTTPException
@ -22,6 +23,7 @@ from hub.utils.hf_errors import hf_error_status
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
iter_destructive_repo_cache_dirs,
repo_cache_dir_name,
)
from hub.utils.gguf import (
extract_quant_label,
@ -233,8 +235,14 @@ def _manifest_variant_blob_hashes(
variant: str,
*,
include_companions: bool = True,
repo_cache_dir: Optional[Path] = None,
) -> frozenset[str]:
manifest = download_manifest.read_manifest("model", repo_id, variant)
manifest = download_manifest.read_manifest(
"model",
repo_id,
variant,
hub_cache = repo_cache_dir.parent if repo_cache_dir is not None else None,
)
if manifest is None:
return frozenset()
variant_key = variant.lower()
@ -257,6 +265,7 @@ def gguf_variant_blob_hashes(
*,
include_companions: bool = True,
allow_remote: bool = True,
repo_cache_dir: Optional[Path] = None,
) -> frozenset[str]:
key = _variant_blob_hash_cache_key(
repo_id,
@ -271,9 +280,9 @@ def gguf_variant_blob_hashes(
repo_id,
variant,
include_companions = include_companions,
repo_cache_dir = repo_cache_dir,
)
if hashes:
_variant_hash_cache_set(key, hashes)
return hashes
requirement_key = _variant_hash_cache_key(repo_id, variant, hf_token)
requirement = _variant_requirement_cache_get(requirement_key)
@ -287,11 +296,22 @@ def gguf_variant_blob_hashes(
return frozenset()
def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]:
return hf_cache_scan.partial_transport_for("model", repo_id, variant)
def _partial_transport_for_variant(
repo_id: str,
variant: str,
repo_cache_dir: Optional[Path] = None,
) -> Optional[str]:
return hf_cache_scan.partial_transport_for(
"model",
repo_id,
variant,
repo_cache_dir,
)
def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str]]]:
def _local_main_gguf_blobs_by_quant(
repo_id: str, repo_cache_dir: Optional[Path] = None
) -> dict[str, dict[str, set[str]]]:
"""Map quant -> repo-relative expected GGUF filename -> cached blob hashes.
Shared companions are copied into each main-quant bucket so update checks can
@ -313,6 +333,14 @@ def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str
continue
if str(getattr(repo_info, "repo_id", "")).lower() != target_lower:
continue
if repo_cache_dir is not None:
try:
if Path(repo_info.repo_path).resolve(strict = False) != repo_cache_dir.resolve(
strict = False
):
continue
except (AttributeError, OSError, RuntimeError, ValueError):
continue
for path, hashes in cache_inventory._repo_gguf_blob_map(
repo_info,
include_companions = True,
@ -388,6 +416,7 @@ def delete_variant_incomplete_blobs_result(
*,
extra_hashes: frozenset[str] = frozenset(),
companions: bool = True,
root: Optional[Path] = None,
) -> VariantIncompleteDeleteResult:
# With a sibling still downloading, ``companions=False`` keeps a shared mmproj
# from being unlinked out from under it; the repo's last delete reclaims it.
@ -409,8 +438,9 @@ def delete_variant_incomplete_blobs_result(
)
deleted = 0
# Destructive iterator: only the exact-case match (or abort if ambiguous),
# so a case-variant sibling repo's partials are never unlinked.
for entry in iter_destructive_repo_cache_dirs("model", repo_id):
# so a case-variant sibling repo's partials are never unlinked. ``root`` scopes
# the purge to one cache so a delete never touches another cache's partials.
for entry in iter_destructive_repo_cache_dirs("model", repo_id, root = root):
blobs_dir = entry / "blobs"
if not blobs_dir.is_dir():
continue
@ -425,15 +455,37 @@ def delete_variant_incomplete_blobs_result(
return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False)
def _repo_cache_dir_for_request(repo_id: str, local_path: Optional[str]) -> Path:
"""Resolve the one Hub repo cache represented by this variant request."""
expected_name = repo_cache_dir_name("model", repo_id).lower()
if local_path:
try:
local = Path(local_path).expanduser().resolve(strict = False)
for candidate in (local, *local.parents):
if candidate.name.lower() == expected_name:
return candidate
except (OSError, RuntimeError, ValueError):
pass
from utils.hf_cache_settings import get_hf_cache_paths
return get_hf_cache_paths().hub_cache / repo_cache_dir_name("model", repo_id)
def _mark_empty_dir_cleanables(
repo_id: str, response: GgufVariantsResponse
repo_id: str,
response: GgufVariantsResponse,
repo_cache_dir: Optional[Path] = None,
) -> GgufVariantsResponse:
"""Surface empty leftover ``<quant>/`` folders (interrupted downloads) as
partial so the UI can delete them -- on local/offline paths too, not just a
remote listing. A listed quant is flipped to partial; an unlisted one is
appended as a zero-byte cleanable entry."""
try:
empty_labels = list_empty_gguf_variant_dirs(repo_id)
empty_labels = (
list_empty_gguf_variant_dirs(repo_id, root = repo_cache_dir.parent)
if repo_cache_dir is not None
else list_empty_gguf_variant_dirs(repo_id)
)
except Exception as e:
logger.warning(f"Failed to scan empty GGUF variant folders for {repo_id}: {e}")
return response
@ -468,6 +520,11 @@ async def get_gguf_variants_response(
"""
def _compute() -> GgufVariantsResponse:
repo_cache_dir = (
None if is_local_path(repo_id) else _repo_cache_dir_for_request(repo_id, local_path)
)
hub_cache = repo_cache_dir.parent if repo_cache_dir is not None else None
def _local_response(
response_repo_id: str, variants, has_vision: bool
) -> GgufVariantsResponse:
@ -511,6 +568,7 @@ async def get_gguf_variants_response(
partial_transport = _partial_transport_for_variant(
response_repo_id,
v.quant,
repo_cache_dir,
),
)
for v in variants
@ -532,7 +590,7 @@ async def get_gguf_variants_response(
local_only = prefer_local_cache or offline
if local_only:
cached = list_gguf_variants_from_hf_cache(repo_id)
cached = list_gguf_variants_from_hf_cache(repo_id, root = hub_cache)
if cached is not None:
variants, has_vision = cached
return _local_response(repo_id, variants, has_vision)
@ -540,7 +598,7 @@ async def get_gguf_variants_response(
variants, has_vision = list_local_gguf_variants(local_path)
if variants or has_vision:
return _local_response(repo_id, variants, has_vision)
partial = list_partial_gguf_variants_from_state(repo_id)
partial = list_partial_gguf_variants_from_state(repo_id, hub_cache = hub_cache)
if partial is not None:
variants, has_vision = partial
return _partial_local_response(repo_id, variants, has_vision)
@ -560,11 +618,11 @@ async def get_gguf_variants_response(
try:
variants, has_vision, siblings = list_gguf_variants(repo_id, hf_token = hf_token)
except Exception:
cached = list_gguf_variants_from_hf_cache(repo_id)
cached = list_gguf_variants_from_hf_cache(repo_id, root = hub_cache)
if cached is not None:
variants, has_vision = cached
return _local_response(repo_id, variants, has_vision)
partial = list_partial_gguf_variants_from_state(repo_id)
partial = list_partial_gguf_variants_from_state(repo_id, hub_cache = hub_cache)
if partial is not None:
variants, has_vision = partial
return _partial_local_response(repo_id, variants, has_vision)
@ -581,7 +639,7 @@ async def get_gguf_variants_response(
cached_filenames_by_snapshot: list[dict[str, int]] = []
cached_quant_bytes_by_snapshot: list[dict[str, int]] = []
if _is_valid_repo_id(repo_id):
for snap in iter_hf_cache_snapshots(repo_id):
for snap in iter_hf_cache_snapshots(repo_id, root = hub_cache):
try:
gguf_paths = list(_iter_gguf_paths(snap))
except (OSError, RuntimeError, ValueError) as e:
@ -694,11 +752,20 @@ async def get_gguf_variants_response(
partial_quants: set[str] = set()
partial_quant_transports: dict[str, Optional[str]] = {}
try:
incomplete_hashes = download_registry.incomplete_blob_hashes("model", repo_id)
incomplete_hashes = download_registry.incomplete_blob_hashes(
"model",
repo_id,
active_only = True,
root = hub_cache,
)
except Exception as e:
logger.warning(f"Failed to compute partial GGUF variants for {repo_id}: {e}")
incomplete_hashes = set()
scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan("model", repo_id)
scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan(
"model",
repo_id,
repo_cache_dir,
)
# Manifest + marker + main incomplete-blob check: catches variants whose
# download was cancelled or whose expected shards are missing/undersized.
for variant in variants:
@ -711,6 +778,7 @@ async def get_gguf_variants_response(
variant.quant,
hf_token,
include_companions = False,
repo_cache_dir = repo_cache_dir,
)
if hf_cache_scan.is_variant_partial(
repo_id,
@ -718,11 +786,13 @@ async def get_gguf_variants_response(
scan_snapshot_dir,
incomplete_blob_hashes = incomplete_hashes,
variant_blob_hashes = variant_hashes,
repo_cache_dir = repo_cache_dir,
):
partial_quants.add(variant.quant)
partial_quant_transports[variant.quant] = _partial_transport_for_variant(
repo_id,
variant.quant,
repo_cache_dir,
)
except Exception as e:
logger.warning(
@ -744,10 +814,14 @@ async def get_gguf_variants_response(
partial_quants.add(variant.quant)
partial_quant_transports.setdefault(
variant.quant,
_partial_transport_for_variant(repo_id, variant.quant),
_partial_transport_for_variant(
repo_id,
variant.quant,
repo_cache_dir,
),
)
local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id)
local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id, repo_cache_dir)
def _variant_detail(v) -> GgufVariantDetail:
is_partial = v.quant in partial_quants
@ -790,14 +864,20 @@ async def get_gguf_variants_response(
if skip:
raise
enriched = _mark_empty_dir_cleanables(
repo_id, GgufVariantsResponse(repo_id = repo_id, variants = [])
repo_id,
GgufVariantsResponse(repo_id = repo_id, variants = []),
_repo_cache_dir_for_request(repo_id, local_path),
)
if enriched.variants:
return enriched
raise
if skip:
return response
return _mark_empty_dir_cleanables(repo_id, response)
return _mark_empty_dir_cleanables(
repo_id,
response,
_repo_cache_dir_for_request(repo_id, local_path),
)
try:
return await asyncio.to_thread(_compute_with_cleanables)

87
studio/backend/hub/services/models/local_inventory.py
View file

@ -106,11 +106,8 @@ def _is_model_directory_for_scan(path: Path, *, entry_limit: int | None) -> bool
def _resolve_hf_cache_dir() -> Path:
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
from utils.hf_cache_settings import get_hf_cache_paths
return get_hf_cache_paths().hub_cache
def _scan_models_dir(
@ -202,7 +199,12 @@ def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
return False
def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
def _scan_hf_cache(
cache_dir: Path,
*,
entry_limit: int | None = None,
active_cache: bool = True,
) -> List[LocalModelInfo]:
if not _safe_is_dir(cache_dir):
return []
@ -240,7 +242,10 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
repo_dir,
)
gguf_partial = hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir)
has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(model_id)
has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(
model_id,
hub_cache = cache_dir,
)
snapshot_partial_transport = (
hf_cache_scan.partial_transport_for(
"model",
@ -252,23 +257,25 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
)
resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_dir)
scan_path = Path(resolved) if resolved else repo_dir
load_path = repo_dir if active_cache else scan_path
# partial=False here; _apply_format_aware_partial below rewrites per-row
# so a hybrid repo's gguf row doesn't taint its safetensors row.
rows = _classify_local_path(
scan_path,
"hf_cache",
load_path = repo_dir,
load_path = load_path,
display_name = model_id.split("/")[-1],
model_id = model_id,
updated_at = updated_at,
partial = False,
active_cache = active_cache,
)
if not rows:
if has_gguf_variant_state and gguf_partial:
rows = [
_local_model_info(
scan_path = repo_dir,
load_path = repo_dir,
load_path = load_path,
source = "hf_cache",
model_format = "gguf",
display_name = model_id.split("/")[-1],
@ -277,6 +284,7 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
partial = True,
requires_variant = True,
size_bytes = gguf_variant_state_size,
active_cache = active_cache,
)
]
else:
@ -285,13 +293,14 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
rows = [
_local_model_info(
scan_path = repo_dir,
load_path = repo_dir,
load_path = load_path,
source = "hf_cache",
model_format = "unknown",
display_name = model_id.split("/")[-1],
model_id = model_id,
updated_at = updated_at,
partial = snapshot_partial or gguf_partial,
active_cache = active_cache,
)
]
elif (
@ -302,7 +311,7 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
rows.append(
_local_model_info(
scan_path = repo_dir,
load_path = repo_dir,
load_path = load_path,
source = "hf_cache",
model_format = "gguf",
display_name = model_id.split("/")[-1],
@ -311,6 +320,7 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
partial = True,
requires_variant = True,
size_bytes = gguf_variant_state_size,
active_cache = active_cache,
)
)
rows = _apply_format_aware_partial(
@ -515,14 +525,39 @@ async def _collect_models_from_default_sources(
local_models += await _scan_source("HF cache", _scan_hf_cache, hf_cache_dir)
if _safe_is_dir(legacy_hf) and legacy_hf.resolve() != hf_cache_dir.resolve():
local_models += await _scan_source("legacy HF cache", _scan_hf_cache, legacy_hf)
local_models += await _scan_source(
"legacy HF cache",
lambda path: _scan_hf_cache(path, active_cache = False),
legacy_hf,
)
if (
_safe_is_dir(hf_default)
and hf_default.resolve() != hf_cache_dir.resolve()
and hf_default.resolve() != legacy_hf.resolve()
):
local_models += await _scan_source("default HF cache", _scan_hf_cache, hf_default)
local_models += await _scan_source(
"default HF cache",
lambda path: _scan_hf_cache(path, active_cache = False),
hf_default,
)
from utils.hf_cache_settings import known_hf_hub_caches
seen_hf = {
os.path.normcase(str(path.resolve(strict = False)))
for path in (hf_cache_dir, legacy_hf, hf_default)
}
for previous_cache in known_hf_hub_caches():
key = os.path.normcase(str(previous_cache.resolve(strict = False)))
if key in seen_hf:
continue
seen_hf.add(key)
local_models += await _scan_source(
"previous HF cache",
lambda path: _scan_hf_cache(path, active_cache = False),
previous_cache,
)
for lm_dir in lm_dirs:
local_models += await _scan_source("LM Studio", _scan_lmstudio_dir, lm_dir)
@ -543,7 +578,11 @@ def _scan_custom_folder(folder_path: Path) -> List[LocalModelInfo]:
limit = _MAX_MODELS_PER_CUSTOM_FOLDER,
entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES,
)
+ _scan_hf_cache(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES)
+ _scan_hf_cache(
folder_path,
entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES,
active_cache = False,
)
+ _scan_lmstudio_dir(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES)
)
if m.model_format in supported_formats
@ -610,12 +649,20 @@ def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelI
row_key = model.inventory_id or model.id
key = f"{row_key}\x00custom" if model.source == "custom" else row_key
existing = deduped.get(key)
if existing is None or _prefer_complete_larger(
model.partial,
model.size_bytes,
existing.partial,
existing.size_bytes,
):
prefer_candidate = existing is None
if existing is not None:
if model.partial != existing.partial:
prefer_candidate = not model.partial
elif (model.active_cache is True) != (existing.active_cache is True):
prefer_candidate = model.active_cache is True
else:
prefer_candidate = _prefer_complete_larger(
model.partial,
model.size_bytes,
existing.partial,
existing.size_bytes,
)
if prefer_candidate:
deduped[key] = model
return sorted(
deduped.values(),

33
studio/backend/hub/services/snapshot_progress.py
View file

@ -86,9 +86,20 @@ def _snapshot_complete_on_disk(
return False
if variant is None and hf_cache_scan.repo_cache_dir_has_incomplete_blobs(entry):
return False
if download_manifest.has_cancel_marker(repo_type, repo_id, variant):
hub_cache = entry.parent
if download_manifest.has_cancel_marker(
repo_type,
repo_id,
variant,
hub_cache = hub_cache,
):
return False
manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
manifest = download_manifest.read_manifest(
repo_type,
repo_id,
variant,
hub_cache = hub_cache,
)
if manifest is None:
return False
return download_manifest.verify_against_disk(manifest, snapshot_dir).ok
@ -118,6 +129,8 @@ def compute_snapshot_progress(
0,
int(getattr(metadata, "completed_baseline_bytes", 0) or 0),
)
metadata_hub_cache = getattr(metadata, "hub_cache", None)
active_root = Path(metadata_hub_cache) if metadata_hub_cache else None
expected_total = max(expected_bytes, 0)
# Always resolve the revision's blob hashes so stale blobs from a superseded
@ -134,11 +147,17 @@ def compute_snapshot_progress(
count_finalized_unscoped = variant is None
readings: list[tuple[int, int, Optional[str], bool]] = []
for entry in preferred_repo_cache_dirs(
repo_type,
repo_id,
force_active = force_active,
):
cache_dirs = (
preferred_repo_cache_dirs(
repo_type,
repo_id,
force_active = force_active,
active_root = active_root,
)
if active_root is not None
else preferred_repo_cache_dirs(repo_type, repo_id, force_active = force_active)
)
for entry in cache_dirs:
completed_bytes = 0
in_progress_bytes = 0
cache_path = hf_cache_scan.resolve_hf_cache_realpath(entry)

120
studio/backend/hub/tests/test_dataset_services.py
View file

@ -72,57 +72,115 @@ def test_dataset_cache_scan_merges_raw_and_processed_rows(monkeypatch):
assert rows[0]["partial"] is False
def test_delete_cached_dataset_attempts_all_roots_before_raising(monkeypatch):
def test_delete_cached_dataset_scopes_delete_to_selected_root(monkeypatch, tmp_path):
"""A dataset present in the active cache and a previously selected cache is
deleted only from the selected root, so the other cache's copy survives."""
calls = []
purged_state = []
target_hub = tmp_path / "active" / "hub"
other_hub = tmp_path / "previous" / "hub"
for hub in (target_hub, other_hub):
(hub / "datasets--Org--Data").mkdir(parents = True)
class _DeleteStrategy:
def __init__(self, label: str, fail: bool):
def __init__(self, label: str):
self.label = label
self.fail = fail
def execute(self):
calls.append(self.label)
if self.fail:
raise RuntimeError(f"{self.label} failed")
class _Cache:
def __init__(self, label: str, fail: bool):
self.cache_dir = label
self.repos = [
def _cache(label: str, hub):
return SimpleNamespace(
cache_dir = label,
repos = [
SimpleNamespace(
repo_type = "dataset",
repo_id = "Org/Data",
repo_path = str(hub / "datasets--Org--Data"),
revisions = [SimpleNamespace(commit_hash = f"{label}-rev")],
)
]
self.fail = fail
def delete_revisions(self, *_revisions):
return _DeleteStrategy(self.cache_dir, self.fail)
],
delete_revisions = lambda *_revs, _label = label: _DeleteStrategy(_label),
)
monkeypatch.setattr(
cache_inventory,
"_collect_hf_cache_scans",
lambda: ([_Cache("first", True), _Cache("second", False)], set()),
lambda: ([_cache("active", target_hub), _cache("previous", other_hub)], set()),
)
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda _repo_id: (True, []),
lambda _repo_id, **_kwargs: (False, []),
)
monkeypatch.setattr(
cache_inventory.download_manifest,
"purge_all_state_for_repo",
lambda *_args: purged_state.append(True) or 1,
lambda *_args, **_kwargs: 0,
)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = target_hub),
)
monkeypatch.setattr(
"hub.utils.hf_cache_state.hf_cache_roots",
lambda: [target_hub, other_hub],
)
with pytest.raises(HTTPException) as exc_info:
cache_inventory._delete_cached_dataset_blocking("Org/Data")
result = cache_inventory._delete_cached_dataset_blocking("Org/Data")
assert exc_info.value.status_code == 500
assert calls == ["first", "second"]
assert purged_state == []
assert result == {"status": "deleted", "repo_id": "Org/Data"}
# Only the selected (active) cache's revision is deleted; the previous
# cache's copy is never touched.
assert calls == ["active"]
assert not (target_hub / "datasets--Org--Data").exists()
assert (other_hub / "datasets--Org--Data").exists()
def test_delete_processed_only_dataset_accepts_processed_cache_path(monkeypatch, tmp_path):
"""A processed-only dataset row sends its Arrow cache path (<owner>___<repo>
under HF_DATASETS_CACHE), which is not a Hub datasets-- dir. The delete must
accept it and run the processed-cache delete instead of raising 400."""
datasets_root = tmp_path / "datasets"
processed_dir = datasets_root / "Org___Data"
processed_dir.mkdir(parents = True)
# No Hub-cache copy exists; only the processed Arrow cache holds this repo.
monkeypatch.setattr(cache_inventory, "_collect_hf_cache_scans", lambda: ([], set()))
monkeypatch.setattr(cache_inventory, "_hf_datasets_cache_roots", lambda: [datasets_root])
processed_calls: list[str] = []
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda repo_id, **_kwargs: (processed_calls.append(repo_id) or True, []),
)
result = cache_inventory._delete_cached_dataset_blocking("Org/Data", str(processed_dir))
assert result == {"status": "deleted", "repo_id": "Org/Data"}
assert processed_calls == ["Org/Data"]
def test_delete_processed_dataset_scopes_to_selected_root(monkeypatch, tmp_path):
"""A dataset processed under two HF_DATASETS_CACHE roots is deleted only from
the selected root; the copy under the other cache home survives (real delete,
not stubbed)."""
selected_root = tmp_path / "selected" / "datasets"
other_root = tmp_path / "other" / "datasets"
for root in (selected_root, other_root):
(root / "Org___Data").mkdir(parents = True)
monkeypatch.setattr(cache_inventory, "_collect_hf_cache_scans", lambda: ([], set()))
monkeypatch.setattr(
cache_inventory, "_hf_datasets_cache_roots", lambda: [selected_root, other_root]
)
result = cache_inventory._delete_cached_dataset_blocking(
"Org/Data", str(selected_root / "Org___Data")
)
assert result == {"status": "deleted", "repo_id": "Org/Data"}
assert not (selected_root / "Org___Data").exists() # the selected copy is deleted
assert (other_root / "Org___Data").exists() # the other cache home is untouched
def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch):
@ -139,22 +197,22 @@ def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch):
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda _repo_id: (False, []),
lambda _repo_id, **_kwargs: (False, []),
)
monkeypatch.setattr(
cache_inventory,
"purge_repo_cache_dirs",
lambda _repo_type, repo_id: purged_dirs.append(repo_id) or True,
lambda _repo_type, repo_id, **_kwargs: purged_dirs.append(repo_id) or True,
)
monkeypatch.setattr(
cache_inventory,
"purge_partial_repo",
lambda *_args: False,
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
cache_inventory.download_manifest,
"purge_all_state_for_repo",
lambda *_args: 0,
lambda *_args, **_kwargs: 0,
)
result = cache_inventory._delete_cached_dataset_blocking("Org/Data")
@ -172,22 +230,22 @@ def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch):
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda _repo_id: (False, []),
lambda _repo_id, **_kwargs: (False, []),
)
monkeypatch.setattr(
cache_inventory,
"purge_repo_cache_dirs",
lambda *_args: False,
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
cache_inventory,
"purge_partial_repo",
lambda *_args: False,
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
cache_inventory.download_manifest,
"purge_all_state_for_repo",
lambda *_args: 0,
lambda *_args, **_kwargs: 0,
)
with pytest.raises(HTTPException) as exc_info:

62
studio/backend/hub/tests/test_download_manifest_scoping.py Normal file
View file

@ -0,0 +1,62 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import json
from types import SimpleNamespace
from hub.utils import download_manifest, state_dir
def _write_manifest(path, payload):
path.parent.mkdir(parents = True, exist_ok = True)
path.write_text(json.dumps(payload), encoding = "utf-8")
def test_purge_state_preserves_active_legacy_when_deleting_inactive_cache(monkeypatch, tmp_path):
"""A scoped delete of an inactive cache must not erase the unscoped legacy
state, which _legacy_state_applies attributes to the active cache."""
active = tmp_path / "active" / "hub"
previous = tmp_path / "previous" / "hub"
for path in (active, previous):
path.mkdir(parents = True)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = str(active)),
)
# Unowned legacy manifest -> belongs to the active cache.
legacy = state_dir.manifest_path("model", "Org/Model")
_write_manifest(legacy, {"version": 1})
# The inactive cache's own scoped copy is the one being deleted.
scoped = state_dir.manifest_path("model", "Org/Model", hub_cache = str(previous))
_write_manifest(scoped, {"version": 1, "hub_cache": str(previous)})
removed = download_manifest.purge_state("model", "Org/Model", hub_cache = str(previous))
assert removed is True
assert not scoped.is_file() # the inactive cache's copy is gone
assert legacy.is_file() # the active cache's legacy state survives
def test_purge_state_removes_legacy_owned_by_the_deleted_cache(monkeypatch, tmp_path):
"""A legacy file that recorded the deleted cache as its owner is purged."""
active = tmp_path / "active" / "hub"
previous = tmp_path / "previous" / "hub"
for path in (active, previous):
path.mkdir(parents = True)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = str(active)),
)
legacy = state_dir.manifest_path("model", "Org/Model")
_write_manifest(legacy, {"version": 1, "hub_cache": str(previous)})
removed = download_manifest.purge_state("model", "Org/Model", hub_cache = str(previous))
assert removed is True
assert not legacy.is_file() # owned by the deleted cache -> purged

22
studio/backend/hub/tests/test_empty_variant_folder.py
View file

@ -120,10 +120,16 @@ def _force_compute_to_raise(monkeypatch):
monkeypatch.setattr(gguf_variants, "list_gguf_variants", _boom, raising = False)
monkeypatch.setattr(
gguf_variants, "list_gguf_variants_from_hf_cache", lambda repo_id: None, raising = False
gguf_variants,
"list_gguf_variants_from_hf_cache",
lambda repo_id, root = None: None,
raising = False,
)
monkeypatch.setattr(
gguf_variants, "list_partial_gguf_variants_from_state", lambda repo_id: None, raising = False
gguf_variants,
"list_partial_gguf_variants_from_state",
lambda repo_id, hub_cache = None: None,
raising = False,
)
@ -133,7 +139,11 @@ def test_get_variants_surfaces_cleanable_when_metadata_fails(monkeypatch):
import asyncio
_force_compute_to_raise(monkeypatch)
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
monkeypatch.setattr(
gguf_variants,
"list_empty_gguf_variant_dirs",
lambda repo_id, root = None: {"UD-IQ1_S"},
)
resp = asyncio.run(
gguf_variants.get_gguf_variants_response(
@ -152,7 +162,11 @@ def test_get_variants_reraises_when_no_cleanable(monkeypatch):
from fastapi import HTTPException
_force_compute_to_raise(monkeypatch)
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set())
monkeypatch.setattr(
gguf_variants,
"list_empty_gguf_variant_dirs",
lambda repo_id, root = None: set(),
)
try:
asyncio.run(

509
studio/backend/hub/tests/test_model_services.py
View file

@ -2,6 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import json
import sys
from pathlib import Path
from types import SimpleNamespace
@ -102,6 +103,236 @@ def test_big_endian_detection_ignores_model_name_be_token():
)
def _cached_model_row(tmp_path: Path, *, partial: bool, active_cache: bool | None, size_bytes: int):
path = tmp_path / f"cache-{active_cache}-{partial}-{size_bytes}"
return model_common._local_model_info(
scan_path = path,
load_path = path,
source = "hf_cache",
model_format = "safetensors",
model_id = "Org/Model",
partial = partial,
active_cache = active_cache,
size_bytes = size_bytes,
)
@pytest.mark.parametrize("reverse", [False, True])
def test_local_inventory_prefers_complete_previous_cache_copy(tmp_path, reverse):
active_partial = _cached_model_row(
tmp_path,
partial = True,
active_cache = True,
size_bytes = 20,
)
previous_complete = _cached_model_row(
tmp_path,
partial = False,
active_cache = False,
size_bytes = 10,
)
rows = [active_partial, previous_complete]
if reverse:
rows.reverse()
result = local_inventory._dedupe_local_models(rows)
assert result == [previous_complete]
def test_local_inventory_compares_all_non_active_cache_copies(tmp_path):
inactive_partial = _cached_model_row(
tmp_path,
partial = True,
active_cache = False,
size_bytes = 20,
)
custom_complete = _cached_model_row(
tmp_path,
partial = False,
active_cache = None,
size_bytes = 10,
)
assert local_inventory._dedupe_local_models([inactive_partial, custom_complete]) == [
custom_complete
]
def test_local_inventory_prefers_active_cache_when_copies_are_equally_complete(tmp_path):
previous = _cached_model_row(
tmp_path,
partial = False,
active_cache = False,
size_bytes = 20,
)
active = _cached_model_row(
tmp_path,
partial = False,
active_cache = True,
size_bytes = 10,
)
assert local_inventory._dedupe_local_models([previous, active]) == [active]
def test_loaded_repo_match_accepts_previous_cache_snapshot_path(monkeypatch, tmp_path):
repo_dir = tmp_path / "old-hub" / "models--Org--Model"
snapshot = repo_dir / "snapshots" / "revision"
snapshot.mkdir(parents = True)
monkeypatch.setattr(deletion, "iter_repo_cache_dirs", lambda *_args: iter([repo_dir]))
assert deletion._loaded_id_matches_repo(str(snapshot), "Org/Model") is True
assert deletion._loaded_id_matches_repo(str(snapshot / "model.gguf"), "Org/Model") is True
assert deletion._loaded_id_matches_repo(str(tmp_path / "other"), "Org/Model") is False
def test_cached_inventory_loads_previous_cache_copy_by_snapshot(monkeypatch, tmp_path):
active_hub = tmp_path / "active-hub"
previous_repo = tmp_path / "previous-hub" / "models--Org--Model"
snapshot = previous_repo / "snapshots" / "revision"
snapshot.mkdir(parents = True)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = active_hub),
)
fields = cache_inventory._cache_inventory_fields(
"Org/Model",
"safetensors",
repo_path = previous_repo,
snapshot_path = snapshot,
)
assert fields["load_id"] == str(snapshot)
def test_cached_inventory_keeps_repo_id_for_active_cache(monkeypatch, tmp_path):
active_hub = tmp_path / "active-hub"
active_repo = active_hub / "models--Org--Model"
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = active_hub),
)
fields = cache_inventory._cache_inventory_fields(
"Org/Model",
"safetensors",
repo_path = active_repo,
)
assert fields["load_id"] == "Org/Model"
def test_cached_inventory_prefers_active_copy_when_completeness_matches():
previous = {"partial": False, "active_cache": False, "size_bytes": 200}
active = {"partial": False, "active_cache": True, "size_bytes": 100}
assert cache_inventory._prefer_cache_row(active, previous) is True
assert cache_inventory._prefer_cache_row(previous, active) is False
def test_cached_inventory_prefers_complete_copy_before_active_cache():
previous = {"partial": False, "active_cache": False, "size_bytes": 100}
active_partial = {"partial": True, "active_cache": True, "size_bytes": 200}
assert cache_inventory._prefer_cache_row(previous, active_partial) is True
assert cache_inventory._prefer_cache_row(active_partial, previous) is False
def test_inventory_scans_every_dynamic_cache_root(monkeypatch, tmp_path):
first = tmp_path / "first-hub"
second = tmp_path / "second-hub"
unreadable = tmp_path / "unreadable-hub"
first.mkdir()
second.mkdir()
unreadable.mkdir()
scanned = []
monkeypatch.setattr(
inventory_scan,
"hf_cache_roots",
lambda: [first, unreadable, second],
)
def scan_cache(cache_dir):
path = Path(cache_dir)
scanned.append(path)
if path == unreadable:
raise PermissionError("unreadable")
return SimpleNamespace(cache_dir = cache_dir)
monkeypatch.setattr("huggingface_hub.scan_cache_dir", scan_cache)
result = inventory_scan._compute_all_hf_cache_scans()
assert scanned == [first, unreadable, second]
assert [Path(scan.cache_dir) for scan in result] == [first, second]
def test_inventory_applies_download_state_to_its_owning_cache(monkeypatch, tmp_path):
state_root = tmp_path / "state"
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
repo_id = "Org/Model"
repo_name = "models--Org--Model"
repo_a = cache_a / repo_name
repo_b = cache_b / repo_name
snapshot_a = repo_a / "snapshots" / "revision"
snapshot_b = repo_b / "snapshots" / "revision"
snapshot_a.mkdir(parents = True)
snapshot_b.mkdir(parents = True)
(snapshot_a / "config.json").write_bytes(b"x")
(snapshot_b / "config.json").write_bytes(b"xx")
monkeypatch.setattr(state_dir, "cache_root", lambda: state_root)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = cache_b),
)
assert download_manifest.write_manifest(
"model",
repo_id,
None,
[download_manifest.ExpectedFile(path = "config.json", size = 2)],
"http",
hub_cache = cache_a,
)
assert inventory_scan.is_snapshot_partial("model", repo_id, repo_a) is True
assert inventory_scan.is_snapshot_partial("model", repo_id, repo_b) is False
assert inventory_scan.partial_transport_for("model", repo_id, None, repo_a) == "http"
assert inventory_scan.partial_transport_for("model", repo_id, None, repo_b) is None
def test_inventory_scopes_cancel_markers_to_their_owning_cache(monkeypatch, tmp_path):
state_root = tmp_path / "state"
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
repo_id = "Org/Model"
repo_name = "models--Org--Model"
repo_a = cache_a / repo_name
repo_b = cache_b / repo_name
repo_a.mkdir(parents = True)
repo_b.mkdir(parents = True)
monkeypatch.setattr(state_dir, "cache_root", lambda: state_root)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = cache_b),
)
assert download_manifest.write_cancel_marker(
"model",
repo_id,
"Q4_K_M",
"xet",
hub_cache = cache_a,
)
assert inventory_scan.is_variant_partial(repo_id, "Q4_K_M", repo_cache_dir = repo_a) is True
assert inventory_scan.is_variant_partial(repo_id, "Q4_K_M", repo_cache_dir = repo_b) is False
def test_list_local_gguf_variants_skips_big_endian_sibling(tmp_path):
(tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 100)
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"y" * 10)
@ -163,8 +394,19 @@ def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path
"http",
)
marker_path = state_dir.marker_path("model", repo_id, variant)
manifest_path = state_dir.manifest_path("model", repo_id, variant)
hub_cache = download_manifest._canonical_hub_cache()
marker_path = state_dir.marker_path(
"model",
repo_id,
variant,
hub_cache = hub_cache,
)
manifest_path = state_dir.manifest_path(
"model",
repo_id,
variant,
hub_cache = hub_cache,
)
assert marker_path is not None
assert manifest_path is not None
@ -181,6 +423,97 @@ def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path
]
def test_download_state_isolated_across_hub_cache_switches(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
selected = SimpleNamespace(hub_cache = cache_a)
from utils import hf_cache_settings
monkeypatch.setattr(hf_cache_settings, "get_hf_cache_paths", lambda: selected)
expected_a = [download_manifest.ExpectedFile(path = "a.gguf", size = 1)]
expected_b = [download_manifest.ExpectedFile(path = "b.gguf", size = 2)]
assert download_manifest.write_manifest("model", "Owner/Repo", "Q4_K_M", expected_a)
assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")
selected.hub_cache = cache_b
assert download_manifest.write_manifest("model", "Owner/Repo", "Q4_K_M", expected_b)
manifest_b = download_manifest.read_manifest("model", "Owner/Repo", "Q4_K_M")
manifest_a = download_manifest.read_manifest(
"model",
"Owner/Repo",
"Q4_K_M",
hub_cache = cache_a,
)
assert manifest_b is not None and manifest_b.expected_files == tuple(expected_b)
assert manifest_a is not None and manifest_a.expected_files == tuple(expected_a)
assert not download_manifest.has_cancel_marker("model", "Owner/Repo", "Q4_K_M")
assert download_manifest.has_cancel_marker(
"model",
"Owner/Repo",
"Q4_K_M",
hub_cache = cache_a,
)
assert len(list((tmp_path / "hub-state" / "manifests").rglob("*.json"))) == 2
def test_legacy_unscoped_download_state_falls_back_only_for_selected_cache(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = cache_a),
)
manifest = state_dir.manifest_path("model", "Owner/Repo", "Q4_K_M")
marker = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M")
assert manifest is not None and marker is not None
manifest.write_text(
json.dumps(
{
"version": 1,
"repo_id": "Owner/Repo",
"variant": "Q4_K_M",
"expected_files": [{"path": "model.gguf", "size": 10}],
"transport": "http",
}
),
encoding = "utf-8",
)
marker.write_text(
json.dumps({"version": 1, "repo_id": "Owner/Repo", "variant": "Q4_K_M"}),
encoding = "utf-8",
)
assert download_manifest.read_manifest("model", "Owner/Repo", "Q4_K_M") is not None
assert download_manifest.has_cancel_marker("model", "Owner/Repo", "Q4_K_M")
assert list(download_manifest.iter_variant_manifests("model", "Owner/Repo")) == [
("Q4_K_M", manifest)
]
assert list(download_manifest.iter_variant_markers("model", "Owner/Repo")) == [
("Q4_K_M", marker)
]
assert (
download_manifest.read_manifest(
"model",
"Owner/Repo",
"Q4_K_M",
hub_cache = cache_b,
)
is None
)
assert not download_manifest.has_cancel_marker(
"model",
"Owner/Repo",
"Q4_K_M",
hub_cache = cache_b,
)
class _RecordingLogger:
def __init__(self):
self.warnings = []
@ -416,8 +749,15 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)],
"http",
hub_cache = repo_path.parent,
)
assert download_manifest.write_cancel_marker(
"model",
"Org/PartialGguf",
"Q4_K_M",
"http",
hub_cache = repo_path.parent,
)
assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
@ -484,6 +824,7 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa
"Q8_0",
[download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)],
"http",
hub_cache = Path(embedder.repo_path).parent,
)
monkeypatch.setattr(
cache_inventory,
@ -1206,6 +1547,7 @@ def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch,
),
],
"http",
hub_cache = entry.parent,
)
requirement = gguf_variants._GgufVariantRequirement(
@ -1292,6 +1634,7 @@ def test_gguf_progress_subtracts_new_job_completed_baseline(monkeypatch, tmp_pat
),
],
"http",
hub_cache = entry.parent,
)
requirement = gguf_variants._GgufVariantRequirement(
@ -1462,6 +1805,7 @@ def test_gguf_progress_complete_on_disk_ignores_full_baseline(monkeypatch, tmp_p
),
],
"http",
hub_cache = entry.parent,
)
requirement = gguf_variants._GgufVariantRequirement(
@ -1861,8 +2205,15 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)],
"http",
hub_cache = cache_dir,
)
assert download_manifest.write_cancel_marker(
"model",
"Org/PartialGguf",
"Q4_K_M",
"http",
hub_cache = cache_dir,
)
assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
monkeypatch.setattr(
local_inventory.hf_cache_scan,
@ -2117,7 +2468,7 @@ def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch
monkeypatch.setattr(
gguf_variants,
"iter_hf_cache_snapshots",
lambda _repo_id: [snapshot],
lambda _repo_id, root = None: [snapshot],
)
monkeypatch.setattr(
gguf_variants,
@ -2136,6 +2487,70 @@ def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch
assert result.variants[0].partial is True
def test_gguf_variants_scopes_partial_state_to_requested_cache(monkeypatch, tmp_path):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
repo_id = "Org/SharedRepo"
repo_name = "models--Org--SharedRepo"
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
repo_a = cache_a / repo_name
snapshot_a = repo_a / "snapshots" / "revision"
snapshot_a.mkdir(parents = True)
(snapshot_a / "model-Q8_0.gguf").write_bytes(b"complete")
blobs_b = cache_b / repo_name / "blobs"
blobs_b.mkdir(parents = True)
(blobs_b / "q8-hash.incomplete").write_bytes(b"partial")
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = cache_b),
)
assert download_manifest.write_cancel_marker(
"model",
repo_id,
"Q8_0",
"http",
hub_cache = cache_b,
)
monkeypatch.setattr(
gguf_variants,
"list_gguf_variants",
lambda *_args, **_kwargs: (
[
SimpleNamespace(
filename = "model-Q8_0.gguf",
quant = "Q8_0",
display_label = None,
size_bytes = 8,
)
],
False,
[
SimpleNamespace(
rfilename = "model-Q8_0.gguf",
size = 8,
lfs = SimpleNamespace(sha256 = "q8-hash"),
)
],
),
)
monkeypatch.setattr(cache_inventory, "all_hf_cache_scans", lambda: [])
result = asyncio.run(
gguf_variants.get_gguf_variants_response(
repo_id,
local_path = str(repo_a),
)
)
assert result.variants[0].downloaded is True
assert result.variants[0].partial is False
def test_download_registry_repo_keys_are_case_insensitive():
registry = download_registry.DownloadRegistry()
@ -2444,6 +2859,34 @@ def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, t
assert (blobs / "shared-mmproj.incomplete").exists()
def test_prepare_cache_for_transport_uses_captured_root(monkeypatch, tmp_path):
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
repo_name = "models--Org--Repo"
partial_a = cache_a / repo_name / "blobs" / "blob.incomplete"
partial_b = cache_b / repo_name / "blobs" / "blob.incomplete"
partial_a.parent.mkdir(parents = True)
partial_b.parent.mkdir(parents = True)
partial_a.write_bytes(b"a")
partial_b.write_bytes(b"b")
monkeypatch.setattr(
download_registry,
"hf_cache_root",
lambda create = False, root = None: root or cache_b,
)
purged = download_registry.prepare_cache_for_transport(
"model",
"Org/Repo",
download_registry.TRANSPORT_HTTP,
root = cache_a,
)
assert purged == 1
assert not partial_a.exists()
assert partial_b.exists()
def _vision_cache_root(monkeypatch, tmp_path):
root = tmp_path / "hub"
blobs = root / "models--Org--Vision" / "blobs"
@ -2802,6 +3245,47 @@ def test_shutdown_skips_marker_for_worker_that_exits_cleanly(monkeypatch):
assert markers == ["Org/Cut"]
def test_orphan_reaper_uses_worker_cache_root_after_setting_changes(monkeypatch, tmp_path):
workers = tmp_path / "workers"
workers.mkdir()
cache_a = tmp_path / "cache-a" / "hub"
cache_b = tmp_path / "cache-b" / "hub"
partial = cache_a / "models--Org--Model" / "blobs" / "abc.incomplete"
partial.parent.mkdir(parents = True)
partial.write_bytes(b"partial")
cache_b.mkdir(parents = True)
monkeypatch.setattr(state_dir, "workers_dir", lambda: workers)
monkeypatch.setattr(download_registry, "_process_alive", lambda _pid: False)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = cache_b),
)
markers = []
monkeypatch.setattr(
download_registry,
"persist_cancel_marker",
lambda *args, **kwargs: markers.append(args),
)
metadata = download_registry.DownloadMetadata(
repo_type = "model",
repo_id = "Org/Model",
variant = None,
transport = download_registry.TRANSPORT_HTTP,
hub_cache = str(cache_a),
xet_cache = str(tmp_path / "cache-a" / "xet"),
)
download_registry.write_worker_breadcrumb("org/model", 1234, metadata)
[breadcrumb] = list(workers.iterdir())
payload = json.loads(breadcrumb.read_text(encoding = "utf-8"))
assert payload["hub_cache"] == str(cache_a)
assert payload["xet_cache"] == str(tmp_path / "cache-a" / "xet")
download_registry.reap_orphan_workers()
assert markers == [("model", "Org/Model", None, "http")]
assert list(workers.iterdir()) == []
def test_model_claim_register_cancel_uses_registry_marker_owner(monkeypatch):
killed = []
@ -3125,12 +3609,19 @@ def _build_variant_cache_repo(repo_dir, blob_specs, snapshot_links):
return repo
def _patch_variant_delete_side_effects(monkeypatch):
def _patch_variant_delete_side_effects(monkeypatch, hub_cache = None):
monkeypatch.setattr(
deletion.download_manifest,
"purge_state",
lambda *_args, **_kwargs: False,
)
# The repo under test lives in this cache; make it the active one so the
# delete scopes to it (default target root is the active hub cache).
if hub_cache is not None:
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = hub_cache),
)
def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path):
@ -3308,7 +3799,7 @@ def test_delete_variant_keeps_blob_shared_with_other_snapshot(monkeypatch, tmp_p
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
_patch_variant_delete_side_effects(monkeypatch)
_patch_variant_delete_side_effects(monkeypatch, tmp_path)
result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)
@ -3335,7 +3826,7 @@ def test_delete_variant_unlinks_unshared_blob(monkeypatch, tmp_path):
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
_patch_variant_delete_side_effects(monkeypatch)
_patch_variant_delete_side_effects(monkeypatch, tmp_path)
result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)
@ -3361,7 +3852,7 @@ def test_delete_variant_surfaces_locked_file_as_conflict(monkeypatch, tmp_path):
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
_patch_variant_delete_side_effects(monkeypatch)
_patch_variant_delete_side_effects(monkeypatch, tmp_path)
real_unlink = Path.unlink

407
studio/backend/hub/utils/download_manifest.py
View file

@ -77,6 +77,7 @@ class Manifest:
started_at: str
expected_files: tuple[ExpectedFile, ...]
transport: Optional[str] = None
hub_cache: Optional[str] = None
@dataclass(frozen = True)
@ -86,6 +87,78 @@ class VerifyResult:
size_mismatched: tuple[str, ...]
def _canonical_hub_cache(hub_cache: Optional[str | Path] = None) -> Optional[str]:
if hub_cache is None:
try:
from utils.hf_cache_settings import get_hf_cache_paths
hub_cache = get_hf_cache_paths().hub_cache
except Exception:
return None
try:
return str(Path(hub_cache).expanduser().resolve(strict = False))
except (OSError, RuntimeError, ValueError):
return str(hub_cache)
def _read_state_payload(path: Path) -> Optional[dict]:
try:
data = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError) as exc:
logger.debug("Could not read Hub state %s: %s", path, exc)
return None
return data if isinstance(data, dict) else None
def _legacy_state_applies(
path: Path,
requested_hub_cache: Optional[str],
*,
fail_closed: bool = False,
) -> bool:
"""Whether an old unscoped state file belongs to the requested cache.
Transitional files that recorded their cache keep that ownership. Older
files with no ownership can only be attributed to the currently selected
cache, which matches the single-cache behavior under which they were
written without leaking them into remembered inactive caches.
"""
data = _read_state_payload(path)
if data is not None:
recorded = data.get("hub_cache")
if isinstance(recorded, str) and recorded:
return _canonical_hub_cache(recorded) == requested_hub_cache
elif not fail_closed:
return False
return requested_hub_cache == _canonical_hub_cache()
def _state_read_path(
path_factory,
repo_type: RepoType,
repo_id: str,
variant: Optional[str],
hub_cache: Optional[str | Path],
*,
fail_closed: bool = False,
) -> Optional[Path]:
requested = _canonical_hub_cache(hub_cache)
scoped = path_factory(repo_type, repo_id, variant, hub_cache = requested)
try:
if scoped is not None and scoped.is_file():
return scoped
except OSError:
pass
legacy = path_factory(repo_type, repo_id, variant)
if legacy is None or legacy == scoped:
return None
try:
if not legacy.is_file():
return None
except OSError:
return None
return legacy if _legacy_state_applies(legacy, requested, fail_closed = fail_closed) else None
def _atomic_write_json(path: Path, payload: dict) -> bool:
# Per-write uuid suffix so a concurrent caller or a stale tmp from a
# previous crash cannot collide with the in-flight write.
@ -124,6 +197,8 @@ def write_manifest(
variant: Optional[str],
expected_files: Sequence[ExpectedFile],
transport: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> bool:
"""Write/overwrite the manifest for this triple. Best-effort.
@ -131,7 +206,13 @@ def write_manifest(
worst-case fallback is the pre-fix scanner behavior (one missed
partial detection), which is no regression.
"""
path = manifest_path(repo_type, repo_id, variant)
recorded_hub_cache = _canonical_hub_cache(hub_cache)
path = manifest_path(
repo_type,
repo_id,
variant,
hub_cache = recorded_hub_cache,
)
if path is None:
return False
payload = {
@ -149,6 +230,7 @@ def write_manifest(
for f in expected_files
],
"transport": transport,
"hub_cache": recorded_hub_cache,
}
return _atomic_write_json(path, payload)
@ -157,6 +239,8 @@ def read_manifest(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> Optional[Manifest]:
"""Return the manifest if present and parseable; ``None`` otherwise.
@ -171,15 +255,17 @@ def read_manifest(
``_MANIFEST_VERSION`` and widen this check) or live under a different
filename, so an incompatible payload can never mis-classify rows.
"""
path = manifest_path(repo_type, repo_id, variant)
path = _state_read_path(
manifest_path,
repo_type,
repo_id,
variant,
hub_cache,
)
if path is None or not path.is_file():
return None
try:
data = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError) as exc:
logger.debug("Could not read manifest %s: %s", path, exc)
return None
if not isinstance(data, dict):
data = _read_state_payload(path)
if data is None:
return None
if data.get("version") != _MANIFEST_VERSION:
logger.debug(
@ -216,6 +302,7 @@ def read_manifest(
started_at = str(data.get("started_at", "")),
expected_files = tuple(expected),
transport = transport if transport in ("http", "xet") else None,
hub_cache = data.get("hub_cache") if isinstance(data.get("hub_cache"), str) else None,
)
@ -289,6 +376,8 @@ def write_cancel_marker(
repo_id: str,
variant: Optional[str] = None,
transport: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> bool:
"""Record that this triple was cancelled. Idempotent across repeated cancels.
@ -296,7 +385,13 @@ def write_cancel_marker(
inventory rows so the UI labels HTTP retries as continuable and XET
retries as full redownloads. None is accepted for forward-compat.
"""
path = marker_path(repo_type, repo_id, variant)
recorded_hub_cache = _canonical_hub_cache(hub_cache)
path = marker_path(
repo_type,
repo_id,
variant,
hub_cache = recorded_hub_cache,
)
if path is None:
return False
payload = {
@ -306,6 +401,7 @@ def write_cancel_marker(
"variant": variant,
"transport": transport,
"cancelled_at": datetime.now(timezone.utc).isoformat(),
"hub_cache": recorded_hub_cache,
}
return _atomic_write_json(path, payload)
@ -314,6 +410,8 @@ def read_cancel_marker_transport(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> Optional[str]:
"""Return the transport recorded in the cancel marker, or ``None`` if no
marker exists or it is unreadable.
@ -330,15 +428,17 @@ def read_cancel_marker_transport(
``None`` keeps the neutral "Retry" label.
* Unknown future versions → ``None`` (unknown layout, unknown transport).
"""
path = marker_path(repo_type, repo_id, variant)
path = _state_read_path(
marker_path,
repo_type,
repo_id,
variant,
hub_cache,
)
if path is None or not path.is_file():
return None
try:
data = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError) as exc:
logger.debug("Could not read cancel marker %s: %s", path, exc)
return None
if not isinstance(data, dict):
data = _read_state_payload(path)
if data is None:
return None
version = data.get("version")
if version == _LEGACY_MARKER_VERSION:
@ -351,10 +451,30 @@ def read_cancel_marker_transport(
return None
def _all_matching_state_paths(
parent: Optional[Path], repo_type: RepoType, repo_id: str, variant: Optional[str]
) -> tuple[Path, ...]:
if parent is None:
return ()
legacy_path = (
manifest_path(repo_type, repo_id, variant)
if parent.name == "manifests"
else marker_path(repo_type, repo_id, variant)
)
if legacy_path is None:
return ()
try:
return tuple(path for path in parent.rglob(legacy_path.name) if path.is_file())
except OSError:
return ()
def clear_cancel_marker(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> None:
"""Remove the cancel marker for this triple if present.
@ -362,31 +482,48 @@ def clear_cancel_marker(
download-start (a fresh attempt supersedes prior cancel state) and
again at successful completion (cleans up if the start clear failed).
"""
path = marker_path(repo_type, repo_id, variant)
if path is None:
return
try:
path.unlink(missing_ok = True)
except OSError as exc:
logger.debug("Could not clear cancel marker %s: %s", path, exc)
requested = _canonical_hub_cache(hub_cache)
path = marker_path(
repo_type,
repo_id,
variant,
hub_cache = requested,
)
legacy = marker_path(repo_type, repo_id, variant)
paths = [path]
if (
legacy is not None
and legacy != path
and _legacy_state_applies(legacy, requested, fail_closed = True)
):
paths.append(legacy)
for target in paths:
if target is None:
continue
try:
target.unlink(missing_ok = True)
except OSError as exc:
logger.debug("Could not clear cancel marker %s: %s", target, exc)
def has_cancel_marker(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> bool:
"""File-existence check only. Body is never read.
Fail-closed: a corrupt marker still returns ``True`` because the
file's existence is the signal (the user once cancelled this
triple, even if the body is unreadable).
"""
path = marker_path(repo_type, repo_id, variant)
if path is None:
return False
"""Return whether a cancel marker applies to the selected cache."""
path = _state_read_path(
marker_path,
repo_type,
repo_id,
variant,
hub_cache,
fail_closed = True,
)
try:
return path.is_file()
return path is not None and path.is_file()
except OSError:
return False
@ -395,48 +532,124 @@ def delete_manifest(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> bool:
path = manifest_path(repo_type, repo_id, variant)
if path is None:
return False
try:
if not path.is_file():
return False
path.unlink()
return True
except OSError as exc:
logger.debug("Could not delete manifest %s: %s", path, exc)
return False
requested = _canonical_hub_cache(hub_cache)
path = manifest_path(
repo_type,
repo_id,
variant,
hub_cache = requested,
)
legacy = manifest_path(repo_type, repo_id, variant)
paths = [path]
if legacy is not None and legacy != path and _legacy_state_applies(legacy, requested):
paths.append(legacy)
removed = False
for target in paths:
if target is None:
continue
try:
if target.is_file():
target.unlink()
removed = True
except OSError as exc:
logger.debug("Could not delete manifest %s: %s", target, exc)
return removed
def purge_state(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> bool:
"""Remove manifest + cancel marker for this triple. Returns ``True``
when anything was present on disk before the call. Idempotent."""
marker_existed = has_cancel_marker(repo_type, repo_id, variant)
manifest_removed = delete_manifest(repo_type, repo_id, variant)
clear_cancel_marker(repo_type, repo_id, variant)
return marker_existed or manifest_removed
when anything was present on disk before the call. Idempotent.
With ``hub_cache`` set, only that cache's scoped state (plus any legacy
unscoped file that belongs to it) is removed, so purging one cache's copy
never clears another cache's resumable/cancel state."""
if hub_cache is None:
paths = (
*_all_matching_state_paths(manifests_dir(), repo_type, repo_id, variant),
*_all_matching_state_paths(cancelled_dir(), repo_type, repo_id, variant),
)
else:
requested = _canonical_hub_cache(hub_cache)
candidates = [
manifest_path(repo_type, repo_id, variant, hub_cache = hub_cache),
marker_path(repo_type, repo_id, variant, hub_cache = hub_cache),
]
# Legacy unscoped state is shared: an unowned file belongs to the active
# cache (per _legacy_state_applies), so only purge it when it belongs to
# the cache being deleted -- else deleting an inactive cache would erase
# the active cache's resume/cancel state.
for path_factory in (manifest_path, marker_path):
legacy = path_factory(repo_type, repo_id, variant)
if legacy is not None and _legacy_state_applies(legacy, requested):
candidates.append(legacy)
paths = tuple(p for p in candidates if p is not None)
removed = False
for path in paths:
try:
if path.is_file():
path.unlink()
removed = True
except OSError as exc:
logger.debug("Could not purge Hub state %s: %s", path, exc)
return removed
def purge_all_state_for_repo(repo_type: RepoType, repo_id: str) -> int:
def purge_all_state_for_repo(
repo_type: RepoType,
repo_id: str,
*,
hub_cache: Optional[str | Path] = None,
) -> int:
"""Remove the snapshot-level manifest + marker AND every variant-keyed
manifest + marker for this repo. Used by the route delete handlers so
scanner state never outlives the cache it described. Returns the count
of (repo, variant) triples that had any state on disk."""
of (repo, variant) triples that had any state on disk.
With ``hub_cache`` set, only that cache's scoped state (plus any legacy
unscoped file) is enumerated and removed, so deleting one cache's copy does
not clear another cache's resumable/cancel state."""
removed = 0
if purge_state(repo_type, repo_id, None):
if purge_state(repo_type, repo_id, None, hub_cache = hub_cache):
removed += 1
variants: set[str] = set()
for variant, _ in iter_variant_manifests(repo_type, repo_id):
variants.add(variant)
for variant, _ in iter_variant_markers(repo_type, repo_id):
variants.add(variant)
prefix = variant_filename_prefix(repo_type, repo_id)
if hub_cache is None:
search = [(p, True) for p in (manifests_dir(), cancelled_dir()) if p is not None]
else:
# This cache's scoped dir (parent of its scoped path) plus the legacy
# unscoped base; glob (not rglob) so other caches' dirs are not swept.
search = []
for scoped, base in (
(manifest_path(repo_type, repo_id, None, hub_cache = hub_cache), manifests_dir()),
(marker_path(repo_type, repo_id, None, hub_cache = hub_cache), cancelled_dir()),
):
if scoped is not None:
search.append((scoped.parent, False))
if base is not None:
search.append((base, False))
for parent, recursive in search:
try:
entries = tuple(
parent.rglob(f"{prefix}*.json") if recursive else parent.glob(f"{prefix}*.json")
)
except OSError:
continue
for entry in entries:
if not entry.is_file():
continue
fallback = entry.stem[len(prefix) :]
variants.add(_variant_from_state_file(entry, fallback))
for variant in variants:
if purge_state(repo_type, repo_id, variant):
if purge_state(repo_type, repo_id, variant, hub_cache = hub_cache):
removed += 1
return removed
@ -453,35 +666,83 @@ def _variant_from_state_file(path: Path, fallback: str) -> str:
def _iter_variant_state_files(
parent: Optional[Path], repo_type: RepoType, repo_id: str
parent: Optional[Path],
repo_type: RepoType,
repo_id: str,
hub_cache: Optional[str | Path],
*,
cancel_markers: bool,
) -> Iterator[tuple[str, Path]]:
if parent is None:
return
prefix = variant_filename_prefix(repo_type, repo_id)
try:
entries = list(parent.iterdir())
except OSError:
path_factory = marker_path if cancel_markers else manifest_path
requested = _canonical_hub_cache(hub_cache)
scoped_probe = path_factory(
repo_type,
repo_id,
None,
hub_cache = requested,
)
if scoped_probe is None:
return
for entry in entries:
if not entry.is_file() or not entry.name.endswith(".json"):
prefix = variant_filename_prefix(repo_type, repo_id)
seen: set[str] = set()
for directory, legacy in ((scoped_probe.parent, False), (parent, True)):
if legacy and directory == scoped_probe.parent:
continue
stem = entry.name[: -len(".json")]
if not stem.lower().startswith(prefix):
try:
entries = list(directory.iterdir())
except OSError:
continue
variant = stem[len(prefix) :]
if variant:
yield _variant_from_state_file(entry, variant), entry
for entry in entries:
if not entry.is_file() or not entry.name.endswith(".json"):
continue
stem = entry.name[: -len(".json")]
if not stem.lower().startswith(prefix) or entry.name in seen:
continue
if legacy and not _legacy_state_applies(
entry,
requested,
fail_closed = cancel_markers,
):
continue
fallback = stem[len(prefix) :]
if fallback:
seen.add(entry.name)
yield _variant_from_state_file(entry, fallback), entry
def iter_variant_manifests(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]:
def iter_variant_manifests(
repo_type: RepoType,
repo_id: str,
*,
hub_cache: Optional[str | Path] = None,
) -> Iterator[tuple[str, Path]]:
"""Yield (variant, manifest_path) for every variant-keyed manifest
written for this repo. Used by is_gguf_repo_partial to enumerate all
variants present on disk so the all-variants-broken gate can run."""
yield from _iter_variant_state_files(manifests_dir(), repo_type, repo_id)
yield from _iter_variant_state_files(
manifests_dir(),
repo_type,
repo_id,
hub_cache,
cancel_markers = False,
)
def iter_variant_markers(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]:
def iter_variant_markers(
repo_type: RepoType,
repo_id: str,
*,
hub_cache: Optional[str | Path] = None,
) -> Iterator[tuple[str, Path]]:
"""Yield (variant, marker_path) for every variant-keyed cancel marker.
Companion to iter_variant_manifests: catches variants cancelled
before download-start ever wrote a manifest (very early failures)."""
yield from _iter_variant_state_files(cancelled_dir(), repo_type, repo_id)
yield from _iter_variant_state_files(
cancelled_dir(),
repo_type,
repo_id,
hub_cache,
cancel_markers = True,
)

119
studio/backend/hub/utils/download_registry.py
View file

@ -129,6 +129,8 @@ def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMeta
"cancel_marker_transport": metadata.cancel_marker_transport
if metadata is not None
else None,
"hub_cache": metadata.hub_cache if metadata is not None else None,
"xet_cache": metadata.xet_cache if metadata is not None else None,
}
tmp = path.with_name(f".{path.name}.tmp-{pid}")
try:
@ -236,6 +238,7 @@ def _settle_orphaned_download(
repo_id: Optional[str],
variant: Optional[str],
transport: Optional[str],
hub_cache: Optional[str] = None,
) -> None:
"""Persist a cancel marker for a reaped orphan still mid-download so the next
launch settles it to a resumable "cancelled" state instead of a phantom-running
@ -251,18 +254,42 @@ def _settle_orphaned_download(
return
from hub.utils import download_manifest
manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
cache_root = Path(hub_cache) if isinstance(hub_cache, str) and hub_cache else None
manifest = download_manifest.read_manifest(
repo_type,
repo_id,
variant,
hub_cache = cache_root,
)
if repo_type == "model" and variant and manifest is None:
return
if manifest is None:
if not has_active_incomplete_blobs(repo_type, repo_id):
if not has_active_incomplete_blobs(repo_type, repo_id, root = cache_root):
return
else:
if _manifest_verifies_against_active_cache(repo_type, repo_id, manifest):
if _manifest_verifies_against_active_cache(
repo_type,
repo_id,
manifest,
root = cache_root,
):
return
if not _manifest_has_active_incomplete_blobs(repo_type, repo_id, manifest):
if not _manifest_has_active_incomplete_blobs(
repo_type,
repo_id,
manifest,
root = cache_root,
):
return
persist_cancel_marker(repo_type, repo_id, variant, transport, logger = logger)
persist_cancel_marker(
repo_type,
repo_id,
variant,
transport,
hub_cache = hub_cache,
logger = logger,
)
def reap_orphan_workers() -> None:
@ -309,6 +336,7 @@ def reap_orphan_workers() -> None:
repo_id,
data.get("variant"),
data.get("cancel_marker_transport") or data.get("transport"),
data.get("hub_cache"),
)
except Exception as exc:
logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc)
@ -355,8 +383,13 @@ def _purge_incomplete_blobs(
return removed
def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]:
for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
def _iter_active_snapshot_dirs(
repo_type: str,
repo_id: str,
*,
root: Optional[Path] = None,
) -> Iterator[Path]:
for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root):
snapshots_dir = entry / "snapshots"
if not snapshots_dir.is_dir():
continue
@ -369,24 +402,41 @@ def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]:
yield snapshot
def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manifest) -> bool:
def _manifest_verifies_against_active_cache(
repo_type: str,
repo_id: str,
manifest,
*,
root: Optional[Path] = None,
) -> bool:
from hub.utils import download_manifest
for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id):
for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id, root = root):
if download_manifest.verify_against_disk(manifest, snapshot_dir).ok:
return True
return False
def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest) -> bool:
def _manifest_has_active_incomplete_blobs(
repo_type: str,
repo_id: str,
manifest,
*,
root: Optional[Path] = None,
) -> bool:
if not getattr(manifest, "variant", None):
return has_active_incomplete_blobs(repo_type, repo_id)
return has_active_incomplete_blobs(repo_type, repo_id, root = root)
expected_hashes = frozenset(
expected.sha256 for expected in manifest.expected_files if expected.sha256
)
if not expected_hashes:
return has_active_incomplete_blobs(repo_type, repo_id)
return has_active_incomplete_blobs(repo_type, repo_id, root = root)
return bool(
incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection(expected_hashes)
incomplete_blob_hashes(
repo_type,
repo_id,
active_only = True,
root = root,
).intersection(expected_hashes)
)
@ -459,6 +509,7 @@ def prepare_cache_for_transport(
only_blob_hashes: Optional[frozenset[str]] = None,
companion_blob_hashes: Optional[frozenset[str]] = None,
protected_blob_hashes: Optional[frozenset[str]] = None,
root: Optional[Path] = None,
) -> int:
"""Guarantee any pre-existing ``.incomplete`` blobs are SAFE to resume under
*mode*. Returns the number of partial blobs purged for untrusted provenance.
@ -485,14 +536,13 @@ def prepare_cache_for_transport(
they are excluded from every purge so a shared companion is never deleted
mid-write.
Scope: only the active ``HF_HUB_CACHE`` root is inspected. That suffices for
resume safety because ``snapshot_download`` runs without a ``cache_dir``
override and so can only read or resume a ``.incomplete`` under this same
active root. Markers are written for the new mode before returning.
Scope: ``root`` selects the cache captured by the caller. It defaults to the
active ``HF_HUB_CACHE`` root for workers that inherit their cache through
the environment. Markers are written for the new mode before returning.
"""
if mode not in VALID_TRANSPORTS:
raise ValueError(f"Invalid transport mode: {mode!r}")
root = hf_cache_root(create = True)
root = hf_cache_root(create = True) if root is None else hf_cache_root(create = True, root = root)
if root is None:
return 0
target = target_dir_name(repo_type, repo_id)
@ -618,10 +668,11 @@ def incomplete_blob_hashes(
repo_id: str,
*,
active_only: bool = False,
root: Optional[Path] = None,
) -> set[str]:
out: set[str] = set()
entries = (
iter_active_repo_cache_dirs(repo_type, repo_id)
iter_active_repo_cache_dirs(repo_type, repo_id, root = root)
if active_only
else iter_repo_cache_dirs(repo_type, repo_id)
)
@ -638,16 +689,24 @@ def incomplete_blob_hashes(
return out
def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int:
"""Sum finalized blob bytes for *blob_hashes* in the active HF cache root.
def completed_blob_bytes(
repo_type: str,
repo_id: str,
blob_hashes: frozenset[str],
*,
root: Optional[Path] = None,
) -> int:
"""Sum finalized blob bytes for *blob_hashes* in a single HF cache root.
A worker only writes to the active ``HF_HUB_CACHE`` root, so a baseline must
ignore legacy/default roots that ``snapshot_download`` won't reuse this run.
A worker only writes to its captured ``HF_HUB_CACHE`` root, so a baseline
must be scoped to that root (``root``), not re-resolved to whatever cache is
active now; otherwise a runtime cache switch makes the retry baseline count
bytes from the wrong disk.
"""
if not blob_hashes:
return 0
total = 0
for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root):
blobs_dir = entry / "blobs"
if not blobs_dir.is_dir():
continue
@ -712,6 +771,8 @@ class DownloadMetadata:
# Bytes already complete before this job started; not counted as this run's
# progress.
completed_baseline_bytes: int = 0
hub_cache: Optional[str] = None
xet_cache: Optional[str] = None
@dataclass(frozen = True)
@ -752,6 +813,7 @@ def persist_cancel_marker(
variant: Optional[str],
transport: Optional[str],
*,
hub_cache: Optional[str] = None,
logger = logger,
) -> None:
if not repo_type or not repo_id:
@ -763,6 +825,7 @@ def persist_cancel_marker(
repo_id,
variant,
transport = transport,
hub_cache = hub_cache,
):
logger.debug("write_cancel_marker returned False for %s", repo_id)
except Exception as exc:
@ -971,6 +1034,7 @@ class DownloadRegistry:
metadata_to_persist.repo_id,
metadata_to_persist.variant,
metadata_to_persist.transport,
hub_cache = metadata_to_persist.hub_cache,
)
return False
@ -1033,6 +1097,8 @@ class DownloadRegistry:
replace_active: bool = False,
metadata_transport: Optional[str] = None,
cancel_marker_transport: Optional[str] = None,
hub_cache: Optional[str] = None,
xet_cache: Optional[str] = None,
) -> tuple[bool, str]:
key = normalize_job_key(key)
repo = _repo_of_key(key)
@ -1106,6 +1172,8 @@ class DownloadRegistry:
0,
int(completed_baseline_bytes or 0),
),
hub_cache = hub_cache,
xet_cache = xet_cache,
)
if cancel_marker_transport is not None:
self._cancel_marker_transports[key] = cancel_marker_transport
@ -1386,6 +1454,7 @@ class DownloadRegistry:
metadata.repo_id,
metadata.variant,
metadata.cancel_marker_transport or metadata.transport,
hub_cache = metadata.hub_cache,
)
reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = []
for key, proc, metadata in live:
@ -1401,6 +1470,7 @@ class DownloadRegistry:
metadata.repo_id,
metadata.variant,
metadata.cancel_marker_transport or metadata.transport,
hub_cache = metadata.hub_cache,
)
continue
reaped.append((key, proc, metadata))
@ -1421,6 +1491,7 @@ class DownloadRegistry:
metadata.repo_id,
metadata.variant,
metadata.cancel_marker_transport or metadata.transport,
hub_cache = metadata.hub_cache,
)

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

@ -253,11 +253,16 @@ def _env_offline() -> bool:
) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
def iter_hf_cache_snapshots(repo_id: str):
from hub.utils.hf_cache_state import iter_repo_cache_dirs
def iter_hf_cache_snapshots(repo_id: str, root: Optional[Path] = None):
from hub.utils.hf_cache_state import iter_active_repo_cache_dirs, iter_repo_cache_dirs
snapshots: list[Path] = []
for repo_dir in iter_repo_cache_dirs("model", repo_id):
repo_dirs = (
iter_active_repo_cache_dirs("model", repo_id, root = root)
if root is not None
else iter_repo_cache_dirs("model", repo_id)
)
for repo_dir in repo_dirs:
snapshots_dir = repo_dir / "snapshots"
if not snapshots_dir.is_dir():
continue
@ -276,12 +281,17 @@ def iter_hf_cache_snapshots(repo_id: str):
yield from snapshots
def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]:
def list_empty_gguf_variant_dirs(repo_id: str, root: Optional[Path] = None) -> set[str]:
"""Quant labels present only as an EMPTY snapshot ``<quant>/`` folder (an
interrupted split download); a quant with shards in any snapshot is excluded."""
empty: dict[str, str] = {}
nonempty: set[str] = set()
for snapshot in iter_hf_cache_snapshots(repo_id):
snapshots = (
iter_hf_cache_snapshots(repo_id, root = root)
if root is not None
else iter_hf_cache_snapshots(repo_id)
)
for snapshot in snapshots:
try:
entries = list(snapshot.iterdir())
except OSError:
@ -303,8 +313,15 @@ def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]:
return {label for key, label in empty.items() if key not in nonempty}
def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]:
for snapshot in iter_hf_cache_snapshots(repo_id):
def list_gguf_variants_from_hf_cache(
repo_id: str, root: Optional[Path] = None
) -> Optional[tuple[list[GgufVariantInfo], bool]]:
snapshots = (
iter_hf_cache_snapshots(repo_id, root = root)
if root is not None
else iter_hf_cache_snapshots(repo_id)
)
for snapshot in snapshots:
variants, has_vision = list_local_gguf_variants(str(snapshot))
if variants or has_vision:
return variants, has_vision
@ -312,7 +329,7 @@ def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVa
def list_partial_gguf_variants_from_state(
repo_id: str,
repo_id: str, hub_cache: Optional[Path] = None
) -> Optional[tuple[list[GgufVariantInfo], bool]]:
"""Reconstruct GGUF variants from download manifests/markers alone.
@ -328,10 +345,26 @@ def list_partial_gguf_variants_from_state(
# original-casing label over a lowercased cancel marker for the same variant.
seen: set[str] = set()
ordered: list[str] = []
for source in (
download_manifest.iter_variant_manifests("model", repo_id),
download_manifest.iter_variant_markers("model", repo_id),
):
sources = (
(
download_manifest.iter_variant_manifests("model", repo_id),
download_manifest.iter_variant_markers("model", repo_id),
)
if hub_cache is None
else (
download_manifest.iter_variant_manifests(
"model",
repo_id,
hub_cache = hub_cache,
),
download_manifest.iter_variant_markers(
"model",
repo_id,
hub_cache = hub_cache,
),
)
)
for source in sources:
for variant, _path in source:
key = variant.lower()
if key not in seen:
@ -343,7 +376,16 @@ def list_partial_gguf_variants_from_state(
variants: list[GgufVariantInfo] = []
has_vision = False
for variant in ordered:
manifest = download_manifest.read_manifest("model", repo_id, variant)
manifest = (
download_manifest.read_manifest("model", repo_id, variant)
if hub_cache is None
else download_manifest.read_manifest(
"model",
repo_id,
variant,
hub_cache = hub_cache,
)
)
main_filename: Optional[str] = None
size_bytes = 0
companion_bytes = 0

127
studio/backend/hub/utils/hf_cache_state.py
View file

@ -29,12 +29,10 @@ def _safe_is_dir(path: Path) -> bool:
return False
def hf_cache_root(*, create: bool = False) -> Optional[Path]:
try:
from huggingface_hub import constants as hf_constants
except ImportError:
return None
root = Path(hf_constants.HF_HUB_CACHE)
def hf_cache_root(*, create: bool = False, root: Optional[Path] = None) -> Optional[Path]:
from utils.hf_cache_settings import get_hf_cache_paths
root = root or get_hf_cache_paths().hub_cache
if create:
try:
root.mkdir(parents = True, exist_ok = True)
@ -46,6 +44,7 @@ def hf_cache_root(*, create: bool = False) -> Optional[Path]:
def hf_cache_roots() -> list[Path]:
from hub.utils.paths import hf_default_cache_dir, legacy_hf_cache_dir
from utils.hf_cache_settings import known_hf_hub_caches
roots: list[Path] = []
seen: set[str] = set()
@ -62,7 +61,8 @@ def hf_cache_roots() -> list[Path]:
seen.add(key)
roots.append(path)
_add(hf_cache_root())
for configured in known_hf_hub_caches():
_add(configured)
_add(legacy_hf_cache_dir())
_add(hf_default_cache_dir())
return roots
@ -181,12 +181,22 @@ def iter_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]:
continue
def iter_destructive_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]:
def iter_destructive_repo_cache_dirs(
repo_type: str,
repo_id: str,
*,
root: Optional[Path] = None,
) -> Iterator[Path]:
target = repo_cache_dir_name(repo_type, repo_id)
folded_target = target.lower()
for root in hf_cache_roots():
if root is not None:
scoped = hf_cache_root(root = root)
bases = [scoped] if scoped is not None else []
else:
bases = hf_cache_roots()
for base in bases:
try:
entries = [entry for entry in root.iterdir() if entry.name.lower() == folded_target]
entries = [entry for entry in base.iterdir() if entry.name.lower() == folded_target]
except OSError:
continue
matched_names = resolve_destructive_case_matches(
@ -200,8 +210,13 @@ def iter_destructive_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[P
yield entry
def iter_active_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]:
root = hf_cache_root()
def iter_active_repo_cache_dirs(
repo_type: str,
repo_id: str,
*,
root: Optional[Path] = None,
) -> Iterator[Path]:
root = hf_cache_root(root = root)
if root is None:
return
target = target_dir_name(repo_type, repo_id)
@ -218,12 +233,13 @@ def preferred_repo_cache_dirs(
repo_id: str,
*,
force_active: bool = False,
active_root: Optional[Path] = None,
) -> list[Path]:
active_entries = list(iter_active_repo_cache_dirs(repo_type, repo_id))
active_entries = list(iter_active_repo_cache_dirs(repo_type, repo_id, root = active_root))
if active_entries:
return active_entries
if force_active:
root = hf_cache_root()
root = hf_cache_root(root = active_root)
if root is not None:
canonical = repo_cache_dir_name(repo_type, repo_id)
return [root / canonical]
@ -237,8 +253,13 @@ def has_incomplete_blobs(repo_type: str, repo_id: str) -> bool:
return False
def has_active_incomplete_blobs(repo_type: str, repo_id: str) -> bool:
for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
def has_active_incomplete_blobs(
repo_type: str,
repo_id: str,
*,
root: Optional[Path] = None,
) -> bool:
for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root):
if repo_cache_dir_has_incomplete_blobs(entry):
return True
return False
@ -273,9 +294,14 @@ def _prune_empty_dirs(root: Path) -> bool:
return removed
def purge_partial_repo(repo_type: str, repo_id: str) -> bool:
def purge_partial_repo(
repo_type: str,
repo_id: str,
*,
root: Optional[Path] = None,
) -> bool:
removed = False
for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id):
for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id, root = root):
blobs_dir = entry / "blobs"
if blobs_dir.is_dir():
for blob in blobs_dir.iterdir():
@ -290,9 +316,14 @@ def purge_partial_repo(repo_type: str, repo_id: str) -> bool:
return removed
def purge_repo_cache_dirs(repo_type: str, repo_id: str) -> bool:
def purge_repo_cache_dirs(
repo_type: str,
repo_id: str,
*,
root: Optional[Path] = None,
) -> bool:
removed = False
for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id):
for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id, root = root):
try:
if entry.is_symlink() or not entry.is_dir():
continue
@ -301,3 +332,59 @@ def purge_repo_cache_dirs(repo_type: str, repo_id: str) -> bool:
except FileNotFoundError:
continue
return removed
def scoped_delete_root(repo_type: str, repo_id: str, cache_path: Optional[str]) -> Optional[Path]:
"""Resolve the single cache root a delete of this repo may touch.
Returns the active hub cache when *cache_path* is falsy, the owning cache
root when *cache_path* points inside a known cache, or ``None`` when
*cache_path* is set but not inside any known cache (caller should reject).
This keeps a delete of one inventory row from removing copies in other,
previously selected caches.
"""
from utils.hf_cache_settings import get_hf_cache_paths
if not cache_path:
return Path(get_hf_cache_paths().hub_cache).resolve(strict = False)
try:
resolved = Path(cache_path).expanduser().resolve(strict = False)
except (OSError, RuntimeError, ValueError):
return None
expected = repo_cache_dir_name(repo_type, repo_id).lower()
repo_dir = next(
(
candidate
for candidate in (resolved, *resolved.parents)
if candidate.name.lower() == expected
),
None,
)
if repo_dir is None:
return None
allowed = {r.resolve(strict = False) for r in hf_cache_roots()}
root = repo_dir.parent.resolve(strict = False)
return root if root in allowed else None
def resolve_delete_target_root(
repo_type: str, repo_id: str, cache_path: Optional[str], owner_roots
) -> Optional[Path]:
"""Pick the single cache root a delete of this repo should target.
An explicit *cache_path* wins (``None`` when it is not a known cache, so the
caller can reject it). Otherwise prefer the active cache when it holds a
copy, else the sole cache that does -- so a model that lives only in a
previously selected cache stays deletable while other caches are untouched.
"""
if cache_path:
return scoped_delete_root(repo_type, repo_id, cache_path)
from utils.hf_cache_settings import get_hf_cache_paths
active = Path(get_hf_cache_paths().hub_cache).resolve(strict = False)
roots = list(owner_roots)
if active in roots:
return active
if len(roots) == 1:
return roots[0]
return active

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

@ -36,7 +36,7 @@ from hub.utils.state_dir import RepoType
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
has_incomplete_blobs,
hf_cache_root,
hf_cache_roots,
iter_repo_cache_dirs,
latest_snapshot_dir,
repo_cache_dir_has_incomplete_blobs,
@ -127,33 +127,13 @@ def all_hf_cache_scans() -> list:
def _compute_all_hf_cache_scans() -> list:
from huggingface_hub import scan_cache_dir
from hub.utils.paths import legacy_hf_cache_dir, hf_default_cache_dir
scans: list = []
seen: set[str] = set()
try:
from huggingface_hub.constants import HF_HUB_CACHE
active = Path(HF_HUB_CACHE).resolve()
seen.add(str(active))
if active.is_dir():
scans.append(scan_cache_dir())
except Exception as exc:
logger.warning("Could not scan active HF cache: %s", exc)
for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir):
for cache_root in hf_cache_roots():
try:
extra = extra_fn()
# is_dir()/resolve() can raise on an inaccessible path; skip it.
if not extra.is_dir():
continue
resolved = str(extra.resolve())
if resolved in seen:
continue
seen.add(resolved)
scans.append(scan_cache_dir(cache_dir = str(extra)))
scans.append(scan_cache_dir(cache_dir = str(cache_root)))
except Exception as exc:
logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc)
logger.warning("Could not scan HF cache %s: %s", cache_root, exc)
return scans
@ -224,16 +204,8 @@ def _compose_partial(*signals: Callable[[], bool]) -> bool:
return any(signal() for signal in signals)
def _state_applies_to_repo_cache_dir(repo_cache_dir: Optional[Path]) -> bool:
if repo_cache_dir is None:
return True
root = hf_cache_root()
if root is None:
return False
try:
return repo_cache_dir.resolve().parent == root.resolve()
except OSError:
return False
def _hub_cache_for_repo_dir(repo_cache_dir: Optional[Path]) -> Optional[Path]:
return repo_cache_dir.parent if repo_cache_dir is not None else None
def _legacy_partial(
@ -285,12 +257,24 @@ def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir: Path)
return False
def _gguf_variant_manifest_blob_hashes(repo_id: str) -> frozenset[str]:
def _gguf_variant_manifest_blob_hashes(
repo_id: str, repo_cache_dir: Optional[Path] = None
) -> frozenset[str]:
from hub.utils import download_manifest
hashes: set[str] = set()
for variant, _path in download_manifest.iter_variant_manifests("model", repo_id):
manifest = download_manifest.read_manifest("model", repo_id, variant)
hub_cache = _hub_cache_for_repo_dir(repo_cache_dir)
for variant, _path in download_manifest.iter_variant_manifests(
"model",
repo_id,
hub_cache = hub_cache,
):
manifest = download_manifest.read_manifest(
"model",
repo_id,
variant,
hub_cache = hub_cache,
)
if manifest is None:
continue
for expected in manifest.expected_files:
@ -315,7 +299,7 @@ def _snapshot_legacy_partial(
) -> bool:
if repo_type != "model":
return _legacy_partial(repo_type, repo_id, repo_cache_dir)
ignored_hashes = _gguf_variant_manifest_blob_hashes(repo_id)
ignored_hashes = _gguf_variant_manifest_blob_hashes(repo_id, repo_cache_dir)
if repo_cache_dir is not None:
return _repo_cache_dir_has_snapshot_legacy_partial(
repo_cache_dir,
@ -375,9 +359,12 @@ def _manifest_partial(
) -> bool:
from hub.utils import download_manifest
if not _state_applies_to_repo_cache_dir(repo_cache_dir):
return False
manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
manifest = download_manifest.read_manifest(
repo_type,
repo_id,
variant,
hub_cache = _hub_cache_for_repo_dir(repo_cache_dir),
)
if manifest is None:
return False
resolved = (
@ -452,10 +439,13 @@ def is_snapshot_partial(
A manifest without a resolvable snapshot is partial: the worker got
far enough to record expectations but did not leave a usable snapshot."""
from hub.utils import download_manifest
state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir)
return _compose_partial(
lambda: state_applies and download_manifest.has_cancel_marker(repo_type, repo_id, None),
lambda: download_manifest.has_cancel_marker(
repo_type,
repo_id,
None,
hub_cache = _hub_cache_for_repo_dir(repo_cache_dir),
),
lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir),
lambda: _manifest_partial(
repo_type,
@ -484,10 +474,13 @@ def is_variant_partial(
caller is checking many variants of the same repo (see
is_gguf_repo_partial for that usage)."""
from hub.utils import download_manifest
state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir)
return _compose_partial(
lambda: state_applies and download_manifest.has_cancel_marker("model", repo_id, variant),
lambda: download_manifest.has_cancel_marker(
"model",
repo_id,
variant,
hub_cache = _hub_cache_for_repo_dir(repo_cache_dir),
),
lambda: bool(
incomplete_blob_hashes
and variant_blob_hashes
@ -526,22 +519,38 @@ def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) ->
from hub.utils import download_manifest
has_legacy_partial = _legacy_partial("model", repo_id, repo_cache_dir)
state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir)
snapshot_dir = resolve_snapshot_dir_for_scan(
"model",
repo_id,
repo_cache_dir,
)
variants: set[str] = set(_completed_gguf_variants(snapshot_dir))
if state_applies:
for variant, _path in download_manifest.iter_variant_manifests(
"model",
repo_id,
hub_cache = _hub_cache_for_repo_dir(repo_cache_dir)
for variant, _path in download_manifest.iter_variant_manifests(
"model",
repo_id,
hub_cache = hub_cache,
):
if (
download_manifest.read_manifest(
"model",
repo_id,
variant,
hub_cache = hub_cache,
)
is not None
):
variants.add(variant)
for variant, _path in download_manifest.iter_variant_markers(
for variant, _path in download_manifest.iter_variant_markers(
"model",
repo_id,
hub_cache = hub_cache,
):
if download_manifest.has_cancel_marker(
"model",
repo_id,
variant,
hub_cache = hub_cache,
):
variants.add(variant)
if not variants:
@ -576,14 +585,19 @@ def partial_transport_for(
available."""
from hub.utils import download_manifest
if not _state_applies_to_repo_cache_dir(repo_cache_dir):
return None
hub_cache = _hub_cache_for_repo_dir(repo_cache_dir)
marker_transport = download_manifest.read_cancel_marker_transport(
repo_type,
repo_id,
variant,
hub_cache = hub_cache,
)
if marker_transport is not None:
return marker_transport
manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
manifest = download_manifest.read_manifest(
repo_type,
repo_id,
variant,
hub_cache = hub_cache,
)
return manifest.transport if manifest is not None else None

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

@ -277,12 +277,8 @@ def _memo_drop(memo_key: tuple[str, str]) -> None:
def _hf_hub_cache_dir() -> Path:
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception as exc:
logger.debug("Could not read huggingface_hub HF_HUB_CACHE, using default: %s", exc)
return Path.home() / ".cache" / "huggingface" / "hub"
from utils.hf_cache_settings import get_hf_cache_paths
return get_hf_cache_paths().hub_cache
def _hf_hub_cache_dirs() -> list[Path]:
@ -300,7 +296,10 @@ def _hf_hub_cache_dirs() -> list[Path]:
seen.add(key)
roots.append(resolved)
_add(_hf_hub_cache_dir())
from utils.hf_cache_settings import known_hf_hub_caches
for configured in known_hf_hub_caches():
_add(configured)
try:
_add(legacy_hf_cache_dir())
_add(hf_default_cache_dir())

31
studio/backend/hub/utils/state_dir.py
View file

@ -8,9 +8,10 @@ so it survives ``huggingface-cli delete-cache`` and any other HF-side
cache lifecycle. Two subdirectories:
<studio cache>/hub-state/
manifests/ <key>.json per-download expected-files manifest
cancelled/ <key>.json per-download cancel marker
manifests/cache-<digest>/<key>.json expected-files manifest
cancelled/cache-<digest>/<key>.json cancel marker
The cache digest isolates state for the same repo across selectable Hub caches.
The ``<key>`` mirrors HF's cache dir naming while the resulting manifest,
cancel-marker, and atomic-write temp filenames fit common filesystem basename
limits. Very long repo IDs use a stable hash in the state key:
@ -29,6 +30,7 @@ configuration failure.
from __future__ import annotations
import hashlib
import os
import re
from pathlib import Path
from typing import Literal, Optional, get_args
@ -55,6 +57,7 @@ _STATE_EXTENSION = ".json"
# _atomic_write_json writes ".<target>.tmp-<8hex>" beside the final file.
_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8
_MAX_VARIANT_FRAGMENT_LENGTH = 64
_CACHE_SCOPE_DIGEST_LENGTH = 32
def state_root() -> Optional[Path]:
@ -130,13 +133,32 @@ def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str
return f"{variant_filename_prefix(repo_type, repo_id)}{variant_fragment}"
def _cache_scope(parent: Path, hub_cache: Optional[str | Path]) -> Optional[Path]:
if hub_cache is None:
return parent
normalized = os.path.normcase(str(Path(hub_cache).expanduser()))
digest = hashlib.sha256(normalized.encode("utf-8")).hexdigest()[:_CACHE_SCOPE_DIGEST_LENGTH]
scoped = parent / f"cache-{digest}"
try:
scoped.mkdir(parents = True, exist_ok = True)
except OSError as exc:
logger.debug("Could not create cache-scoped Hub state dir %s: %s", scoped, exc)
return None
return scoped
def manifest_path(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> Optional[Path]:
"""Path to the manifest file for this triple. May or may not exist."""
parent = _subdir(_MANIFESTS_SUBDIR)
if parent is None:
return None
parent = _cache_scope(parent, hub_cache)
if parent is None:
return None
return parent / f"{_entry_key(repo_type, repo_id, variant)}.json"
@ -146,9 +168,14 @@ def marker_path(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
*,
hub_cache: Optional[str | Path] = None,
) -> Optional[Path]:
"""Path to the cancel-marker file for this triple. May or may not exist."""
parent = _subdir(_CANCELLED_SUBDIR)
if parent is None:
return None
parent = _cache_scope(parent, hub_cache)
if parent is None:
return None
return parent / f"{_entry_key(repo_type, repo_id, variant)}.json"

1
studio/backend/hub/workers/hf_download.py
View file

@ -661,6 +661,7 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod
variant,
plan.main_hashes,
hf_token,
hub_cache = Path(snapshot_path).parents[2],
)
except Exception as e:
print(

64
studio/backend/main.py
View file

@ -289,6 +289,7 @@ from fastapi import Depends, FastAPI, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, Response
from starlette.middleware.gzip import GZipMiddleware
from pathlib import Path
from datetime import datetime
@ -308,12 +309,14 @@ from routes import (
training_router,
)
from routes.llama import router as llama_router
from routes.whisper import router as whisper_router
from routes.preview import router as preview_router
from hub.routes import (
inventory_router as hub_inventory_router,
datasets_router as hub_datasets_router,
token_router as hub_token_router,
)
from picker.routes import templates_router as picker_templates_router
from hub.schemas.downloads import TransportCapabilities
from hub.utils.download_registry import (
get_download_transport_capabilities,
@ -753,6 +756,8 @@ app.add_middleware(SecurityHeadersMiddleware)
# headroom; non-upload routes keep the default body cap.
import json as _json_for_413 # noqa: E402
from utils.upload_limits import ( # noqa: E402
STT_AUDIO_JSON_MAX_BYTES,
STT_AUDIO_RAW_MAX_BYTES,
UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES,
default_request_body_limit_bytes,
upload_request_limit_bytes,
@ -763,6 +768,7 @@ _BODY_PROTECTED_PREFIXES = (
"/v1/completions",
"/p/",
"/api/inference",
"/api/picker",
"/api/data-recipe",
"/api/datasets",
"/api/hub",
@ -790,6 +796,14 @@ def _get_upload_passthrough_request_max_bytes(path: str) -> int:
return default_request_body_limit_bytes()
def _get_request_body_max_bytes(path: str) -> int:
if path.startswith("/api/inference/audio/transcribe/raw"):
return STT_AUDIO_RAW_MAX_BYTES
if path.startswith("/api/inference/audio/transcribe"):
return STT_AUDIO_JSON_MAX_BYTES
return default_request_body_limit_bytes()
async def _send_411(send) -> None:
payload = _json_for_413.dumps(
{"detail": "Content-Length required for upload requests."},
@ -832,12 +846,14 @@ class MaxBodyMiddleware:
app,
max_bytes_getter,
protected_prefixes: tuple,
request_max_bytes_getter = None,
upload_passthrough_prefixes: tuple = (),
upload_passthrough_max_bytes_getter = None,
):
self.app = app
self.max_bytes_getter = max_bytes_getter
self.protected_prefixes = protected_prefixes
self.request_max_bytes_getter = request_max_bytes_getter
self.upload_passthrough_prefixes = upload_passthrough_prefixes
self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter
@ -854,6 +870,14 @@ class MaxBodyMiddleware:
except Exception:
return int(self.max_bytes_getter())
def _request_max_bytes(self, path: str) -> int:
if self.request_max_bytes_getter is None:
return int(self.max_bytes_getter())
try:
return int(self.request_max_bytes_getter(path))
except Exception:
return int(self.max_bytes_getter())
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.app(scope, receive, send)
@ -866,7 +890,7 @@ class MaxBodyMiddleware:
await self.app(scope, receive, send)
return
max_bytes = int(self.max_bytes_getter())
max_bytes = self._request_max_bytes(path)
declared = None
for name, value in scope.get("headers", []):
if name == b"content-length":
@ -931,6 +955,7 @@ app.add_middleware(
MaxBodyMiddleware,
max_bytes_getter = default_request_body_limit_bytes,
protected_prefixes = _BODY_PROTECTED_PREFIXES,
request_max_bytes_getter = _get_request_body_max_bytes,
upload_passthrough_prefixes = _BODY_UPLOAD_PASSTHROUGH_PREFIXES,
upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes,
)
@ -989,11 +1014,13 @@ app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"])
app.include_router(whisper_router, prefix = "/api/whisper", tags = ["whisper"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"])
app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"])
app.include_router(picker_templates_router, prefix = "/api/picker", tags = ["picker"])
app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"])
# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic
@ -1509,6 +1536,34 @@ def _should_inject_bootstrap(request: Request) -> bool:
return _is_local_bootstrap_request(request)
_IMMUTABLE_ASSET_CACHE_CONTROL = "public, max-age=31536000, immutable"
class ImmutableStaticFiles(StaticFiles):
"""Serve Vite's content-hashed assets without browser revalidation."""
def file_response(
self,
full_path,
stat_result,
scope,
status_code = 200,
):
response = super().file_response(full_path, stat_result, scope, status_code)
response.headers["Cache-Control"] = _IMMUTABLE_ASSET_CACHE_CONTROL
return response
class _AssetGZipMiddleware(GZipMiddleware):
"""Serve range requests uncompressed; gzip + 206 mislabels Content-Range."""
async def __call__(self, scope, receive, send):
if scope["type"] == "http" and any(key == b"range" for key, _ in scope["headers"]):
await self.app(scope, receive, send)
return
await super().__call__(scope, receive, send)
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@ -1516,7 +1571,12 @@ def setup_frontend(app: FastAPI, build_path: Path):
assets_dir = build_path / "assets"
if assets_dir.exists():
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
assets_app = _AssetGZipMiddleware(
ImmutableStaticFiles(directory = assets_dir),
minimum_size = 1024,
compresslevel = 6,
)
app.mount("/assets", assets_app, name = "assets")
def _build_index_response(request: Request) -> Response:
content = (build_path / "index.html").read_bytes()

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

@ -18,6 +18,8 @@ from pydantic import (
model_validator,
)
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
@ -54,8 +56,16 @@ class LoadRequest(BaseModel):
@field_validator("chat_template_override")
@classmethod
def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]:
if value is not None and value.strip() == "":
if value is None:
return None
# Char count is a lower bound on UTF-8 byte length: reject an oversized
# template before spending work encoding it.
if len(value) > MAX_CHAT_TEMPLATE_BYTES:
raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
if value.strip() == "":
return None
if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
return value
cache_type_kv: Optional[str] = Field(
@ -177,6 +187,32 @@ class UnloadRequest(BaseModel):
model_path: str = Field(..., description = "Model identifier to unload")
class TranscribeRequest(BaseModel):
"""Speech-to-text request for the dictation STT sidecar."""
audio: str = Field(..., description = "Base64-encoded audio (any common format)")
model: Optional[str] = Field(None, description = "STT model id; defaults server-side")
language: Optional[str] = Field(None, description = "BCP-47 language, or 'auto'/None to detect")
fast: bool = Field(
False,
description = "Use low-latency single-candidate decoding for dictation",
)
engine: Optional[str] = Field(
None,
description = "STT engine: 'transformers' (default) or 'gguf' (whisper.cpp)",
)
class SttLoadRequest(BaseModel):
"""Warm the STT sidecar with a model without transcribing."""
model: Optional[str] = Field(None, description = "STT model id; defaults server-side")
engine: Optional[str] = Field(
None,
description = "STT engine: 'transformers' (default) or 'gguf' (whisper.cpp)",
)
class ValidateModelRequest(BaseModel):
"""Check whether an identifier resolves to a ModelConfig; does NOT load weights."""
@ -206,6 +242,13 @@ class ValidateModelRequest(BaseModel):
description = "Also read the native context length from the local GGUF header. "
"Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.",
)
include_chat_template: bool = Field(
False,
description = "Also read the embedded chat template from the local GGUF header, so a "
"native (picked / drag-drop) file's default template can be shown before it is loaded. "
"Opt-in and, like include_context_length, a metadata-only probe that skips the training "
"guard. Only the leased file's own embedded template is read, never sibling sidecars.",
)
class TransformersUpgradeInfo(BaseModel):
@ -266,6 +309,11 @@ class ValidateModelResponse(BaseModel):
description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF "
"header alongside context_length; 0 for dense models, None when not read.",
)
chat_template: Optional[str] = Field(
None,
description = "Embedded GGUF chat template, read from the header when include_chat_template "
"is set (native lease-backed picks); None for non-GGUF, over-cap, or not-read templates.",
)
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,

8
studio/backend/models/models.py
View file

@ -178,6 +178,14 @@ class LocalModelInfo(BaseModel):
None,
description = "HF repo id for cached models, e.g. org/model",
)
active_cache: Optional[bool] = Field(
None,
description = "Whether an HF model belongs to the current download cache.",
)
partial: bool = Field(
False,
description = "Whether the cached model has an incomplete download.",
)
model_format: Optional[str] = Field(
None,
description = "Detected weights format ('gguf' when known). Lets the UI "

7
studio/backend/models/training.py
View file

@ -505,6 +505,13 @@ class TrainingStartRequest(BaseModel):
description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.",
)
@field_validator("target_modules", mode = "before")
@classmethod
def _normalize_target_modules(cls, value: Any) -> Any:
# Sanitized non-LoRA history stores the unused value as null; treat it as a
# fresh request's omitted/default empty list on resume.
return [] if value is None else value
@model_validator(mode = "after")
def _validate_streaming_splits(self) -> "TrainingStartRequest":
# Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]"

2
studio/backend/picker/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

6
studio/backend/picker/routes/__init__.py Normal file
View file

@ -0,0 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from .templates import router as templates_router
__all__ = ["templates_router"]

45
studio/backend/picker/routes/templates.py Normal file
View file

@ -0,0 +1,45 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import asyncio
from typing import Optional
from fastapi import APIRouter, Body, Depends, Query
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
from ..schemas import (
MAX_CHAT_TEMPLATE_BYTES,
ModelTemplateResponse,
ValidateChatTemplateRequest,
ValidateChatTemplateResponse,
)
from ..service import read_default_chat_template, validate_chat_template
router = APIRouter()
@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse)
async def validate_chat_template_route(
body: ValidateChatTemplateRequest = Body(...),
current_subject: str = Depends(get_current_subject),
) -> ValidateChatTemplateResponse:
return await asyncio.to_thread(validate_chat_template, body.template)
@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse)
async def get_default_chat_template_route(
model_name: str,
gguf_variant: Optional[str] = Query(None),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
) -> ModelTemplateResponse:
template = await asyncio.to_thread(
read_default_chat_template, model_name, hf_token, gguf_variant
)
if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
template = None
return ModelTemplateResponse(model_name = model_name, chat_template = template)

32
studio/backend/picker/schemas.py Normal file
View file

@ -0,0 +1,32 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from typing import Optional
from pydantic import BaseModel, Field, field_validator
# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at
# the API boundary so a direct caller cannot make Jinja parse an oversized
# template. MaxBodyMiddleware only caps the whole request body, not this field.
MAX_CHAT_TEMPLATE_BYTES = 65_536
class ValidateChatTemplateRequest(BaseModel):
template: str = Field(default = "")
@field_validator("template")
@classmethod
def _enforce_template_size(cls, value: str) -> str:
if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
return value
class ValidateChatTemplateResponse(BaseModel):
valid: bool
error: Optional[str] = None
class ModelTemplateResponse(BaseModel):
model_name: str
chat_template: Optional[str] = None

432
studio/backend/picker/service.py Normal file
View file

@ -0,0 +1,432 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import json
import logging
import os
import re
from pathlib import Path
from typing import Optional
from hub.services.models.folder_browser import (
_build_browse_allowlist,
_is_path_inside_allowlist,
)
from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots
from utils.models.gguf_metadata import read_gguf_chat_template
from utils.models.model_config import (
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mmproj,
_is_mtp_drafter,
)
from utils.hf_cache_settings import active_hf_hub_cache
from utils.paths.path_utils import (
is_local_path,
normalize_path,
resolve_cached_repo_id_case,
)
from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse
logger = logging.getLogger(__name__)
_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
def _is_valid_repo_id(repo_id: str) -> bool:
return bool(_VALID_REPO_ID.fullmatch(repo_id))
_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json")
_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja")
_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json")
# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory
# before its template is size-checked. The JSON envelope may exceed a bare template
# (it carries other tokenizer metadata); the extracted template is still bounded by
# MAX_CHAT_TEMPLATE_BYTES downstream.
MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024
def _read_bounded_text(path: Path, limit: int) -> Optional[str]:
"""Read at most `limit` bytes of UTF-8 text; None if larger or unreadable."""
try:
with path.open("rb") as f:
data = f.read(limit + 1)
except OSError:
return None
if len(data) > limit:
return None
try:
return data.decode("utf-8")
except UnicodeError:
return None
def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool:
# Block symlinked children from escaping the validated directory (realpath-checked).
# None = trusted caller (HF cache / remote download).
return allow_roots is None or _is_path_inside_allowlist(path, allow_roots)
def validate_chat_template(template: str) -> ValidateChatTemplateResponse:
text = (template or "").strip()
if not text:
return ValidateChatTemplateResponse(valid = True, error = None)
# Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a
# missing dependency must not crash API startup.
try:
from jinja2 import TemplateError
from jinja2.ext import Extension
from jinja2.sandbox import ImmutableSandboxedEnvironment
except ImportError:
return ValidateChatTemplateResponse(valid = True, error = None)
class _GenerationTag(Extension):
# Accept Transformers' {% generation %} assistant-mask tag so a pasted HF
# chat template validates (we only parse it).
tags = {"generation"}
def parse(self, parser):
next(parser.stream)
return parser.parse_statements(["name:endgeneration"], drop_needle = True)
try:
env = ImmutableSandboxedEnvironment(
trim_blocks = True,
lstrip_blocks = True,
extensions = ["jinja2.ext.loopcontrols", _GenerationTag],
)
env.parse(text)
return ValidateChatTemplateResponse(valid = True, error = None)
except TemplateError as exc:
message = getattr(exc, "message", None) or str(exc)
lineno = getattr(exc, "lineno", None)
if lineno:
message = f"Line {lineno}: {message}"
return ValidateChatTemplateResponse(valid = False, error = message)
except Exception as exc:
return ValidateChatTemplateResponse(valid = False, error = str(exc))
def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]:
if not isinstance(config, dict):
return None
raw = config.get("chat_template")
if isinstance(raw, str) and raw.strip():
return raw
if isinstance(raw, list):
fallback: Optional[str] = None
for entry in raw:
if not isinstance(entry, dict):
continue
template = entry.get("template")
if not isinstance(template, str):
continue
if entry.get("name") == "default":
return template
if fallback is None:
fallback = template
return fallback
return None
def _chat_template_from_jinja_file(
dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
for rel in _JINJA_TEMPLATE_PATHS:
template_file = dir_path / rel
if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots):
continue
try:
if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES:
continue
template = template_file.read_text(encoding = "utf-8")
except Exception:
continue
if template.strip():
return template
return None
def _chat_template_from_processor_payload(payload: object) -> Optional[str]:
# processor chat_template.json may be the template string itself or a
# {name: template} map, not only a tokenizer_config-shaped object.
if isinstance(payload, str):
return payload if payload.strip() else None
template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type]
if template:
return template
if isinstance(payload, dict):
# Named-template map: prefer "default", else the first non-empty entry
# (mirrors the tokenizer-config list fallback).
default = payload.get("default")
if isinstance(default, str) and default.strip():
return default
for value in payload.values():
if isinstance(value, str) and value.strip():
return value
return None
def _chat_template_from_processor_json(
dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
for rel in _PROCESSOR_TEMPLATE_PATHS:
config_file = dir_path / rel
if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
continue
raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
if raw is None:
continue
try:
payload = json.loads(raw)
except Exception:
continue
template = _chat_template_from_processor_payload(payload)
if template:
return template
return None
def _chat_template_from_tokenizer_dir(
dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
jinja = _chat_template_from_jinja_file(dir_path, allow_roots)
if jinja:
return jinja
for rel in _TOKENIZER_CONFIG_PATHS:
config_file = dir_path / rel
if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
continue
raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
if raw is None:
continue
try:
config = json.loads(raw)
except Exception:
continue
template = _chat_template_from_tokenizer_config(config)
if template:
return template
return _chat_template_from_processor_json(dir_path, allow_roots)
_GGUF_SCAN_MAX_DEPTH = 2
def _iter_ggufs(dir_path: Path) -> list[Path]:
if dir_path == dir_path.parent:
return []
root = str(dir_path)
found: list[Path] = []
for current, dirs, files in os.walk(root, followlinks = False):
rel = os.path.relpath(current, root)
depth = 0 if rel == os.curdir else rel.count(os.sep) + 1
if depth >= _GGUF_SCAN_MAX_DEPTH:
dirs[:] = []
for name in files:
if not name.lower().endswith(".gguf") or _is_mmproj(name):
continue
path = Path(current) / name
try:
rel = path.relative_to(dir_path).as_posix()
except ValueError:
rel = name
quant = _extract_quant_label(rel)
if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
continue
found.append(path)
return found
def _variant_matches(relative_path: str, needle: str) -> bool:
quant = _extract_quant_label(relative_path).lower()
if quant == needle:
return True
if extract_quant_label(relative_path).lower() == needle:
return True
prefix = f"{needle}-"
if not quant.startswith(prefix):
return False
suffix = quant[len(prefix) :]
if not suffix.endswith("bpw"):
return False
value = suffix[:-3]
return bool(value) and value.replace(".", "", 1).isdigit()
_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE)
def _is_nonfirst_gguf_split(path: Path) -> bool:
match = _GGUF_SPLIT_INDEX_RE.search(path.stem)
return match is not None and int(match.group(1)) != 1
def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]:
try:
ggufs = sorted(_iter_ggufs(dir_path))
except OSError:
return None
if not ggufs:
return None
needle = (gguf_variant or "").strip().lower()
if needle:
for path in ggufs:
try:
relative = path.relative_to(dir_path).as_posix()
except ValueError:
relative = path.name
if _variant_matches(relative, needle):
return path
return None
candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs
try:
return max(candidates, key = lambda path: path.stat().st_size)
except OSError:
return candidates[0]
def _chat_template_from_dir(
dir_path: Path,
gguf_variant: Optional[str] = None,
allow_roots: Optional[list[Path]] = None,
) -> Optional[str]:
def from_gguf() -> Optional[str]:
gguf = _find_gguf_in_dir(dir_path, gguf_variant)
if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots):
return None
return read_gguf_chat_template(str(gguf))
# Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the
# author's maintained template and supersede the GGUF's possibly-stale embedded
# copy. The variant only picks the GGUF fallback, so tokenizer-first precedence
# holds whether or not a variant is given.
return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf()
def read_default_chat_template(
model_name: str,
hf_token: Optional[str] = None,
gguf_variant: Optional[str] = None,
) -> Optional[str]:
if not isinstance(model_name, str) or not model_name.strip():
return None
name = model_name.strip()
if is_local_path(name):
try:
target = Path(normalize_path(name)).expanduser()
allow_roots = _build_browse_allowlist()
if not _is_path_inside_allowlist(target, allow_roots):
logger.debug("Refused chat template read outside allowed folders: %s", name)
return None
if name.lower().endswith(".gguf"):
# Prefer a maintained sidecar next to the file over the GGUF's
# embedded copy (tokenizer-first precedence, as elsewhere).
sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots)
if sidecar:
return sidecar
return read_gguf_chat_template(str(target))
return _chat_template_from_dir(target, gguf_variant, allow_roots)
except Exception as exc:
logger.debug("Could not read local chat template for %s: %s", name, exc)
return None
if not _is_valid_repo_id(name):
return None
resolved = resolve_cached_repo_id_case(name)
try:
# Resolve within each cached revision, newest first. A revision's sidecar
# supersedes its own embedded GGUF copy, but must not override a newer
# revision, so precedence stays per-snapshot rather than global.
for snapshot in iter_hf_cache_snapshots(resolved):
template = _chat_template_from_dir(snapshot, gguf_variant)
if template:
return template
except Exception as exc:
logger.debug("Could not read cached chat template for %s: %s", resolved, exc)
try:
from huggingface_hub import HfApi, hf_hub_download
_api = HfApi()
def _remote_exceeds_cap(rel: str) -> bool:
# Best-effort: skip the download when the remote's advertised size
# exceeds the cap, so a maliciously large sidecar is never fetched.
try:
infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token)
except Exception:
return False
for info in infos:
size = getattr(info, "size", None)
if (
getattr(info, "path", None) == rel
and isinstance(size, int)
and size > MAX_TEMPLATE_METADATA_BYTES
):
return True
return False
def _download_text(rel: str) -> Optional[str]:
if _remote_exceeds_cap(rel):
return None
try:
path = hf_hub_download(
resolved,
rel,
token = hf_token,
cache_dir = active_hf_hub_cache(),
)
return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES)
except Exception:
return None
for rel in _JINJA_TEMPLATE_PATHS:
template = _download_text(rel)
if not template or not template.strip():
continue
# A raw Jinja sidecar is the whole template, so it must fit the route's
# response cap (the local path skips oversized .jinja too). Download stays
# bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small
# template still extracts below, but an over-cap Jinja is dropped so the
# search falls through to the tokenizer/processor template.
if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
continue
return template
for rel in _TOKENIZER_CONFIG_PATHS:
raw = _download_text(rel)
if not raw:
continue
try:
config = json.loads(raw)
except Exception:
continue
template = _chat_template_from_tokenizer_config(config)
if template:
return template
for rel in _PROCESSOR_TEMPLATE_PATHS:
raw = _download_text(rel)
if not raw:
continue
try:
payload = json.loads(raw)
except Exception:
continue
template = _chat_template_from_processor_payload(payload)
if template:
return template
return None
except Exception as exc:
logger.debug("Could not fetch chat template for %s: %s", resolved, exc)
return None

1
studio/backend/requirements/extras.txt
View file

@ -10,6 +10,7 @@ omegaconf
einx
pyloudnorm
openai-whisper
av # PyAV: decode dictation audio (webm/opus/mp3/…) for the Whisper STT sidecar
uroman # 4.0 MB - used for Outetts.
MeCab # 19.9 MB - used for Outetts.
inflect # number-to-words, required by OuteTTS

2
studio/backend/requirements/no-torch-runtime.txt
View file

@ -7,7 +7,7 @@
# (current PyPI metadata still declares torch as a hard dep).
# unsloth direct deps (from pyproject.toml [project].dependencies)
typer
typer>=0.12.0
# typer's full runtime dep tree. Required explicitly because this
# file is installed with --no-deps. On Linux/Mac CI runners these
# are often cached transitively; on a fresh windows-latest venv they

5
studio/backend/routes/auth.py
View file

@ -494,6 +494,11 @@ async def change_password(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Current password is incorrect",
)
if any(ch.isspace() for ch in payload.new_password):
raise HTTPException(
status_code = status.HTTP_400_BAD_REQUEST,
detail = "New password cannot contain spaces",
)
if payload.current_password == payload.new_password:
raise HTTPException(
status_code = status.HTTP_400_BAD_REQUEST,

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

@ -23,40 +23,6 @@ def _is_valid_repo_id(repo_id: str) -> bool:
return bool(_VALID_REPO_ID.fullmatch(repo_id))
_dataset_size_cache: dict[str, int] = {}
def _get_dataset_size_cached(repo_id: str) -> int:
if repo_id in _dataset_size_cache:
return _dataset_size_cache[repo_id]
try:
from huggingface_hub import dataset_info as hf_dataset_info
info = hf_dataset_info(repo_id, token = None, files_metadata = True)
total = sum(s.size for s in info.siblings if getattr(s, "size", None))
_dataset_size_cache[repo_id] = total
return total
except Exception:
return 0
def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
"""Resolved realpath for a HF cache repo dir: most-recent snapshot, else cache root.
Mirrors routes/models.py; duplicated here to keep this module self-contained.
"""
try:
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.is_dir():
snaps = [s for s in snapshots_dir.iterdir() if s.is_dir()]
if snaps:
latest = max(snaps, key = lambda s: s.stat().st_mtime)
return str(latest.resolve())
return str(repo_dir.resolve())
except Exception:
return None
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
@ -64,6 +30,7 @@ if str(backend_path) not in sys.path:
from utils.datasets import check_dataset_format
from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
router = APIRouter()
logger = get_logger(__name__)
@ -292,11 +259,13 @@ def _download_hf_metadata(*, repo_id: str, repo_files: list[str], token: str | N
try:
from huggingface_hub import hf_hub_download
from utils.hf_cache_settings import active_hf_hub_cache
local_path = hf_hub_download(
repo_id = repo_id,
filename = metadata_file,
repo_type = "dataset",
token = token,
cache_dir = active_hf_hub_cache(),
)
except Exception as exc:
logger.warning(f"Could not read HF dataset metadata for {repo_id}: {exc}")
@ -525,77 +494,15 @@ def list_local_datasets(
@router.get("/download-progress")
async def get_dataset_download_progress(
repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Return download progress for a HuggingFace dataset repo.
Mirrors ``GET /api/models/download-progress`` but scans the
``datasets--owner--name`` cache dir under HF_HUB_CACHE, where in-progress
download bytes are visible. Returns ``cache_path`` so the UI can show it.
"""
_empty = {
"downloaded_bytes": 0,
"expected_bytes": 0,
"progress": 0,
"cache_path": None,
}
try:
if not _is_valid_repo_id(repo_id):
return _empty
from huggingface_hub import constants as hf_constants
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"datasets--{repo_id.replace('/', '--')}".lower()
completed_bytes = 0
in_progress_bytes = 0
cache_path: Optional[str] = None
if cache_dir.is_dir():
for entry in cache_dir.iterdir():
if entry.name.lower() != target:
continue
cache_path = _resolve_hf_cache_realpath(entry)
blobs_dir = entry / "blobs"
if not blobs_dir.is_dir():
break
for f in blobs_dir.iterdir():
if not f.is_file():
continue
if f.name.endswith(".incomplete"):
in_progress_bytes += f.stat().st_size
else:
completed_bytes += f.stat().st_size
break
downloaded_bytes = completed_bytes + in_progress_bytes
if downloaded_bytes == 0:
return {**_empty, "cache_path": cache_path}
expected_bytes = _get_dataset_size_cached(repo_id)
if expected_bytes <= 0:
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": 0,
"progress": 0,
"cache_path": cache_path,
}
# 95% threshold (as in the model endpoint): HF blob dedup makes
# completed_bytes drift under expected_bytes; inter-file gaps look "done".
if completed_bytes >= expected_bytes * 0.95:
progress = 1.0
else:
progress = min(downloaded_bytes / expected_bytes, 0.99)
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": expected_bytes,
"progress": round(progress, 3),
"cache_path": cache_path,
}
except Exception as e:
logger.warning(f"Error checking dataset download progress for {repo_id}: {e}")
return _empty
"""Compatibility route backed by the shared multi-cache progress service."""
from hub.services.datasets import downloads
return await downloads.get_dataset_download_progress_response(
repo_id,
hf_token = hf_token,
)
@router.post("/check-format", response_model = CheckFormatResponse)

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

@ -29,6 +29,8 @@ import re as _re
from utils.models import extract_model_size_b as _extract_model_size_b
from utils.api_errors import openai_error_body, anthropic_error_body
from utils.upload_limits import STT_AUDIO_B64_MAX_CHARS, STT_AUDIO_RAW_MAX_BYTES
from hub.dependencies import get_hf_token
from core.inference.orchestrator import GenStreamError, GenStreamErrorRaised
from core.inference.llama_admission import (
LlamaAdmissionCancelled,
@ -1692,6 +1694,8 @@ async def _aiter_llama_stream_items(
from models.inference import (
LoadRequest,
UnloadRequest,
TranscribeRequest,
SttLoadRequest,
GenerateRequest,
LoadResponse,
LoadProgressResponse,
@ -3406,9 +3410,8 @@ async def _acquire_swap_gate() -> None:
await asyncio.sleep(0.02)
# Counts in-flight auto-switch requests per (target, variant). The busy guard
# subtracts same-target waiters so concurrent requests for one model load once
# instead of each 409-ing the other.
# Counts auto-switch requests queued to load each (target, variant). They are not
# generating, so the drain wait below excludes them from the active inference count.
_auto_switch_waiters: dict[tuple[str, str], int] = {}
_auto_switch_waiters_guard = threading.Lock()
@ -3426,35 +3429,31 @@ def _note_switch_waiter(key: tuple[str, str], delta: int) -> None:
_auto_switch_waiters.pop(key, None)
def _same_target_waiters(key: tuple[str, str]) -> int:
def _switch_waiter_count() -> int:
with _auto_switch_waiters_guard:
return _auto_switch_waiters.get(key, 0)
return sum(max(0, count) for count in _auto_switch_waiters.values())
# A second waiter map keyed by the raw requested model, registered before the
# (slow) resolve. The middleware counts a concurrent same-model request as
# in-flight before it resolves and joins _auto_switch_waiters, so without this
# the first request would see it as an unrelated request and 409.
_auto_switch_request_waiters: dict[str, int] = {}
_auto_switch_request_waiters_guard = threading.Lock()
async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None:
"""Wait until a model replacement cannot interrupt active inference.
def _request_waiter_key(requested_model: str) -> str:
return requested_model.strip().lower()
def _note_request_waiter(key: str, delta: int) -> None:
with _auto_switch_request_waiters_guard:
n = _auto_switch_request_waiters.get(key, 0) + delta
if n > 0:
_auto_switch_request_waiters[key] = n
else:
_auto_switch_request_waiters.pop(key, None)
def _same_request_waiters(key: str) -> int:
with _auto_switch_request_waiters_guard:
return _auto_switch_request_waiters.get(key, 0)
The caller holds ``inference_lifecycle_gate``, which prevents new inference
from starting while existing requests drain. Auto-switch requests that have
resolved their targets are scheduler waiters, not active generations, so
exclude them to avoid a queue deadlock.
"""
from core.inference.llama_keepwarm import other_inference_request_count
while True:
queued_switches = _switch_waiter_count()
if current_request_counted and queued_switches > 0:
queued_switches -= 1
active_others = other_inference_request_count(
current_request_counted = current_request_counted,
include_pending = False,
)
if active_others <= queued_switches:
return
await asyncio.sleep(0.02)
def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]:
@ -3582,7 +3581,6 @@ async def _maybe_auto_switch_model(
from core.inference.local_model_resolver import resolve_local_gguf
from core.inference.llama_keepwarm import (
get_last_unloaded_model,
other_inference_request_count,
inference_lifecycle_gate,
)
@ -3603,12 +3601,7 @@ async def _maybe_auto_switch_model(
if not auto_switch_on and get_auto_unload_idle_seconds() <= 0:
return
# Register by the raw requested model before resolving (which can be slow):
# the middleware already counts a concurrent same-model request as in-flight,
# so the busy guard must know it shares this target even while it resolves.
request_key = _request_waiter_key(requested_model)
_note_request_waiter(request_key, 1)
try:
async def _resolve_and_switch() -> None:
# Off the loop: a cold-cache rebuild walks several model dirs + HF caches.
# With auto-switch off (or an omitted-model reload-only request), skip the
# resolve so only the reload-stash path runs and no name is ever matched.
@ -3706,6 +3699,7 @@ async def _maybe_auto_switch_model(
)
key = _switch_key(override_id, variant)
_note_switch_waiter(key, 1)
waiter_noted = True
try:
async with _auto_switch_lock():
# The asyncio lock is per loop; add a process-wide gate so a swap on
@ -3718,31 +3712,6 @@ async def _maybe_auto_switch_model(
if _already_serving():
_record_serving_alias()
return
# Single slot: refuse a cross-model swap while another inference
# request is active rather than killing its response. Requests
# heading to this same target (by resolved id or raw name) are
# excluded, so concurrent requests for one model load once. A
# pending request is still in the middleware, not generating, so
# it is not counted here.
same_others = max(
_same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0
)
others = other_inference_request_count(
current_request_counted = True, include_pending = False
)
# Not gated on the GGUF being loaded: _load_model_impl also
# tears down an active Unsloth backend before loading a GGUF,
# so refuse whenever any other inference request is in flight.
if others > same_others:
raise HTTPException(
status_code = 409,
detail = openai_error_body(
"Cannot switch models while another inference request is in progress.",
status = 409,
code = "model_switch_busy",
param = "model",
),
)
# Apply this model's saved launch flags so the swap honors the config.
override = get_model_override(override_id)
load_kwargs = {"model_path": target_id, "gguf_variant": variant}
@ -3757,16 +3726,22 @@ async def _maybe_auto_switch_model(
LoadRequest(**load_kwargs),
fastapi_request,
current_subject,
current_request_counted = True,
)
# Advertise the repo id (not the concrete load path) as the loaded
# model's public id and override key for /v1/models and idle stash.
get_llama_cpp_backend()._openai_advertised_id = override_id
finally:
# Deregister before releasing the gate: otherwise a swap on another
# loop counts this finished request as queued and unloads its model.
_note_switch_waiter(key, -1)
waiter_noted = False
_auto_switch_process_lock.release()
finally:
_note_switch_waiter(key, -1)
finally:
_note_request_waiter(request_key, -1)
if waiter_noted:
_note_switch_waiter(key, -1)
await _resolve_and_switch()
async def _auto_switch_from_request_body(request: Request, current_subject: str):
@ -4186,6 +4161,15 @@ def _maybe_unsupported_message(msg: str) -> str:
return msg
def _raise_if_sidecar_swap_in_progress() -> None:
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
raise HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
@router.post("/load", response_model = LoadResponse)
async def load_model(
request: LoadRequest,
@ -4206,24 +4190,23 @@ async def load_model(
# install can reserve while this request queues on the gate, so the pre-gate
# check alone is only a fast path.
from core.inference.llama_keepwarm import inference_lifecycle_gate
from utils.transformers_version import sidecar_swap_in_progress
_swap_409 = HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
if sidecar_swap_in_progress():
raise _swap_409
_raise_if_sidecar_swap_in_progress()
# Hold the lifecycle gate across the load so idle auto-unload can't unload the
# model mid-load. Auto-switch calls _load_model_impl directly since it already
# holds this gate.
async with inference_lifecycle_gate():
if sidecar_swap_in_progress():
raise _swap_409
_raise_if_sidecar_swap_in_progress()
return await _load_model_impl(request, fastapi_request, current_subject)
async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str):
async def _load_model_impl(
request: LoadRequest,
fastapi_request: Request,
current_subject: str,
*,
current_request_counted: bool = False,
):
from core.inference.llama_cpp import LlamaServerNotFoundError
# A new load starts here; arm the progress throttle so this load's first
@ -4557,6 +4540,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
),
)
# Keep the resident model alive until every active generation finishes;
# the caller's lifecycle gate blocks new starts.
await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
# A sidecar install can reserve the gate while inference drains, after the
# route-level checks above, so recheck before replacing either backend.
_raise_if_sidecar_swap_in_progress()
# Unload any active Unsloth model only after every hub conflict check.
if unsloth_backend.active_model_name:
logger.info(
@ -4767,6 +4757,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Unload any active GGUF model first
llama_backend = get_llama_cpp_backend()
await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
_raise_if_sidecar_swap_in_progress()
if llama_backend.is_loaded:
logger.info("Unloading GGUF model before loading Unsloth model")
llama_backend.unload_model()
@ -5144,10 +5136,10 @@ async def validate_model(
latest_tier_active_for, config.identifier, request.hf_token
):
effective_load_in_4bit = False
# A metadata-only probe just reads the GGUF header and allocates no VRAM,
# so it must not be refused by the training guard. Real loads validate
# without include_context_length and /load applies the guard again.
if not request.include_context_length:
# A metadata-only probe reads the GGUF header and allocates no VRAM, so the
# training guard must not refuse it. Real loads omit include_context_length /
# include_chat_template, and /load applies the guard again.
if not (request.include_context_length or request.include_chat_template):
# Match /load's inherited llama.cpp extras and parallel slot count so
# validation cannot pass a smaller estimate than the subsequent load.
effective_extra_args = _resolve_inherited_extra_args(
@ -5189,9 +5181,15 @@ async def validate_model(
context_length: Optional[int] = None
layer_count: Optional[int] = None
moe_layer_count: Optional[int] = None
if request.include_context_length and is_gguf:
chat_template: Optional[str] = None
# Both header probes read the same local GGUF, so resolve it once.
if (request.include_context_length or request.include_chat_template) and is_gguf:
from hub.utils.gguf import resolve_local_gguf_path
from utils.models.gguf_metadata import read_gguf_staged_dims
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
from utils.models.gguf_metadata import (
read_gguf_chat_template,
read_gguf_staged_dims,
)
# Best-effort: a header-read failure must never fail validation of an
# otherwise-valid model (the outer except turns it into a 400).
@ -5207,13 +5205,24 @@ async def validate_model(
model_identifier, request.gguf_variant
)
if local_gguf:
# Header walk reads tokenizer arrays for dense models (tens of
# ms); keep it off the event loop.
dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf)
if dims:
context_length = dims["context_length"]
layer_count = dims["layer_count"]
moe_layer_count = dims["moe_layer_count"]
if request.include_context_length:
# Header walk reads tokenizer arrays (tens of ms); keep it
# off the event loop.
dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf)
if dims:
context_length = dims["context_length"]
layer_count = dims["layer_count"]
moe_layer_count = dims["moe_layer_count"]
if request.include_chat_template:
# Read only the leased GGUF's own embedded template (the copy
# llama.cpp loads), never a sibling sidecar: the native grant
# authorizes just this path, so neighbours would be scope escalation.
raw_template = await asyncio.to_thread(read_gguf_chat_template, local_gguf)
if (
raw_template is not None
and len(raw_template.encode("utf-8")) <= MAX_CHAT_TEMPLATE_BYTES
):
chat_template = raw_template
except Exception as e:
logger.debug("Header probe failed for %s: %s", model_log_label, e)
@ -5232,6 +5241,7 @@ async def validate_model(
context_length = context_length,
layer_count = layer_count,
moe_layer_count = moe_layer_count,
chat_template = chat_template,
requires_transformers_upgrade = transformers_upgrade is not None,
transformers_upgrade = transformers_upgrade,
)
@ -6107,6 +6117,342 @@ async def generate_audio(
)
# =====================================================================
# Speech-to-text (STT) sidecar (/audio/transcribe, /audio/stt/*)
# =====================================================================
def _resolve_stt_engine(engine: Optional[str]) -> str:
"""Normalize the requested STT engine name; default is Transformers."""
normalized = (engine or "transformers").strip().lower()
if normalized in ("", "transformers", "whisper"):
return "transformers"
if normalized in ("gguf", "ggml", "whisper_cpp", "whisper.cpp"):
return "gguf"
raise HTTPException(
status_code = 422,
detail = f"Unknown STT engine '{engine}'. Use 'transformers' or 'gguf'.",
)
def _resolve_serving_stt_engine(engine: Optional[str]) -> str:
"""Resolve the engine that will actually serve a model.
whisper.cpp (gguf) only accepts curated ids, which Transformers serves too,
so when whisper-server is not installed (the common case: `unsloth studio
update` does not yet build it) fall back to Transformers instead of 501-ing
on every recording. Used for download/load/transcribe; unload targets a
specific engine via _resolve_stt_engine.
"""
resolved = _resolve_stt_engine(engine)
if resolved == "gguf":
from core.inference import stt_ggml_sidecar
if not stt_ggml_sidecar.is_available():
return "transformers"
return resolved
def _stt_sidecar_for(engine: str):
if engine == "gguf":
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
return get_ggml_stt_sidecar()
from core.inference.stt_sidecar import get_stt_sidecar
return get_stt_sidecar()
@studio_router.get("/audio/stt/status")
async def stt_status(
model: Optional[str] = None, current_subject: str = Depends(get_current_subject)
):
"""Report STT availability and which model, if any, is resident.
``model`` extends the Transformers ``downloaded_models`` check to a
custom Hugging Face repository beyond the curated defaults.
"""
from core.inference import stt_ggml_sidecar, stt_sidecar
from core.inference.stt_sidecar import (
DEFAULT_STT_MODEL,
STT_MODELS,
get_stt_sidecar,
is_available,
)
sidecar = get_stt_sidecar()
ggml = stt_ggml_sidecar.get_ggml_stt_sidecar()
transformers_downloaded = [
model_id for model_id in STT_MODELS if stt_sidecar.is_model_downloaded(model_id)
]
if model and model not in STT_MODELS and stt_sidecar.is_model_downloaded(model):
transformers_downloaded.append(model)
return JSONResponse(
content = {
"available": is_available(),
"loaded_model": sidecar.loaded_model,
"loading": sidecar.is_loading(),
"device": sidecar.device,
"keep_alive_seconds": sidecar.keep_alive_seconds,
"default_model": DEFAULT_STT_MODEL,
"models": list(STT_MODELS.keys()),
# Transformers engine, same shape as "gguf" below so clients read
# either generically. Top-level fields above kept for old clients.
"transformers": {
"available": is_available(),
"loaded_model": sidecar.loaded_model,
"loading": sidecar.is_loading(),
"device": sidecar.device,
"keep_alive_seconds": sidecar.keep_alive_seconds,
"default_model": DEFAULT_STT_MODEL,
"models": list(STT_MODELS.keys()),
"downloaded_models": transformers_downloaded,
"download": stt_sidecar.download_status(),
},
# whisper.cpp (GGUF) engine.
"gguf": {
"available": stt_ggml_sidecar.is_available(),
"loaded_model": ggml.loaded_model,
"loading": ggml.is_loading(),
"device": ggml.device,
"keep_alive_seconds": ggml.keep_alive_seconds,
"default_model": stt_ggml_sidecar.DEFAULT_GGML_STT_MODEL,
"models": list(stt_ggml_sidecar.GGML_STT_MODELS.keys()),
"downloaded_models": [
model_id
for model_id in stt_ggml_sidecar.GGML_STT_MODELS
if stt_ggml_sidecar._cached_model_path(model_id) is not None
],
"download": stt_ggml_sidecar.download_status(),
},
}
)
@studio_router.post("/audio/stt/download")
async def stt_download(
payload: SttLoadRequest,
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""Start a background download of a dictation model.
Both engines download directly (a GGML checkpoint is a single file the Model
Hub's GGUF variant planner cannot express; a Transformers checkpoint is a
whole snapshot). Progress is reported by /audio/stt/status.
"""
from core.inference import stt_ggml_sidecar, stt_sidecar
from core.inference.stt_sidecar import (
SttModelCompatibilityError,
SttModelIdError,
validate_remote_model,
)
engine = _resolve_serving_stt_engine(payload.engine)
module = stt_ggml_sidecar if engine == "gguf" else stt_sidecar
try:
# Transformers accepts custom `owner/model` repos, so confirm the repo is
# a Whisper checkpoint (metadata-only) before snapshot_download pulls a
# possibly-large non-STT repo into the shared cache. Curated ids
# short-circuit; GGUF only accepts curated ids, so it needs no check.
if engine != "gguf":
validated = await asyncio.to_thread(validate_remote_model, payload.model, hf_token)
# Pin the download to the commit that was just validated so the
# repo cannot be swapped between validation and snapshot_download.
await asyncio.to_thread(
module.start_model_download,
payload.model,
hf_token,
validated.get("revision"),
)
else:
await asyncio.to_thread(module.start_model_download, payload.model, hf_token)
except SttModelIdError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttModelCompatibilityError as e:
raise HTTPException(status_code = 422, detail = str(e))
return JSONResponse(content = module.download_status())
@studio_router.post("/audio/stt/load")
async def stt_load(payload: SttLoadRequest, current_subject: str = Depends(get_current_subject)):
"""Load the selected STT model after the user starts local dictation."""
from core.inference.stt_sidecar import (
SttLoadCancelledError,
SttModelCompatibilityError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
get_stt_sidecar,
)
sidecar = _stt_sidecar_for(_resolve_serving_stt_engine(payload.engine))
try:
await asyncio.to_thread(sidecar.load, payload.model)
except SttModelNotDownloadedError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttUnavailableError as e:
raise HTTPException(status_code = 501, detail = str(e))
except SttLoadCancelledError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttModelIdError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttModelCompatibilityError as e:
raise HTTPException(status_code = 422, detail = str(e))
except Exception as e:
logger.error(f"STT load error: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
return JSONResponse(content = {"loaded_model": sidecar.loaded_model, "device": sidecar.device})
@studio_router.post("/audio/stt/validate")
async def stt_validate(
payload: SttLoadRequest,
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""Verify a Hub repository is a Whisper checkpoint before downloading it."""
from core.inference.stt_sidecar import (
SttModelCompatibilityError,
SttModelIdError,
validate_remote_model,
)
try:
result = await asyncio.to_thread(validate_remote_model, payload.model, hf_token)
except (SttModelIdError, SttModelCompatibilityError) as e:
raise HTTPException(status_code = 422, detail = str(e))
return JSONResponse(content = result)
@studio_router.post("/audio/stt/unload")
async def stt_unload(
engine: Optional[str] = None, current_subject: str = Depends(get_current_subject)
):
"""Release the local STT model when dictation is idle.
Without an engine, both sidecars unload so an engine switch in Voice
settings always frees whichever backend was resident.
"""
if engine is None:
engines = ["transformers", "gguf"]
else:
# Use the serving resolver: a "gguf" pick without whisper-server is
# actually served by the Transformers fallback, so unload must target
# that same engine or the resident model is never freed.
engines = [_resolve_serving_stt_engine(engine)]
# Attempt every engine even if one raises, so failing to unload one never
# skips freeing the other (both can be resident after a switch).
failed: list[str] = []
for name in engines:
try:
await asyncio.to_thread(_stt_sidecar_for(name).unload)
except Exception as exc: # noqa: BLE001 - report after attempting all engines
logger.warning("Failed to unload STT engine '%s': %s", name, exc)
failed.append(name)
if failed:
raise HTTPException(
status_code = 500,
detail = f"Failed to unload STT engine(s): {', '.join(failed)}",
)
return JSONResponse(content = {"loaded_model": None, "device": None})
async def _transcribe_audio_bytes(
raw: bytes,
model: Optional[str],
language: Optional[str],
fast: bool,
engine: Optional[str] = None,
) -> JSONResponse:
"""Run STT for already-decoded request bytes."""
from core.inference.stt_sidecar import (
SttAudioDecodeError,
SttAudioTooLongError,
SttLanguageError,
SttLoadCancelledError,
SttModelCompatibilityError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
)
if not raw:
raise HTTPException(status_code = 400, detail = "Audio is empty.")
if len(raw) > _MAX_AUDIO_RAW_BYTES:
raise HTTPException(status_code = 413, detail = "Audio is too large.")
sidecar = _stt_sidecar_for(_resolve_serving_stt_engine(engine))
try:
result = await asyncio.to_thread(
sidecar.transcribe,
raw,
model,
language,
fast,
)
except SttUnavailableError as e:
raise HTTPException(status_code = 501, detail = str(e))
except SttLoadCancelledError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttModelNotDownloadedError as e:
raise HTTPException(status_code = 409, detail = str(e))
except SttModelIdError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttModelCompatibilityError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttLanguageError as e:
raise HTTPException(status_code = 422, detail = str(e))
except SttAudioTooLongError as e:
raise HTTPException(status_code = 413, detail = str(e))
except SttAudioDecodeError as e:
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
logger.error(f"Transcription error: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
return JSONResponse(content = result)
@studio_router.post("/audio/transcribe")
async def transcribe_audio(
payload: TranscribeRequest, current_subject: str = Depends(get_current_subject)
):
"""Transcribe dictation audio to text via the STT sidecar.
Runs alongside the chat model without evicting it, so any model (including
text-only ones) can be driven by voice.
"""
b64 = payload.audio or ""
if not b64:
raise HTTPException(status_code = 400, detail = "No audio provided.")
if len(b64) > _MAX_AUDIO_B64_CHARS:
raise HTTPException(status_code = 413, detail = "Audio is too large.")
try:
raw = base64.b64decode(b64, validate = True)
except Exception:
raise HTTPException(status_code = 400, detail = "Audio is not valid base64.")
return await _transcribe_audio_bytes(
raw, payload.model, payload.language, payload.fast, payload.engine
)
@studio_router.post("/audio/transcribe/raw")
async def transcribe_audio_raw(
request: Request,
model: Optional[str] = None,
language: Optional[str] = None,
fast: bool = False,
engine: Optional[str] = None,
current_subject: str = Depends(get_current_subject),
):
"""Transcribe a raw audio body without base64 or JSON conversion overhead."""
chunks: list[bytes] = []
size = 0
async for chunk in request.stream():
size += len(chunk)
if size > _MAX_AUDIO_RAW_BYTES:
raise HTTPException(status_code = 413, detail = "Audio is too large.")
chunks.append(chunk)
return await _transcribe_audio_bytes(b"".join(chunks), model, language, fast, engine)
# =====================================================================
# OpenAI-Compatible Chat Completions (/chat/completions)
# =====================================================================
@ -6148,8 +6494,8 @@ def _decode_audio_base64(b64: str) -> np.ndarray:
# cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally
# bounds the *decoded* length, since a small compressed file (opus/flac/etc.)
# can expand to a far larger PCM array than the encoded-size cap implies.
_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024
_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3
_MAX_AUDIO_RAW_BYTES = STT_AUDIO_RAW_MAX_BYTES
_MAX_AUDIO_B64_CHARS = STT_AUDIO_B64_MAX_CHARS
_MAX_AUDIO_SECONDS = 30 * 60
_WAV_HEADER_BYTES = 44
_MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000
@ -7078,7 +7424,7 @@ async def openai_chat_completions(
if payload.provider_id or payload.provider_type:
# External provider: this request won't touch the local GGUF, so drop it
# from the keep-warm count or its in-flight stream would falsely block a
# concurrent local auto-switch with model_switch_busy.
# concurrent local model switch from proceeding.
from core.inference.llama_keepwarm import untrack_current_request
untrack_current_request(request.scope)

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

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""llama.cpp prebuilt update endpoints.
"""llama.cpp prebuilt update endpoints -- the single main update item.
GET /api/llama/update-status -> is a newer prebuilt available + job state
POST /api/llama/update -> download + atomically swap to the latest
@ -9,13 +9,19 @@ POST /api/llama/update -> download + atomically swap to the latest
Detection reuses utils.llama_cpp_freshness; the swap reuses
install_llama_prebuilt.py via utils.llama_cpp_update. Both fail open so the UI
never blocks on a missing marker / offline GitHub.
whisper.cpp updates piggyback here: the status payload carries a whisper
sub-status (update_available is the llama OR whisper union) and the apply job
chains a whisper phase after the llama phase when whisper is behind, with a
per-phase breakdown in job.phases. All pre-existing top-level fields keep
their shape, so older clients keep working unchanged.
"""
from __future__ import annotations
import asyncio
import threading
from typing import Optional
from typing import Literal, Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
@ -38,6 +44,31 @@ class LlamaUpdateJob(BaseModel):
progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.")
started_at: Optional[str] = None
finished_at: Optional[str] = None
phases: Optional[dict] = Field(
None,
description = (
"Per-phase breakdown of a chained llama+whisper job "
"(name -> state/progress/to_tag/...); None for pre-chaining jobs."
),
)
class WhisperSubStatus(BaseModel):
"""The whisper piggyback inside the llama update item."""
update_available: bool = Field(
False, description = "True when the chained apply would run a whisper phase."
)
installed_tag: Optional[str] = None
latest_tag: Optional[str] = None
update_size_bytes: Optional[int] = None
skip_reason: Optional[str] = Field(
None,
description = (
"Why the whisper phase would be skipped "
"(up_to_date | local_link | source_build | not_installed | ...)."
),
)
class LlamaUpdateStatusResponse(BaseModel):
@ -46,7 +77,18 @@ class LlamaUpdateStatusResponse(BaseModel):
description = "True when the install came from an Unsloth prebuilt (has a marker).",
)
update_available: bool = Field(
False, description = "True when the latest release is genuinely newer than the install."
False,
description = (
"True when an update would do something: llama.cpp is behind OR the "
"whisper piggyback is behind."
),
)
llama_update_available: bool = Field(
False, description = "True when the latest llama.cpp release is newer than the install."
)
update_component: Optional[Literal["llama", "whisper"]] = Field(
None,
description = "Component whose versions the combined update banner should display.",
)
stale: bool = Field(
False, description = "Update available AND install older than the staleness threshold."
@ -62,6 +104,9 @@ class LlamaUpdateStatusResponse(BaseModel):
update_size_bytes: Optional[int] = Field(
None, description = "Download size of the prebuilt Update would fetch, in bytes."
)
whisper: Optional[WhisperSubStatus] = Field(
None, description = "Whisper piggyback sub-status; None when the probe is unavailable."
)
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)

724
studio/backend/routes/models.py
View file

@ -60,13 +60,52 @@ def _safe_is_dir(path) -> bool:
# Shared with the hub inventory scans; keep the private aliases so existing
# importers (core.inference.local_model_resolver, tests) stay valid.
# importers stay valid. ``_HF_REPO_ID_RE`` is the Hub repo id shape ("owner/name");
# anything else is treated as a local filesystem path.
from utils.hidden_models import (
_HF_REPO_ID_RE,
_existing_resolved_path,
_safe_resolve,
is_hidden_model as _is_hidden_model,
)
def hidden_model_matchers() -> tuple[list[str], list[str], list[str]]:
"""Substring needles, exact repo ids, and exact resolved paths identifying
infra models (the RAG embedder and the llama.cpp install validation probe)
that pickers hide. Served by the ``/api/hub/hidden-models`` endpoint. A
configured HF-repo embedder is published as its exact lowercased repo id
(mirroring ``utils.hidden_models.is_hidden_model``) and a local-path
embedder as its exact resolved path only: a generic basename like "model"
must not substring-hide unrelated chat models."""
from core.rag import config as rag_config
needles = [
# The validation probe's repo and its exact filename. The filename carries
# .gguf so it won't hide unrelated repos like ``user/stories260K-finetune-GGUF``.
"ggml-org/models",
"stories260k.gguf",
]
exact_ids: list[str] = []
exact_paths: list[str] = []
for model in (
rag_config.effective_embedding_model(),
rag_config.effective_gguf_repo(),
):
# Resolve an existing local path before the repo-id regex: a local embedder
# shaped like "models/embedder" is an exact path, not a Hub repo id.
existing_path = _existing_resolved_path(model)
if existing_path:
exact_paths.append(existing_path.lower())
elif _HF_REPO_ID_RE.match(model):
exact_ids.append(model.lower())
else:
resolved = _safe_resolve(Path(model).expanduser())
if resolved:
exact_paths.append(resolved.lower())
return needles, exact_ids, exact_paths
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
@ -91,6 +130,7 @@ try:
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mtp_drafter,
is_audio_input_type,
)
from core.inference import get_inference_backend
@ -123,6 +163,7 @@ except ImportError:
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mtp_drafter,
is_audio_input_type,
)
from core.inference import get_inference_backend
@ -183,11 +224,8 @@ def derive_model_type(
def _resolve_hf_cache_dir() -> Path:
"""Resolve local HF cache root used by hub downloads."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
from utils.hf_cache_settings import get_hf_cache_paths
return get_hf_cache_paths().hub_cache
def _is_model_directory(d: Path) -> bool:
@ -329,10 +367,12 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
return found
def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
def _scan_hf_cache(cache_dir: Path, *, active_cache: bool = True) -> List[LocalModelInfo]:
if not cache_dir.exists() or not cache_dir.is_dir():
return []
from hub.utils import inventory_scan as hf_cache_scan
found: List[LocalModelInfo] = []
for repo_dir in cache_dir.glob("models--*"):
if not repo_dir.is_dir():
@ -348,13 +388,21 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
except OSError:
updated_at = None
partial = hf_cache_scan.is_snapshot_partial("model", model_id, repo_dir)
partial = partial or hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir)
load_id = model_id
if not active_cache:
load_id = _resolve_hf_cache_realpath(repo_dir) or str(repo_dir.resolve())
found.append(
LocalModelInfo(
id = model_id,
id = load_id,
model_id = model_id,
display_name = model_id.split("/")[-1],
path = str(repo_dir),
path = load_id if not active_cache else str(repo_dir),
source = "hf_cache",
active_cache = active_cache,
partial = partial,
updated_at = updated_at,
),
)
@ -735,26 +783,34 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
legacy_hf_cache_dir,
lmstudio_model_dirs,
)
from utils.hf_cache_settings import known_hf_hub_caches
hf_cache_dir = _resolve_hf_cache_dir()
legacy_hf = legacy_hf_cache_dir()
hf_default = hf_default_cache_dir()
lm_dirs = lmstudio_model_dirs()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real:
local_models += _scan_hf_cache(hf_default)
local_models = _scan_models_dir(models_root)
active_cache_real = _safe_resolve(hf_cache_dir)
active_cache_key = os.path.normcase(active_cache_real) if active_cache_real else None
seen_hf: set[str] = set()
for cache_dir in (
hf_cache_dir,
*known_hf_hub_caches(),
legacy_hf,
hf_default,
):
cache_real = _safe_resolve(cache_dir)
if cache_real is None:
continue
cache_key = os.path.normcase(str(cache_real))
if cache_key in seen_hf:
continue
seen_hf.add(cache_key)
local_models += _scan_hf_cache(
cache_dir,
active_cache = cache_key == active_cache_key,
)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
@ -776,7 +832,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_hf_cache(folder_path, active_cache = False)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
@ -797,13 +853,23 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
semantic_id = model.model_id if model.source == "hf_cache" and model.model_id else model.id
key = f"{semantic_id}\x00custom" if model.source == "custom" else semantic_id
existing = deduped.get(key)
prefer_model = existing is None
if existing is not None and model.source == existing.source == "hf_cache":
if model.partial != existing.partial:
prefer_model = not model.partial
elif bool(model.active_cache) != bool(existing.active_cache):
prefer_model = bool(model.active_cache)
else:
prefer_model = (model.updated_at or 0) > (existing.updated_at or 0)
if prefer_model:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
key = lambda item: item.updated_at or 0,
reverse = True,
)
return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)]
@ -1161,10 +1227,7 @@ def _build_browse_allowlist(
legacy_hf_cache_dir,
well_known_model_dirs,
)
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from utils.paths import external_media
from storage.studio_db import list_scan_folders
candidates: list[Path] = []
@ -1181,9 +1244,12 @@ def _build_browse_allowlist(
_add(Path.home())
if media_roots is None:
media_roots = linux_run_media_mount_roots()
media_roots = [
*external_media.linux_run_media_mount_roots(),
*external_media.macos_volume_roots(),
]
if drive_roots is None:
drive_roots = windows_drive_roots()
drive_roots = external_media.windows_drive_roots()
for p in media_roots:
_add(p)
for p in drive_roots:
@ -1461,10 +1527,7 @@ def browse_folders(
then hidden (if ``show_hidden=true``).
"""
from utils.paths import hf_default_cache_dir, well_known_model_dirs
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from utils.paths import external_media
from storage.studio_db import (
contains_sensitive_path_component,
is_denied_system_path,
@ -1473,8 +1536,11 @@ def browse_folders(
# Probe removable-media and Windows drive roots once; the allowlist and
# chips reuse the result so a disconnected mapped drive isn't scanned twice.
media_roots = linux_run_media_mount_roots()
drive_roots = windows_drive_roots()
media_roots = [
*external_media.linux_run_media_mount_roots(),
*external_media.macos_volume_roots(),
]
drive_roots = external_media.windows_drive_roots()
# Build once; the sandbox check and suggestion chips share it.
allowed_roots = _build_browse_allowlist(media_roots, drive_roots)
@ -1750,9 +1816,11 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op
async def get_model_config(
model_name: str,
hf_token: Optional[str] = Query(None),
header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Get configuration for a specific model (wraps load_model_defaults)."""
hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
if not is_local_path(model_name):
resolved = resolve_cached_repo_id_case(model_name)
@ -1991,19 +2059,19 @@ async def discard_remote_code_download(
# Never delete a model that is loaded for inference.
try:
from hub.services.models.deletion import _loaded_id_matches_repo
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and llama_backend.model_identifier:
loaded = llama_backend.model_identifier.lower()
if loaded == model_name.lower() or loaded.startswith(model_name.lower()):
if _loaded_id_matches_repo(llama_backend.model_identifier, model_name):
return {"deleted": False, "reason": "loaded"}
except Exception:
pass
try:
inference_backend = get_inference_backend()
if inference_backend.active_model_name:
active = inference_backend.active_model_name.lower()
if active == model_name.lower() or active.startswith(model_name.lower()):
if _loaded_id_matches_repo(inference_backend.active_model_name, model_name):
return {"deleted": False, "reason": "loaded"}
except Exception:
pass
@ -2471,6 +2539,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get
async def check_vision_model(
model_name: str,
hf_token: Optional[str] = Query(None),
header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""
@ -2478,6 +2547,7 @@ async def check_vision_model(
This endpoint wraps the backend is_vision_model function.
"""
hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
logger.info(f"Checking if vision model: {model_name}")
# Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision).
@ -2503,6 +2573,7 @@ async def check_vision_model(
async def check_embedding_model(
model_name: str,
hf_token: Optional[str] = Query(None),
header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""
@ -2510,6 +2581,7 @@ async def check_embedding_model(
This endpoint wraps the backend is_embedding_model function.
"""
hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
logger.info(f"Checking if embedding model: {model_name}")
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
@ -2541,13 +2613,10 @@ def _read_native_context_length(repo_id: str, is_local: bool) -> Optional[int]:
if is_local:
roots = [Path(repo_id)]
else:
from huggingface_hub import constants as hf_constants
from hub.utils.hf_cache_state import iter_repo_cache_dirs
if not _is_valid_repo_id(repo_id):
return None
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
roots = [e for e in cache_dir.iterdir() if e.name.lower() == target]
roots = list(iter_repo_cache_dirs("model", repo_id))
for root in roots:
for f in _iter_gguf_paths(root):
@ -2573,47 +2642,32 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio
Q8_0 weights). Never raises.
"""
try:
from utils.models.model_config import (
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mtp_drafter,
)
if is_local:
roots = [Path(repo_id)]
else:
from huggingface_hub import constants as hf_constants
from hub.utils.hf_cache_state import iter_repo_cache_dirs
if not _is_valid_repo_id(repo_id):
return None, 0
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
roots = []
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
snaps = entry / "snapshots"
if snaps.is_dir():
roots.extend(s for s in snaps.iterdir() if s.is_dir())
for entry in iter_repo_cache_dirs("model", repo_id):
snaps = entry / "snapshots"
if snaps.is_dir():
roots.extend(s for s in snaps.iterdir() if s.is_dir())
want = quant.lower().replace("-", "").replace("_", "")
want = _normalized_quant_label(quant)
best_total = 0
best_first: Optional[str] = None
for root in roots:
matches: list[tuple[str, Path]] = []
total = 0
for f in _iter_gguf_paths(root):
if _is_mmproj_filename(f.name):
continue
try:
rel = f.relative_to(root).as_posix()
except ValueError:
rel = f.name
if _is_mtp_drafter(rel):
continue
q = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, q):
continue
if q.lower().replace("-", "").replace("_", "") != want:
q = _main_variant_gguf_label(rel)
if q is None or _normalized_quant_label(q) != want:
continue
try:
total += f.stat().st_size
@ -2699,6 +2753,8 @@ async def get_gguf_variants(
repo_id: str = Query(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
prefer_local_cache: bool = False,
local_path: Optional[str] = None,
hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"),
hf_token_header: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
@ -2710,9 +2766,16 @@ async def get_gguf_variants(
response = await hub_gguf_variants.get_gguf_variants_response(
repo_id,
prefer_local_cache = prefer_local_cache,
local_path = local_path,
hf_token = hf_token,
)
local = is_local_path(repo_id)
context_model = (
local_path
if prefer_local_cache and local_path and is_local_path(local_path)
else repo_id
)
local = is_local_path(context_model)
return GgufVariantsResponse(
repo_id = response.repo_id,
@ -2734,7 +2797,7 @@ async def get_gguf_variants(
# The header walk reads tokenizer arrays on dense models (tens of
# ms per uncached file); keep it off the event loop.
context_length = await asyncio.to_thread(
_read_native_context_length, repo_id, is_local = local
_read_native_context_length, context_model, is_local = local
),
)
except HTTPException:
@ -2752,69 +2815,17 @@ async def get_gguf_download_progress(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
variant: str = Query("", description = "Quantization variant (e.g. UD-TQ1_0)"),
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Download progress from cached GGUF files for a specific variant.
Tracks completed shards in snapshots and in-progress (.incomplete)
downloads in the blobs directory.
"""
try:
if not _is_valid_repo_id(repo_id):
return {
"downloaded_bytes": 0,
"expected_bytes": expected_bytes,
"progress": 0,
}
from huggingface_hub import constants as hf_constants
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
variant_lower = variant.lower().replace("-", "").replace("_", "")
downloaded_bytes = 0
in_progress_bytes = 0
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
# Completed .gguf files for this variant in snapshots.
# Exclude mmproj so a vision adapter can't satisfy a same-label
# main variant (e.g. mmproj-F16 vs an F16 weight).
for f in _iter_gguf_paths(entry):
if _is_mmproj_filename(f.name):
continue
rel = f.relative_to(entry).as_posix()
quant = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, quant):
continue
rel_key = rel.lower().replace("-", "").replace("_", "")
if not variant_lower or variant_lower in rel_key:
try:
downloaded_bytes += f.stat().st_size
except OSError:
continue # broken symlink / unreadable: skip
# In-progress (.incomplete) downloads in blobs.
blobs_dir = entry / "blobs"
if blobs_dir.is_dir():
for f in blobs_dir.iterdir():
if f.is_file() and f.name.endswith(".incomplete"):
try:
in_progress_bytes += f.stat().st_size
except OSError:
continue
break
total_progress_bytes = downloaded_bytes + in_progress_bytes
progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0
# Report 1.0 only when all bytes are in completed files.
if expected_bytes > 0 and downloaded_bytes >= expected_bytes:
progress = 1.0
return {
"downloaded_bytes": total_progress_bytes,
"expected_bytes": expected_bytes,
"progress": round(progress, 3),
}
except Exception:
return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0}
"""Compatibility route backed by the shared multi-cache progress service."""
from hub.services.models import downloads
return await downloads.get_gguf_download_progress_response(
repo_id,
variant = variant,
expected_bytes = expected_bytes,
hf_token = hf_token,
)
def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
@ -2839,98 +2850,12 @@ def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
@router.get("/download-progress")
async def get_download_progress(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Return download progress for any HuggingFace model repo.
Checks the local HF cache for completed blobs and in-progress
(.incomplete) downloads. Gets the expected total size from the HF API
on the first call, then caches it for later polls. Also returns
``cache_path``: the realpath of the snapshot dir (or cache repo root
if no snapshot yet) so the UI can show where weights live on disk.
"""
_empty = {
"downloaded_bytes": 0,
"expected_bytes": 0,
"progress": 0,
"cache_path": None,
}
try:
if not _is_valid_repo_id(repo_id):
return _empty
from huggingface_hub import constants as hf_constants
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
completed_bytes = 0
in_progress_bytes = 0
cache_path: Optional[str] = None
for entry in cache_dir.iterdir():
if entry.name.lower() != target:
continue
cache_path = _resolve_hf_cache_realpath(entry)
blobs_dir = entry / "blobs"
if not blobs_dir.is_dir():
break
for f in blobs_dir.iterdir():
if not f.is_file():
continue
if f.name.endswith(".incomplete"):
in_progress_bytes += f.stat().st_size
else:
completed_bytes += f.stat().st_size
break
downloaded_bytes = completed_bytes + in_progress_bytes
if downloaded_bytes == 0:
return {**_empty, "cache_path": cache_path}
expected_bytes = _get_repo_size_cached(repo_id)
if expected_bytes <= 0:
# Total unknown; report bytes only, no percentage.
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": 0,
"progress": 0,
"cache_path": cache_path,
}
# 95% threshold (blob dedup can skew completed_bytes). Do NOT
# treat "no .incomplete files" as done: HF downloads sequentially,
# so none exist between files even when far from finished.
if completed_bytes >= expected_bytes * 0.95:
progress = 1.0
else:
progress = min(downloaded_bytes / expected_bytes, 0.99)
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": expected_bytes,
"progress": round(progress, 3),
"cache_path": cache_path,
}
except Exception as e:
logger.warning(f"Error checking download progress for {repo_id}: {e}")
return _empty
_repo_size_cache: dict[str, int] = {}
def _get_repo_size_cached(repo_id: str) -> int:
if repo_id in _repo_size_cache:
return _repo_size_cache[repo_id]
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token = None, files_metadata = True)
total = sum(s.size for s in info.siblings if s.size)
_repo_size_cache[repo_id] = total
return total
except Exception as e:
logger.warning(f"Failed to get repo size for {repo_id}: {e}")
return 0
"""Compatibility route backed by the shared multi-cache progress service."""
from hub.services.models import downloads
return await downloads.get_download_progress_response(repo_id, hf_token = hf_token)
def _repo_in_any_hf_cache(model_name: str) -> bool:
@ -2943,25 +2868,13 @@ def _repo_in_any_hf_cache(model_name: str) -> bool:
would delete a model they did not download via the scan. Mirrors the cache set in
``_all_hf_cache_scans`` but only probes for the one repo dir (cheap, no full scan).
"""
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
resolve_cached_repo_id_case,
)
from utils.paths import resolve_cached_repo_id_case
dirname = f"models--{resolve_cached_repo_id_case(model_name).replace('/', '--')}"
dirname_lower = dirname.lower()
candidates = []
try:
from huggingface_hub.constants import HF_HUB_CACHE
candidates.append(Path(HF_HUB_CACHE))
except Exception:
pass
for fn in (legacy_hf_cache_dir, hf_default_cache_dir):
try:
candidates.append(fn())
except Exception:
continue
from hub.utils.hf_cache_state import hf_cache_roots
candidates = hf_cache_roots()
# resolve_cached_repo_id_case only normalizes the ACTIVE cache, but discard deletes
# case-insensitively across all caches, so detect case-insensitively too -- else a
# pre-existing case-variant repo is misreported as scan-created and deleted on decline.
@ -2985,38 +2898,8 @@ def _all_hf_cache_scans():
broken symlink, OS-redirected ~/.cache) is skipped, not fatal, so the
Downloaded list never blanks out and downloads never leak into Recommended.
"""
from huggingface_hub import scan_cache_dir
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir
scans = []
# Guard the active cache too: degrade to "no downloads" instead of raising.
try:
scans.append(scan_cache_dir())
except Exception as exc:
logger.warning("Could not scan active HF cache: %s", exc)
seen: set[str] = set()
try:
# Resolve the active cache dir for dedup.
from huggingface_hub.constants import HF_HUB_CACHE
seen.add(str(Path(HF_HUB_CACHE).resolve()))
except Exception:
pass
for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir):
try:
extra = extra_fn()
# is_dir()/resolve() can raise on an inaccessible path; skip it.
if not extra.is_dir():
continue
resolved = str(extra.resolve())
if resolved in seen:
continue
seen.add(resolved)
scans.append(scan_cache_dir(cache_dir = str(extra)))
except Exception as exc:
logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc)
return scans
from hub.utils.inventory_scan import all_hf_cache_scans
return all_hf_cache_scans()
def _is_gguf_filename(name: str) -> bool:
@ -3035,6 +2918,22 @@ def _is_main_gguf_filename(name: str) -> bool:
return _is_gguf_filename(name) and not _is_mmproj_filename(name)
def _main_variant_gguf_label(rel_path: str) -> Optional[str]:
name = rel_path.rsplit("/", 1)[-1]
if not _is_main_gguf_filename(name):
return None
if _is_mtp_drafter(rel_path):
return None
label = _extract_quant_label(rel_path)
if _is_big_endian_gguf_path(rel_path, label):
return None
return label
def _normalized_quant_label(label: str) -> str:
return label.lower().replace("-", "").replace("_", "")
def _repo_has_mmproj(repo_info) -> bool:
"""True if the repo ships a GGUF vision adapter (mmproj), so it can
take image inputs. Cheap: scans already-listed file names only."""
@ -3127,7 +3026,9 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
if repo_info.repo_type != "model":
continue
repo_id = repo_info.repo_id
if _is_hidden_model(repo_id):
# Pass the snapshot path too so the config check also hides
# custom Whisper checkpoints, not just curated repo ids.
if _is_hidden_model(repo_id, str(repo_info.repo_path)):
continue
total_size = _repo_gguf_size_bytes(repo_info)
if total_size == 0:
@ -3184,7 +3085,9 @@ async def list_cached_models(
if repo_info.repo_type != "model":
continue
repo_id = repo_info.repo_id
if _is_hidden_model(repo_id):
# Pass the snapshot path too so the config check also hides
# custom Whisper checkpoints, not just curated repo ids.
if _is_hidden_model(repo_id, str(repo_info.repo_path)):
continue
if _repo_has_gguf_files(repo_info):
continue
@ -3242,124 +3145,177 @@ async def list_cached_models(
async def delete_cached_model(
repo_id: str = Body(...),
variant: Optional[str] = Body(None),
cache_path: Optional[str] = Body(None),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Delete a cached model repo (or a specific GGUF variant) from the HF cache.
"""Compatibility route backed by the shared multi-cache deletion service."""
from hub.services.models import deletion
return await deletion.delete_cached_model_response(repo_id, variant, hf_token, cache_path)
With *variant*, only GGUF files matching that quant label are removed
(e.g. ``UD-Q4_K_XL``); otherwise the whole repo is deleted. Refuses
if the model is currently loaded for inference.
"""
def _resolve_cached_model_path(repo_id: str, variant: Optional[str]) -> Path:
"""Absolute path of a cached repo (newest snapshot dir) or, with *variant*,
that quant's main GGUF file (first split of a sharded quant). Paths come
from the HF cache scan only, so callers can't probe arbitrary paths."""
cache_scans = _all_hf_cache_scans()
matching_repos = []
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
if repo_info.repo_type != "model":
continue
if repo_info.repo_id.lower() == repo_id.lower():
matching_repos.append(repo_info)
if not matching_repos:
raise HTTPException(status_code = 404, detail = "Model not found in cache")
if variant:
want = _normalized_quant_label(variant)
candidate_revisions = sorted(
(rev for repo_info in matching_repos for rev in repo_info.revisions),
key = lambda rev: getattr(rev, "last_modified", 0) or 0,
reverse = True,
)
for rev in candidate_revisions:
snapshot = getattr(rev, "snapshot_path", None)
matches = []
for f in rev.files:
p = Path(f.file_path)
rel = f.file_name
if snapshot:
try:
rel = p.relative_to(snapshot).as_posix()
except ValueError:
pass
label = _main_variant_gguf_label(rel)
if label is None or _normalized_quant_label(label) != want:
continue
if p.exists() or p.is_symlink():
matches.append((rel, p))
if matches:
# Path-sorted so a sharded quant deterministically yields its first split.
return sorted(matches, key = lambda m: m[0].lower())[0][1]
raise HTTPException(
status_code = 404,
detail = f"Variant {variant} not found in cache for {repo_id}",
)
def repo_size(repo_info) -> int:
gguf_size = _repo_gguf_size_bytes(repo_info)
if gguf_size > 0:
return gguf_size
return sum(
(getattr(f, "size_on_disk", None) or 0)
for rev in repo_info.revisions
for f in rev.files
)
def repo_last_modified(repo_info) -> float:
return max(
(getattr(rev, "last_modified", 0) or 0 for rev in repo_info.revisions),
default = 0,
)
target_repo = max(
matching_repos,
key = lambda repo_info: (repo_size(repo_info), repo_last_modified(repo_info)),
)
# Whole repo: the newest revision's snapshot dir holds the visible files.
revisions = sorted(
(rev for rev in target_repo.revisions if getattr(rev, "snapshot_path", None)),
key = lambda rev: getattr(rev, "last_modified", 0) or 0,
reverse = True,
)
for rev in revisions:
p = Path(rev.snapshot_path)
if p.exists():
return p
p = Path(target_repo.repo_path)
if p.exists():
return p
raise HTTPException(status_code = 404, detail = "Cached model path not found")
def _wsl_reveal_in_explorer(path: Path) -> bool:
import subprocess
from utils.paths.path_utils import _IS_WSL
if not _IS_WSL:
return False
try:
windows_path = subprocess.run(
["wslpath", "-w", str(path)],
capture_output = True,
text = True,
check = True,
timeout = 10,
).stdout.strip()
if not windows_path:
return False
argument = f"/select,{windows_path}" if path.is_file() else windows_path
subprocess.Popen(["explorer.exe", argument])
return True
except (OSError, subprocess.SubprocessError):
return False
def _reveal_in_file_manager(path: Path) -> None:
"""Open the OS file manager with *path* selected (best effort per platform)."""
import subprocess
target = str(path)
if sys.platform == "darwin":
cmd = ["open", "-R", target] if path.is_file() else ["open", target]
subprocess.Popen(cmd)
elif os.name == "nt":
if path.is_file():
subprocess.Popen(["explorer", f"/select,{target}"])
else:
os.startfile(target) # noqa: S606 - local user's own file manager
elif not _wsl_reveal_in_explorer(path):
# No cross-desktop "select file" standard on Linux; open the directory.
directory = target if path.is_dir() else str(path.parent)
subprocess.Popen(["xdg-open", directory])
class CachedModelPathResponse(BaseModel):
path: str
is_dir: bool
@router.get("/cached-model-path", response_model = CachedModelPathResponse)
async def get_cached_model_path(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
variant: str = Query("", description = "Quantization variant (empty for whole repo)"),
current_subject: str = Depends(get_current_subject),
):
"""Absolute on-disk path of a cached repo or one of its GGUF variants."""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant.strip() or None)
return {"path": str(path), "is_dir": path.is_dir()}
# Refuse if the model is currently loaded.
@router.post("/reveal-cached-model")
async def reveal_cached_model(
repo_id: str = Body(...),
variant: Optional[str] = Body(None),
current_subject: str = Depends(get_current_subject),
):
"""Reveal a cached repo (or one GGUF variant's file) in the OS file manager."""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
variant = (variant or "").strip() or None
path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant)
try:
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and llama_backend.model_identifier:
loaded_id = llama_backend.model_identifier.lower()
if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "Unload the model before deleting",
)
except HTTPException:
raise
except Exception:
pass
try:
inference_backend = get_inference_backend()
if inference_backend.active_model_name:
active = inference_backend.active_model_name.lower()
if active == repo_id.lower() or active.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "Unload the model before deleting",
)
except HTTPException:
raise
except Exception:
pass
try:
cache_scans = _all_hf_cache_scans()
target_repo = None
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
if repo_info.repo_type != "model":
continue
if repo_info.repo_id.lower() == repo_id.lower():
target_repo = repo_info
break
if target_repo is not None:
break
if target_repo is None:
raise HTTPException(status_code = 404, detail = "Model not found in cache")
# ── Per-variant GGUF deletion ────────────────────────────
if variant:
deleted_bytes = 0
deleted_count = 0
for rev in target_repo.revisions:
for f in rev.files:
if not _is_gguf_filename(f.file_name):
continue
quant = _extract_quant_label(f.file_name)
if quant.lower() != variant.lower():
continue
# Delete the blob (data) and the snapshot symlink.
try:
blob = Path(f.blob_path)
snap = Path(f.file_path)
size = blob.stat().st_size if blob.exists() else 0
if snap.exists() or snap.is_symlink():
snap.unlink()
if blob.exists():
blob.unlink()
deleted_bytes += size
deleted_count += 1
except Exception as e:
logger.warning(f"Failed to delete {f.file_name}: {e}")
if deleted_count == 0:
raise HTTPException(
status_code = 404,
detail = f"Variant {variant} not found in cache for {repo_id}",
)
freed_mb = deleted_bytes / (1024 * 1024)
logger.info(
f"Deleted {deleted_count} file(s) for {repo_id} variant {variant}: "
f"{freed_mb:.1f} MB freed"
)
return {"status": "deleted", "repo_id": repo_id, "variant": variant}
# ── Full repo deletion ───────────────────────────────────
revision_hashes = [rev.commit_hash for rev in target_repo.revisions]
if not revision_hashes:
raise HTTPException(status_code = 404, detail = "No revisions found for model")
delete_strategy = hf_cache.delete_revisions(*revision_hashes)
logger.info(
f"Deleting cached model {repo_id}: "
f"{delete_strategy.expected_freed_size_str} will be freed"
)
delete_strategy.execute()
return {"status": "deleted", "repo_id": repo_id}
except HTTPException:
raise
await asyncio.to_thread(_reveal_in_file_manager, path)
except Exception as e:
logger.error(f"Error deleting cached model {repo_id}: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to delete cached model",
)
logger.error(f"Failed to reveal {path}: {e}")
raise HTTPException(status_code = 500, detail = "Failed to open file manager")
return {"status": "ok", "path": str(path)}
@router.get("/checkpoints", response_model = CheckpointListResponse)

119
studio/backend/routes/settings.py
View file

@ -60,6 +60,7 @@ from utils.embedding_model_settings import (
set_rag_embedding_model,
validate_embedding_model,
)
from utils.hf_cache_settings import cache_status, get_hf_cache_paths, set_hf_cache_home
router = APIRouter()
@ -89,6 +90,23 @@ class HelperPrecacheResponse(BaseModel):
disabled_by_env: bool
class HuggingFaceCachePayload(BaseModel):
cache_home: Optional[str] = Field(default = None, max_length = 4096)
class HuggingFaceCacheResponse(BaseModel):
cache_home: str
hub_cache: str
xet_cache: str
source: Literal["default", "studio", "environment"]
editable: bool
is_custom: bool
available: bool
writable: bool
free_bytes: Optional[int] = None
environment_variable: Optional[str] = None
class OpenAIAutoSwitchPayload(BaseModel):
enabled: bool
# None leaves the stored value untouched (partial updates can't clobber it).
@ -135,6 +153,30 @@ def _helper_precache_response(enabled: bool | None = None) -> HelperPrecacheResp
)
def _hugging_face_cache_response() -> HuggingFaceCacheResponse:
return HuggingFaceCacheResponse(**cache_status(get_hf_cache_paths()))
@router.get("/hugging-face-cache", response_model = HuggingFaceCacheResponse)
def get_hugging_face_cache(
current_subject: str = Depends(get_current_subject),
) -> HuggingFaceCacheResponse:
return _hugging_face_cache_response()
@router.put("/hugging-face-cache", response_model = HuggingFaceCacheResponse)
def update_hugging_face_cache(
payload: HuggingFaceCachePayload, current_subject: str = Depends(get_current_subject)
) -> HuggingFaceCacheResponse:
try:
set_hf_cache_home(payload.cache_home)
except RuntimeError as exc:
raise HTTPException(status_code = 409, detail = str(exc)) from exc
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
return _hugging_face_cache_response()
@router.get("/upload-limit", response_model = UploadLimitResponse)
def get_upload_limit(current_subject: str = Depends(get_current_subject)) -> UploadLimitResponse:
return _upload_limit_response(get_upload_limit_mb())
@ -416,6 +458,11 @@ def update_embedding_model(
log = logger,
) from exc
hf_token = (payload.hf_token or "").strip() or None
from utils.utils import hf_env_offline
# Offline, both the Hub malware scan and the is-embedding check are unreachable and degrade
# to the local cache below; capture the state once.
local_only_load = hf_env_offline()
# The env/default model needs no verification; saving it is a no-op override.
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
# what the backend loads, and HF metadata cannot verify a local path.
@ -439,26 +486,41 @@ def update_embedding_model(
# Fall back to the loader's own token so a gated/private repo is actually scanned
# (a token-less scan fails open for exactly the repo that would still load).
scan_token = hf_token or _ambient_hf_token()
# Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
# one blocks instead of passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys(
(
*security_load_subdirs(model, scan_token),
*_st_module_subdirs(model, scan_token),
# Offline: subdir probes would hit the network and hang; the offline gate walks the
# whole cached snapshot, so no load-subdir hints are needed.
if local_only_load:
load_subdirs = ()
else:
# Include ST module dirs (0_Transformer/) so a flagged pickle directly under one
# blocks instead of passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys(
(
*security_load_subdirs(model, scan_token),
*_st_module_subdirs(model, scan_token),
)
)
)
)
if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
if evaluate_file_security(
model,
hf_token = scan_token,
load_subdirs = load_subdirs,
local_only_load = local_only_load,
).blocked:
# 403, not 409: the client routes every 409 into the forceable "save anyway"
# flow, but this block is a hard, non-forceable security refusal.
raise HTTPException(
status_code = 403,
if local_only_load:
detail = (
f"{model!r} has cached pickle weights that cannot be security-scanned "
"offline and no safetensors alternative, so it cannot be used as the "
"embedding model. Re-download it with safetensors weights while online."
)
else:
detail = (
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
"cannot be used as the embedding model."
),
)
)
raise HTTPException(status_code = 403, detail = detail)
if model != default_embedding_model() and not payload.force and not is_local_gguf:
from core.rag import config as rag_config
@ -468,15 +530,28 @@ def update_embedding_model(
# which would wrongly 409 a valid online GGUF embedder.
gguf_named = _llama_backend_active() and rag_config._names_gguf(model)
if not gguf_named and not is_embedding_model(model, hf_token = hf_token):
raise HTTPException(
status_code = 409,
detail = (
f"Could not verify {model!r} as an embedding model on "
"Hugging Face (it may be the wrong model type, gated, or "
"you may be offline)."
),
)
gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token)
# Offline, is_embedding_model can only confirm the ST layout (modules.json); a
# transformers-native embedder (e.g. gte-modernbert) is unverifiable without Hub
# metadata. If already cached and loadable, accept it rather than raising a 409 that
# online would not (ST can load any cached encoder). Uncached -> 409.
from utils.utils import hf_cache_snapshot_is_loadable
# Require a genuinely loadable cache (config + weights), not just a resolved refs/main,
# so a metadata-only partial cache still gets the forceable 409.
offline_cached = local_only_load and hf_cache_snapshot_is_loadable(model)
if not offline_cached:
raise HTTPException(
status_code = 409,
detail = (
f"Could not verify {model!r} as an embedding model on "
"Hugging Face (it may be the wrong model type, gated, or "
"you may be offline)."
),
)
# The Hub GGUF probe (list_repo_files) can hang offline; skip it. Local check stays.
gguf_error = _local_gguf_backend_error(model)
if gguf_error is None and not local_only_load:
gguf_error = _hf_gguf_backend_error(model, hf_token)
if gguf_error:
raise HTTPException(status_code = 409, detail = gguf_error)
set_rag_embedding_model(model)

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

@ -109,7 +109,9 @@ async def get_hardware_utilization(current_subject: str = Depends(get_current_su
@router.get("/hardware/visible")
async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)):
from utils.hardware import get_visible_gpu_utilization
return get_visible_gpu_utilization()
# Off the event loop: the ROCm fallbacks shell out (Windows perf counters, sysfs) and the System view polls this route.
return await asyncio.to_thread(get_visible_gpu_utilization)
@router.post("/start")
@ -196,6 +198,7 @@ async def start_training(
request.local_eval_datasets, "Local eval dataset"
)
resume_output_dir: Optional[str] = None
resume_run: Optional[dict] = None
if request.resume_from_checkpoint:
try:
resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint)
@ -208,7 +211,7 @@ async def start_training(
if not resume_run or not can_resume_run(resume_run):
raise HTTPException(
status_code = 400,
detail = "Resume checkpoint must belong to a stopped run with saved trainer state.",
detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.",
)
resume_checkpoint = get_resume_checkpoint_path(resume_output_dir)
if not resume_checkpoint:
@ -412,53 +415,39 @@ async def start_training(
try:
from routes.training_vram import (
can_keep_chat_during_training,
free_chat_models_for_training,
summarize_resident_chat,
coordinate_models_for_training,
)
resident = summarize_resident_chat()
if not resident["any"]:
return
if resident.get("loading"):
# In-flight load can't be sized -> free rather than risk OOM.
freed = free_chat_models_for_training(reason = "chat model still loading")
logger.info("Freed in-flight chat load for training: %s", freed)
return
keep, info = can_keep_chat_during_training(
model_name = training_kwargs["model_name"],
hf_token = training_kwargs["hf_token"],
training_type = training_kwargs["training_type"],
load_in_4bit = training_kwargs["load_in_4bit"],
batch_size = training_kwargs["batch_size"],
max_seq_length = training_kwargs["max_seq_length"],
lora_rank = training_kwargs["lora_r"],
target_modules = training_kwargs["target_modules"],
gradient_checkpointing = training_kwargs["gradient_checkpointing"],
optimizer = training_kwargs["optim"],
gpu_ids = training_kwargs["gpu_ids"],
)
if keep:
logger.info(
"Keeping chat model(s) loaded during training "
"(free ~%s GB, needs ~%s GB): %s",
info.get("usable_gb"),
info.get("required_gb"),
resident,
def _can_keep_resident_models():
return can_keep_chat_during_training(
model_name = training_kwargs["model_name"],
hf_token = training_kwargs["hf_token"],
training_type = training_kwargs["training_type"],
load_in_4bit = training_kwargs["load_in_4bit"],
batch_size = training_kwargs["batch_size"],
max_seq_length = training_kwargs["max_seq_length"],
lora_rank = training_kwargs["lora_r"],
target_modules = training_kwargs["target_modules"],
gradient_checkpointing = training_kwargs["gradient_checkpointing"],
optimizer = training_kwargs["optim"],
gpu_ids = training_kwargs["gpu_ids"],
)
else:
freed = free_chat_models_for_training(
reason = "insufficient VRAM to run training alongside chat",
)
logger.info("Freed chat model(s) for training: %s", freed)
freed = coordinate_models_for_training(_can_keep_resident_models)
if freed:
logger.info("Freed models for training: %s", freed)
except Exception as e:
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
logger.warning("Inference/training memory coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
from utils.transformers_version import SidecarSwapInProgress
try:
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
job_id = job_id,
before_spawn = _free_vram_for_training,
resume_source_run_id = resume_run["id"] if resume_run else None,
**training_kwargs,
)
except SidecarSwapInProgress as exc:
# Expected loss of the race against a sidecar install: a retryable
@ -521,7 +510,10 @@ async def stop_training(
status = "idle", message = "No training job is currently running"
)
backend.stop_training(save = body.save)
if not backend.stop_training(save = body.save):
return TrainingStopResponse(
status = "idle", message = "No training job is currently running"
)
return TrainingStopResponse(
status = "stopped",
@ -637,9 +629,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject
"loss": getattr(progress, "loss", None),
"learning_rate": getattr(progress, "learning_rate", None),
}
output_dir = getattr(backend, "_output_dir", None)
if output_dir:
details["output_dir"] = output_dir
# Always present: an explicit null tells the client to drop a cached
# path (stop without save clears the run's output_dir).
details["output_dir"] = getattr(backend, "_output_dir", None) or None
# Metric history for chart recovery after SSE reconnection.
metric_history = None

148
studio/backend/routes/training_vram.py
View file

@ -1,15 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""VRAM coordination between chat/inference and training.
"""Memory coordination between inference and training.
Decides, from live free VRAM, whether a resident chat model can stay loaded
during training or must be unloaded, and unloads it across all backends
(HF/MLX orchestrator + llama.cpp GGUF server). In the route layer because the
GGUF accessor lives in routes/inference.py; backends are imported lazily.
Uses live free VRAM to keep resident chat and STT models when they fit. STT is
evicted before chat when training needs memory.
"""
from typing import Any, Dict, List, Optional, Tuple
from typing import Any, Callable, Dict, List, Optional, Tuple
from loggers import get_logger
@ -77,6 +75,37 @@ def summarize_resident_chat() -> Dict[str, Any]:
}
def summarize_resident_stt() -> Dict[str, Any]:
"""Report the resident dictation model (either engine). Never raises."""
try:
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
from core.inference.stt_sidecar import get_stt_sidecar
sidecar = get_stt_sidecar()
model = sidecar.loaded_model
device = sidecar.device
loading = sidecar.is_loading()
# whisper.cpp holds GPU memory via its subprocess, and both engines can be
# live at once (engine switch or direct /audio/stt/load). Always fold the
# GGUF sidecar in: a resident Transformers model must not mask a GGUF
# server still binding its backend, or admission lets training launch into
# that startup and OOM.
ggml = get_ggml_stt_sidecar()
if not model:
model = ggml.loaded_model
device = device or ggml.device
loading = loading or ggml.is_loading()
return {
"model": model,
"device": device,
"loading": loading,
"any": bool(model or loading),
}
except Exception as e:
logger.warning("Could not inspect STT sidecar: %s", e)
return {"model": None, "device": None, "loading": False, "any": False}
def can_keep_chat_during_training(
*,
model_name: str,
@ -366,3 +395,110 @@ def free_chat_models_for_training(reason: str) -> List[str]:
logger.warning("Could not unload GGUF chat model: %s", e)
return freed
def free_stt_model_for_training(reason: str) -> List[str]:
"""Unload the dictation model(s) before training. Never raises.
The Transformers and GGUF sidecars are freed under independent exception
boundaries so a failure unloading one backend never skips freeing the other
(both can hold accelerator memory at once after an engine switch).
"""
freed: List[str] = []
try:
from core.inference.stt_sidecar import get_stt_sidecar
sidecar = get_stt_sidecar()
if sidecar.is_loading() and sidecar.cancel_pending_load():
logger.info("Cancelling STT model load for training (%s)", reason)
# The loader may still be in from_pretrained()/.to(device) holding
# VRAM; wait for it to observe the cancel and release first.
sidecar.wait_for_load_to_settle()
# A load that finished before seeing the cancel leaves a resident
# model; unload it so training gets the memory back.
if sidecar.loaded_model:
sidecar.unload()
freed.append("stt:loading")
else:
model = sidecar.loaded_model
if model:
logger.info("Unloading STT model '%s' for training (%s)", model, reason)
sidecar.unload()
freed.append(f"stt:{model}")
except Exception as e:
logger.warning("Could not unload Transformers STT model: %s", e)
# Check the GGUF sidecar even after a cancelled/failed Transformers unload;
# both engines can hold memory at once (engine switch or direct load).
try:
from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar
ggml = get_ggml_stt_sidecar()
if ggml.is_loading() and ggml.cancel_pending_load():
logger.info("Cancelling GGUF STT model load for training (%s)", reason)
# whisper-server may still be binding its backend; wait for the
# cancelled startup to be killed and reaped before training claims
# the memory (loaded_model stays unset until it is ready).
ggml.wait_for_load_to_settle()
if ggml.loaded_model:
ggml.unload()
freed.append("stt:gguf-loading")
else:
ggml_model = ggml.loaded_model
if ggml_model:
logger.info("Unloading GGUF STT model '%s' for training (%s)", ggml_model, reason)
ggml.unload()
freed.append(f"stt:{ggml_model}")
except Exception as e:
logger.warning("Could not unload GGUF STT model: %s", e)
return freed
def coordinate_models_for_training(
can_keep: Callable[[], Tuple[bool, Dict[str, Any]]],
) -> List[str]:
"""Keep resident models when they fit, evicting STT before chat."""
resident_chat = summarize_resident_chat()
resident_stt = summarize_resident_stt()
if not resident_chat["any"] and not resident_stt["any"]:
return []
if resident_chat.get("loading"):
freed = free_stt_model_for_training(reason = "chat model still loading")
freed += free_chat_models_for_training(reason = "chat model still loading")
return freed
freed: List[str] = []
if resident_stt.get("loading"):
released_stt = free_stt_model_for_training(reason = "STT model still loading")
freed += released_stt
resident_stt = (
{"model": None, "device": None, "loading": False, "any": False}
if released_stt
else summarize_resident_stt()
)
if not resident_chat["any"] and not resident_stt["any"]:
return freed
keep, info = can_keep()
if keep:
logger.info(
"Keeping resident models loaded during training (free ~%s GB, needs ~%s GB): %s",
info.get("usable_gb"),
info.get("required_gb"),
{"chat": resident_chat, "stt": resident_stt},
)
return freed
if resident_stt["any"]:
freed += free_stt_model_for_training(reason = "insufficient training memory")
if not resident_chat["any"]:
return freed
keep, _info = can_keep()
if keep:
logger.info("Keeping chat model loaded after freeing STT: %s", resident_chat)
return freed
freed += free_chat_models_for_training(
reason = "insufficient VRAM to run training alongside chat",
)
return freed

74
studio/backend/routes/whisper.py Normal file
View file

@ -0,0 +1,74 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""whisper.cpp prebuilt status endpoint.
GET /api/whisper/update-status -> is a newer prebuilt available + job state
Detection reuses utils.whisper_cpp_freshness and fails open so the UI never
blocks on a missing marker / offline GitHub. There is no whisper-only update
trigger: whisper updates piggyback on the single main update item
(POST /api/llama/update chains a whisper phase when whisper is behind).
"""
from __future__ import annotations
import asyncio
from typing import Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject
from utils.whisper_cpp_update import get_update_status
router = APIRouter()
class WhisperUpdateJob(BaseModel):
state: str = Field("idle", description = "idle | running | success | error")
message: str = ""
from_tag: Optional[str] = None
to_tag: Optional[str] = None
reload_required: Optional[bool] = None
error: Optional[str] = None
progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.")
started_at: Optional[str] = None
finished_at: Optional[str] = None
class WhisperUpdateStatusResponse(BaseModel):
supported: bool = Field(
False,
description = "True when the install came from an Unsloth prebuilt (has a marker).",
)
update_available: bool = Field(
False, description = "True when the latest release is genuinely newer than the install."
)
stale: bool = Field(
False, description = "Update available AND install older than the staleness threshold."
)
installed_tag: Optional[str] = None
latest_tag: Optional[str] = None
published_repo: Optional[str] = None
installed_at_utc: Optional[str] = None
age_days: Optional[int] = None
source_build: bool = Field(
False, description = "True when there is no marker (source build) but a prebuilt is offered."
)
update_size_bytes: Optional[int] = Field(
None, description = "Download size of the prebuilt an update would fetch, in bytes."
)
job: WhisperUpdateJob = Field(default_factory = WhisperUpdateJob)
@router.get("/update-status", response_model = WhisperUpdateStatusResponse)
async def whisper_update_status(
force_refresh: bool = Query(
False, description = "Bypass the 24h release cache for an explicit check."
),
current_subject: str = Depends(get_current_subject),
) -> WhisperUpdateStatusResponse:
# Off the event loop: detection may probe the host and read GitHub.
status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh)
return WhisperUpdateStatusResponse(**status)

7
studio/backend/run.py
View file

@ -1244,6 +1244,13 @@ def _apply_supplied_password(password_value: "Optional[str]") -> None:
flush = True,
)
sys.exit(1)
if any(ch.isspace() for ch in supplied):
print(
"Error: password cannot contain spaces; not starting.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
if _is_current_password(supplied):
print(
"Error: the new password must differ from the current bootstrap "

114
studio/backend/storage/studio_db.py
View file

@ -192,13 +192,18 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
error_message TEXT,
duration_seconds REAL,
loss_sparkline TEXT,
display_name TEXT
display_name TEXT,
resume_blocked INTEGER NOT NULL DEFAULT 0
)
"""
)
existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()}
if "display_name" not in existing_cols:
conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
if "resume_blocked" not in existing_cols:
conn.execute(
"ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS training_metrics (
@ -734,16 +739,43 @@ def create_run(
config_json: str,
started_at: str,
total_steps: Optional[int],
*,
output_dir: Optional[str] = None,
cancel_requested: bool = False,
resumed_from_run_id: Optional[str] = None,
) -> None:
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps)
VALUES (?, ?, ?, ?, ?, ?)
INSERT INTO training_runs (
id, model_name, dataset_name, config_json, started_at, total_steps,
output_dir, resume_blocked
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(id, model_name, dataset_name, config_json, started_at, total_steps),
(
id,
model_name,
dataset_name,
config_json,
started_at,
total_steps,
None if cancel_requested else output_dir,
int(cancel_requested),
),
)
if resumed_from_run_id:
claimed = conn.execute(
"""
UPDATE training_runs SET resume_blocked = 1
WHERE id = ? AND status IN ('stopped', 'error')
AND output_dir = ? AND resume_blocked = 0
""",
(resumed_from_run_id, output_dir),
)
if claimed.rowcount != 1:
raise RuntimeError("Resume source is no longer available")
conn.commit()
finally:
conn.close()
@ -786,6 +818,8 @@ def finish_run(
loss_sparkline: Optional[str] = None,
output_dir: Optional[str] = None,
error_message: Optional[str] = None,
clear_output_dir: bool = False,
resume_blocked: bool = False,
) -> None:
conn = get_connection()
try:
@ -793,9 +827,16 @@ def finish_run(
"""
UPDATE training_runs
SET status = ?, ended_at = ?, final_step = ?, final_loss = ?,
duration_seconds = ?, loss_sparkline = ?, output_dir = ?,
error_message = ?
WHERE id = ?
duration_seconds = ?, loss_sparkline = ?,
output_dir = CASE
WHEN resume_blocked = 1 OR ? = 1 THEN NULL
WHEN ? IS NOT NULL THEN ?
WHEN ? IN ('error', 'stopped') THEN output_dir
ELSE NULL
END,
error_message = ?,
resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END
WHERE id = ? AND status = 'running'
""",
(
status,
@ -804,8 +845,13 @@ def finish_run(
final_loss,
duration_seconds,
loss_sparkline,
int(clear_output_dir),
output_dir,
output_dir,
status,
error_message,
int(clear_output_dir),
int(resume_blocked),
id,
),
)
@ -865,6 +911,38 @@ def update_run_display_name(id: str, display_name: Optional[str]) -> None:
conn.close()
def update_run_output_dir(id: str, output_dir: Optional[str]) -> None:
conn = get_connection()
try:
conn.execute(
"""
UPDATE training_runs SET output_dir = ?
WHERE id = ? AND status = 'running' AND resume_blocked = 0
""",
(output_dir, id),
)
conn.commit()
finally:
conn.close()
def mark_run_cancel_requested(id: str) -> bool:
"""Clear resume/export state only while the exact run is still active."""
conn = get_connection()
try:
cursor = conn.execute(
"""
UPDATE training_runs SET output_dir = NULL, resume_blocked = 1
WHERE id = ? AND status = 'running'
""",
(id,),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def list_runs(limit: int = 50, offset: int = 0) -> dict:
conn = get_connection()
try:
@ -874,15 +952,15 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
r.ended_at, r.total_steps, r.final_step, r.final_loss,
r.output_dir, r.duration_seconds, r.error_message,
r.loss_sparkline, r.display_name, r.config_json,
r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked,
CASE
WHEN r.status = 'stopped'
WHEN r.status IN ('stopped', 'error')
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
@ -917,13 +995,13 @@ def get_run(id: str) -> Optional[dict]:
"""
SELECT r.*,
CASE
WHEN r.status = 'stopped'
WHEN r.status IN ('stopped', 'error')
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
@ -958,12 +1036,12 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
0 AS resumed_later
FROM training_runs r
WHERE r.output_dir = ?
AND r.status = 'stopped'
AND r.status IN ('stopped', 'error')
AND NOT EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
ORDER BY r.started_at DESC
@ -1066,8 +1144,12 @@ def cleanup_orphaned_runs() -> None:
conn.execute(
"""
UPDATE training_runs
SET status = 'error',
error_message = 'Server restarted during training',
SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END,
error_message = CASE
WHEN resume_blocked = 1 THEN NULL
ELSE 'Server restarted during training'
END,
output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END,
ended_at = ?
WHERE status = 'running'
""",

319
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -66,6 +66,66 @@ def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path):
assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"]
def test_legacy_hf_scan_uses_snapshot_path_for_inactive_cache(tmp_path):
repo = tmp_path / "models--Org--Model"
snapshot = repo / "snapshots" / "revision"
snapshot.mkdir(parents = True)
[row] = models_route._scan_hf_cache(tmp_path, active_cache = False)
assert row.model_id == "Org/Model"
assert row.id == str(snapshot.resolve())
assert row.path == str(snapshot.resolve())
def test_collect_local_models_scans_previous_cache(monkeypatch, tmp_path):
active = tmp_path / "active"
previous = tmp_path / "previous"
active.mkdir()
snapshot = previous / "models--Org--Previous" / "snapshots" / "revision"
snapshot.mkdir(parents = True)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
monkeypatch.setattr("utils.paths.legacy_hf_cache_dir", lambda: tmp_path / "legacy")
monkeypatch.setattr("utils.paths.hf_default_cache_dir", lambda: tmp_path / "default")
monkeypatch.setattr("utils.paths.lmstudio_model_dirs", lambda: [])
monkeypatch.setattr("utils.hf_cache_settings.known_hf_hub_caches", lambda: [active, previous])
monkeypatch.setattr("storage.studio_db.list_scan_folders", lambda: [])
rows = models_route.collect_local_models(tmp_path / "models")
previous_row = next(row for row in rows if row.model_id == "Org/Previous")
assert previous_row.id == str(snapshot.resolve())
def test_collect_local_models_prefers_complete_previous_copy(monkeypatch, tmp_path):
active = tmp_path / "active"
previous = tmp_path / "previous"
active_partial = active / "models--Org--Model" / "blobs" / "abc.incomplete"
active_partial.parent.mkdir(parents = True)
active_partial.write_bytes(b"partial")
snapshot = previous / "models--Org--Model" / "snapshots" / "revision"
snapshot.mkdir(parents = True)
(snapshot / "model.safetensors").write_bytes(b"complete")
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
monkeypatch.setattr("utils.paths.legacy_hf_cache_dir", lambda: tmp_path / "legacy")
monkeypatch.setattr("utils.paths.hf_default_cache_dir", lambda: tmp_path / "default")
monkeypatch.setattr("utils.paths.lmstudio_model_dirs", lambda: [])
monkeypatch.setattr(
"utils.hf_cache_settings.known_hf_hub_caches",
lambda: [active, previous],
)
monkeypatch.setattr("storage.studio_db.list_scan_folders", lambda: [])
rows = models_route.collect_local_models(tmp_path / "models")
[row] = [row for row in rows if row.model_id == "Org/Model"]
assert row.id == str(snapshot.resolve())
assert row.partial is False
assert row.active_cache is False
def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monkeypatch, tmp_path):
repo = _repo(
"HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive",
@ -131,6 +191,72 @@ def test_is_hidden_model_hides_validation_probe_everywhere():
assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF")
def test_is_hidden_model_hides_dictation_models(tmp_path):
assert models_route._is_hidden_model("unsloth/whisper-tiny")
assert models_route._is_hidden_model("unsloth/whisper-base")
assert models_route._is_hidden_model("unsloth/whisper-small")
assert models_route._is_hidden_model("unsloth/whisper-large-v3-turbo")
assert models_route._is_hidden_model(
"/hf/models--unsloth--whisper-large-v3/snapshots/abc/model.safetensors"
)
assert not models_route._is_hidden_model("user/whisper-finetune")
assert not models_route._is_hidden_model(
"C:\\cache\\models--unsloth--whisper-small-finetune\\model.safetensors"
)
custom = tmp_path / "custom-whisper"
custom.mkdir()
(custom / "config.json").write_text(
'{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}'
)
(custom / "model.safetensors").write_bytes(b"weights")
assert models_route._is_hidden_model(
"user/custom-checkpoint",
str(custom / "model.safetensors"),
)
named_only = tmp_path / "whisper-finetune"
named_only.mkdir()
(named_only / "config.json").write_text('{"model_type": "llama"}')
assert not models_route._is_hidden_model("user/whisper-finetune", str(named_only))
def test_list_cached_models_hides_custom_whisper_by_config(monkeypatch, tmp_path):
# Regression: the legacy /cached-models picker must pass the snapshot path so
# the config check hides a custom (non-curated) Whisper checkpoint; a bare
# repo id cannot ("user/whisper-finetune" is not in the curated set).
repo_path = tmp_path / "models--user--whisper-finetune"
snap = repo_path / "snapshots" / "abc"
snap.mkdir(parents = True)
(snap / "config.json").write_text(
'{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}'
)
(snap / "model.safetensors").write_bytes(b"weights")
captured: list = []
real_hidden = models_route._is_hidden_model
def spy(*values):
captured.append(values)
return real_hidden(*values)
monkeypatch.setattr(models_route, "_is_hidden_model", spy)
repo = _repo(
"user/whisper-finetune",
[SimpleNamespace(file_name = "model.safetensors", size_on_disk = 10)],
repo_path,
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
)
result = asyncio.run(
models_route.list_cached_models(current_subject = "test-user", hf_token = None)
)
# The route passed the snapshot path (not just the repo id) ...
assert any(str(repo_path) in values for values in captured)
# ... so the custom Whisper checkpoint is hidden from the chat picker.
assert result["cached"] == []
def test_is_hidden_model_matches_repo_ids_exactly(monkeypatch):
"""A custom embedder with a generic basename is hidden by EXACT repo-id
match only, so unrelated cached repos that merely contain the basename stay
@ -573,33 +699,14 @@ def _gfile(name: str, size: int, mtime: float) -> SimpleNamespace:
)
def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_path):
"""An unreadable auxiliary cache (e.g. an inaccessible
``~/.cache/huggingface/hub``) must be skipped, not abort the scan.
Regression guard for ``extra.is_dir()`` raising and wiping the response.
"""
import huggingface_hub
import utils.paths as paths_mod
def test_all_hf_cache_scans_uses_shared_inventory(monkeypatch, tmp_path):
from hub.utils import inventory_scan
active = SimpleNamespace(
repos = [_repo("Org/Active", [_file("Q4_K_M.gguf", 5_000)], tmp_path / "active")]
)
def _fake_scan(cache_dir = None):
if cache_dir is None:
return active
raise AssertionError("auxiliary scan should have been skipped")
class _Boom:
def is_dir(self):
raise PermissionError(13, "Permission denied")
def resolve(self):
raise PermissionError(13, "Permission denied")
monkeypatch.setattr(huggingface_hub, "scan_cache_dir", _fake_scan)
monkeypatch.setattr(paths_mod, "legacy_hf_cache_dir", lambda: _Boom())
monkeypatch.setattr(paths_mod, "hf_default_cache_dir", lambda: _Boom())
monkeypatch.setattr(inventory_scan, "all_hf_cache_scans", lambda: [active])
scans = models_route._all_hf_cache_scans()
assert scans == [active]
@ -686,13 +793,17 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
"list_gguf_variants",
lambda repo_id, hf_token = None: (variants, True, []),
)
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
monkeypatch.setattr(
GV,
"_local_main_gguf_blobs_by_quant",
lambda _repo_id, repo_cache_dir = None: {},
)
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16"
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id, root = None: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -705,6 +816,52 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
assert flags["F16"] is False
def test_gguf_variants_route_scopes_local_probe_to_selected_cache(monkeypatch, tmp_path):
snapshot = tmp_path / "inactive" / "models--org--repo" / "snapshots" / "rev"
snapshot.mkdir(parents = True)
calls = []
async def scoped_variants(repo_id, **kwargs):
calls.append((repo_id, kwargs))
return SimpleNamespace(
repo_id = repo_id,
variants = [],
has_vision = False,
default_variant = None,
)
context_calls = []
monkeypatch.setattr(GV, "get_gguf_variants_response", scoped_variants)
monkeypatch.setattr(
models_route,
"_read_native_context_length",
lambda model, *, is_local: context_calls.append((model, is_local)) or 8192,
)
result = asyncio.run(
models_route.get_gguf_variants(
repo_id = "org/repo",
prefer_local_cache = True,
local_path = str(snapshot),
hf_token = None,
current_subject = "test-user",
)
)
assert calls == [
(
"org/repo",
{
"prefer_local_cache": True,
"local_path": str(snapshot),
"hf_token": None,
},
)
]
assert context_calls == [(str(snapshot), True)]
assert result.context_length == 8192
def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
siblings = [
SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100),
@ -726,12 +883,16 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
siblings,
),
)
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
monkeypatch.setattr(
GV,
"_local_main_gguf_blobs_by_quant",
lambda _repo_id, repo_cache_dir = None: {},
)
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id, root = None: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -758,12 +919,16 @@ def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, t
"list_gguf_variants",
lambda repo_id, hf_token = None: (variants, False, []),
)
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
monkeypatch.setattr(
GV,
"_local_main_gguf_blobs_by_quant",
lambda _repo_id, repo_cache_dir = None: {},
)
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id, root = None: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -774,66 +939,82 @@ def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, t
assert result.variants[0].downloaded is False
def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path):
"""A cached mmproj adapter must not count toward a same-label main
variant's download progress (mmproj-F16 vs an F16 weight)."""
import huggingface_hub.constants as hf_constants
def test_legacy_gguf_progress_delegates_to_shared_service(monkeypatch):
calls = []
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # only the adapter on disk
async def shared(repo_id, *, variant, expected_bytes, hf_token):
calls.append((repo_id, variant, expected_bytes, hf_token))
return {"downloaded_bytes": 10, "expected_bytes": 20, "progress": 0.5}
result = asyncio.run(
models_route.get_gguf_download_progress(
repo_id = "org/repo",
variant = "F16",
expected_bytes = 20_000,
current_subject = "test-user",
)
monkeypatch.setattr(
"hub.services.models.downloads.get_gguf_download_progress_response",
shared,
)
assert result["downloaded_bytes"] == 0
assert result["progress"] == 0
def test_gguf_download_progress_excludes_big_endian_sibling(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M-be.gguf").write_bytes(b"y" * 20_000)
result = asyncio.run(
models_route.get_gguf_download_progress(
repo_id = "org/repo",
variant = "Q4_K_M",
expected_bytes = 20_000,
expected_bytes = 20,
hf_token = "token",
current_subject = "test-user",
)
)
assert result["downloaded_bytes"] == 0
assert result["progress"] == 0
assert result["progress"] == 0.5
assert calls == [("org/repo", "Q4_K_M", 20, "token")]
def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
def test_legacy_model_progress_delegates_to_shared_service(monkeypatch):
calls = []
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
snap = tmp_path / "models--org--repo" / "snapshots" / "rev" / "Q4_K_M"
snap.mkdir(parents = True)
(snap / "foo.gguf").write_bytes(b"x" * 20_000)
async def shared(repo_id, *, hf_token):
calls.append((repo_id, hf_token))
return {"downloaded_bytes": 10, "expected_bytes": 20, "progress": 0.5}
monkeypatch.setattr(
"hub.services.models.downloads.get_download_progress_response",
shared,
)
result = asyncio.run(
models_route.get_gguf_download_progress(
models_route.get_download_progress(
repo_id = "org/repo",
variant = "Q4_K_M",
expected_bytes = 20_000,
hf_token = "token",
current_subject = "test-user",
)
)
assert result["downloaded_bytes"] == 20_000
assert result["progress"] == 1.0
assert result["progress"] == 0.5
assert calls == [("org/repo", "token")]
def test_legacy_delete_delegates_to_shared_service(monkeypatch):
calls = []
async def shared(
repo_id,
variant,
hf_token,
cache_path = None,
):
calls.append((repo_id, variant, hf_token, cache_path))
return {"status": "deleted", "repo_id": repo_id}
monkeypatch.setattr(
"hub.services.models.deletion.delete_cached_model_response",
shared,
)
result = asyncio.run(
models_route.delete_cached_model(
repo_id = "org/repo",
variant = None,
cache_path = "/data/hf/hub",
hf_token = "token",
current_subject = "test-user",
)
)
assert result == {"status": "deleted", "repo_id": "org/repo"}
assert calls == [("org/repo", None, "token", "/data/hf/hub")]

75
studio/backend/tests/test_change_password_policy.py Normal file
View file

@ -0,0 +1,75 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import importlib.util
import sys
from pathlib import Path
import pytest
from fastapi import HTTPException
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
from models.auth import ChangePasswordRequest # noqa: E402
# Load routes/auth.py directly so collection does not execute routes/__init__.py,
# which pulls in the heavy training/models/inference routers.
_route_path = _BACKEND_ROOT / "routes" / "auth.py"
_spec = importlib.util.spec_from_file_location("_change_password_route", _route_path)
assert _spec is not None and _spec.loader is not None
auth_routes = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(auth_routes)
@pytest.fixture
def _user(monkeypatch):
monkeypatch.setattr(
auth_routes.storage,
"get_user_and_secret",
lambda username: ("salt", "hash", "jwt-secret", False),
)
monkeypatch.setattr(
auth_routes.hashing,
"verify_password",
lambda password, salt, pwd_hash: password == "bootstrap-pw",
)
def _change(new_password):
payload = ChangePasswordRequest(
current_password = "bootstrap-pw",
new_password = new_password,
)
return asyncio.run(auth_routes.change_password(payload, None, "unsloth"))
def test_rejects_whitespace_only_password(_user):
with pytest.raises(HTTPException) as excinfo:
_change(" " * 8)
assert excinfo.value.status_code == 400
assert "spaces" in excinfo.value.detail
def test_rejects_tabs_and_spaces_password(_user):
with pytest.raises(HTTPException) as excinfo:
_change(" \t \t \t \t ")
assert excinfo.value.status_code == 400
def test_rejects_password_containing_spaces(_user):
with pytest.raises(HTTPException) as excinfo:
_change("correct horse battery")
assert excinfo.value.status_code == 400
assert "spaces" in excinfo.value.detail
def test_allows_password_without_spaces(_user, monkeypatch):
monkeypatch.setattr(auth_routes.storage, "update_password", lambda *args, **kwargs: True)
monkeypatch.setattr(auth_routes, "create_access_token", lambda subject: "at")
monkeypatch.setattr(auth_routes, "create_refresh_token", lambda subject: "rt")
token = _change("correct-horse-battery")
assert token.access_token == "at"
assert token.must_change_password is False

74
studio/backend/tests/test_chat_load_during_training.py
View file

@ -801,6 +801,80 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(guard_called, [])
def _validate_gguf_template(
self,
*,
template,
canonical_path = "/picked/model.gguf",
):
# Drive validate_model for a native lease-backed GGUF template probe and
# capture what the embedded-template reader was called with.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
model_path = "model.gguf",
gguf_variant = "Q4_K_M",
native_path_lease = "signed-lease",
include_chat_template = True,
)
cfg = SimpleNamespace(
identifier = canonical_path,
display_name = "model.gguf",
is_gguf = True,
is_lora = False,
is_vision = False,
gguf_file = canonical_path,
path = None,
base_model = None,
)
import utils.models.gguf_metadata as gguf_meta
seen = {}
def _fake_read(path):
seen["path"] = path
return template
guard_called = []
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = (canonical_path, "model.gguf", True),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
patch.object(gguf_meta, "read_gguf_chat_template", _fake_read),
patch.object(
self.route,
"_guard_chat_load_against_training",
lambda *a, **kw: guard_called.append(True),
),
):
resp = asyncio.run(self.route.validate_model(request, current_subject = "u"))
return resp, seen, guard_called
def test_include_chat_template_reads_leased_gguf_embedded_template(self):
# The picker chat-template GET has no lease plumbing, so a native picked
# GGUF surfaces its default template through this lease-aware probe: the
# embedded template is read from the granted canonical path and returned.
resp, seen, _ = self._validate_gguf_template(template = "{{ messages }}")
self.assertEqual(resp.chat_template, "{{ messages }}")
# Read strictly the leased file's own embedded template, never a sibling
# sidecar: the grant authorizes just this one path.
self.assertEqual(seen["path"], "/picked/model.gguf")
def test_include_chat_template_skips_training_guard(self):
# A template-only probe allocates no VRAM, so like include_context_length
# it must not be refused by the training guard.
_, _, guard_called = self._validate_gguf_template(template = "{{ messages }}")
self.assertEqual(guard_called, [])
def test_include_chat_template_over_cap_is_dropped(self):
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
resp, _, _ = self._validate_gguf_template(template = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1))
self.assertIsNone(resp.chat_template)
# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────

223
studio/backend/tests/test_cloudflare_tunnel.py
View file

@ -403,11 +403,234 @@ def test_reader_ignores_api_endpoint_failure_line():
assert t.error == "cloudflared exited before emitting a tunnel URL"
# ── public reachability probe ────────────────────────────────────────
class _FakeResponse:
def __init__(self, body):
self._body = body
def read(self, size = -1):
return self._body
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def _patch_urlopen(monkeypatch, handler):
import urllib.request
monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: handler(req))
@pytest.fixture(autouse = True)
def _stub_dns_wait(monkeypatch, request):
if request.node.name.startswith("test_verify_public_url"):
monkeypatch.setattr(ct, "_wait_for_dns", lambda *a, **kw: None)
def test_wait_for_dns_polls_until_answer(monkeypatch):
calls = []
def handler(req):
calls.append(req.full_url)
if len(calls) < 3:
return _FakeResponse(b'{"Status":3}')
return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}')
_patch_urlopen(monkeypatch, handler)
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5)
assert len(calls) == 3
assert "name=words.trycloudflare.com" in calls[0]
def test_wait_for_dns_gives_up_at_deadline(monkeypatch):
_patch_urlopen(monkeypatch, lambda req: _FakeResponse(b'{"Status":3}'))
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 0.05)
def test_wait_for_dns_retries_transient_doh_error(monkeypatch):
calls = []
def handler(req):
calls.append(req.full_url)
if len(calls) < 3:
raise OSError("transient")
return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}')
_patch_urlopen(monkeypatch, handler)
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5)
assert len(calls) == 3
def test_wait_for_dns_bails_on_persistent_doh_errors(monkeypatch):
calls = []
def handler(req):
calls.append(req.full_url)
raise OSError("blocked")
_patch_urlopen(monkeypatch, handler)
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5)
assert len(calls) == ct._DNS_MAX_DOH_ERRORS
def test_verify_public_url_accepts_studio_marker(monkeypatch):
seen = {}
def handler(req):
seen["url"] = req.full_url
return _FakeResponse(b'{"status":"healthy","service":"Unsloth UI Backend"}')
_patch_urlopen(monkeypatch, handler)
assert ct.verify_public_url("https://words.trycloudflare.com") is True
assert seen["url"] == "https://words.trycloudflare.com/api/health"
def test_verify_public_url_waits_for_dns_first(monkeypatch):
order = []
monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: order.append(("dns", host)))
def handler(req):
order.append(("probe", req.full_url))
return _FakeResponse(b'{"service":"Unsloth UI Backend"}')
_patch_urlopen(monkeypatch, handler)
assert ct.verify_public_url("https://words.trycloudflare.com") is True
assert order[0] == ("dns", "words.trycloudflare.com")
assert order[1][0] == "probe"
def test_verify_public_url_dns_wait_and_probe_share_deadline(monkeypatch):
# An exhausted DNS wait leaves the probe a single attempt, not a fresh window.
calls = []
monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: None)
def handler(req):
calls.append(req.full_url)
raise OSError("unreachable")
_patch_urlopen(monkeypatch, handler)
assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0) is False
assert len(calls) == 1
def test_verify_public_url_retries_then_succeeds(monkeypatch):
calls = []
def handler(req):
calls.append(req.full_url)
if len(calls) < 3:
raise OSError("Name or service not known")
return _FakeResponse(b'{"service":"Unsloth UI Backend"}')
_patch_urlopen(monkeypatch, handler)
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
assert ct.verify_public_url("https://words.trycloudflare.com") is True
assert len(calls) == 3
def test_verify_public_url_rejects_unreachable_host(monkeypatch):
def handler(req):
raise OSError("Name or service not known")
_patch_urlopen(monkeypatch, handler)
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False
def test_verify_public_url_rejects_foreign_responder(monkeypatch):
# e.g. a Cloudflare error page: no service marker in the body.
_patch_urlopen(monkeypatch, lambda req: _FakeResponse(b"<html>error 1033</html>"))
monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False
@pytest.fixture(autouse = True)
def _stub_public_probe(monkeypatch, request):
# start_studio_tunnel tests use fake hostnames; keep them off the network.
if not request.node.name.startswith("test_start_studio_tunnel"):
return
monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: True)
def test_start_studio_tunnel_no_binary(monkeypatch):
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None)
assert ct.start_studio_tunnel(8080) is None
def test_start_studio_tunnel_drops_url_that_is_not_publicly_reachable(monkeypatch):
attempts = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
attempts.append(protocol)
def start(self):
self.url = "https://words.trycloudflare.com"
def wait_for_ready(self, timeout):
return self.url
def stop(self):
pass
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: False)
assert ct.start_studio_tunnel(8080) is None
assert attempts == [None]
assert ct._active_tunnel is None
def test_start_studio_tunnel_returns_url_once_probe_passes(monkeypatch):
probed = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
self.protocol = protocol
def start(self):
self.url = "https://words.trycloudflare.com"
def wait_for_ready(self, timeout):
return self.url
def stop(self):
pass
def _probe(url, **kw):
probed.append(url)
return True
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
monkeypatch.setattr(ct, "verify_public_url", _probe)
try:
assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com"
assert probed == ["https://words.trycloudflare.com"]
finally:
ct.stop_studio_tunnel()
def test_start_studio_tunnel_registers_before_wait(monkeypatch):
# The tunnel must be visible to stop_studio_tunnel() during the readiness
# wait, else a shutdown in that window orphans cloudflared.

735
studio/backend/tests/test_combined_update.py Normal file
View file

@ -0,0 +1,735 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hermetic tests for the combined llama+whisper update item.
llama.cpp is the single main update item; whisper.cpp piggybacks on it. These
pin the union status (update_available = llama behind OR whisper behind), the
chained apply (llama phase first, whisper phase only when behind), the failure
policy (llama failure aborts; whisper failure keeps the llama partial success),
the silent whisper skips, and the backward-compatible payload shape.
"""
from __future__ import annotations
import json
import sys
import time
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import utils.llama_cpp_freshness as freshness # noqa: E402
import utils.llama_cpp_update as upd # noqa: E402
import utils.whisper_cpp_freshness as wfresh # noqa: E402
import utils.whisper_cpp_update as wupd # noqa: E402
MARKER = "UNSLOTH_PREBUILT_INFO.json"
WHISPER_MARKER = "UNSLOTH_WHISPER_PREBUILT_INFO.json"
# The top-level status and job fields that predate the whisper piggyback; the
# combined payload must stay an exact superset so current UI code keeps working.
LEGACY_STATUS_FIELDS = {
"supported",
"update_available",
"stale",
"installed_tag",
"latest_tag",
"published_repo",
"installed_at_utc",
"age_days",
"source_build",
"update_size_bytes",
"job",
}
LEGACY_JOB_FIELDS = {
"state",
"message",
"from_tag",
"to_tag",
"reload_required",
"error",
"progress",
"started_at",
"finished_at",
}
class _FakeInstallerPopen:
"""Stands in for the streamed llama installer process."""
def __init__(
self,
cmd,
*,
returncode = 0,
lines = None,
on_start = None,
**kwargs,
):
if on_start is not None:
on_start(list(cmd))
self.returncode = returncode
self.stdout = iter(lines or [])
def wait(self):
return self.returncode
def kill(self):
pass
def _patch_llama_installer(
monkeypatch,
*,
returncode = 0,
lines = None,
on_start = None,
):
# Only intercept the installer invocation: importing routes.inference inside
# the worker can Popen unrelated host probes (ldconfig etc).
def _popen(cmd, **kw):
is_installer = any("install_llama_prebuilt" in str(part) for part in cmd)
return _FakeInstallerPopen(
cmd,
returncode = returncode if is_installer else 0,
lines = lines if is_installer else None,
on_start = on_start if is_installer else None,
)
monkeypatch.setattr(upd.subprocess, "Popen", _popen)
def _write_llama_install(dir_: Path, tag: str) -> str:
"""Create a fake llama prebuilt install and return the llama-server path."""
bin_dir = dir_ / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
binary = bin_dir / "llama-server"
binary.write_text("stub")
(dir_ / MARKER).write_text(
json.dumps(
{
"tag": tag,
"release_tag": tag,
"published_repo": "unslothai/llama.cpp",
"installed_at_utc": "2020-01-01T00:00:00Z",
}
)
)
return str(binary)
def _write_whisper_install(
dir_: Path,
tag: str,
backend: str = "cpu",
) -> str:
"""Create a fake whisper prebuilt install and return the whisper-server path."""
bin_dir = dir_ / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
binary = bin_dir / "whisper-server"
binary.write_text("stub")
(dir_ / WHISPER_MARKER).write_text(
json.dumps(
{
"release_tag": tag,
"upstream_tag": tag.split("-")[0],
"published_repo": "unslothai/whisper.cpp",
"backend": backend,
"installed_at_utc": "2020-01-01T00:00:00Z",
}
)
)
return str(binary)
@pytest.fixture(autouse = True)
def _clean_state(monkeypatch, tmp_path):
freshness.reset_caches()
wfresh.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
wupd._resolve_memo.clear()
monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".llama_cache")
monkeypatch.setattr(wfresh, "_cache_dir", lambda: tmp_path / ".whisper_cache")
for var in (
"LLAMA_SERVER_PATH",
"UNSLOTH_LLAMA_CPP_PATH",
"WHISPER_SERVER_PATH",
"UNSLOTH_WHISPER_CPP_PATH",
):
monkeypatch.delenv(var, raising = False)
# Never hit the network in these tests.
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
monkeypatch.setattr(wfresh, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
yield
freshness.reset_caches()
wfresh.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
wupd._resolve_memo.clear()
def _setup_llama(
monkeypatch,
tmp_path,
*,
installed = "b9493",
latest = "b9518",
):
"""Marker-managed llama install; behind when installed != latest."""
install_dir = tmp_path / "llama.cpp"
binary = _write_llama_install(install_dir, installed)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest)
return install_dir
def _setup_whisper(
monkeypatch,
tmp_path,
*,
installed = "v1.9.1-unsloth.1",
latest = "v1.9.2-unsloth.1",
):
"""Marker-managed whisper install; behind when latest is newer."""
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, installed)
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
monkeypatch.setattr(wupd, "_installer_script", lambda: tmp_path / "install_whisper_prebuilt.py")
monkeypatch.setattr(wfresh, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest)
return install_dir
def _patch_whisper_phase(
monkeypatch,
events,
*,
to_tag = "v1.9.2-unsloth.1",
error = None,
):
"""Record whisper phase runs without touching a real installer."""
def _run(phase, set_progress):
events.append("whisper")
if error is not None:
raise RuntimeError(error)
set_progress(0.5)
return {
"to_tag": to_tag,
"reload_required": False,
"message": f"Updated whisper.cpp to {to_tag}.",
}
monkeypatch.setattr(wupd, "run_chained_phase", _run)
def _wait_for_job():
deadline = time.time() + 10
while time.time() < deadline:
with upd._job_lock:
job = dict(upd._job)
if job["state"] in ("success", "error"):
return job
time.sleep(0.05)
with upd._job_lock:
return dict(upd._job)
# --- status: the single item folds whisper in ---
def test_status_payload_is_exact_superset_of_legacy_fields(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
st = upd.get_update_status(force_refresh = True)
assert LEGACY_STATUS_FIELDS <= set(st)
assert LEGACY_JOB_FIELDS <= set(st["job"])
# The new fields ride alongside, never replacing the legacy ones.
assert st["llama_update_available"] is True
assert st["whisper"]["update_available"] is True
assert st["whisper"]["latest_tag"] == "v1.9.2-unsloth.1"
assert st["update_component"] == "llama"
def test_status_union_whisper_only_surfaces_update(monkeypatch, tmp_path):
# llama current, whisper behind: the single item still shows an update.
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path)
st = upd.get_update_status(force_refresh = True)
assert st["llama_update_available"] is False
assert st["whisper"]["update_available"] is True
assert st["update_available"] is True
assert st["update_component"] == "whisper"
assert st["installed_tag"] == "b9518"
assert st["latest_tag"] == "b9518"
assert st["whisper"]["installed_tag"] == "v1.9.1-unsloth.1"
assert st["whisper"]["latest_tag"] == "v1.9.2-unsloth.1"
def test_status_whisper_current_does_not_flip_union(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1")
st = upd.get_update_status(force_refresh = True)
assert st["update_available"] is False
assert st["whisper"]["skip_reason"] == "up_to_date"
assert st["update_component"] is None
def test_status_survives_whisper_probe_failure(monkeypatch, tmp_path):
# The piggyback fails open: llama status still works without a whisper probe.
_setup_llama(monkeypatch, tmp_path)
def _boom(*, force_refresh = False):
raise RuntimeError("probe exploded")
monkeypatch.setattr(wupd, "chained_phase_plan", _boom)
st = upd.get_update_status(force_refresh = True)
assert st["update_available"] is True
assert st["whisper"] is None
# --- whisper chained_phase_plan: silent skips ---
def test_whisper_plan_skips_local_link(monkeypatch, tmp_path):
monkeypatch.setattr(wupd, "_find_binary", lambda: str(tmp_path / "whisper-server"))
monkeypatch.setattr(wupd, "_active_install_is_local_link", lambda b: True)
plan = wupd.chained_phase_plan()
assert plan["update_available"] is False
assert plan["skip_reason"] == "local_link"
assert plan["phase"] is None
def test_whisper_plan_skips_source_build(monkeypatch, tmp_path):
binary = tmp_path / "whisper.cpp" / "build" / "bin" / "whisper-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # no marker
monkeypatch.setattr(wupd, "_find_binary", lambda: str(binary))
plan = wupd.chained_phase_plan()
assert plan["skip_reason"] == "source_build"
assert plan["phase"] is None
def test_whisper_update_targets_canonical_root_when_inner_marker_exists(tmp_path):
install_dir = tmp_path / "whisper.cpp"
binary = install_dir / "build" / "bin" / "whisper-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
(install_dir / WHISPER_MARKER).write_text("{}")
(binary.parent / WHISPER_MARKER).write_text("{}")
assert wupd._install_dir_for(str(binary)) == install_dir
def test_whisper_plan_skips_when_not_installed(monkeypatch):
monkeypatch.setattr(wupd, "_find_binary", lambda: None)
plan = wupd.chained_phase_plan()
assert plan["skip_reason"] == "not_installed"
assert plan["phase"] is None
def test_whisper_plan_eligible_when_behind(monkeypatch, tmp_path):
install_dir = _setup_whisper(monkeypatch, tmp_path)
script = tmp_path / "install_whisper_prebuilt.py"
script.write_text("stub")
plan = wupd.chained_phase_plan(force_refresh = True)
assert plan["update_available"] is True
assert plan["skip_reason"] is None
assert plan["phase"]["install_dir"] == install_dir
assert plan["phase"]["repo"] == "unslothai/whisper.cpp"
assert plan["phase"]["backend"] == "cpu"
# Pin to the exact release the freshness check offered: unpinned, the
# installer's download-host /releases/latest pointer can lag published_at
# and reinstall an older build in a loop.
assert plan["phase"]["pin_release_tag"] == "v1.9.2-unsloth.1"
def test_whisper_plan_requires_a_repairable_pair_for_slim_installs(monkeypatch, tmp_path):
install_dir = _setup_whisper(monkeypatch, tmp_path)
marker_path = install_dir / WHISPER_MARKER
marker = json.loads(marker_path.read_text())
marker["install_kind"] = "slim"
marker_path.write_text(json.dumps(marker))
wfresh.reset_caches()
monkeypatch.setattr(
wupd,
"_resolve_prebuilt_for_host",
lambda **kwargs: {"prebuilt_available": False},
)
plan = wupd.chained_phase_plan(force_refresh = True)
assert plan["update_available"] is False
assert plan["skip_reason"] == "paired_llama_unavailable"
repaired = wupd.chained_phase_plan(
force_refresh = True,
paired_llama_will_update = True,
)
assert repaired["update_available"] is True
assert repaired["phase"] is not None
def test_whisper_phase_pins_installer_to_checked_release(monkeypatch, tmp_path):
calls = []
monkeypatch.setattr(
wupd._flow,
"stream_installer",
lambda cmd, env, **kw: calls.append(cmd),
)
monkeypatch.setattr(wupd, "reset_caches", lambda **kw: None)
monkeypatch.setattr(wupd, "latest_published_release", lambda repo, **kw: "v9")
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, "v9")
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
wupd.run_chained_phase(
{
"install_dir": install_dir,
"repo": "unslothai/whisper.cpp",
"asset": None,
"backend": "cpu",
"script": tmp_path / "install_whisper_prebuilt.py",
"pin_release_tag": "v9",
},
lambda f: None,
)
cmd = calls[0]
assert "--published-release-tag" in cmd
assert cmd[cmd.index("--published-release-tag") + 1] == "v9"
def test_whisper_phase_exit_2_is_a_failed_phase(monkeypatch, tmp_path):
# No install occurred, so incompatibility must remain an actionable job
# error instead of producing a false success toast and hiding the banner.
def _raise_exit_2(cmd, env, **kw):
raise wupd._flow.InstallerExit(2, "installer exited 2: incompatible release")
monkeypatch.setattr(wupd._flow, "stream_installer", _raise_exit_2)
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, "v1")
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
with pytest.raises(wupd._flow.InstallerExit) as exc_info:
wupd.run_chained_phase(
{
"install_dir": install_dir,
"repo": "unslothai/whisper.cpp",
"asset": None,
"backend": "cpu",
"script": tmp_path / "install_whisper_prebuilt.py",
"pin_release_tag": None,
},
lambda f: None,
)
assert exc_info.value.returncode == 2
def test_llama_update_survives_unavailable_whisper_module(monkeypatch, tmp_path):
import builtins
llama_dir = _setup_llama(monkeypatch, tmp_path)
monkeypatch.setattr(upd, "_whisper_chain_status", lambda **kw: None)
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: _write_llama_install(llama_dir, "b9518"),
)
real_import = builtins.__import__
def guarded_import(
name,
globals = None,
locals = None,
fromlist = (),
level = 0,
):
if name == "utils" and "whisper_cpp_update" in fromlist:
raise AssertionError("whisper module was re-imported after its failed probe")
return real_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(builtins, "__import__", guarded_import)
# A failed optional whisper probe must not be followed by an unconditional
# import. The valid llama phase still starts and completes.
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert job["phases"]["llama"]["state"] == "success"
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "unavailable"
def test_macos_status_uses_compatible_resolver_release(monkeypatch, tmp_path):
_setup_whisper(
monkeypatch,
tmp_path,
installed = "v1.9.1-unsloth.1",
latest = "v1.9.2-unsloth.1",
)
monkeypatch.setattr(wupd.sys, "platform", "darwin")
monkeypatch.setattr(
wupd,
"_resolve_prebuilt_for_host",
lambda **kw: {
"prebuilt_available": True,
"release_tag": "v1.9.1-unsloth.1",
},
)
status = wupd.get_update_status(force_refresh = True)
assert status["latest_tag"] == "v1.9.1-unsloth.1"
assert status["update_available"] is False
assert status["stale"] is False
def test_whisper_phase_integrity_failure_is_not_swallowed(monkeypatch, tmp_path):
def _raise_exit_1(cmd, env, **kw):
raise wupd._flow.InstallerExit(1, "installer exited 1: checksum mismatch")
monkeypatch.setattr(wupd._flow, "stream_installer", _raise_exit_1)
install_dir = tmp_path / "whisper.cpp"
binary = _write_whisper_install(install_dir, "v1")
monkeypatch.setattr(wupd, "_find_binary", lambda: binary)
with pytest.raises(wupd._flow.InstallerExit, match = "checksum mismatch"):
wupd.run_chained_phase(
{
"install_dir": install_dir,
"repo": "unslothai/whisper.cpp",
"asset": None,
"backend": "cpu",
"script": tmp_path / "install_whisper_prebuilt.py",
"pin_release_tag": None,
},
lambda f: None,
)
# --- apply: the chained job ---
def test_apply_runs_llama_then_whisper(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events)
res = upd.start_update()
assert res["started"] is True, res
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["llama", "whisper"] # llama phase strictly first
assert job["phases"]["llama"]["state"] == "success"
assert job["phases"]["llama"]["to_tag"] == "b9518"
assert job["phases"]["whisper"]["state"] == "success"
assert job["phases"]["whisper"]["to_tag"] == "v1.9.2-unsloth.1"
# Legacy top-level fields keep their llama meaning.
assert job["from_tag"] == "b9493"
assert job["to_tag"] == "b9518"
assert "Updated llama.cpp to b9518." in job["message"]
assert "Updated whisper.cpp to v1.9.2-unsloth.1." in job["message"]
assert job["progress"] == 1.0
assert LEGACY_JOB_FIELDS <= set(job)
def test_apply_llama_only_when_whisper_current(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1")
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["llama"]
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "up_to_date"
def test_apply_whisper_only_noops_llama(monkeypatch, tmp_path):
# llama current + whisper behind: the same single apply runs, with the llama
# phase a cheap already-matches no-op and the whisper phase doing the work.
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
events = []
_patch_llama_installer(monkeypatch, on_start = lambda cmd: events.append("llama"))
_patch_whisper_phase(monkeypatch, events)
res = upd.start_update()
assert res["started"] is True, res
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["whisper"] # the llama installer never ran
# The legacy job-level to_tag means "llama tag"; a whisper-only round
# leaves it unset so the UI never reports a llama update that never ran.
assert job["to_tag"] is None
assert job["phases"]["llama"]["state"] == "skipped"
assert job["phases"]["llama"]["reason"] == "up_to_date"
assert job["phases"]["whisper"]["state"] == "success"
assert "Updated whisper.cpp to v1.9.2-unsloth.1." in job["message"]
def test_whisper_reload_never_raises_job_reload_flag(monkeypatch, tmp_path):
# A whisper-only update that had to unload a warm sidecar reports
# reload_required on its phase, but the JOB flag stays down: the chat
# frontend resyncs (and clears the local checkpoint) off the job flag,
# which must mean "the llama server changed", not "the sidecar restarted".
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
def _whisper_phase(phase, set_progress):
return {
"to_tag": "v1.9.2-unsloth.1",
"reload_required": True,
"message": "Updated whisper.cpp to v1.9.2-unsloth.1.",
}
monkeypatch.setattr(wupd, "run_chained_phase", _whisper_phase)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert job["phases"]["whisper"]["reload_required"] is True
assert not job["reload_required"]
def test_apply_refuses_when_both_current(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518")
_setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1")
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"
def test_apply_llama_failure_aborts_before_whisper(monkeypatch, tmp_path):
_setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
events = []
_patch_llama_installer(monkeypatch, returncode = 2, lines = ["boom: disk full\n"])
_patch_whisper_phase(monkeypatch, events)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "error", job
assert "boom" in (job["error"] or "")
assert events == [] # whisper never attempted
assert job["phases"]["llama"]["state"] == "error"
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "aborted"
assert job["message"] == "llama.cpp update failed."
def test_apply_whisper_failure_keeps_llama_partial_success(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
# An active model makes the llama phase report reload_required.
import threading
from types import ModuleType
class _FakeBackend:
def __init__(self):
self._serial_load_lock = threading.Lock()
self._llama_update_in_progress = False
self.is_active = True
def unload_model(self):
self.is_active = False
backend = _FakeBackend()
routes_pkg = ModuleType("routes")
routes_pkg.__path__ = []
inference_mod = ModuleType("routes.inference")
inference_mod.get_llama_cpp_backend = lambda: backend
monkeypatch.setitem(sys.modules, "routes", routes_pkg)
monkeypatch.setitem(sys.modules, "routes.inference", inference_mod)
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events, error = "whisper installer exploded")
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "error", job
assert events == ["llama", "whisper"]
# The message says both halves: llama landed, whisper did not.
assert "Updated llama.cpp to b9518." in job["message"]
assert "whisper.cpp update failed." in job["message"]
assert "whisper installer exploded" in (job["error"] or "")
# The llama phase's reload_required survives the whisper failure.
assert job["reload_required"] is True
assert job["to_tag"] == "b9518"
assert job["phases"]["llama"]["state"] == "success"
assert job["phases"]["whisper"]["state"] == "error"
def test_apply_skips_whisper_local_link_silently(monkeypatch, tmp_path):
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
monkeypatch.setattr(wupd, "_active_install_is_local_link", lambda b: True)
events = []
_patch_llama_installer(
monkeypatch,
on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")),
)
_patch_whisper_phase(monkeypatch, events)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert events == ["llama"]
assert job["phases"]["whisper"]["state"] == "skipped"
assert job["phases"]["whisper"]["reason"] == "local_link"
assert job["message"] == "Updated llama.cpp to b9518."
def test_chained_progress_windows(monkeypatch, tmp_path):
# The llama phase fills roughly the first 0.7 slice and whisper the rest.
llama_dir = _setup_llama(monkeypatch, tmp_path)
_setup_whisper(monkeypatch, tmp_path)
(tmp_path / "install_whisper_prebuilt.py").write_text("stub")
seen = {}
def _whisper_phase(phase, set_progress):
with upd._job_lock:
seen["at_whisper_start"] = upd._job["progress"]
set_progress(0.5)
with upd._job_lock:
seen["mid_whisper"] = upd._job["progress"]
return {"to_tag": "v1.9.2-unsloth.1", "reload_required": False, "message": "ok"}
monkeypatch.setattr(wupd, "run_chained_phase", _whisper_phase)
_patch_llama_installer(
monkeypatch,
lines = ["Downloading app.tar.gz: 100.0% (35.0 MiB/35.0 MiB) at 9.0 MiB/s\n"],
on_start = lambda cmd: _write_llama_install(llama_dir, "b9518"),
)
assert upd.start_update()["started"] is True
job = _wait_for_job()
assert job["state"] == "success", job
assert seen["at_whisper_start"] == pytest.approx(0.7)
assert seen["mid_whisper"] == pytest.approx(0.7 + 0.5 * 0.3)
assert job["progress"] == 1.0

9
studio/backend/tests/test_consent_gate.py
View file

@ -873,6 +873,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
if fn == "config.json":
import json
@ -899,6 +900,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -932,6 +934,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -972,6 +975,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -1008,6 +1012,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -1037,6 +1042,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -1079,6 +1085,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -1120,6 +1127,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile
@ -1182,6 +1190,7 @@ class TestScannerCoversAllExecutableCode:
repo,
fn,
token = None,
cache_dir = None,
):
import json
import tempfile

50
studio/backend/tests/test_embedding_model_security_gate.py
View file

@ -106,6 +106,56 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch):
assert unverified.status_code == 409
def test_offline_cached_non_st_model_is_accepted(client, monkeypatch):
# Offline, a cached transformers-native embedder (no modules.json) is unverifiable via HF
# metadata, but ST can load any cached encoder, so accept it (no 409).
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
import utils.models as _models
import utils.utils as _uu
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: True)
r = c.put("/embedding-model", json = {"embedding_model": "acme/gte-modernbert"})
assert r.status_code == 200
assert saved.get("model") == "acme/gte-modernbert"
def test_offline_partial_or_uncached_model_still_409(client, monkeypatch):
# Offline but not loadable (uncached or metadata-only partial cache): keep the forceable
# 409, since the cache-only load would fail anyway.
c, _saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
import utils.models as _models
import utils.utils as _uu
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: False)
r = c.put("/embedding-model", json = {"embedding_model": "acme/uncached-embedder"})
assert r.status_code == 409
def test_offline_skips_remote_gguf_probe(client, monkeypatch):
# Offline + llama backend: the remote GGUF probe (list_repo_files) must be skipped so a
# dead-DNS session cannot hang.
c, _saved = client
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
monkeypatch.setattr(settings, "_local_gguf_backend_error", lambda model: None)
def _boom(*a, **k):
raise AssertionError("hit the network for the GGUF probe")
monkeypatch.setattr(settings, "_hf_gguf_backend_error", _boom)
import utils.models as _models
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: True)
r = c.put("/embedding-model", json = {"embedding_model": "acme/embedder"})
assert r.status_code == 200
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.

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

@ -158,6 +158,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None
utils_model_config._extract_quant_label = lambda value: value
utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False
utils_model_config._is_mtp_drafter = lambda *args, **kwargs: False
utils_model_config.is_audio_input_type = lambda *args, **kwargs: None
monkeypatch.setitem(
sys.modules,

59
studio/backend/tests/test_gguf_load_cache_reuse.py
View file

@ -103,6 +103,10 @@ def _build_cache(
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: _types.SimpleNamespace(hub_cache = tmp_path),
)
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
return tmp_path
@ -117,6 +121,61 @@ def _fail_get_paths_info(*_args, **_kwargs):
class TestLoadReusesCachedCopy:
def test_download_uses_selected_cache_for_lookup_preflight_and_write(
self, tmp_path, monkeypatch
):
backend = LlamaCppBackend()
selected = tmp_path / "selected" / "hub"
startup = tmp_path / "startup" / "hub"
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(startup))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: _types.SimpleNamespace(hub_cache = selected),
)
seen = {"lookups": [], "disk": [], "downloads": []}
def cached_lookup(
repo_id,
filename,
*,
cache_dir = None,
**_kwargs,
):
seen["lookups"].append((repo_id, filename, cache_dir))
return None
def disk_usage(path):
seen["disk"].append(str(path))
return _types.SimpleNamespace(free = 1024)
def download(repo_id, filename, _token, **kwargs):
seen["downloads"].append((repo_id, filename, kwargs.get("cache_dir")))
return str(selected / filename)
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch(
"huggingface_hub.get_paths_info",
lambda _repo, paths, **_kwargs: [
_types.SimpleNamespace(path = path, size = 4) for path in paths
],
),
patch("huggingface_hub.try_to_load_from_cache", cached_lookup),
patch("core.inference.llama_cpp.shutil.disk_usage", disk_usage),
patch(
"core.inference.llama_cpp.hf_hub_download_with_xet_fallback",
download,
),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(selected / MAIN)
assert seen == {
"lookups": [(REPO, MAIN, str(selected))],
"disk": [str(selected)],
"downloads": [(REPO, MAIN, str(selected))],
}
def test_online_reuse_after_revision_bump(self, hf_cache):
"""A new repo revision does not replace a complete cached model."""
backend = LlamaCppBackend()

205
studio/backend/tests/test_gpu_memory_mode.py
View file

@ -733,20 +733,211 @@ def test_split_pin_without_mask_only_sets_pci_order(monkeypatch):
def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch):
# ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR
# mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP
# would index into the already-reduced set).
# ROCm with the mask sourced from HIP: the pin must land in
# HIP_VISIBLE_DEVICES too, and an inherited ROCR mask is cleared so the
# mask can't apply twice (ROCR re-indexes, then HIP would index into the
# already-reduced set).
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = "6.0")
monkeypatch.setitem(sys.modules, "torch", torch_stub)
env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"}
_rocm_torch_stub(monkeypatch)
env = {
"CUDA_VISIBLE_DEVICES": "3,1",
"HIP_VISIBLE_DEVICES": "3,1",
"ROCR_VISIBLE_DEVICES": "3,1",
}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
assert env["HIP_VISIBLE_DEVICES"] == "1,3"
assert "ROCR_VISIBLE_DEVICES" not in env
def test_split_pin_preserves_inherited_rocr_mask(monkeypatch):
# Mask sourced from ROCR alone (e.g. an AMD SDK parent): the pin must
# re-emit at the ROCr layer, not swap to HIP -- clearing ROCR re-exposes
# every agent to HSA enumeration, which can segfault at startup on an
# unsupported GPU the parent mask was hiding (#7272 review). CUDA carries
# the post-ROCR ordinals, mirroring the prefer_rocr emission.
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
_rocm_torch_stub(monkeypatch)
env = {"ROCR_VISIBLE_DEVICES": "3,1"}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["ROCR_VISIBLE_DEVICES"] == "1,3"
assert env["CUDA_VISIBLE_DEVICES"] == "0,1"
assert "HIP_VISIBLE_DEVICES" not in env
def test_split_pin_keeps_hip_on_windows_despite_stray_rocr(monkeypatch):
# On Windows the ROCR var is dead (no ROCr layer) and the resolver never
# reads it, so a stray value must not flip the pin to the ROCR emission:
# the HIP mask is the only effective selector there.
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = "6.0")
monkeypatch.setitem(sys.modules, "torch", torch_stub)
monkeypatch.setattr(sys, "platform", "win32")
env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "9"}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
assert env["HIP_VISIBLE_DEVICES"] == "1,3"
assert "ROCR_VISIBLE_DEVICES" not in env
def _rocm_torch_stub(monkeypatch):
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = "6.0")
monkeypatch.setitem(sys.modules, "torch", torch_stub)
# prefer_rocr is Linux-only (ROCR is an ROCr variable); pin the platform so
# these Linux-behaviour tests also pass on a Windows dev box.
monkeypatch.setattr(sys, "platform", "linux")
def test_subset_pin_masks_via_rocr_on_rocm(monkeypatch):
# A GPU-subset pin must exclude the rest at the ROCr/HSA layer: HIP masking
# still enumerates every agent first, which segfaults the build on an
# unsupported deselected GPU (e.g. a gfx1103 iGPU under a gfx110X prebuilt).
# ROCR drops it at the driver layer; only one mask is set (HIP cleared).
_rocm_torch_stub(monkeypatch)
env = {"HIP_VISIBLE_DEVICES": "9"} # stale/inherited HIP mask must not survive
LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True)
assert env["ROCR_VISIBLE_DEVICES"] == "0"
assert env["CUDA_VISIBLE_DEVICES"] == "0"
assert "HIP_VISIBLE_DEVICES" not in env
def test_prefer_rocr_remaps_cuda_to_post_rocr_ordinals(monkeypatch):
# ROCR re-indexes the visible agents from 0, and HIP (cleared here) falls back
# to CUDA_VISIBLE_DEVICES -- so on the prefer_rocr path CUDA must carry the
# post-ROCR ordinals, not the physical ids, else a non-zero pick indexes out
# of range and the child sees no GPU and drops to CPU (#7272 review).
_rocm_torch_stub(monkeypatch)
# Single non-zero GPU: ROCR keeps the physical id, CUDA becomes ordinal 0.
env = {}
LlamaCppBackend._emit_child_gpu_visibility(env, "1", prefer_rocr = True)
assert env["ROCR_VISIBLE_DEVICES"] == "1"
assert env["CUDA_VISIBLE_DEVICES"] == "0"
assert "HIP_VISIBLE_DEVICES" not in env
# Multi-GPU subset: ROCR keeps the physical ids, CUDA is the 0-based ordinals.
env = {}
LlamaCppBackend._emit_child_gpu_visibility(env, "1,3", prefer_rocr = True)
assert env["ROCR_VISIBLE_DEVICES"] == "1,3"
assert env["CUDA_VISIBLE_DEVICES"] == "0,1"
assert "HIP_VISIBLE_DEVICES" not in env
def test_subset_pin_default_still_uses_hip_and_clears_rocr(monkeypatch):
# Without prefer_rocr the masking is unchanged: HIP narrows, inherited ROCR
# is cleared so the two can't double-mask.
_rocm_torch_stub(monkeypatch)
env = {"ROCR_VISIBLE_DEVICES": "0,1"}
LlamaCppBackend._emit_child_gpu_visibility(env, "1")
assert env["HIP_VISIBLE_DEVICES"] == "1"
assert "ROCR_VISIBLE_DEVICES" not in env
def test_cpu_only_pin_keeps_hip_even_with_prefer_rocr(monkeypatch):
# The CPU-only sentinel never routes through ROCR (no portable "hide all"
# spelling); it hides every GPU via HIP.
_rocm_torch_stub(monkeypatch)
env = {}
LlamaCppBackend._emit_child_gpu_visibility(env, "-1", prefer_rocr = True)
assert env["HIP_VISIBLE_DEVICES"] == "-1"
assert "ROCR_VISIBLE_DEVICES" not in env
def _amd_sdk_torch_stub(monkeypatch):
# AMD SDK wheel: torch.version.hip is None but __version__ encodes rocm.
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = None)
torch_stub.__version__ = "2.9.1+rocm7.2.1"
monkeypatch.setitem(sys.modules, "torch", torch_stub)
monkeypatch.setattr(sys, "platform", "linux")
def test_prefer_rocr_falls_back_to_hip_on_windows(monkeypatch):
# ROCR_VISIBLE_DEVICES is a Linux ROCr variable (Windows HIP has no ROCr
# layer), so on Windows ROCm prefer_rocr must keep the HIP mask or a nonzero
# pick loses its only effective selector (#7272 review).
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = "6.0")
monkeypatch.setitem(sys.modules, "torch", torch_stub)
monkeypatch.setattr(sys, "platform", "win32")
env = {"ROCR_VISIBLE_DEVICES": "9"}
LlamaCppBackend._emit_child_gpu_visibility(env, "1", prefer_rocr = True)
assert env["HIP_VISIBLE_DEVICES"] == "1"
assert env["CUDA_VISIBLE_DEVICES"] == "1"
assert "ROCR_VISIBLE_DEVICES" not in env
def test_amd_sdk_wheel_hip_none_still_masks_rocr(monkeypatch):
# An AMD SDK wheel leaves torch.version.hip unset but has "rocm" in __version__.
# It must still get the ROCR mask, else only CUDA_VISIBLE_DEVICES is set and an
# unsupported iGPU keeps enumerating and can crash llama-server.
_amd_sdk_torch_stub(monkeypatch)
env = {"HIP_VISIBLE_DEVICES": "9"}
LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True)
assert env["ROCR_VISIBLE_DEVICES"] == "0"
assert "HIP_VISIBLE_DEVICES" not in env
def test_cuda_wheel_hip_none_gets_no_rocm_mask(monkeypatch):
# A CUDA wheel (hip=None, no "rocm" in __version__) must NOT get a HIP/ROCR mask
# -- only CUDA_VISIBLE_DEVICES -- so the version-string check can't false-positive.
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = None)
torch_stub.__version__ = "2.9.1+cu124"
monkeypatch.setitem(sys.modules, "torch", torch_stub)
env = {}
LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True)
assert env["CUDA_VISIBLE_DEVICES"] == "0"
assert "ROCR_VISIBLE_DEVICES" not in env
assert "HIP_VISIBLE_DEVICES" not in env
def test_resolve_physical_ids_reads_rocr_on_amd_sdk_wheel(monkeypatch):
# _resolve_visible_physical_ids must use the same ROCm detection as
# _emit_child_gpu_visibility: on an AMD SDK wheel (hip=None, rocm in
# __version__) an inherited ROCR mask IS the ordinal->physical mapping.
# Reading it as "no mask" labels ordinal 0 as physical 0 and the child's
# ROCR pin then re-exposes the GPU the mask was hiding (#7272 review).
_amd_sdk_torch_stub(monkeypatch)
for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
monkeypatch.delenv(var, raising = False)
monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1")
assert LlamaCppBackend._resolve_visible_physical_ids() == [1]
def test_resolve_physical_ids_ignores_rocr_on_cuda_wheel(monkeypatch):
# A CUDA wheel (hip=None, no "rocm") keeps CUDA-only semantics: a stray
# ROCR var must not be read as the mask.
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = None)
torch_stub.__version__ = "2.9.1+cu124"
monkeypatch.setitem(sys.modules, "torch", torch_stub)
for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
monkeypatch.delenv(var, raising = False)
monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1")
assert LlamaCppBackend._resolve_visible_physical_ids() is None
def test_resolve_physical_ids_ignores_rocr_on_windows(monkeypatch):
# ROCR_VISIBLE_DEVICES is a Linux ROCr variable: Windows HIP has no ROCr
# layer, so a stray ROCR var there does not mask the runtime. Reading it as
# the ordinal->physical mapping would label ordinal 0 with a stale ROCR id
# while the runtime still enumerates every adapter, so auto-selection could
# budget one card and pin another (#7272 review). HIP must still be honoured.
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = None)
torch_stub.__version__ = "2.9.1+rocm7.2.1" # AMD SDK wheel
monkeypatch.setitem(sys.modules, "torch", torch_stub)
monkeypatch.setattr(sys, "platform", "win32")
for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
monkeypatch.delenv(var, raising = False)
monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1")
assert LlamaCppBackend._resolve_visible_physical_ids() is None
# HIP precedence is unchanged on Windows.
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
assert LlamaCppBackend._resolve_visible_physical_ids() == [1]
# ── Diffusion single-device selection ───────────────────────────────────────

290
studio/backend/tests/test_hf_cache_settings.py Normal file
View file

@ -0,0 +1,290 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import os
import sys
import threading
import time
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from hub.services.models.common import _local_model_info
from utils import hf_cache_settings
from utils import native_path_leases
@pytest.fixture()
def settings_store(monkeypatch, tmp_path):
store = {}
monkeypatch.setattr(hf_cache_settings, "_EXPLICIT_CACHE_ENV", {})
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
monkeypatch.setattr(
"storage.studio_db.get_app_setting",
lambda key, fallback = None: store.get(key, fallback),
)
monkeypatch.setattr(
"storage.studio_db.upsert_app_settings",
lambda values: store.update(values) or values,
)
return store
def test_studio_cache_switch_is_live_and_keeps_history(settings_store, tmp_path):
first = tmp_path / "external-a" / "huggingface"
second = tmp_path / "external-b" / "huggingface"
first.parent.mkdir()
second.parent.mkdir()
selected = hf_cache_settings.set_hf_cache_home(str(first))
assert selected.hub_cache == first / "hub"
assert selected.xet_cache == first / "xet"
assert selected.child_env({}) == {
"HF_HUB_CACHE": str(first / "hub"),
"HF_XET_CACHE": str(first / "xet"),
}
hf_cache_settings.set_hf_cache_home(str(second))
assert settings_store[hf_cache_settings.CACHE_HISTORY_SETTING_KEY] == [str(first)]
assert first / "hub" in hf_cache_settings.known_hf_hub_caches()
reset = hf_cache_settings.set_hf_cache_home(None)
assert reset.source == "default"
assert second in hf_cache_settings.known_hf_cache_homes()
def test_environment_cache_is_read_only(monkeypatch, tmp_path):
custom = tmp_path / "managed"
monkeypatch.setattr(
hf_cache_settings,
"_EXPLICIT_CACHE_ENV",
{"HF_HOME": str(custom)},
)
paths = hf_cache_settings.get_hf_cache_paths()
assert paths.source == "environment"
assert paths.editable is False
assert paths.hub_cache == custom / "hub"
with pytest.raises(RuntimeError, match = "environment variable"):
hf_cache_settings.set_hf_cache_home(str(tmp_path / "other"))
def test_explicit_hub_cache_is_the_displayed_location(monkeypatch, tmp_path):
custom_hub = tmp_path / "models-cache"
custom_hub.mkdir()
monkeypatch.setattr(
hf_cache_settings,
"_EXPLICIT_CACHE_ENV",
{"HF_HUB_CACHE": str(custom_hub)},
)
paths = hf_cache_settings.get_hf_cache_paths()
status = hf_cache_settings.cache_status(paths)
assert paths.cache_home == custom_hub
assert paths.hub_cache == custom_hub
assert status["cache_home"] == str(custom_hub)
assert status["available"] is True
assert custom_hub / "hub" not in hf_cache_settings.known_hf_hub_caches()
def test_explicit_hub_cache_display_wins_over_hf_home(monkeypatch, tmp_path):
hf_home = tmp_path / "hf-home"
custom_hub = tmp_path / "other-disk" / "models-cache"
hf_home.mkdir()
custom_hub.mkdir(parents = True)
monkeypatch.setattr(
hf_cache_settings,
"_EXPLICIT_CACHE_ENV",
{"HF_HOME": str(hf_home), "HF_HUB_CACHE": str(custom_hub)},
)
paths = hf_cache_settings.get_hf_cache_paths()
assert paths.cache_home == custom_hub
assert paths.hub_cache == custom_hub
assert paths.xet_cache == hf_home / "xet"
assert custom_hub / "hub" not in hf_cache_settings.known_hf_hub_caches()
assert hf_home / "hub" in hf_cache_settings.known_hf_hub_caches()
def test_xet_only_override_keeps_model_cache_editable(settings_store, monkeypatch, tmp_path):
xet_cache = tmp_path / "chunks"
stored = tmp_path / "stored-cache"
settings_store[hf_cache_settings.CACHE_HOME_SETTING_KEY] = str(stored)
monkeypatch.setattr(
hf_cache_settings,
"_EXPLICIT_CACHE_ENV",
{"HF_XET_CACHE": str(xet_cache)},
)
paths = hf_cache_settings.get_hf_cache_paths()
assert paths.cache_home == stored
assert paths.hub_cache == stored / "hub"
assert paths.xet_cache == xet_cache
assert paths.editable is True
selected = tmp_path / "selected-cache"
selected.parent.mkdir(exist_ok = True)
updated = hf_cache_settings.set_hf_cache_home(str(selected))
assert updated.hub_cache == selected / "hub"
assert updated.xet_cache == xet_cache
def test_worker_environment_is_applied_before_import(monkeypatch, tmp_path):
hub = str(tmp_path / "hub")
xet = str(tmp_path / "xet")
observed = {}
class Module:
@staticmethod
def run():
import os
return os.environ["HF_HUB_CACHE"], os.environ["HF_XET_CACHE"]
def fake_import(name):
import os
observed["name"] = name
observed["hub"] = os.environ.get("HF_HUB_CACHE")
return Module
monkeypatch.setattr(native_path_leases.importlib, "import_module", fake_import)
result = native_path_leases.run_without_native_path_secret(
"fake.worker",
"run",
{"HF_HUB_CACHE": hub, "HF_XET_CACHE": xet},
)
assert observed == {"name": "fake.worker", "hub": hub}
assert result == (hub, xet)
def test_spawn_environment_is_applied_then_restored(monkeypatch, tmp_path):
hub = str(tmp_path / "hub")
xet = str(tmp_path / "xet")
monkeypatch.setenv("HF_HUB_CACHE", "parent-hub")
monkeypatch.delenv("HF_XET_CACHE", raising = False)
with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": hub, "HF_XET_CACHE": xet}):
import os
assert os.environ["HF_HUB_CACHE"] == hub
assert os.environ["HF_XET_CACHE"] == xet
assert os.environ["HF_HUB_CACHE"] == "parent-hub"
assert "HF_XET_CACHE" not in os.environ
def test_spawn_environment_supports_nested_contexts(monkeypatch):
monkeypatch.setenv("HF_HUB_CACHE", "parent")
with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "outer"}):
assert os.environ["HF_HUB_CACHE"] == "outer"
with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "inner"}):
assert os.environ["HF_HUB_CACHE"] == "inner"
assert os.environ["HF_HUB_CACHE"] == "outer"
assert os.environ["HF_HUB_CACHE"] == "parent"
def test_spawn_environment_serializes_threads(monkeypatch):
monkeypatch.setenv("HF_HUB_CACHE", "parent")
first_entered = threading.Event()
release_first = threading.Event()
observations: list[tuple[str, str]] = []
def first():
with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "first"}):
observations.append(("first", os.environ["HF_HUB_CACHE"]))
first_entered.set()
assert release_first.wait(timeout = 2)
def second():
assert first_entered.wait(timeout = 2)
with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "second"}):
observations.append(("second", os.environ["HF_HUB_CACHE"]))
first_thread = threading.Thread(target = first)
second_thread = threading.Thread(target = second)
first_thread.start()
second_thread.start()
assert first_entered.wait(timeout = 2)
time.sleep(0.02)
assert observations == [("first", "first")]
release_first.set()
first_thread.join(timeout = 2)
second_thread.join(timeout = 2)
assert observations == [("first", "first"), ("second", "second")]
assert os.environ["HF_HUB_CACHE"] == "parent"
def test_cache_switch_invalidates_inventory(settings_store, tmp_path, monkeypatch):
invalidations = []
monkeypatch.setattr(
"hub.utils.inventory_scan.invalidate_hf_cache_scans",
lambda: invalidations.append(True),
)
selected = tmp_path / "external" / "huggingface"
selected.parent.mkdir()
hf_cache_settings.set_hf_cache_home(str(selected))
assert invalidations == [True]
def test_cache_validation_write_tests_hub_and_xet(settings_store, tmp_path, monkeypatch):
selected = tmp_path / "external" / "huggingface"
selected.parent.mkdir()
tested = []
real_named_temporary_file = hf_cache_settings.tempfile.NamedTemporaryFile
def recording_write_test(*args, **kwargs):
tested.append(Path(kwargs["dir"]))
return real_named_temporary_file(*args, **kwargs)
monkeypatch.setattr(
hf_cache_settings.tempfile,
"NamedTemporaryFile",
recording_write_test,
)
hf_cache_settings.set_hf_cache_home(str(selected))
assert tested == [selected / "hub", selected / "xet"]
def test_cache_validation_rejects_unwritable_child(settings_store, tmp_path, monkeypatch):
selected = tmp_path / "external" / "huggingface"
selected.parent.mkdir()
def reject_hub(*args, **kwargs):
if Path(kwargs["dir"]).name == "hub":
raise PermissionError("read-only")
raise AssertionError("xet should not be tested after hub fails")
monkeypatch.setattr(hf_cache_settings.tempfile, "NamedTemporaryFile", reject_hub)
with pytest.raises(ValueError, match = "permission"):
hf_cache_settings.set_hf_cache_home(str(selected))
def test_inactive_cache_model_loads_from_snapshot_path(tmp_path):
snapshot = tmp_path / "snapshots" / "revision"
snapshot.mkdir(parents = True)
row = _local_model_info(
scan_path = snapshot,
load_path = snapshot,
source = "hf_cache",
model_format = "safetensors",
model_id = "org/model",
active_cache = False,
)
assert row.model_id == "org/model"
assert row.active_cache is False
assert row.load_id == str(snapshot)

32
studio/backend/tests/test_hf_xet_fallback.py
View file

@ -101,13 +101,23 @@ def test_shim_injects_studio_prepare_on_http_retry(monkeypatch):
prepared = []
monkeypatch.setattr(
"hub.utils.download_registry.prepare_cache_for_transport",
lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)),
lambda repo_type, repo_id, mode, *a, **k: prepared.append(
(repo_type, repo_id, mode, k.get("root"))
),
)
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
selected_cache = "/captured/hub"
out = xf.hf_hub_download_with_xet_fallback(
DL_REPO,
FILE,
None,
cache_dir = selected_cache,
)
assert out == "/cache/model.gguf"
assert seen_disable_xet == [False, True] # Xet first, then HTTP
assert prepared == [("model", DL_REPO, "http")], "shim must run Unsloth's marker-aware prep"
assert prepared == [
("model", DL_REPO, "http", Path(selected_cache))
], "shim must prepare the cache captured by the download"
def test_shim_snapshot_injects_studio_prepare(monkeypatch):
@ -120,10 +130,22 @@ def test_shim_snapshot_injects_studio_prepare(monkeypatch):
return "/tmp/snap-dir"
monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot)
out = xf.snapshot_download_with_xet_fallback("org/model")
selected_cache = "/captured/hub"
out = xf.snapshot_download_with_xet_fallback(
"org/model",
cache_dir = selected_cache,
)
assert out == "/tmp/snap-dir"
assert captured["repo_id"] == "org/model"
assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http
prepared = []
monkeypatch.setattr(
"hub.utils.download_registry.prepare_cache_for_transport",
lambda repo_type, repo_id, mode, *a, **k: prepared.append(
(repo_type, repo_id, mode, k.get("root"))
),
)
captured["prepare_for_http_fn"]("model", "org/model")
assert prepared == [("model", "org/model", "http", Path(selected_cache))]
def test_degrades_gracefully_without_shared_helper(monkeypatch):

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

@ -200,14 +200,9 @@ def _gpu_linux_host(caps):
)
def test_host_is_blackwell_includes_datacenter_parts():
assert ilp._host_is_blackwell(_gpu_linux_host(["10.0"])) is True # B200 sm_100
assert ilp._host_is_blackwell(_gpu_linux_host(["10.3"])) is True # B300 sm_103
assert ilp._host_is_blackwell(_gpu_linux_host(["12.0"])) is True # RTX 50 sm_120
assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121
assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper
assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere
assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins
# _host_is_blackwell / _blackwell_min_toolkit_for_host are prebuilt_core
# re-exports; their value tables moved verbatim to
# tests/studio/install/test_prebuilt_core.py.
def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile):
@ -285,16 +280,6 @@ def test_drop_blackwell_incapable_windows_cuda_applies_to_datacenter():
assert [a.name for a in kept] == [cuda13.name]
def test_blackwell_min_toolkit_is_sm_aware():
# Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it.
f = ilp._blackwell_min_toolkit_for_host
assert f(_gpu_linux_host(["10.0"])) == (12, 8) # B200
assert f(_gpu_linux_host(["12.0"])) == (12, 8) # RTX 50
assert f(_gpu_linux_host(["10.3"])) == (12, 9) # B300
assert f(_gpu_linux_host(["12.1"])) == (12, 9) # DGX Spark
assert f(_gpu_linux_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins
def test_sm103_host_drops_cuda128_windows_build():
# B300 (sm_103) needs cuda-12.9: a legacy win-cuda-12.8 build must be dropped.
host = _host(

231
studio/backend/tests/test_install_whisper_prebuilt_checksums.py Normal file
View file

@ -0,0 +1,231 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Trust-anchor tests for install_whisper_prebuilt.py.
Whisper verifies each download against the release's own
whisper-prebuilt-sha256.json checksum index (the same model as
install_llama_prebuilt.py), not a committed pins file. These pin the index
parser, the fail-closed behaviour when an asset is not covered, the
tampered-manifest guard, and the newest-release resolution.
"""
from __future__ import annotations
import importlib
import sys
from pathlib import Path
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
iwp = importlib.import_module("install_whisper_prebuilt")
if not hasattr(iwp, "parse_release_checksums"):
pytest.skip("checksum-model symbols not present - check branch", allow_module_level = True)
_A = "0" * 64
_B = "1" * 64
_TAG = "v1.9.1-unsloth.1"
_REPO = "unslothai/whisper.cpp"
def _index(**overrides) -> dict:
payload = {
"schema_version": 1,
"component": "whisper.cpp",
"release_tag": _TAG,
"upstream_tag": "v1.9.1",
"artifacts": {
"whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz": {"sha256": _A},
"whisper-v1.9.1-unsloth.1-linux-x64-cuda12-portable.tar.gz": {"sha256": _B},
},
}
payload.update(overrides)
return payload
# parse_release_checksums / expected_sha256_for are prebuilt_core re-exports;
# their valid/fail-closed matrix is asserted against the real whisper
# descriptor in tests/studio/install/test_prebuilt_core.py. The download-host
# fast-path tests below still route through this module's parse wrapper.
# release tag resolution.
def test_resolve_release_tag_explicit_override_passthrough():
assert iwp.resolve_release_tag(_REPO, published_release_tag = "v1.9.1-unsloth.2") == (
"v1.9.1-unsloth.2"
)
def test_resolve_release_tag_resolves_newest_when_no_override(monkeypatch):
monkeypatch.setattr(iwp, "resolve_newest_release_tag", lambda repo: "v9.9.9-unsloth.9")
assert iwp.resolve_release_tag(_REPO, published_release_tag = None) == "v9.9.9-unsloth.9"
def test_resolve_newest_release_tag_picks_latest_published(monkeypatch):
releases = [
{"tag_name": "v1.9.1-unsloth.1", "published_at": "2026-01-01T00:00:00Z"},
{"tag_name": "v1.9.1-unsloth.3", "published_at": "2026-03-01T00:00:00Z"},
{"tag_name": "v1.9.1-unsloth.2", "published_at": "2026-02-01T00:00:00Z"},
{"tag_name": "draft", "published_at": "2026-09-01T00:00:00Z", "draft": True},
{"tag_name": "pre", "published_at": "2026-09-01T00:00:00Z", "prerelease": True},
]
monkeypatch.setattr(iwp, "fetch_json", lambda url: releases)
assert iwp.resolve_newest_release_tag(_REPO) == "v1.9.1-unsloth.3"
def test_resolve_newest_release_tag_none_published_fails_closed(monkeypatch):
monkeypatch.setattr(iwp, "fetch_json", lambda url: [{"tag_name": "d", "draft": True}])
with pytest.raises(iwp.PrebuiltFallback):
iwp.resolve_newest_release_tag(_REPO)
def test_pins_symbols_are_gone():
# The committed-pins trust model was removed in favour of llama's runtime index.
for gone in ("load_pins", "pins_path", "resolve_expected_sha256", "PINS_FILENAME"):
assert not hasattr(iwp, gone), f"{gone} should have been removed"
# Download-host fast path (resolve + fetch the JSON assets with no GitHub API).
_CPU_ASSET = "whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz"
def _manifest() -> dict:
return {
"schema_version": 1,
"component": "whisper.cpp",
"upstream_tag": "v1.9.1",
"artifacts": [{"asset": _CPU_ASSET, "os": "linux", "arch": "x64", "backend": "cpu"}],
}
def _no_api(monkeypatch):
"""Fail loudly if any code path touches api.github.com."""
def _boom(*a, **k):
raise AssertionError("api.github.com was used on the fast path")
monkeypatch.setattr(iwp, "fetch_json", _boom)
monkeypatch.setattr(iwp, "github_release", _boom)
monkeypatch.setattr(iwp, "fetch_release_bundle", _boom)
def test_fetch_release_for_install_prefers_download_host(monkeypatch):
_no_api(monkeypatch)
monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG)
def _dhj(url):
if url.endswith(iwp.SHA256_ASSET_NAME):
return _index()
if url.endswith(iwp.MANIFEST_ASSET_NAME):
return _manifest()
raise AssertionError(f"unexpected url {url}")
monkeypatch.setattr(iwp, "_download_host_json", _dhj)
bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = None)
assert bundle.release_tag == _TAG
assert checks[_CPU_ASSET] == _A
# asset_urls point at the download host (github.com), not the API.
assert bundle.asset_urls[iwp.SHA256_ASSET_NAME].startswith(
f"https://github.com/{_REPO}/releases/"
)
assert bundle.asset_urls[_CPU_ASSET].startswith(
f"https://github.com/{_REPO}/releases/download/"
)
walked = iwp._fetch_release_candidate(_REPO, _TAG)
assert iwp.SHA256_ASSET_NAME in walked.asset_urls
assert _CPU_ASSET in walked.asset_urls
def test_fetch_release_for_install_explicit_tag_skips_the_head(monkeypatch):
# An explicit tag needs no /releases/latest HEAD: resolving it must not call it.
monkeypatch.setattr(
iwp,
"_download_host_latest_release_tag",
lambda repo: (_ for _ in ()).throw(AssertionError("HEAD used for an explicit tag")),
)
monkeypatch.setattr(
iwp,
"_download_host_json",
lambda url: _index() if url.endswith(iwp.SHA256_ASSET_NAME) else _manifest(),
)
_no_api(monkeypatch)
bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = _TAG)
assert bundle.release_tag == _TAG
def test_fetch_release_for_install_falls_back_to_api(monkeypatch):
# Fast path returns None (e.g. a 404) -> the API path resolves the release.
monkeypatch.setattr(iwp, "_resolve_release_via_download_host", lambda repo, tag: None)
sentinel = iwp.ReleaseBundle(repo = _REPO, release_tag = _TAG, manifest = _manifest(), asset_urls = {})
monkeypatch.setattr(iwp, "resolve_release_tag", lambda repo, *, published_release_tag: _TAG)
monkeypatch.setattr(iwp, "fetch_release_bundle", lambda repo, tag: sentinel)
monkeypatch.setattr(iwp, "fetch_release_checksums", lambda bundle: {_CPU_ASSET: _A})
bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = None)
assert bundle is sentinel
assert checks == {_CPU_ASSET: _A}
def test_resolve_via_download_host_sha_404_returns_none(monkeypatch):
import urllib.error
monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG)
def _dhj(url):
raise urllib.error.HTTPError(url, 404, "not found", {}, None)
monkeypatch.setattr(iwp, "_download_host_json", _dhj)
assert iwp._resolve_release_via_download_host(_REPO, None) is None
def test_resolve_via_download_host_tag_mismatch_returns_none(monkeypatch):
# A checksum index whose self-reported release_tag disagrees is rejected (None).
monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG)
monkeypatch.setattr(
iwp, "_download_host_json", lambda url: _index(release_tag = "v1.9.1-unsloth.2")
)
assert iwp._resolve_release_via_download_host(_REPO, None) is None
def test_download_host_latest_release_tag_parses_redirect(monkeypatch):
class _Resp:
def __enter__(self):
return self
def __exit__(self, *a):
return False
def geturl(self):
return f"https://github.com/{_REPO}/releases/tag/{_TAG}"
class _Opener:
def open(
self,
req,
timeout = None,
):
return _Resp()
monkeypatch.setattr(iwp, "_URL_OPENER", _Opener())
assert iwp._download_host_latest_release_tag(_REPO) == _TAG
def test_download_host_latest_release_tag_404_returns_none(monkeypatch):
import urllib.error
class _Opener:
def open(
self,
req,
timeout = None,
):
raise urllib.error.HTTPError(req.full_url, 404, "nf", {}, None)
monkeypatch.setattr(iwp, "_URL_OPENER", _Opener())
assert iwp._download_host_latest_release_tag(_REPO) is None

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

@ -254,8 +254,10 @@ def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tm
)
fake_external_media = SimpleNamespace(
linux_run_media_mount_roots = lambda: [media_root],
macos_volume_roots = lambda: [],
windows_drive_roots = lambda: [],
)
fake_paths.external_media = fake_external_media
fake_studio_db = SimpleNamespace(
list_scan_folders = lambda: [],
contains_sensitive_path_component = studio_db.contains_sensitive_path_component,

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

@ -119,6 +119,9 @@ def _clean_state(monkeypatch, tmp_path):
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
# Never hit the network in these tests.
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
# Keep the whisper piggyback out of the llama-only tests: no host probe, no
# whisper phase (test_combined_update.py covers the chained flow).
monkeypatch.setattr(upd, "_whisper_chain_status", lambda **kwargs: None)
yield
freshness.reset_caches()
upd._reset_job_for_tests()

15
studio/backend/tests/test_llama_route.py
View file

@ -100,6 +100,21 @@ def test_status_response_exposes_update_size_bytes():
assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None
def test_status_response_exposes_update_component():
model = rl.LlamaUpdateStatusResponse(
supported = True,
update_available = True,
llama_update_available = False,
update_component = "whisper",
whisper = {
"update_available": True,
"installed_tag": "v1",
"latest_tag": "v2",
},
)
assert model.model_dump()["update_component"] == "whisper"
def test_status_handler_runs_off_event_loop(monkeypatch):
seen = {}

7
studio/backend/tests/test_local_llama_cpp_link.py
View file

@ -21,6 +21,13 @@ from utils import llama_cpp_update as u
from core.inference.llama_cpp import LlamaCppBackend
@pytest.fixture(autouse = True)
def _no_whisper_piggyback(monkeypatch):
# Keep the whisper piggyback probe off the host: these tests exercise the
# llama local-link contract only.
monkeypatch.setattr(u, "_whisper_chain_status", lambda **kwargs: None)
def _make_link(link: Path, target: Path) -> None:
"""Create a directory junction (Windows) / symlink (POSIX); neither needs
elevation."""

116
studio/backend/tests/test_middleware.py
View file

@ -14,6 +14,7 @@ import pytest
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import Response
from fastapi.testclient import TestClient
from starlette.middleware.gzip import GZipMiddleware
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
@ -33,6 +34,7 @@ def main_module():
def _make_protected_app(
max_bytes: int,
main_module,
request_max_bytes_getter = None,
upload_passthrough_prefixes: tuple = (),
upload_passthrough_max_bytes_getter = None,
):
@ -40,7 +42,13 @@ def _make_protected_app(
app.add_middleware(
main_module.MaxBodyMiddleware,
max_bytes_getter = lambda: max_bytes,
protected_prefixes = ("/v1/chat/completions", "/api/settings", "/api/train"),
protected_prefixes = (
"/v1/chat/completions",
"/api/inference",
"/api/settings",
"/api/train",
),
request_max_bytes_getter = request_max_bytes_getter,
upload_passthrough_prefixes = upload_passthrough_prefixes,
upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter,
)
@ -67,6 +75,10 @@ def _make_protected_app(
total += len(chunk)
return {"ok": True, "chunks": chunks, "total": total}
@app.post("/api/inference/audio/transcribe/raw")
async def transcribe_raw(request: Request):
return {"ok": True, "total": len(await request.body())}
@app.get("/api/train/status")
async def status_get():
return {"ok": True, "get": True}
@ -96,6 +108,43 @@ class TestMaxBodyMiddleware:
assert r.status_code == 200
assert r.json()["unprotected"] is True
def test_route_specific_cap_overrides_default(self, main_module):
app = _make_protected_app(
4096,
main_module,
request_max_bytes_getter = lambda path: (
128 if path.endswith("/transcribe/raw") else 4096
),
)
c = TestClient(app)
rejected = c.post(
"/api/inference/audio/transcribe/raw",
content = b"x" * 129,
)
accepted = c.post(
"/api/inference/audio/transcribe/raw",
content = b"x" * 128,
)
assert rejected.status_code == 413
assert accepted.status_code == 200
assert accepted.json()["total"] == 128
def test_stt_routes_use_audio_specific_caps(self, main_module):
from utils.upload_limits import (
STT_AUDIO_JSON_MAX_BYTES,
STT_AUDIO_RAW_MAX_BYTES,
)
assert (
main_module._get_request_body_max_bytes("/api/inference/audio/transcribe/raw")
== STT_AUDIO_RAW_MAX_BYTES
)
assert (
main_module._get_request_body_max_bytes("/api/inference/audio/transcribe")
== STT_AUDIO_JSON_MAX_BYTES
)
def test_settings_put_body_over_cap_rejected(self, main_module):
app = _make_protected_app(1024, main_module)
c = TestClient(app)
@ -471,6 +520,71 @@ class TestSecurityHeadersMiddleware:
assert b"server" in names
class TestFrontendAssets:
def test_hashed_assets_are_compressed_and_cached(self, tmp_path, main_module):
content = b"export const value = 'responsive';\n" * 200
(tmp_path / "page-abc123.js").write_bytes(content)
app = FastAPI()
assets_app = GZipMiddleware(
main_module.ImmutableStaticFiles(directory = tmp_path),
minimum_size = 1024,
compresslevel = 6,
)
app.mount("/assets", assets_app, name = "assets")
response = TestClient(app).get(
"/assets/page-abc123.js",
headers = {"Accept-Encoding": "gzip"},
)
assert response.status_code == 200
assert response.content == content
assert response.headers["content-encoding"] == "gzip"
assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
assert "accept-encoding" in response.headers["vary"].lower()
def test_asset_revalidation_keeps_immutable_cache_header(self, tmp_path, main_module):
(tmp_path / "page-abc123.js").write_text("export {};", encoding = "utf-8")
app = FastAPI()
app.mount(
"/assets",
main_module.ImmutableStaticFiles(directory = tmp_path),
name = "assets",
)
client = TestClient(app)
first = client.get("/assets/page-abc123.js")
response = client.get(
"/assets/page-abc123.js",
headers = {"If-None-Match": first.headers["etag"]},
)
assert response.status_code == 304
assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
def test_range_request_is_not_compressed(self, tmp_path, main_module):
content = b"export const value = 'responsive';\n" * 200
(tmp_path / "page-abc123.js").write_bytes(content)
app = FastAPI()
assets_app = main_module._AssetGZipMiddleware(
main_module.ImmutableStaticFiles(directory = tmp_path),
minimum_size = 1024,
compresslevel = 6,
)
app.mount("/assets", assets_app, name = "assets")
response = TestClient(app).get(
"/assets/page-abc123.js",
headers = {"Accept-Encoding": "gzip", "Range": "bytes=0-99"},
)
assert response.status_code == 206
assert response.headers.get("content-encoding") != "gzip"
assert response.headers["content-range"] == f"bytes 0-99/{len(content)}"
assert response.content == content[:100]
assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
# /api/health auth gate

410
studio/backend/tests/test_mlx_inference_backend.py
View file

@ -922,3 +922,413 @@ def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch):
"<think>vision</think>",
"<think>vision</think> answer",
]
class _FakeLRUPromptCache:
def __init__(
self,
max_size = 10,
max_bytes = 1 << 63,
):
self.max_size = max_size
self.max_bytes = max_bytes
self.entries = {}
def fetch_nearest_cache(self, key, tokens):
import copy
stored = self.entries.get(key, {})
exact = stored.get(tuple(tokens))
if exact is not None:
return copy.deepcopy(exact), []
best = None
for candidate, cache in stored.items():
if len(candidate) < len(tokens) and tuple(tokens[: len(candidate)]) == candidate:
if best is None or len(candidate) > len(best[0]):
best = (candidate, cache)
if best is not None:
return copy.deepcopy(best[1]), list(tokens[len(best[0]) :])
return None, list(tokens)
def insert_cache(
self,
key,
tokens,
prompt_cache,
*,
cache_type = "assistant",
):
import copy
self.entries.setdefault(key, {})[tuple(tokens)] = copy.deepcopy(prompt_cache)
class _FakeCacheEntry:
def __init__(
self,
offset = 0,
nbytes = 1,
):
self.offset = offset
self.nbytes = nbytes
def _install_fake_prompt_cache_api(monkeypatch, trimmable = True):
from core.inference import mlx_inference
def _make_prompt_cache(_model):
return [_FakeCacheEntry()]
def _can_trim_prompt_cache(_cache):
return trimmable
def _trim_prompt_cache(cache, num):
cache[0].offset = max(cache[0].offset - num, 0)
return num
monkeypatch.setattr(
mlx_inference,
"_mlx_prompt_cache_api",
lambda: (
_FakeLRUPromptCache,
_make_prompt_cache,
_can_trim_prompt_cache,
_trim_prompt_cache,
),
)
def test_mlx_prompt_cache_max_bytes_budget(monkeypatch):
from core.inference.mlx_inference import (
PROMPT_CACHE_FALLBACK_BYTES,
PROMPT_CACHE_MEMORY_FRACTION,
_prompt_cache_max_bytes,
)
monkeypatch.delenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", raising = False)
assert _prompt_cache_max_bytes(None) == PROMPT_CACHE_FALLBACK_BYTES
assert _prompt_cache_max_bytes(20.0) == int(20.0 * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "4096")
assert _prompt_cache_max_bytes(20.0) == 4096
monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "0")
assert _prompt_cache_max_bytes(20.0) == 0
monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "not-a-number")
assert _prompt_cache_max_bytes(20.0) == int(20.0 * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
def test_mlx_prompt_cache_never_returns_empty_remainder(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
from core.inference.mlx_inference import _MLXPromptCacheHistory
history = _MLXPromptCacheHistory(6, 1 << 30)
tokens = list(range(10))
cache, rest = history.fetch(object(), "key", tokens)
assert len(rest) == 10
cache[0].offset = len(tokens)
history.insert("key", tokens, cache)
_cache, rest = history.fetch(object(), "key", tokens)
assert rest == tokens[-1:]
longer = tokens + [99, 100]
_cache, rest = history.fetch(object(), "key", longer)
assert rest == [99, 100]
_install_fake_prompt_cache_api(monkeypatch, trimmable = False)
history = _MLXPromptCacheHistory(6, 1 << 30)
cache, _rest = history.fetch(object(), "key", tokens)
cache[0].offset = len(tokens)
history.insert("key", tokens, cache)
_cache, rest = history.fetch(object(), "key", tokens)
assert rest == tokens, "untrimmable entry must not be reused"
def test_mlx_prompt_cache_key_isolates_adapter_state(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import MLXInferenceBackend
class _Tok:
bos_token = None
def encode(
self,
text,
add_special_tokens = True,
):
return [ord(c) for c in text]
backend = MLXInferenceBackend()
backend._model = object()
backend._tokenizer = _Tok()
backend.active_model_name = "model-a"
prompt = "shared prefix"
_rest, cache, key, tokens, cached = backend._prepare_prompt_cache(prompt, True)
assert cached == 0
cache[0].offset = len(tokens)
backend._prompt_cache_history.insert(key, tokens, cache)
_rest, _cache, _key, _tokens, cached_same = backend._prepare_prompt_cache(prompt, True)
assert cached_same > 0
_rest, _cache, _key, _tokens, cached_flipped = backend._prepare_prompt_cache(prompt, False)
assert cached_flipped == 0
def _install_fake_text_stack(
monkeypatch,
token_map,
captured,
markers = None,
):
import types as _types
from core.inference import mlx_inference
_install_fake_mlx(monkeypatch)
monkeypatch.setattr(
mlx_inference,
"_temporary_mlx_adapter_state",
lambda _model, _state: __import__("contextlib").nullcontext(),
)
monkeypatch.setattr(
"core.inference.chat_template_helpers.apply_chat_template_for_generation",
lambda _tok, messages, **_kw: messages[-1]["content"],
)
monkeypatch.setattr(
"core.inference.chat_template_helpers.render_with_native_template_fallback",
lambda formatted_prompt, **_kw: SimpleNamespace(
prompt = formatted_prompt,
reasoning_channel_markers = markers,
),
)
monkeypatch.setattr(
"core.inference.chat_template_helpers.detect_think_prefill",
lambda *_a, **_kw: "",
)
class _Resp:
def __init__(self, token, processed):
self.token = token
self.text = f"<{token}>"
self.prompt_tokens = processed
self.prompt_tps = 10.0
self.generation_tokens = 1
self.generation_tps = 5.0
def _stream_generate(_model, _tokenizer, **kwargs):
captured.append(kwargs)
processed = len(kwargs["prompt"])
cache = kwargs.get("prompt_cache")
if cache is not None:
cache[0].offset += processed
for token in token_map["generated"]:
if cache is not None:
cache[0].offset += 1
yield _Resp(token, processed)
mlx_lm_pkg = _types.ModuleType("mlx_lm")
mlx_lm_pkg.stream_generate = _stream_generate
mlx_lm_sample = _types.ModuleType("mlx_lm.sample_utils")
mlx_lm_sample.make_sampler = lambda **_kw: object()
mlx_lm_sample.make_logits_processors = lambda **_kw: []
monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg)
monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample)
class _Tok:
bos_token = None
chat_template = "x"
def encode(
self,
text,
add_special_tokens = True,
):
return list(token_map[text])
def decode(
self,
ids,
skip_special_tokens = False,
):
return "".join(str(i) for i in ids)
from core.inference.mlx_inference import MLXInferenceBackend
backend = MLXInferenceBackend()
backend._model = object()
backend._tokenizer = _Tok()
backend._is_vlm = False
backend.active_model_name = "model-a"
return backend
def _run_turn(backend, prompt):
list(
backend.generate_chat_response(
messages = [{"role": "user", "content": prompt}],
max_new_tokens = 4,
)
)
def test_mlx_text_reuses_prompt_cache_on_the_next_turn(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
captured = []
token_map = {
"P1": [1, 2, 3],
"P2": [1, 2, 3, 7, 8, 9, 10],
"generated": [7, 8],
}
backend = _install_fake_text_stack(monkeypatch, token_map, captured)
_run_turn(backend, "P1")
assert captured[0]["prompt"] == [1, 2, 3]
assert "prompt_cache" in captured[0]
assert backend.last_generation_stats["timings"]["cache_n"] == 0
_run_turn(backend, "P2")
assert captured[1]["prompt"] == [9, 10], "turn two should prefill only the new tail"
stats = backend.last_generation_stats
assert stats["timings"]["cache_n"] == 5
assert stats["timings"]["prompt_n"] == 2
assert stats["usage"]["prompt_tokens"] == 7
def test_mlx_text_without_lru_prompt_cache_prefills_the_full_prompt(monkeypatch):
from core.inference import mlx_inference
monkeypatch.setattr(mlx_inference, "_mlx_prompt_cache_api", lambda: None)
captured = []
token_map = {"P1": [1, 2, 3], "generated": [7]}
backend = _install_fake_text_stack(monkeypatch, token_map, captured)
_run_turn(backend, "P1")
assert captured[0]["prompt"] == "P1"
assert "prompt_cache" not in captured[0]
assert backend.last_generation_stats["timings"]["cache_n"] == 0
def test_mlx_text_tracks_tokens_on_the_native_reasoning_path(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
captured = []
token_map = {"P1": [1, 2, 3], "P2": [1, 2, 3, 7, 8, 9], "generated": [7, 8]}
backend = _install_fake_text_stack(monkeypatch, token_map, captured, markers = ("<a>", "</a>"))
_run_turn(backend, "P1")
_run_turn(backend, "P2")
assert captured[1]["prompt"] == [9]
def test_mlx_presence_penalty_latches_the_first_decode_step():
mx = pytest.importorskip("mlx.core")
import numpy as np
from core.inference.mlx_inference import _make_mlx_presence_penalty_processor
processor = _make_mlx_presence_penalty_processor(2.0)
logits = mx.zeros((1, 5))
out = processor(mx.array([3]), logits)
assert np.array_equal(np.array(out), np.zeros((1, 5))), "prompt must not be penalized"
out = processor(mx.array([3, 1]), mx.zeros((1, 5)))
penalized = np.array(out)[0]
assert penalized[1] == -2.0
assert penalized[3] == 0.0
def test_mlx_prompt_cache_survives_reset_but_not_unload(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
_install_fake_mlx(monkeypatch)
sys.modules["mlx.core"].clear_cache = lambda: None
from core.inference.mlx_inference import MLXInferenceBackend
backend = MLXInferenceBackend()
backend.active_model_name = "model-a"
history = backend._prompt_cache()
assert history is not None
backend.reset_generation_state()
assert backend._prompt_cache_history is history
backend.unload_model("model-a")
assert backend._prompt_cache_history is None
def test_mlx_prompt_cache_skips_entries_over_budget(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
from core.inference.mlx_inference import _MLXPromptCacheHistory
history = _MLXPromptCacheHistory(6, 1000)
history.insert("key", [1, 2, 3], [_FakeCacheEntry(offset = 3, nbytes = 400)])
assert len(history._lru.entries.get("key", {})) == 1
history.insert("key", list(range(50)), [_FakeCacheEntry(offset = 50, nbytes = 5000)])
stored = history._lru.entries.get("key", {})
assert tuple([1, 2, 3]) in stored
assert tuple(range(50)) not in stored
def test_mlx_prompt_cache_keys_on_what_the_kv_covers(monkeypatch):
_install_fake_prompt_cache_api(monkeypatch)
from core.inference.mlx_inference import _MLXPromptCacheHistory
class _Entry:
def __init__(
self,
offset,
nbytes = 1,
):
self.offset = offset
self.nbytes = nbytes
history = _MLXPromptCacheHistory(6, 1 << 30)
history.insert("key", list(range(10)), [_Entry(offset = 8)])
assert tuple(range(8)) in history._lru.entries["key"]
assert tuple(range(10)) not in history._lru.entries["key"]
history.insert("other", list(range(4)), [_Entry(offset = 9)])
assert "other" not in history._lru.entries
def test_mlx_prompt_cache_only_stores_verifiable_prefix_coverage(monkeypatch):
mx = pytest.importorskip("mlx.core")
from mlx_lm.models.cache import CacheList, ChunkedKVCache, KVCache, RotatingKVCache
_install_fake_prompt_cache_api(monkeypatch)
from core.inference.mlx_inference import _kv_prefix_coverage, _MLXPromptCacheHistory
def feed(entry, n):
for _ in range(n):
block = mx.zeros((1, 2, 1, 4), dtype = mx.float16)
entry.update_and_fetch(block, block)
mx.eval(entry.state)
return entry
plain = feed(KVCache(), 30)
unwrapped = feed(RotatingKVCache(max_size = 100, keep = 2), 30)
wrapped = feed(RotatingKVCache(max_size = 10, keep = 2), 30)
chunked = feed(ChunkedKVCache(chunk_size = 8), 30)
slid = feed(ChunkedKVCache(chunk_size = 8), 30)
slid.maybe_trim_front()
assert _kv_prefix_coverage([plain]) == 30
assert _kv_prefix_coverage([unwrapped]) == 30
assert _kv_prefix_coverage([chunked]) == 30
assert wrapped.offset == 30 and wrapped.state[0].shape[2] == 10
assert _kv_prefix_coverage([wrapped]) is None
assert slid.start_position > 0
assert _kv_prefix_coverage([slid]) is None
assert _kv_prefix_coverage([CacheList(feed(KVCache(), 30), feed(KVCache(), 30))]) == 30
assert _kv_prefix_coverage([CacheList(feed(KVCache(), 30), wrapped)]) is None
assert _kv_prefix_coverage([feed(KVCache(), 30), feed(KVCache(), 29)]) is None
assert _kv_prefix_coverage([]) is None
history = _MLXPromptCacheHistory(6, 1 << 40)
for unsafe in (wrapped, slid):
history.insert("key", list(range(30)), [unsafe])
assert "key" not in history._lru.entries
history.insert("key", list(range(30)), [plain])
assert tuple(range(30)) in history._lru.entries["key"]

137
studio/backend/tests/test_mlx_stop_checkpoint.py Normal file
View file

@ -0,0 +1,137 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression tests for MLX stop-and-save checkpoint handling."""
import importlib.util
import json
import sys
import types
from pathlib import Path
import numpy as np
from safetensors.numpy import save_file
_BACKEND = Path(__file__).resolve().parents[1]
def _load_worker_module():
spec = importlib.util.spec_from_file_location(
"training_worker_under_test",
_BACKEND / "core" / "training" / "worker.py",
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
worker = _load_worker_module()
class _FakeTrainer:
def __init__(self, step: int):
self._global_step = step
self._train_loss_history = []
self.model = object()
def _write_checkpoint(out: Path, step: int) -> Path:
checkpoint = out / f"checkpoint-{step}"
checkpoint.mkdir(parents = True, exist_ok = True)
(checkpoint / "trainer_state.json").write_text(
json.dumps({"global_step": step}), encoding = "utf-8"
)
save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors")
save_file(
{"state": np.ones(1, dtype = np.float32)},
checkpoint / "optimizer_state.safetensors",
)
return checkpoint
def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path):
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
assert worker._mlx_has_checkpoint_at_step(out, 5) is True
def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path):
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True
def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists(
tmp_path, monkeypatch
):
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
saved_steps: list[int] = []
def _save_state(_value, path, name):
save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name))
def _save_trainer_state(state, ckpt_dir, **_kwargs):
Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8")
saved_steps.append(int(state["global_step"]))
fake_utils = types.SimpleNamespace(
save_trainable_adapters = lambda model, path: _save_state(
model, path, "adapters.safetensors"
),
save_optimizer_state = lambda optimizer, path: _save_state(
optimizer, path, "optimizer_state.safetensors"
),
save_trainer_state = _save_trainer_state,
)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True
assert saved_steps == [10]
assert (out / "checkpoint-10" / "trainer_state.json").is_file()
def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path):
out = tmp_path / "outputs" / "run_x"
out.mkdir(parents = True)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False
def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path):
out = tmp_path / "outputs" / "run_x"
ckpt = out / "checkpoint-5"
ckpt.mkdir(parents = True)
(ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8")
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False
def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path):
# An older checkpoint does not cover the current step, so this still fails.
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False
def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch):
out = tmp_path / "outputs" / "run_x"
out.mkdir(parents = True)
def _boom(*_args, **_kwargs):
raise RuntimeError("save failed")
fake_utils = types.SimpleNamespace(
save_trainable_adapters = _boom,
save_optimizer_state = lambda *_a, **_k: None,
save_trainer_state = lambda *_a, **_k: None,
)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False

232
studio/backend/tests/test_model_picker_regression.py Normal file
View file

@ -0,0 +1,232 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression guards for the model-picker per-model-config feature (the set of
bugs that got the predecessor PR reverted). Pure-function / validation checks
only, so they run on CPU in the backend pytest job with no model download.
Covers, at the backend layer:
- infra-model hiding: the RAG embedder (bge-small-en-v1.5) and the llama.cpp
install-validation probe (ggml-org/models / stories260K) stay hidden, while
normal chat repos are not hidden;
- the HF token is honored from the dedicated header with the query string as a
fallback, never the other way around;
- the chat-template byte caps reject oversized overrides (both the char-count
fast path and the UTF-8 byte path) and the sidecar reader is size-bounded.
"""
from __future__ import annotations
import sys
import types
import pytest
# Keep this test runnable without the optional structlog dependency (mirrors
# tests/test_cached_gguf_routes.py), since importing routes.models pulls it in.
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules["structlog"] = types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
)
import routes.models as models_route
from core.rag import config as rag_config
from hub.dependencies import get_hf_token
from models.inference import LoadRequest
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
from picker.service import _read_bounded_text
from utils.hidden_models import is_hidden_model
@pytest.fixture(autouse = True)
def _pin_default_embedder(monkeypatch):
"""Pin the effective embedder to Studio's static default so hiding is
deterministic and cannot depend on ambient RAG config / env."""
default = "unsloth/bge-small-en-v1.5"
monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", default, raising = False)
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: default)
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: default)
monkeypatch.setattr(rag_config, "default_gguf_repo", lambda: default)
# --------------------------------------------------------------------------- #
# Infra-model hiding (the "infra models resurfaced in the picker" regression) #
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize(
"value",
[
"ggml-org/models", # the probe repo id
"unsloth/bge-small-en-v1.5", # the RAG embedder repo
"unsloth/bge-small-en-v1.5-GGUF", # its GGUF companion
"/root/.cache/huggingface/hub/x/stories260K.gguf", # probe on disk
"/root/.cache/x/Stories260K.GGUF", # case-insensitive
r"C:\\models\\stories260K.gguf", # windows-style path
"/opt/models/bge-small-en-v1.5", # embedder basename folder
"/opt/models/bge-small-en-v1.5-Q8_0.gguf", # suffixed local weight
],
)
def test_infra_models_are_hidden(value):
assert is_hidden_model(value) is True
@pytest.mark.parametrize(
"value",
[
"unsloth/gemma-3-270m-it-GGUF", # a normal small chat GGUF
"unsloth/Qwen3-0.6B", # a normal non-GGUF chat model
"user/stories260K-finetune-GGUF", # repo id merely contains "stories260k"
"user/model-chat", # generic repo must not be hidden
"meta-llama/Llama-3.1-8B-Instruct",
],
)
def test_normal_models_are_not_hidden(value):
assert is_hidden_model(value) is False
def test_is_hidden_model_ignores_empty_values():
assert is_hidden_model(None) is False
assert is_hidden_model("") is False
assert is_hidden_model(None, "", "unsloth/gemma-3-270m-it-GGUF") is False
def test_hidden_model_matchers_expose_probe_needles():
needles, exact_ids, _exact_paths = models_route.hidden_model_matchers()
lowered = [n.lower() for n in needles]
assert "ggml-org/models" in lowered
assert "stories260k.gguf" in lowered
# The configured embedder is exposed as an exact repo id, never as a
# basename needle that would substring-hide unrelated chat models.
assert "bge-small-en-v1.5" not in lowered
assert "unsloth/bge-small-en-v1.5" in exact_ids
def test_hidden_model_matchers_custom_repo_publishes_exact_ids(monkeypatch):
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
needles, exact_ids, exact_paths = models_route.hidden_model_matchers()
assert needles == ["ggml-org/models", "stories260k.gguf"]
assert "org/model" in exact_ids
assert "org/model-gguf" in exact_ids
assert exact_paths == []
def test_hidden_model_matchers_local_owner_name_path_is_exact_path(monkeypatch, tmp_path):
# A local embedder shaped like owner/name that exists on disk must be an
# exact resolved path, not a Hub repo id (mirroring is_hidden_model), so the
# local row stays hidden instead of showing as a chat model.
(tmp_path / "models" / "embedder").mkdir(parents = True)
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "ggml-org/models")
_needles, exact_ids, exact_paths = models_route.hidden_model_matchers()
resolved = str((tmp_path / "models" / "embedder").resolve()).lower()
assert resolved in exact_paths
assert "models/embedder" not in exact_ids
# --------------------------------------------------------------------------- #
# HF token via header, query string only as a fallback (the token-leak fix) #
# --------------------------------------------------------------------------- #
def test_get_hf_token_strips_and_returns():
assert get_hf_token(" hf_abc ") == "hf_abc"
@pytest.mark.parametrize("value", [None, "", " ", "\n\t"])
def test_get_hf_token_blank_is_none(value):
assert get_hf_token(value) is None
@pytest.mark.parametrize(
"value,expected",
[(" hf_x ", "hf_x"), ("", None), (" ", None), (None, None), (1234, None)],
)
def test_normalize_hf_token(value, expected):
assert models_route._normalize_hf_token(value) == expected
def test_header_token_wins_over_query():
header, query = "hf_header", "hf_query"
resolved = models_route._normalize_hf_token(header) or models_route._normalize_hf_token(query)
assert resolved == "hf_header"
def test_query_token_is_fallback_when_header_absent():
resolved = models_route._normalize_hf_token(None) or models_route._normalize_hf_token(
"hf_query"
)
assert resolved == "hf_query"
# --------------------------------------------------------------------------- #
# Chat-template byte caps (the unbounded-template hardening) #
# --------------------------------------------------------------------------- #
def _load_request(**overrides):
data = {"model_path": "unsloth/test-model-GGUF", "gguf_variant": "Q4_K_M"}
data.update(overrides)
return LoadRequest.model_validate(data)
def test_blank_chat_template_override_normalizes_to_none():
assert _load_request(chat_template_override = " \n\t").chat_template_override is None
def test_nonblank_chat_template_override_preserved_verbatim():
template = " {{ messages }} "
assert _load_request(chat_template_override = template).chat_template_override == template
def test_chat_template_at_byte_limit_is_accepted():
template = "a" * MAX_CHAT_TEMPLATE_BYTES # exactly the limit, 1 byte/char
assert (
len(_load_request(chat_template_override = template).chat_template_override)
== MAX_CHAT_TEMPLATE_BYTES
)
def test_chat_template_over_char_limit_is_rejected():
with pytest.raises(Exception): # pydantic ValidationError wrapping ValueError
_load_request(chat_template_override = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1))
def test_chat_template_over_byte_limit_is_rejected():
# Char count stays under the limit but UTF-8 bytes exceed it (3 bytes/char),
# so only the byte-count branch can catch this.
multibyte = "€" * (MAX_CHAT_TEMPLATE_BYTES // 2) # euro sign, 3 bytes each
assert len(multibyte) <= MAX_CHAT_TEMPLATE_BYTES
assert len(multibyte.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES
with pytest.raises(Exception):
_load_request(chat_template_override = multibyte)
def test_read_bounded_text_reads_within_limit(tmp_path):
p = tmp_path / "t.json"
p.write_text("hello", encoding = "utf-8")
assert _read_bounded_text(p, 16) == "hello"
def test_read_bounded_text_rejects_over_limit(tmp_path):
p = tmp_path / "big.json"
p.write_bytes(b"x" * 100)
assert _read_bounded_text(p, 50) is None
def test_read_bounded_text_at_limit_is_read(tmp_path):
p = tmp_path / "exact.json"
p.write_bytes(b"x" * 50)
assert _read_bounded_text(p, 50) == "x" * 50
def test_read_bounded_text_missing_file_is_none(tmp_path):
assert _read_bounded_text(tmp_path / "nope.json", 50) is None

199
studio/backend/tests/test_model_update_robustness.py
View file

@ -112,13 +112,21 @@ def patch_hub_gguf(monkeypatch):
blob_ids = [local_blob],
gguf_files = {"model-Q4_K_M.gguf": 1000},
)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = tmp_path),
)
monkeypatch.setattr(
GV,
"list_gguf_variants",
lambda r, hf_token = None: (_variants(), False, [remote_sibling]),
raising = True,
)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
monkeypatch.setattr(
GV,
"iter_hf_cache_snapshots",
lambda _repo_id, root = None: [snap],
)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
@ -217,6 +225,10 @@ def test_variant_update_check_detects_companion_only_update(
companion_path: 100,
},
)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = tmp_path),
)
siblings = [
patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "mainsha"),
patch_hub_gguf.sibling(companion_path, 100, "new-companion"),
@ -227,7 +239,11 @@ def test_variant_update_check_detects_companion_only_update(
lambda r, hf_token = None: (_variants(), has_vision, siblings),
raising = True,
)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
monkeypatch.setattr(
GV,
"iter_hf_cache_snapshots",
lambda _repo_id, root = None: [snap],
)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
@ -314,6 +330,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
file_name = "model.safetensors",
size_on_disk = 100,
blob_path = str(repo_path / "blobs" / "modelsha"),
blob_last_modified = 3_000.0,
),
]
)
@ -336,6 +353,98 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
assert rows[0]["repo_id"] == "Org/SafeTensorRepo"
assert rows[0]["model_format"] == "safetensors"
assert rows[0]["size_bytes"] == 100
assert rows[0]["last_modified"] == 3_000.0
def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path):
repo_path = tmp_path / "models--Org--GgufRepo"
repo = SimpleNamespace(
repo_id = "Org/GgufRepo",
repo_type = "model",
repo_path = repo_path,
revisions = [
SimpleNamespace(
files = [
SimpleNamespace(
file_name = "model-Q4_K_M.gguf",
size_on_disk = 100,
blob_path = None,
blob_last_modified = 5_000.0,
),
]
)
],
)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
monkeypatch.setattr(
CI.hf_cache_scan,
"is_gguf_repo_partial",
lambda *args, **kwargs: False,
)
monkeypatch.setattr(
CI,
"_gguf_variant_state_summary",
lambda _repo_id, **_kwargs: (False, 0),
)
rows = CI._scan_cached_gguf()
assert len(rows) == 1
assert rows[0]["repo_id"] == "Org/GgufRepo"
assert rows[0]["model_format"] == "gguf"
assert rows[0]["size_bytes"] == 100
assert rows[0]["last_modified"] == 5_000.0
def test_cached_model_scan_hides_custom_whisper_repo(monkeypatch, tmp_path):
repo_path = tmp_path / "models--Org--CustomWhisper"
snapshot = repo_path / "snapshots" / ("a" * 40)
snapshot.mkdir(parents = True)
(snapshot / "config.json").write_text(
'{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}'
)
repo = SimpleNamespace(
repo_id = "Org/CustomWhisper",
repo_type = "model",
repo_path = repo_path,
revisions = [
SimpleNamespace(
files = [
SimpleNamespace(
file_name = "config.json",
size_on_disk = 10,
blob_path = None,
),
SimpleNamespace(
file_name = "model.safetensors",
size_on_disk = 100,
blob_path = str(repo_path / "blobs" / "modelsha"),
),
]
)
],
)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
monkeypatch.setattr(
CI,
"_cached_model_snapshot_path",
lambda _repo_path: snapshot,
)
monkeypatch.setattr(
CI.hf_cache_scan,
"is_snapshot_partial",
lambda *args, **kwargs: False,
)
assert CI._scan_cached_models() == []
# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───
@ -584,7 +693,12 @@ def test_reclaim_replaced_gguf_variant_prunes_old_revision_only(monkeypatch, tmp
invalidated = []
monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: invalidated.append(True))
result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"NEWsha"}))
result = D.reclaim_replaced_gguf_variant(
repo_id,
"Q4_K_M",
frozenset({"NEWsha"}),
hub_cache = tmp_path,
)
assert result["removed_snapshots"] == 1
assert result["deleted_blobs"] == 1
@ -631,8 +745,85 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch
monkeypatch.setattr(CI, "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo_info])])
monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: None)
result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"REMOTEsha256"}))
result = D.reclaim_replaced_gguf_variant(
repo_id,
"Q4_K_M",
frozenset({"REMOTEsha256"}),
hub_cache = tmp_path,
)
assert snap.exists() is True # the current file must survive
assert result["removed_snapshots"] == 0
assert result["deleted_blobs"] == 0
def test_reclaim_replaced_gguf_variant_only_mutates_worker_cache(monkeypatch, tmp_path):
repo_id = "org/repo-GGUF"
cache_a = tmp_path / "cache-a"
cache_b = tmp_path / "cache-b"
def cached_repo(cache_dir, revision):
repo_path = cache_dir / "models--org--repo-GGUF"
snap = repo_path / "snapshots" / revision / "model-Q4_K_M.gguf"
blob = repo_path / "blobs" / "OLDsha"
snap.parent.mkdir(parents = True, exist_ok = True)
blob.parent.mkdir(parents = True, exist_ok = True)
blob.write_bytes(b"old")
snap.symlink_to(blob)
return (
SimpleNamespace(
repo_id = repo_id,
repo_type = "model",
repo_path = repo_path,
revisions = [
SimpleNamespace(
files = [
SimpleNamespace(
file_name = snap.name,
file_path = str(snap),
blob_path = str(blob),
)
]
)
],
),
snap,
blob,
)
repo_a, snap_a, blob_a = cached_repo(cache_a, "a" * 40)
repo_b, snap_b, blob_b = cached_repo(cache_b, "b" * 40)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo_a]), SimpleNamespace(repos = [repo_b])],
)
monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: None)
result = D.reclaim_replaced_gguf_variant(
repo_id,
"Q4_K_M",
frozenset({"NEWsha"}),
hub_cache = cache_b,
)
assert result["removed_snapshots"] == 1
assert snap_b.exists() is False
assert blob_b.exists() is False
assert snap_a.exists() is True
assert blob_a.exists() is True
def _mmproj_repo(*file_names: str):
return SimpleNamespace(
revisions = [SimpleNamespace(files = [SimpleNamespace(file_name = n) for n in file_names])]
)
def test_repo_has_mmproj_requires_gguf_projector():
# A non-GGUF sidecar whose name merely contains "mmproj" must NOT mark the
# repo vision-capable; the runtime's projector detection is GGUF-only.
assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj_config.json")) is False
assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "README-mmproj.md")) is False
# A real GGUF projector still marks the repo vision-capable.
assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj-F16.gguf")) is True

11
studio/backend/tests/test_models_get_model_config_case_resolution.py
View file

@ -84,7 +84,7 @@ def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, mon
# covers the active cache; discard deletes case-insensitively, so detection must too,
# else a decline deletes a pre-existing user repo).
import utils.paths as paths_pkg
import huggingface_hub.constants as hf_constants
import hub.utils.paths as hub_paths
active = tmp_path / "active"
legacy = tmp_path / "legacy"
@ -96,9 +96,12 @@ def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, mon
# No active-cache variant; case resolution is a no-op here.
monkeypatch.setattr(paths_pkg, "resolve_cached_repo_id_case", lambda name: name)
monkeypatch.setattr(paths_pkg, "legacy_hf_cache_dir", lambda: legacy)
monkeypatch.setattr(paths_pkg, "hf_default_cache_dir", lambda: default)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(active))
monkeypatch.setattr(hub_paths, "legacy_hf_cache_dir", lambda: legacy)
monkeypatch.setattr(hub_paths, "hf_default_cache_dir", lambda: default)
monkeypatch.setattr(
"utils.hf_cache_settings.known_hf_hub_caches",
lambda: [active],
)
assert models_route._repo_in_any_hf_cache("unsloth/foo") is True
# Absent from every cache -> reported absent.

598
studio/backend/tests/test_offline_embedding_minimal.py Normal file
View file

@ -0,0 +1,598 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Offline RAG embedding-model handling (issue #6817).
Offline the studio must never call the Hub (a DNS-dead session hangs on retries). Using a fake
HF cache under a temp HF_HUB_CACHE, assert that offline: is_embedding_model classifies from the
cached modules.json without the Hub; the file-security gate fails CLOSED on an unscanned pickle
weight with no safetensors alternative and allows an inert cache; the embedder threads
local_files_only into the load. Online behavior is unchanged (bounded timeout + cache fallback).
"""
import sys
import types
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import pytest
from utils.security import evaluate_file_security
from utils.utils import (
hf_cache_snapshot_dir,
hf_cache_snapshot_is_loadable,
hf_env_offline,
st_repo_id_candidates,
)
# Minimal sentence-transformers modules.json (the marker the gate keys on).
MODULES_JSON = (
'[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}]'
)
def _modules_json(*paths):
"""modules.json listing one Transformer module per path (a load root)."""
import json
return json.dumps(
[
{
"idx": i,
"name": str(i),
"path": p,
"type": "sentence_transformers.models.Transformer",
}
for i, p in enumerate(paths)
]
)
_COMMIT = "0123456789abcdef0123456789abcdef01234567"
def _make_cache(
root,
repo_id,
files,
commit = _COMMIT,
):
"""Build a canonical HF-cache snapshot (refs/main + snapshots/<commit>/) for repo_id under
root from {relpath: contents}; returns the snapshot dir."""
from huggingface_hub.file_download import repo_folder_name
repo_dir = Path(root) / repo_folder_name(repo_id = repo_id, repo_type = "model")
(repo_dir / "refs").mkdir(parents = True, exist_ok = True)
(repo_dir / "refs" / "main").write_text(commit)
snapshot = repo_dir / "snapshots" / commit
snapshot.mkdir(parents = True, exist_ok = True)
for rel, contents in files.items():
path = snapshot / rel
path.parent.mkdir(parents = True, exist_ok = True)
path.write_text(contents)
return snapshot
def _no_network():
"""Patch model_info to fail loudly if any offline path reaches the network."""
return patch("huggingface_hub.model_info", side_effect = AssertionError("hit the network"))
def _is_embedding_model(*args, **kwargs):
from utils.models.model_config import is_embedding_model
return is_embedding_model(*args, **kwargs)
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
"""Point the HF cache at a fresh temp dir.
get_hf_cache_paths() reads an import-time env snapshot, not live os.environ,
so point it (and thus active_hf_hub_cache + the snapshot lookup's selected
root) at this temp cache too."""
root = tmp_path / "hub"
root.mkdir()
monkeypatch.setenv("HF_HOME", str(tmp_path))
monkeypatch.setenv("HF_HUB_CACHE", str(root))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = root),
)
return root
@pytest.fixture(autouse = True)
def _clean_env(monkeypatch):
"""Start each test online with an empty detection cache; offline tests opt in."""
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
from utils.models import model_config as mc
mc._embedding_detection_cache.clear()
yield
mc._embedding_detection_cache.clear()
# ── hf_env_offline ───────────────────────────────────────────────
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "])
def test_hf_env_offline_true(monkeypatch, value):
monkeypatch.setenv("HF_HUB_OFFLINE", value)
assert hf_env_offline() is True
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
def test_hf_env_offline_false(monkeypatch, value):
monkeypatch.setenv("HF_HUB_OFFLINE", value)
assert hf_env_offline() is False
def test_hf_env_offline_honors_transformers_flag(monkeypatch):
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
assert hf_env_offline() is True
def test_hf_env_offline_default_false():
assert hf_env_offline() is False
# ── st_repo_id_candidates ────────────────────────────────────────
def test_candidates_slashless_adds_st_alias():
assert st_repo_id_candidates("all-MiniLM-L6-v2") == [
"all-MiniLM-L6-v2",
"sentence-transformers/all-MiniLM-L6-v2",
]
def test_candidates_with_org_is_verbatim():
assert st_repo_id_candidates("org/model") == ["org/model"]
def test_candidates_empty_name():
assert st_repo_id_candidates(" ") == []
# ── hf_cache_snapshot_dir ────────────────────────────────────────
def test_snapshot_dir_resolves_active_commit(hf_cache):
snapshot = _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON})
assert hf_cache_snapshot_dir("org/emb") == snapshot
def test_snapshot_dir_none_when_uncached(hf_cache):
assert hf_cache_snapshot_dir("org/missing") is None
def test_snapshot_dir_uses_st_alias_for_slashless(hf_cache):
snapshot = _make_cache(
hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON}
)
assert hf_cache_snapshot_dir("all-MiniLM-L6-v2") == snapshot
def test_snapshot_dir_none_when_snapshot_missing(hf_cache):
from huggingface_hub.file_download import repo_folder_name
repo_dir = hf_cache / repo_folder_name(repo_id = "org/broken", repo_type = "model")
(repo_dir / "refs").mkdir(parents = True)
(repo_dir / "refs" / "main").write_text("deadbeef") # no snapshots/deadbeef dir
assert hf_cache_snapshot_dir("org/broken") is None
def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch):
# An unexpanded $VAR in HF_HUB_CACHE must resolve where the loader looks.
real = tmp_path / "hub"
real.mkdir()
monkeypatch.setenv("MY_HF_CACHE", str(real))
monkeypatch.setenv("HF_HUB_CACHE", "$MY_HF_CACHE")
monkeypatch.delenv("HF_HOME", raising = False)
monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False)
snapshot = _make_cache(real, "org/emb", {"modules.json": MODULES_JSON})
assert hf_cache_snapshot_dir("org/emb") == snapshot
def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch):
# ST uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must inspect it too.
st_home = tmp_path / "st_home"
st_home.mkdir()
monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
monkeypatch.delenv("HF_HUB_CACHE", raising = False)
monkeypatch.delenv("HF_HOME", raising = False)
snapshot = _make_cache(st_home, "org/emb", {"modules.json": MODULES_JSON})
assert hf_cache_snapshot_dir("org/emb") == snapshot
def test_snapshot_dir_prefers_selected_cache_over_st_home(tmp_path, monkeypatch):
# The RAG loader passes cache_folder=active_hf_hub_cache(), which overrides
# SENTENCE_TRANSFORMERS_HOME, so the snapshot + offline security lookup must
# search the selected cache even when ST_HOME points elsewhere. Otherwise the
# gate scans a cache the model never loads from and a pickle weight in the
# selected cache slips through.
st_home = tmp_path / "st_home"
st_home.mkdir()
selected = tmp_path / "hub"
selected.mkdir()
monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
monkeypatch.delenv("HF_HUB_CACHE", raising = False)
monkeypatch.delenv("HF_HOME", raising = False)
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: SimpleNamespace(hub_cache = selected),
)
snapshot = _make_cache(selected, "org/emb", {"modules.json": MODULES_JSON}) # only in selected
assert hf_cache_snapshot_dir("org/emb") == snapshot
def test_snapshot_is_loadable_with_config_and_weights(hf_cache):
_make_cache(hf_cache, "org/emb", {"config.json": "{}", "model.safetensors": "x"})
assert hf_cache_snapshot_is_loadable("org/emb") is True
def test_snapshot_is_not_loadable_when_metadata_only(hf_cache):
# A partial cache (refs/main resolves but no weights) is not loadable.
_make_cache(hf_cache, "org/partial", {"config.json": "{}", "modules.json": MODULES_JSON})
assert hf_cache_snapshot_is_loadable("org/partial") is False
def test_snapshot_is_not_loadable_when_uncached(hf_cache):
assert hf_cache_snapshot_is_loadable("org/missing") is False
def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch):
# A pickle under SENTENCE_TRANSFORMERS_HOME must still fail closed offline.
st_home = tmp_path / "st_home"
st_home.mkdir()
monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
monkeypatch.delenv("HF_HUB_CACHE", raising = False)
monkeypatch.delenv("HF_HOME", raising = False)
_make_cache(st_home, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
with _no_network():
assert evaluate_file_security("org/pk", local_only_load = True).blocked is True
# ── is_embedding_model: offline (no network) ─────────────────────
def test_offline_true_for_cached_st_model(hf_cache, monkeypatch):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
_make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON, "config.json": "{}"})
with _no_network():
assert _is_embedding_model("org/emb") is True
def test_offline_false_for_cached_non_st_model(hf_cache, monkeypatch):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
_make_cache(hf_cache, "org/plain", {"config.json": "{}", "model.safetensors": "x"})
with _no_network():
assert _is_embedding_model("org/plain") is False
def test_offline_false_when_uncached(hf_cache, monkeypatch):
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
with _no_network():
assert _is_embedding_model("org/missing") is False
def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
_make_cache(hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON})
with _no_network():
assert _is_embedding_model("all-MiniLM-L6-v2") is True
def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch):
# An online lookup memoizes True for an UNCACHED repo (tags say embedding, no weights). Once
# offline, is_embedding_model must reclassify from the empty cache and return False, not the
# stale online True that would make settings accept a repo _get() cannot load.
with patch(
"huggingface_hub.model_info",
side_effect = lambda *a, **k: SimpleNamespace(
tags = ["sentence-transformers"], pipeline_tag = None
),
):
assert _is_embedding_model("org/uncached-emb") is True # memoized True online
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
with _no_network():
assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache
def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch):
# Because the offline branch never records a memo, once an uncached repo's snapshot
# materializes (another process populates the cache) the next call re-reports True.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
with _no_network():
assert _is_embedding_model("org/later") is False # uncached
_make_cache(hf_cache, "org/later", {"modules.json": MODULES_JSON})
assert _is_embedding_model("org/later") is True # cache now present, no stale negative
# ── is_embedding_model: online (bounded + fallback) ──────────────
def test_online_passes_bounded_timeout(hf_cache):
seen = {}
def _mi(
name,
token = None,
timeout = None,
**kw,
):
seen["timeout"] = timeout
return SimpleNamespace(tags = ["sentence-transformers"], pipeline_tag = None)
with patch("huggingface_hub.model_info", side_effect = _mi):
assert _is_embedding_model("org/emb") is True
assert seen["timeout"] == 15.0
def test_online_error_falls_back_to_cache_marker(hf_cache):
_make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON})
with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")):
assert _is_embedding_model("org/emb") is True
def test_online_error_without_cache_returns_false(hf_cache):
with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")):
assert _is_embedding_model("org/missing") is False
# ── evaluate_file_security: offline fail-closed gate ─────────────
def _offline_decision(name):
return evaluate_file_security(name, local_only_load = True)
def test_gate_allows_safetensors_only(hf_cache):
_make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"})
with _no_network():
assert _offline_decision("org/st").blocked is False
def test_gate_blocks_pickle_without_safetensors(hf_cache):
_make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
with _no_network():
decision = _offline_decision("org/pk")
assert decision.blocked is True
assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files)
def test_gate_allows_pickle_with_safetensors_sibling(hf_cache):
_make_cache(hf_cache, "org/both", {"pytorch_model.bin": "x", "model.safetensors": "y"})
with _no_network():
assert _offline_decision("org/both").blocked is False
def test_gate_blocks_sharded_pickle(hf_cache):
_make_cache(
hf_cache,
"org/shard",
{
"pytorch_model-00001-of-00002.bin": "a",
"pytorch_model-00002-of-00002.bin": "b",
},
)
with _no_network():
assert _offline_decision("org/shard").blocked is True
def test_gate_allows_nothing_cached(hf_cache):
with _no_network():
assert _offline_decision("org/missing").blocked is False
def test_gate_allows_gguf_only(hf_cache):
_make_cache(hf_cache, "org/gg", {"model.gguf": "x"})
with _no_network():
assert _offline_decision("org/gg").blocked is False
def test_gate_blocks_pickle_in_module_subdir(hf_cache):
# 0_Transformer is a module load root (listed in modules.json), so its pickle blocks.
_make_cache(
hf_cache,
"org/mod",
{"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"},
)
with _no_network():
assert _offline_decision("org/mod").blocked is True
def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache):
_make_cache(
hf_cache,
"org/mod2",
{
"modules.json": _modules_json("0_Transformer"),
"0_Transformer/pytorch_model.bin": "x",
"0_Transformer/model.safetensors": "y",
},
)
with _no_network():
assert _offline_decision("org/mod2").blocked is False
def test_gate_allows_unreferenced_nested_pickle(hf_cache):
# A pickle in a dir NOT referenced by modules.json (e.g. nemo/) is never deserialized, so it
# must not block the offline load (matches the online gate).
_make_cache(
hf_cache,
"org/aux",
{
"modules.json": MODULES_JSON, # Transformer at the root only
"model.safetensors": "w",
"nemo/pytorch_model.bin": "x",
},
)
with _no_network():
assert _offline_decision("org/aux").blocked is False
def test_gate_blocks_adapter_pickle_without_safetensors(hf_cache):
_make_cache(hf_cache, "org/ad", {"config.json": "{}", "adapter_model.bin": "x"})
with _no_network():
decision = _offline_decision("org/ad")
assert decision.blocked is True
assert any(u["path"] == "adapter_model.bin" for u in decision.unsafe_files)
def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache):
_make_cache(hf_cache, "org/ad2", {"adapter_model.bin": "x", "adapter_model.safetensors": "y"})
with _no_network():
assert _offline_decision("org/ad2").blocked is False
def test_gate_blocks_base_pickle_with_only_adapter_safetensors_decoy(hf_cache):
# A decoy adapter_model.safetensors must NOT suppress a base pytorch_model.bin (the base
# loader would still deserialize the unscanned pickle).
_make_cache(hf_cache, "org/decoy", {"pytorch_model.bin": "x", "adapter_model.safetensors": "y"})
with _no_network():
assert _offline_decision("org/decoy").blocked is True
def test_gate_blocks_adapter_pickle_with_only_base_safetensors_decoy(hf_cache):
# Symmetric: a base model.safetensors must NOT suppress an adapter_model.bin.
_make_cache(hf_cache, "org/decoy2", {"adapter_model.bin": "x", "model.safetensors": "y"})
with _no_network():
assert _offline_decision("org/decoy2").blocked is True
def test_gate_reports_snapshot_relative_path(hf_cache):
_make_cache(
hf_cache,
"org/mod3",
{"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"},
)
with _no_network():
decision = _offline_decision("org/mod3")
assert decision.blocked is True
assert any(u["path"] == "0_Transformer/pytorch_model.bin" for u in decision.unsafe_files)
# ── evaluate_file_security: online path unchanged ────────────────
def test_online_default_blocks_unsafe():
status = {
"scansDone": True,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
}
with patch(
"huggingface_hub.model_info",
side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status),
):
assert evaluate_file_security("org/x").blocked is True
def test_online_default_allows_clean():
status = {"scansDone": True, "filesWithIssues": []}
with patch(
"huggingface_hub.model_info",
side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status),
):
assert evaluate_file_security("org/x").blocked is False
# ── embeddings guard + loader ────────────────────────────────────
def test_guard_offline_blocks_pickle_only(hf_cache):
from core.rag.embeddings import UnsafeEmbeddingModelError, _guard_model_security
_make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
with _no_network():
with pytest.raises(UnsafeEmbeddingModelError):
_guard_model_security("org/pk", local_only = True)
def test_guard_offline_allows_safetensors(hf_cache):
from core.rag.embeddings import _guard_model_security
_make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"})
with _no_network():
_guard_model_security("org/st", local_only = True) # must not raise
def _install_fake_sentence_transformers(monkeypatch, captured):
class FakeSentenceTransformer:
def __init__(
self,
name,
*,
device = None,
model_kwargs = None,
local_files_only = False,
**kw,
):
captured["name"] = name
captured["device"] = device
captured["local_files_only"] = local_files_only
module = types.ModuleType("sentence_transformers")
module.SentenceTransformer = FakeSentenceTransformer
monkeypatch.setitem(sys.modules, "sentence_transformers", module)
def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch):
from core.rag import embeddings
snapshot = _make_cache(
hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}
)
# TRANSFORMERS_OFFLINE only: a cached model loads from its local snapshot dir (a local path,
# never the Hub), offline-safe on ANY sentence-transformers version.
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.setattr(embeddings, "_model", None, raising = False)
monkeypatch.setattr(embeddings, "_name", None, raising = False)
monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
captured = {}
_install_fake_sentence_transformers(monkeypatch, captured)
with _no_network():
embeddings._get("org/st")
assert captured["name"] == str(snapshot)
def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch):
from core.rag import embeddings
empty = tmp_path / "hub"
empty.mkdir()
monkeypatch.setenv("HF_HUB_CACHE", str(empty))
monkeypatch.delenv("HF_HOME", raising = False)
monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False)
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.setattr(embeddings, "_model", None, raising = False)
monkeypatch.setattr(embeddings, "_name", None, raising = False)
monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
# No cache -> repo-id load forced cache-only (fails fast offline, not a hang).
monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None)
captured = {}
_install_fake_sentence_transformers(monkeypatch, captured)
embeddings._get("org/uncached-xyz")
assert captured["name"] == "org/uncached-xyz"
assert captured["local_files_only"] is True
def test_get_online_omits_local_files_only(monkeypatch):
from core.rag import embeddings
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setattr(embeddings, "_model", None, raising = False)
monkeypatch.setattr(embeddings, "_name", None, raising = False)
monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
# Isolate the loader wiring from the online guard's network calls.
monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None)
captured = {}
_install_fake_sentence_transformers(monkeypatch, captured)
embeddings._get("org/online")
assert captured["local_files_only"] is False

106
studio/backend/tests/test_offline_gguf_cache_fallback.py
View file

@ -119,10 +119,21 @@ def _build_cache(
return snap
def _symlink_or_skip(link: Path, target: Path) -> None:
try:
link.symlink_to(target)
except OSError as exc:
pytest.skip(f"symlinks unavailable: {exc}")
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
"""Point ``huggingface_hub.constants.HF_HUB_CACHE`` at a temp dir."""
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: _types.SimpleNamespace(hub_cache = tmp_path),
)
return tmp_path
@ -220,6 +231,10 @@ class TestGgufVariantFileResolution:
return f"/fake/{repo_id}/{filename}"
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: _types.SimpleNamespace(hub_cache = tmp_path),
)
with (
patch(
"huggingface_hub.list_repo_files",
@ -427,6 +442,40 @@ class TestGgufVariantFileResolution:
assert out == str(snap / "mmproj-F16.gguf")
def test_download_companion_uses_selected_cache_not_import_time_default(
self, monkeypatch, tmp_path
):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
import_time_cache = tmp_path / "import-time-cache"
selected_cache = tmp_path / "selected-cache"
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(import_time_cache))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: _types.SimpleNamespace(hub_cache = selected_cache),
)
repo = "unsloth/vision-GGUF"
snap = _build_cache(selected_cache, repo, {"mmproj-F16.gguf": 4})
backend = LlamaCppBackend()
offline_error = type("OfflineModeIsEnabled", (Exception,), {})
def fail_list(*_args, **_kwargs):
raise offline_error("offline")
def fail_download(*_args, **_kwargs):
raise AssertionError("selected-cache companion must not download")
with (
patch("huggingface_hub.list_repo_files", fail_list),
patch(
"core.inference.llama_cpp.hf_hub_download_with_xet_fallback",
fail_download,
),
):
out = backend._download_mmproj(hf_repo = repo)
assert out == str(snap / "mmproj-F16.gguf")
def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path):
backend = LlamaCppBackend()
downloaded: list[str] = []
@ -453,6 +502,10 @@ class TestGgufVariantFileResolution:
return f"/fake/{repo_id}/{filename}"
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(
"utils.hf_cache_settings.get_hf_cache_paths",
lambda: _types.SimpleNamespace(hub_cache = tmp_path),
)
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
@ -1084,7 +1137,7 @@ class TestListLocalGgufVariantsSubdir:
target.write_bytes(b"\0" * 20)
out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M")
assert out == str(target.resolve())
assert out == str(target.absolute())
def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant
@ -1094,6 +1147,57 @@ class TestListLocalGgufVariantsSubdir:
assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None
def test_find_local_gguf_by_variant_keeps_split_symlink_name(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant
blobs = tmp_path / "blobs"
blobs.mkdir()
snap = tmp_path / "snapshots" / "rev" / "BF16"
snap.mkdir(parents = True)
(tmp_path / "snapshots" / "rev" / "config.json").write_text("{}")
for i, sha in enumerate(("aa" * 32, "bb" * 32), start = 1):
(blobs / sha).write_bytes(b"\0" * 10)
_symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha)
out = _find_local_gguf_by_variant(str(tmp_path / "snapshots" / "rev"), "BF16")
assert out is not None
assert Path(out).name == "model-BF16-00001-of-00002.gguf"
def test_detect_gguf_model_keeps_split_symlink_name(self, tmp_path):
from utils.models.model_config import detect_gguf_model
blobs = tmp_path / "blobs"
blobs.mkdir()
snap = tmp_path / "snapshots" / "rev"
snap.mkdir(parents = True)
for i, (sha, size) in enumerate((("cc" * 32, 10), ("dd" * 32, 20)), start = 1):
(blobs / sha).write_bytes(b"\0" * size)
_symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha)
out = detect_gguf_model(str(snap))
assert out is not None
assert Path(out).name == "model-BF16-00001-of-00002.gguf"
def test_lone_split_symlink_uses_colocated_target_shards(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant, detect_gguf_model
target_dir = tmp_path / "external" / "BF16"
target_dir.mkdir(parents = True)
target = target_dir / "model-BF16-00001-of-00002.gguf"
target.write_bytes(b"\0" * 10)
(target_dir / "model-BF16-00002-of-00002.gguf").write_bytes(b"\0" * 10)
local = tmp_path / "local"
local.mkdir()
(local / "config.json").write_text("{}")
link = local / target.name
_symlink_or_skip(link, target)
expected = str(target.absolute())
assert _find_local_gguf_by_variant(str(local), "BF16") == expected
assert detect_gguf_model(str(local)) == expected
assert detect_gguf_model(str(link)) == expected
def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path):
from utils.models.model_config import ModelConfig

140
studio/backend/tests/test_openai_auto_switch.py
View file

@ -68,7 +68,13 @@ class _LoadRecorder:
request,
fastapi_request,
current_subject = None,
*,
current_request_counted = False,
):
# Mirror the production load boundary before recording any replacement.
await inference_route._wait_for_model_switch_idle(
current_request_counted = current_request_counted
)
self.calls.append(request)
if self.fail:
from fastapi import HTTPException
@ -94,7 +100,6 @@ def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder):
# gate that auto-switch already owns, so it calls the impl directly).
monkeypatch.setattr(inference_route, "_load_model_impl", recorder)
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
def _run_hook(model = "some/model"):
@ -1091,6 +1096,7 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch
from pathlib import Path
import routes.models as models_route
from utils import paths as upaths
from utils import hf_cache_settings
import storage.studio_db as studio_db
scanned = []
@ -1111,13 +1117,18 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch
)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active")
monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False)
monkeypatch.setattr(
hf_cache_settings,
"known_hf_hub_caches",
lambda: [tmp_path / "active", tmp_path / "previous"],
)
monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy")
monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default")
monkeypatch.setattr(upaths, "lmstudio_model_dirs", lambda: [tmp_path / "lmstudio"])
monkeypatch.setattr(
studio_db, "list_scan_folders", lambda: [{"path": str(tmp_path / "custom")}]
)
for sub in ("active", "legacy", "default", "lmstudio", "custom"):
for sub in ("active", "previous", "legacy", "default", "lmstudio", "custom"):
(tmp_path / sub).mkdir()
resolver._build_index()
@ -1126,6 +1137,7 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch
lm = {p for k, p in scanned if k == "lm"}
assert str((tmp_path / "legacy").resolve()) in hf
assert str((tmp_path / "default").resolve()) in hf
assert str((tmp_path / "previous").resolve()) in hf
assert str((tmp_path / "custom").resolve()) in hf
assert str((tmp_path / "lmstudio").resolve()) in lm
@ -1205,10 +1217,9 @@ def test_middleware_ignores_non_post(monkeypatch):
# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ──
def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
# A cross-model swap must 409 (not kill) while another inference request is in
# flight; the requesting call itself is excluded from the count.
from fastapi import HTTPException
def test_auto_switch_waits_for_another_inference_to_finish(monkeypatch):
# A cross-model swap queues while another request is generating, then loads
# after that request drains. The requesting call itself is excluded.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M")
@ -1222,10 +1233,18 @@ def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
)
monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one
monkeypatch.setattr(kw, "_pending", 0)
with pytest.raises(HTTPException) as exc:
_run_hook("org/B-GGUF:Q8_0")
assert exc.value.status_code == 409
assert rec.calls == []
async def _drive():
task = asyncio.create_task(
inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
)
await asyncio.sleep(0.05)
assert rec.calls == []
kw._note_end() # the other generation finishes; this request remains counted
await asyncio.wait_for(task, timeout = 1)
asyncio.run(_drive())
assert len(rec.calls) == 1
def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch):
@ -1411,13 +1430,12 @@ def test_concurrent_same_target_requests_load_once(monkeypatch):
monkeypatch.setattr(kw, "_pending", 0)
inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1 # loads once, no 409
assert len(rec.calls) == 1
def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch):
# A concurrent request heading to a different target still blocks the swap: the
# same-target exclusion must not swallow a genuinely conflicting request.
from fastapi import HTTPException
def test_queued_different_target_does_not_deadlock_current_swap(monkeypatch):
# A concurrent request already queued for another target is not generating,
# so it must not prevent the current serialized swap from proceeding.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@ -1432,10 +1450,8 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat
monkeypatch.setattr(kw, "_inflight", 2)
monkeypatch.setattr(kw, "_pending", 0)
inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1)
with pytest.raises(HTTPException) as exc:
_run_hook("org/B-GGUF:Q8_0")
assert exc.value.status_code == 409
assert rec.calls == []
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1
def test_v1_models_advertises_repo_id_not_load_path(monkeypatch):
@ -1481,6 +1497,37 @@ def test_load_route_holds_lifecycle_gate(monkeypatch):
assert "_load_model_impl" in src
def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded():
# Both replacement directions drain active inference, then recheck whether a
# sidecar install reserved the lifecycle gate during that wait. Exact-model
# reuse exits earlier, so an already-loaded model never waits on unrelated inference.
import inspect
src = inspect.getsource(inference_route._load_model_impl)
gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:"))
gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait)
unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait)
standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1)
standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
already_loaded = src.index('status = "already_loaded"')
assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth
assert standard_wait < standard_sidecar_check < unload_gguf
def test_switch_waiter_deregisters_before_swap_gate_release():
# A waiter left registered after the swap gate is released would let a swap on
# another event loop count the finished request as still queued, pass the drain
# early, and unload the model that request is about to generate against.
import inspect
src = inspect.getsource(inference_route._maybe_auto_switch_model)
deregister = src.index("_note_switch_waiter(key, -1)")
release = src.index("_auto_switch_process_lock.release()")
assert deregister < release
def _anthropic_payload(max_tokens = None):
from models.inference import AnthropicMessagesRequest, AnthropicMessage
return AnthropicMessagesRequest(
@ -1519,9 +1566,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch):
# ── review round 6: concurrency ordering, external untrack, unload gate, ids ──
def test_pending_same_target_request_does_not_force_409(monkeypatch):
def test_pending_same_target_request_does_not_block_swap(monkeypatch):
# A second same-target request blocked in the middleware (pending, not yet
# generating) must not make the first request 409: pending is excluded.
# generating) must not block the first request: pending is excluded.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@ -1536,13 +1583,13 @@ def test_pending_same_target_request_does_not_force_409(monkeypatch):
monkeypatch.setattr(kw, "_inflight", 1) # just the caller
monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1 # loads once, no 409
assert len(rec.calls) == 1
def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch):
def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch):
# The real middleware counts a concurrent same-model request as in-flight
# before it resolves and registers a target waiter. The raw-request waiter,
# registered before resolve, must still exclude it so the first request loads.
# before it resolves and registers a target waiter. Treat it as active until
# its target is known, then recognize it as another queued switch request.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@ -1556,10 +1603,20 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat
)
monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin
monkeypatch.setattr(kw, "_pending", 0)
# The twin has only registered its raw requested model (not yet a target waiter).
inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1)
_run_hook("org/B-GGUF:Q8_0")
assert len(rec.calls) == 1 # loads once, no 409
# The twin is still resolving, so it is counted in-flight but has not joined
# the concrete target queue yet.
async def _drive():
task = asyncio.create_task(
inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
)
await asyncio.sleep(0.05)
assert rec.calls == []
inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
await asyncio.wait_for(task, timeout = 1)
asyncio.run(_drive())
assert len(rec.calls) == 1
def test_external_untrack_decrements_inflight_and_is_idempotent():
@ -1595,11 +1652,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch):
assert not backend.is_loaded # torn down despite the active request
def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
def test_auto_switch_waits_when_unsloth_stream_active(monkeypatch):
# The GGUF slot is empty but an Unsloth model is streaming (counted in-flight).
# _load_model_impl would unload it, so auto-switch must 409, not only when a
# GGUF is loaded.
from fastapi import HTTPException
# The replacement waits for it just as it does for a GGUF generation.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None) # no GGUF loaded
@ -1613,10 +1668,18 @@ def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
)
monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request
monkeypatch.setattr(kw, "_pending", 0)
with pytest.raises(HTTPException) as exc:
_run_hook("org/B-GGUF:Q8_0")
assert exc.value.status_code == 409
assert rec.calls == [] # the active Unsloth model is not torn down
async def _drive():
task = asyncio.create_task(
inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
)
await asyncio.sleep(0.05)
assert rec.calls == []
kw._note_end()
await asyncio.wait_for(task, timeout = 1)
asyncio.run(_drive())
assert len(rec.calls) == 1
def test_public_model_id_prefers_advertised_over_path():
@ -3097,6 +3160,8 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
request,
fastapi_request,
current_subject = None,
*,
current_request_counted = False,
):
with slock:
state["cur"] += 1
@ -3114,7 +3179,6 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load)
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
barrier = threading.Barrier(2)

16
studio/backend/tests/test_password_prompt.py
View file

@ -183,6 +183,22 @@ def test_loop_short_password_reprompts(monkeypatch):
assert "at least 8 characters" in out
def test_loop_whitespace_only_reprompts(monkeypatch):
ok, applied, out = _run_loop(monkeypatch, _keys(" " * 8, "long-enough-pw", "long-enough-pw"))
assert ok is True
assert applied == ["long-enough-pw"]
assert "contain spaces" in out
def test_loop_password_with_inner_space_reprompts(monkeypatch):
ok, applied, out = _run_loop(
monkeypatch, _keys("has space pw", "long-enough-pw", "long-enough-pw")
)
assert ok is True
assert applied == ["long-enough-pw"]
assert "contain spaces" in out
def test_loop_rejects_current_password(monkeypatch):
ok, applied, out = _run_loop(
monkeypatch, _keys("bootstrap-pw", "fresh-password", "fresh-password")

272
studio/backend/tests/test_picker_service.py Normal file
View file

@ -0,0 +1,272 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import json
from types import SimpleNamespace
from picker.service import (
MAX_TEMPLATE_METADATA_BYTES,
_chat_template_from_dir,
_chat_template_from_processor_json,
_chat_template_from_tokenizer_config,
_chat_template_from_tokenizer_dir,
_find_gguf_in_dir,
_iter_ggufs,
read_default_chat_template,
validate_chat_template,
)
def test_iter_ggufs_skips_gguf_companions(tmp_path):
mtp_dir = tmp_path / "MTP"
mtp_dir.mkdir()
main = tmp_path / "model-Q8_0.gguf"
main.write_bytes(b"")
(tmp_path / "mmproj-F16.gguf").write_bytes(b"")
(tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"")
(mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
(tmp_path / "model-Q8_0-be.gguf").write_bytes(b"")
assert _iter_ggufs(tmp_path) == [main]
def test_find_gguf_in_dir_matches_quant_label(tmp_path):
mtp_dir = tmp_path / "MTP"
mtp_dir.mkdir()
main = tmp_path / "model-Q8_0.gguf"
main.write_bytes(b"")
(mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
assert _find_gguf_in_dir(tmp_path, "Q8_0") == main
assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path):
smaller = tmp_path / "a-model-Q4_K_M.gguf"
larger = tmp_path / "z-model-Q8_0.gguf"
smaller.write_bytes(b"0")
larger.write_bytes(b"00")
assert _find_gguf_in_dir(tmp_path, None) == larger
def test_find_gguf_in_dir_without_variant_prefers_first_split(tmp_path):
first = tmp_path / "model-Q4_K_M-00001-of-00003.gguf"
second = tmp_path / "model-Q4_K_M-00002-of-00003.gguf"
third = tmp_path / "model-Q4_K_M-00003-of-00003.gguf"
first.write_bytes(b"0")
second.write_bytes(b"000")
third.write_bytes(b"00")
assert _find_gguf_in_dir(tmp_path, None) == first
first.unlink()
assert _find_gguf_in_dir(tmp_path, None) == second
def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path):
target = tmp_path / "model-IQ4_XS-3.53bpw.gguf"
target.write_bytes(b"")
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target
assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target
assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
def test_validate_chat_template_accepts_valid_and_empty():
assert validate_chat_template("{{ messages[0].content }}").valid is True
assert validate_chat_template("").valid is True
assert validate_chat_template(" ").valid is True
def test_validate_chat_template_reports_syntax_error_with_line():
result = validate_chat_template("{% if %}{% endif %}")
assert result.valid is False
assert result.error is not None
assert result.error.startswith("Line ")
def test_chat_template_from_tokenizer_config_reads_string():
assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO"
assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None
assert _chat_template_from_tokenizer_config({}) is None
def test_chat_template_from_tokenizer_config_prefers_named_default():
config = {
"chat_template": [
{"name": "tool_use", "template": "TOOL"},
{"name": "default", "template": "DEFAULT"},
]
}
assert _chat_template_from_tokenizer_config(config) == "DEFAULT"
def test_chat_template_from_tokenizer_config_falls_back_to_first_entry():
config = {
"chat_template": [
{"name": "tool_use", "template": "TOOL"},
{"name": "other", "template": "OTHER"},
]
}
assert _chat_template_from_tokenizer_config(config) == "TOOL"
def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path):
(tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8")
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA"
def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path):
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path):
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG"
def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch):
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# Selecting a variant must not flip precedence to the embedded GGUF template.
assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG"
def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch):
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# With no tokenizer sidecar, the embedded GGUF template is still the fallback.
assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF"
def test_chat_template_from_dir_returns_none_when_absent(tmp_path):
assert _chat_template_from_dir(tmp_path) is None
def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch):
gguf = tmp_path / "model-Q4_K_M.gguf"
gguf.write_bytes(b"")
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# A directly selected .gguf must prefer a maintained sidecar over its embedded copy.
assert read_default_chat_template(str(gguf)) == "FROM_CONFIG"
def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch):
gguf = tmp_path / "model-Q4_K_M.gguf"
gguf.write_bytes(b"")
monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# With no sidecar next to the file, the embedded GGUF template is the fallback.
assert read_default_chat_template(str(gguf)) == "FROM_GGUF"
def test_tokenizer_config_over_size_limit_is_skipped_not_parsed(tmp_path):
# An oversized tokenizer_config.json must be skipped before json.loads so a
# hostile sidecar cannot exhaust memory.
padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024)
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "HELLO", "_pad": padding}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) is None
def test_processor_json_over_size_limit_is_skipped_not_parsed(tmp_path):
padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024)
(tmp_path / "chat_template.json").write_text(
json.dumps({"default": "HELLO", "_pad": padding}), encoding = "utf-8"
)
assert _chat_template_from_processor_json(tmp_path) is None
def test_tokenizer_config_at_size_limit_is_still_read(tmp_path):
# A normal-sized config is unaffected by the bound (regression guard).
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
def test_remote_template_over_size_limit_is_skipped_before_download(monkeypatch):
# An uncached Hub repo whose template exceeds the cap must be skipped via the
# remote size pre-check, never downloaded.
import huggingface_hub
monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name)
monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: [])
def _fail_download(*args, **kwargs):
raise AssertionError("oversized remote template must not be downloaded")
def _fake_get_paths_info(self, repo_id, paths, **kwargs):
return [SimpleNamespace(path = p, size = MAX_TEMPLATE_METADATA_BYTES + 1) for p in paths]
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fail_download)
monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info)
assert read_default_chat_template("org/oversized-model") is None
def test_remote_oversized_jinja_falls_through_to_tokenizer_template(tmp_path, monkeypatch):
# A raw chat_template.jinja between the response cap (MAX_CHAT_TEMPLATE_BYTES)
# and the download bound (MAX_TEMPLATE_METADATA_BYTES) must not be returned: the
# route drops it, so the remote path must skip the oversized Jinja and fall
# through to the smaller tokenizer_config.json.
import huggingface_hub
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
big_jinja = tmp_path / "chat_template.jinja"
big_jinja.write_text("{{ x }}" * (MAX_CHAT_TEMPLATE_BYTES // 4), encoding = "utf-8")
assert MAX_CHAT_TEMPLATE_BYTES < big_jinja.stat().st_size < MAX_TEMPLATE_METADATA_BYTES
tokenizer_config = tmp_path / "tokenizer_config.json"
tokenizer_config.write_text(json.dumps({"chat_template": "SMALL_TEMPLATE"}), encoding = "utf-8")
files = {
"chat_template.jinja": big_jinja,
"tokenizer_config.json": tokenizer_config,
}
selected_cache = tmp_path / "selected-cache" / "hub"
observed_cache_dirs = []
monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name)
monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: [])
monkeypatch.setattr("picker.service.active_hf_hub_cache", lambda: str(selected_cache))
def _fake_download(repo_id, rel, **kwargs):
observed_cache_dirs.append(kwargs.get("cache_dir"))
target = files.get(rel)
if target is None:
raise FileNotFoundError(rel)
return str(target)
def _fake_get_paths_info(self, repo_id, paths, **kwargs):
return [
SimpleNamespace(
path = p,
size = files[p].stat().st_size if p in files else 0,
)
for p in paths
]
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download)
monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info)
assert read_default_chat_template("org/big-jinja-model") == "SMALL_TEMPLATE"
assert observed_cache_dirs
assert set(observed_cache_dirs) == {str(selected_cache)}

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