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

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

Browse source 
# Conflicts:
#	unsloth/models/_utils.py
This commit is contained in:
Daniel Han 2026-06-18 22:54:30 -07:00
commit 73bf482f65
496 changed files with 49331 additions and 10331 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -434,7 +434,7 @@ jobs:
# ─────────────────────────────────────────────────────────────
# Semgrep: design-flaw detection (catches what regex-pattern
# scanning of malicious authors cannot — first-party logic bugs
# scanning of malicious authors cannot, e.g. first-party logic bugs
# like langchain-core CVE-2025-68664 dumps/dumpd injection,
# n8n CVE-2025-68668 _pyodide.eval_code sandbox escape, marimo
# CVE-2026-39987 unauth WebSocket).
@ -849,10 +849,13 @@ jobs:
grep -q "Standalone pre-install package scanner" scripts/scan_packages.py
- name: Scan declared + transitive Python deps
# scan_packages.py exits 1 on CRITICAL/HIGH findings, 0 on
# clean. We swallow the exit because the baseline isn't
# triaged yet; surface the findings in the workflow summary.
# Drop continue-on-error after the first clean run on main.
# scan_packages.py exits 1 on NON-baselined CRITICAL/HIGH
# findings, 0 otherwise. It scans code-only (docstrings and
# comments are blanked first) and suppresses reviewed
# known-good findings via scripts/scan_packages_baseline.json,
# so legitimate-library noise no longer red-fails the gate.
# The step stays advisory until SCAN_ENFORCE=1 (see env below);
# then PIPESTATUS propagates the scanner's exit code.
#
# `--with-deps` walks PyPI metadata to enumerate every
# transitive dep the declared set would install, then scans
@ -869,6 +872,14 @@ jobs:
# downloads in exchange for wall-clock parallelism.
env:
SHARD_FILES: ${{ matrix.shard.files }}
# Enforcement switch. "1" = blocking: a non-baselined CRITICAL/HIGH
# fails the build. scan_packages.py scans code-only (docstrings/comments
# stripped), fetches sdist-only packages directly from PyPI (no build)
# so every shard resolves, and honors the reviewed allowlist at
# scripts/scan_packages_baseline.json, so only NON-baselined
# CRITICAL/HIGH cause its exit 1. The committed baseline makes all three
# shards exit 0 today; set this back to "0" to return to advisory.
SCAN_ENFORCE: "1"
run: |
set +e
mkdir -p logs
@ -884,12 +895,14 @@ jobs:
fi
done
echo "::endgroup::"
rc=0
if [ ${#REQ_ARGS[@]} -eq 0 ]; then
echo "[security-audit] shard ${{ matrix.shard.id }}: no PyPI specs, nothing to scan" \
| tee "$LOG"
else
python scripts/scan_packages.py --with-deps "${REQ_ARGS[@]}" \
2>&1 | tee "$LOG"
rc=${PIPESTATUS[0]}
fi
{
echo "## scan_packages :: shard ${{ matrix.shard.id }}"
@ -897,11 +910,19 @@ jobs:
echo "### Files in this shard"
for f in $SHARD_FILES; do echo "- audit-reqs/$f.txt"; done
echo
echo "scan_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)"
echo
echo '### Findings (tail)'
echo '```'
tail -200 "$LOG"
echo '```'
} >> "$GITHUB_STEP_SUMMARY"
# Advisory by default; blocking once SCAN_ENFORCE=1 and the baseline
# is committed. PIPESTATUS is captured above so `tee` does not mask the
# scanner's exit code.
if [ "$SCAN_ENFORCE" = "1" ]; then
exit "$rc"
fi
- uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
if: always()
@ -975,24 +996,37 @@ jobs:
python3 -c "import ast; ast.parse(open('scripts/scan_npm_packages.py').read())"
- name: Scan npm tarballs (declared + transitive, no install)
# The script exits 1 on HIGH/CRITICAL findings; we capture the
# full log and surface it in the step summary either way. It
# never runs `npm install`, never executes anything from a
# downloaded tarball, and only fetches from registry.npmjs.org.
# Initially non-blocking so the baseline can settle; drop
# continue-on-error once the baseline is clean for a week.
# scan_npm_packages.py exits 1 on NON-baselined HIGH/CRITICAL
# findings, 0 otherwise. It scans code-only (JS/TS comments are
# blanked first) and honors a reviewed allowlist at
# scripts/scan_npm_packages_baseline.json. It never runs
# `npm install`, never executes anything from a downloaded
# tarball, and only fetches from registry.npmjs.org. The npm
# corpus is clean (the baseline is empty), so the gate is
# enforcing (SCAN_ENFORCE=1) and any new finding fails the build.
env:
SCAN_ENFORCE: "1"
run: |
set -o pipefail
set +e
LOG=logs-scan-npm.txt
python3 scripts/scan_npm_packages.py 2>&1 | tee "$LOG"
rc=${PIPESTATUS[0]}
{
echo "## scan_npm_packages"
echo
echo "scan_npm_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)"
echo
echo '### Findings (tail)'
echo '```'
tail -300 "$LOG"
echo '```'
} >> "$GITHUB_STEP_SUMMARY"
# Blocking: the npm corpus is clean, so any non-baselined
# HIGH/CRITICAL is new and should fail the build. PIPESTATUS is
# captured above so `tee` does not mask the scanner's exit code.
if [ "$SCAN_ENFORCE" = "1" ]; then
exit "$rc"
fi
- uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
if: always()

3
.github/workflows/studio-backend-ci.yml vendored
View file

@ -224,7 +224,8 @@ jobs:
tests/sh/test_mac_intel_compat.sh \
tests/sh/test_nvcc_meets_llama_minimum.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh; do
tests/sh/test_torch_constraint.sh \
tests/sh/test_torch_flavor.sh; do
echo "::group::$s"
bash "$s"
echo "::endgroup::"

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

@ -317,7 +317,7 @@ jobs:
timeout-minutes: 25
env:
# Tool calling is the highest-volume GGUF in this workflow
# (Qwen3.5-2B at IQ3_XXS = ~890 MiB). Caching HF_HOME would
# (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB). Caching HF_HOME would
# store xet chunks + blobs + snapshots = ~4 GiB compressed --
# 4-5x file-size inflation, dominated by xet chunks. Use main's
# `--local-dir gguf-cache` pattern to cache the flat .gguf only.
@ -326,8 +326,11 @@ jobs:
# path keeps the test off HF_HOME entirely so the cache size
# tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images
# jobs still cover the gguf_variant resolution path.
# Q4_K_XL, not IQ3_XXS: at IQ3_XXS this model emits malformed
# tool calls that llama-server's peg-native parser rejects with a
# 500. Mac/Windows already use Q4_K_XL for the same reason.
GGUF_REPO: unsloth/Qwen3.5-2B-GGUF
GGUF_FILE: Qwen3.5-2B-UD-IQ3_XXS.gguf
GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf
STUDIO_PORT: '18889'
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
@ -772,6 +775,9 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
ss -tln | grep ":${STUDIO_PORT}" || true
# Capture backend + llama-server logs so a 500 has a server-side traceback.
mkdir -p logs/server-logs
cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true
- name: Upload logs
# Always upload so green runs are still reviewable.
@ -784,6 +790,7 @@ jobs:
path: |
logs/studio.log
logs/install.log
logs/server-logs/
retention-days: 7
# ─────────────────────────────────────────────────────────────────────

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

@ -276,6 +276,10 @@ jobs:
run: |
kill "${STUDIO_IME_PID}" 2>/dev/null || true
sleep 2
# Capture backend + llama-server logs (all three Studios share this
# dir) so a stray 500 has a server-side traceback.
mkdir -p logs/server-logs
cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true
- name: Upload Playwright artifacts
# Always upload so a green run's screenshots stay reviewable --
@ -289,6 +293,7 @@ jobs:
logs/studio_extra.log
logs/studio_ime.log
logs/install.log
logs/server-logs/
logs/playwright
logs/playwright_extra
logs/playwright_ime

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

@ -65,17 +65,20 @@ jobs:
with:
persist-credentials: false
# Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit
# test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke.
- name: setup.ps1 unit test (Resolve-CudaToolkit)
# Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell
# unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke.
- name: PowerShell installer unit tests
shell: pwsh
run: |
$errs = $null
[void][System.Management.Automation.Language.Parser]::ParseFile(
(Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs)
if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
Write-Host "setup.ps1 parsed with no errors"
foreach ($f in @('install.ps1', 'studio/setup.ps1')) {
$errs = $null
[void][System.Management.Automation.Language.Parser]::ParseFile(
(Resolve-Path $f).Path, [ref]$null, [ref]$errs)
if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
Write-Host "$f parsed with no errors"
}
pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1
pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:

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

@ -136,6 +136,17 @@ jobs:
}
}
- name: Seed a legacy launch-studio.vbs (upgrade-cleanup check)
# Simulate a pre-hardening install so the post-install assertion below
# proves the installer DELETES an existing launch-studio.vbs (the exact
# Kaspersky-flagged file), not merely stops generating it.
shell: pwsh
run: |
$appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio'
New-Item -ItemType Directory -Force -Path $appDir | Out-Null
Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode
Write-Host "seeded legacy launch-studio.vbs at $appDir"
- name: Install Studio (--local, --no-torch)
# install.ps1 is the supported Windows installer. install.sh
# has no Windows branch (apt-get / brew calls). The PS1
@ -192,6 +203,69 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut)
# The shortcut launch path is otherwise untested here (the steps below
# boot `unsloth studio` directly). Guard against re-introducing the VBS
# that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk
# pointing anywhere other than hidden PowerShell over launch-studio.ps1.
shell: pwsh
run: |
$appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio'
if (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.vbs')) {
throw "regression: launch-studio.vbs exists (the Kaspersky VBS-FP shape)"
}
if (-not (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.ps1'))) {
throw "missing launch-studio.ps1 in $appDir"
}
$lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk'
if (-not (Test-Path -LiteralPath $lnk)) {
$lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk'
}
if (-not (Test-Path -LiteralPath $lnk)) { throw "no Unsloth Studio.lnk on Desktop or Start Menu" }
$sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk)
Write-Host "shortcut target: $($sc.TargetPath)"
Write-Host "shortcut args: $($sc.Arguments)"
if ($sc.TargetPath -match 'wscript\.exe$') { throw "shortcut still targets wscript.exe (VBS host)" }
if ($sc.TargetPath -notmatch 'powershell\.exe$') { throw "unexpected shortcut target: $($sc.TargetPath)" }
if ($sc.Arguments -notmatch '-WindowStyle Hidden') {
throw "shortcut must launch windowless (-WindowStyle Hidden)"
}
Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)"
- name: Launch Studio via the shortcut and assert health
# Run the exact command the .lnk stores (hidden PowerShell over
# launch-studio.ps1) and confirm it brings the backend up. This is the
# only step that proves the shortcut launch is not silently broken.
# Default port range is 8888-8908; the later UI tests use 18896/18897, so
# there is no conflict, and we tear this server down before they boot.
shell: pwsh
run: |
$lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk'
if (-not (Test-Path -LiteralPath $lnk)) {
$lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk'
}
$sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk)
Write-Host "launching: $($sc.TargetPath) $($sc.Arguments)"
Start-Process -FilePath $sc.TargetPath -ArgumentList $sc.Arguments -WorkingDirectory $sc.WorkingDirectory
$foundPort = 0
foreach ($i in 1..180) {
foreach ($port in 8888..8908) {
try {
$r = Invoke-RestMethod -Uri "http://127.0.0.1:$port/api/health" -TimeoutSec 1
if ($r.status -eq 'healthy' -and $r.service -eq 'Unsloth UI Backend') { $foundPort = $port; break }
} catch {}
}
if ($foundPort) { break }
Start-Sleep -Seconds 1
}
# Tear down the shortcut-launched server before the main UI tests boot.
try {
$owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess
if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null }
} catch {}
if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" }
Write-Host "Studio healthy on port $foundPort (launched via the shortcut)"
- name: Add Studio shim to GITHUB_PATH
# install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe
# and adds that dir to the User PATH via the Windows registry.

2
.gitignore vendored
View file

@ -235,3 +235,5 @@ package-lock.json
!studio/backend/core/data_recipe/oxc-validator/package-lock.json
!studio/package-lock.json
llama.cpp/
# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo.
/~/

2
.pre-commit-config.yaml
View file

@ -1,6 +1,6 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.16
rev: v0.15.17
hooks:
- id: ruff
args:

313
install.ps1
View file

@ -494,6 +494,37 @@ function Install-UnslothStudio {
}
}
# Retry Invoke-InstallCommand on transient uv download failures with backoff.
# Returns the last exit code on permanent failure so rollback still fires.
function Invoke-InstallCommandRetry {
param(
[Parameter(Mandatory = $true, Position = 0)][ScriptBlock]$Command,
[string]$Label = "install step"
)
# Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables).
# TryParse with bounds avoids an Int32 overflow throw. Bounds: 1..100 retries, 0..3600s.
$maxAttempts = 3
$parsedAttempts = 0
if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRIES, [ref]$parsedAttempts) -and $parsedAttempts -ge 1 -and $parsedAttempts -le 100) {
$maxAttempts = $parsedAttempts
}
$delay = 3
$parsedDelay = 0
if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRY_DELAY, [ref]$parsedDelay) -and $parsedDelay -ge 0 -and $parsedDelay -le 3600) {
$delay = $parsedDelay
}
$attempt = 1
while ($true) {
$code = Invoke-InstallCommand $Command
if ($code -eq 0) { return 0 }
if ($attempt -ge $maxAttempts) { return $code }
substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow"
Start-Sleep -Seconds $delay
$attempt++
$delay = $delay * 2
}
}
function New-StudioShortcuts {
param(
[Parameter(Mandatory = $true)][string]$UnslothExePath
@ -518,7 +549,6 @@ function Install-UnslothStudio {
}
$appDir = $StudioDataDir
$launcherPs1 = Join-Path $appDir "launch-studio.ps1"
$launcherVbs = Join-Path $appDir "launch-studio.vbs"
$desktopDir = [Environment]::GetFolderPath("Desktop")
$desktopLink = if ($desktopDir -and $desktopDir.Trim()) {
Join-Path $desktopDir "Unsloth Studio.lnk"
@ -811,19 +841,30 @@ exit 0
# even when install.ps1 is executed from PowerShell 7.
$utf8Bom = New-Object System.Text.UTF8Encoding($true)
[System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom)
# shell.Run(cmd, 0, ...) already hides the window, so -WindowStyle Hidden
# is redundant; omitting it trims an AV-heuristic token (Kaspersky FP).
$vbsContent = @"
Set shell = CreateObject("WScript.Shell")
cmd = "powershell -NoProfile -ExecutionPolicy Bypass -File ""$launcherPs1"""
shell.Run cmd, 0, False
"@
# WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths.
Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force
# No .vbs launcher is written. A WScript.Shell .vbs that spawns a hidden
# ExecutionPolicy-Bypass PowerShell is exactly the shape VBS-dropper
# heuristics score (e.g. Kaspersky HEUR:Trojan.VBS.Agent.gen). The .lnk
# shortcuts instead point straight at powershell.exe running
# launch-studio.ps1 with a hidden window (selected below).
# Delete any launch-studio.vbs left by a pre-hardening install. New
# installs no longer generate it, but an upgrade that merely stopped
# generating it would leave the exact file AV flags on disk, so remove
# it explicitly. Covers default and env-mode installs (same $appDir).
$legacyLauncherVbs = Join-Path $appDir "launch-studio.vbs"
if (Test-Path -LiteralPath $legacyLauncherVbs) {
Remove-Item -LiteralPath $legacyLauncherVbs -Force -ErrorAction SilentlyContinue
}
# Prefer bundled icon from local clone/dev installs.
# If not available, best-effort download from raw GitHub.
# We only attach the icon if the resulting file has a valid ICO header.
# Snapshot the existing icon first so we can tell whether it actually
# changed and gate the heavier icon-cache refresh on a real change.
$preIconHash = $null
if (Test-Path -LiteralPath $iconPath) {
try { $preIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash } catch {}
}
$hasValidIcon = $false
if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) {
try {
@ -859,6 +900,24 @@ shell.Run cmd, 0, False
}
}
# Did the icon content actually change vs the previous install?
# Only a real change (or a first/removed icon) should trigger the heavy
# refresh; a no-op reinstall with no icon at all must not.
$iconChanged = $false
if ($hasValidIcon) {
if (-not $preIconHash) {
$iconChanged = $true
} else {
try {
$postIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash
$iconChanged = ($postIconHash -ne $preIconHash)
} catch { $iconChanged = $true }
}
} elseif ($preIconHash) {
# A previously present icon was removed or invalidated.
$iconChanged = $true
}
# Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts
# that may point at a deleted workspace; launcher + icon stay.
if ($StudioRedirectMode -eq 'env') {
@ -866,8 +925,22 @@ shell.Run cmd, 0, False
return
}
$wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe"
$shortcutArgs = "//B //Nologo `"$launcherVbs`""
# Whether this is effectively a first install (no pre-existing .lnk).
# Used to gate the heavier icon-cache refresh below so a no-op reinstall
# does not repeatedly clear caches / restart StartMenuExperienceHost --
# a behavioral cluster AV heuristics can score as dropper-like.
$firstInstall = -not (
($desktopLink -and (Test-Path -LiteralPath $desktopLink)) -or
($startMenuLink -and (Test-Path -LiteralPath $startMenuLink))
)
# Launch transport for the shortcuts: powershell.exe runs
# launch-studio.ps1 with a hidden window. We deliberately avoid a
# .vbs/WScript.Shell wrapper -- that script-engine shape is what AV
# VBS-dropper heuristics score (Kaspersky HEUR:Trojan.VBS.Agent.gen).
$powershellForLnk = Join-Path $env:SystemRoot "System32\WindowsPowerShell\v1.0\powershell.exe"
$shortcutTarget = $powershellForLnk
$shortcutArgs = "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$launcherPs1`""
try {
$wshell = New-Object -ComObject WScript.Shell
@ -877,9 +950,11 @@ shell.Run cmd, 0, False
if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue }
try {
$shortcut = $wshell.CreateShortcut($linkPath)
$shortcut.TargetPath = $wscriptExe
$shortcut.TargetPath = $shortcutTarget
$shortcut.Arguments = $shortcutArgs
$shortcut.WorkingDirectory = $appDir
# Start minimized so the brief PowerShell console flash is muted.
$shortcut.WindowStyle = 7
$shortcut.Description = "Launch Unsloth Studio"
if ($hasValidIcon) {
$shortcut.IconLocation = "$iconPath,0"
@ -893,15 +968,13 @@ shell.Run cmd, 0, False
}
if ($createdShortcutCount -gt 0) {
substep "Created Unsloth Studio shortcut"
# Force Explorer to re-read each new shortcut's icon so it renders
# immediately instead of a stale/generic entry (a same-name .lnk
# recreated across reinstalls keeps Explorer's cached per-item icon).
# The reliable, non-disruptive fix (no explorer restart) is a per-item
# SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW per .lnk; the global
# SHCNE_ASSOCCHANGED broadcast alone does NOT recover a stale item.
# Also clear the on-disk icon cache (covers heavier staleness).
try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {}
try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {}
# Always do the cheap, non-disruptive per-item refresh so a
# rewritten same-name .lnk renders with its new target/icon
# immediately (a same-name .lnk recreated across reinstalls keeps
# Explorer's cached per-item icon). The reliable fix (no explorer
# restart) is a per-item SHChangeNotify SHCNE_UPDATEITEM +
# SHCNF_PATHW per .lnk; the global SHCNE_ASSOCCHANGED broadcast
# alone does NOT recover a stale item.
try {
Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);' -ErrorAction SilentlyContinue
# SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut
@ -911,21 +984,31 @@ shell.Run cmd, 0, False
# SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders)
[UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero)
} catch {}
# Win11's Start Menu (StartMenuExperienceHost) keeps its OWN
# pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT
# invalidate, so a rewritten same-name shortcut shows the old tile
# until the host restarts. Drop only the render caches (NEVER
# start2.bin -- the pinned layout) and let the host rebuild.
# Best-effort; Win10 has no such host (Test-Path skips it).
try {
$smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
if (Test-Path -LiteralPath $smehTemp) {
Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch {}
# Heavier on-disk icon-cache clear + StartMenuExperienceHost tile
# rebuild only when the icon actually changed or this is a first
# install. Running "clear icon cache + kill StartMenuExperienceHost"
# on every no-op reinstall is a dropper-like behavioral cluster and
# is unnecessary when the icon is unchanged (the per-item notify
# above already refreshes the rewritten shortcut).
if ($firstInstall -or $iconChanged) {
try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {}
try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {}
# Win11's Start Menu (StartMenuExperienceHost) keeps its OWN
# pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT
# invalidate, so a rewritten same-name shortcut shows the old tile
# until the host restarts. Drop only the render caches (NEVER
# start2.bin -- the pinned layout) and let the host rebuild.
# Best-effort; Win10 has no such host (Test-Path skips it).
try {
$smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
if (Test-Path -LiteralPath $smehTemp) {
Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch {}
}
} else {
substep "no Unsloth Studio shortcuts were created" "Yellow"
}
@ -1193,7 +1276,7 @@ shell.Run cmd, 0, False
# ── Install uv ──
Write-TauriLog "STEP" "Installing uv package manager"
$UvMinVersion = "0.7.22"
$UvMinVersion = "0.8.16"
function Test-UvVersionOk {
$cmd = Get-Command uv -ErrorAction SilentlyContinue
if (-not $cmd) { return $false }
@ -1264,6 +1347,15 @@ shell.Run cmd, 0, False
$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"
}
# uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read
# timeout for large wheel downloads. User-provided values are preserved.
if (-not $env:UV_HTTP_RETRIES) {
$env:UV_HTTP_RETRIES = "5"
}
if (-not $env:UV_HTTP_TIMEOUT) {
$env:UV_HTTP_TIMEOUT = "180"
}
# ── Create venv (migrate old layout if possible, otherwise fresh) ──
# Pass the resolved executable path to uv so it does not re-resolve
# a version string back to a conda interpreter.
@ -1565,7 +1657,9 @@ shell.Run cmd, 0, False
$HipSdkInstalled = $true # binary found → SDK is installed regardless of device state
try {
$hipOut = & $hipinfoExe.Source 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
if ($hipOut -match "(?i)gcnArchName") {
# hipinfo can crash after printing gcnArchName (#6043).
# Once the arch is printed, keep the ROCm wheel path.
$HasROCm = $true
$_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() })
$_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
@ -1575,8 +1669,13 @@ shell.Run cmd, 0, False
} else {
$ROCmGpuLabel = "AMD ROCm"
}
if ($LASTEXITCODE -ne 0) {
Write-Host " [INFO] hipinfo exited with code $LASTEXITCODE but reported gcnArchName -- treating as ROCm-capable (see #6043)" -ForegroundColor Cyan
}
} elseif ($LASTEXITCODE -ne 0) {
# hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected")
# hipinfo ran but returned a HIP runtime error without any gcnArchName
# output (e.g. "no ROCm-capable device detected"), or crashed before
# printing device info.
$firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1)
Write-Host " [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow
Write-Host " $firstLine" -ForegroundColor Yellow
@ -1816,6 +1915,64 @@ shell.Run cmd, 0, False
substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow"
return "$baseUrl/cu126"
}
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
# matching setup.ps1's stale-venv parse.
function ConvertTo-TorchFlavorTag {
param([string]$TorchVersion)
if (-not $TorchVersion) { return $null }
if ($TorchVersion -match '\+(cu\d+)') { return $Matches[1] }
if ($TorchVersion -match '\+rocm') { return 'rocm' }
if ($TorchVersion -match '\+cpu') { return 'cpu' }
return 'cpu'
}
# Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a
# gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops.
function Get-ExpectedTorchFlavorTag {
param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if ($leaf -match '^cu\d+$') { return $leaf }
if ($leaf -eq 'cpu') { return 'cpu' }
if ($leaf -match '^rocm') { return 'rocm' }
if ($leaf -match '^gfx') { return 'rocm' }
return $null
}
# Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses
# ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference.
function Get-InstalledTorchTag {
param([string]$PythonExe)
if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null }
try {
$psi = New-Object System.Diagnostics.ProcessStartInfo
$psi.FileName = $PythonExe
$psi.Arguments = '-c "import torch; print(torch.__version__)"'
$psi.RedirectStandardOutput = $true
$psi.RedirectStandardError = $true
$psi.UseShellExecute = $false
$psi.CreateNoWindow = $true
$proc = [System.Diagnostics.Process]::Start($psi)
# Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd()
# before the wait would block forever if a wedged "import torch" never
# closes stdout; leaving the redirected stderr undrained would deadlock a
# child that floods it past the pipe buffer. Async reads let a noisy-but-
# exiting probe finish, while a truly hung one still hits the 30s timeout
# and is killed -- bounded either way.
$outTask = $proc.StandardOutput.ReadToEndAsync()
$errTask = $proc.StandardError.ReadToEndAsync()
$finished = $proc.WaitForExit(30000)
if (-not $finished) { try { $proc.Kill() } catch {}; return $null }
$torchVer = $outTask.GetAwaiter().GetResult().Trim()
[void]$errTask.GetAwaiter().GetResult()
if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null }
return ConvertTo-TorchFlavorTag $torchVer
} catch { return $null }
}
$TorchIndexUrl = Get-TorchIndexUrl
# ===== Windows-on-ARM + NVIDIA GPU -> automatic WSL2 fallback (N1X "RTX Spark" / DGX Spark-class) =====
@ -2286,21 +2443,21 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
$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.6.8" "unsloth-zoo>=2026.6.6" }
if ($baseInstallExit -eq 0) {
# 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.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
}
if ($baseInstallExit -eq 0) {
$NoTorchReq = Find-NoTorchRuntimeFile
if ($NoTorchReq) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
}
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2314,7 +2471,7 @@ shell.Run cmd, 0, False
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
}
substep "overlaying unsloth-zoo from git main..."
$zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
$zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
if ($zooOverlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
@ -2327,7 +2484,7 @@ shell.Run cmd, 0, False
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch from $ROCmIndexUrl..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2335,7 +2492,7 @@ shell.Run cmd, 0, False
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2347,21 +2504,21 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
}
if ($baseInstallExit -eq 0) {
$NoTorchReq = Find-NoTorchRuntimeFile
if ($NoTorchReq) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2376,7 +2533,7 @@ shell.Run cmd, 0, False
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
}
substep "overlaying unsloth-zoo from git main..."
$zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
$zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
if ($zooOverlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
@ -2387,7 +2544,7 @@ shell.Run cmd, 0, False
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --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)
@ -2399,13 +2556,13 @@ shell.Run cmd, 0, False
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
}
substep "overlaying unsloth-zoo from git main..."
$zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
$zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
if ($zooOverlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" }
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)
@ -2413,6 +2570,50 @@ shell.Run cmd, 0, False
}
}
# ── 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
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
if (-not $SkipTorch) {
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') {
$installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) {
# AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path
# would have force-reinstalled. Repair from the same repo.amd.com index.
$rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec torchvision torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
}
$installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
}
$installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
}
}
# Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly.
if ($installedTorchTag -eq 'cpu') {
Write-Host ""
Write-Host " [WARN] PyTorch is CPU-only but a $expectedTorchTag GPU build was expected for this machine." -ForegroundColor Yellow
Write-Host " [WARN] Training and GPU inference will run on CPU until this is fixed." -ForegroundColor Yellow
Write-Host " [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." -ForegroundColor Yellow
}
}
}
# Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped
# for --local: the editable install above already makes _PACKAGE_ROOT in
# unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__).

231
install.sh
View file

@ -163,6 +163,40 @@ run_install_cmd() {
return $_rc
}
# Retry run_install_cmd on transient uv download failures with backoff. Returns
# the last exit code on permanent failure so the set -e rollback trap still fires.
: "${UNSLOTH_INSTALL_RETRIES:=3}"
: "${UNSLOTH_INSTALL_RETRY_DELAY:=3}"
run_install_cmd_retry() {
_ricr_label="$1"
# Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables).
# Length guard precedes the numeric test so a huge value can't overflow `[ -ge ]`.
# 0?* rejects leading-zero delays ("08"/"09" break the later $((delay*2)) as octal);
# bare "0" stays valid. Bounds: 1..100 retries, 0..3600s base delay.
case "$UNSLOTH_INSTALL_RETRIES" in
''|*[!0-9]*|0) _ricr_max=3 ;;
*) if [ "${#UNSLOTH_INSTALL_RETRIES}" -le 3 ] && [ "$UNSLOTH_INSTALL_RETRIES" -ge 1 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRIES" -le 100 ] 2>/dev/null; then _ricr_max=$UNSLOTH_INSTALL_RETRIES; else _ricr_max=3; fi ;;
esac
case "$UNSLOTH_INSTALL_RETRY_DELAY" in
''|*[!0-9]*|0?*) _ricr_delay=3 ;;
*) if [ "${#UNSLOTH_INSTALL_RETRY_DELAY}" -le 4 ] && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -ge 0 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -le 3600 ] 2>/dev/null; then _ricr_delay=$UNSLOTH_INSTALL_RETRY_DELAY; else _ricr_delay=3; fi ;;
esac
_ricr_attempt=1
while :; do
# AND-OR (not `if`) preserves the real failure code: $? after a non-taken
# `if` is 0 in sh/dash/bash, which would break the rollback path.
run_install_cmd "$@" && return 0
_ricr_rc=$?
if [ "$_ricr_attempt" -ge "$_ricr_max" ]; then
return "$_ricr_rc"
fi
substep "retrying \"$_ricr_label\" after transient failure (attempt $((_ricr_attempt + 1))/$_ricr_max, waiting ${_ricr_delay}s)..." "$C_WARN"
sleep "$_ricr_delay" || true
_ricr_attempt=$((_ricr_attempt + 1))
_ricr_delay=$((_ricr_delay * 2))
done
}
# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
@ -1229,6 +1263,10 @@ if (-not \$targetExe) { exit 1 }
# AppData\Local (renders blank); a profile-path icon renders everywhere.
\$iconDir = Join-Path \$env:USERPROFILE '.unsloth'
\$iconPath = Join-Path \$iconDir 'unsloth.ico'
\$preIconHash = \$null
if (Test-Path -LiteralPath \$iconPath) {
try { \$preIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash } catch {}
}
if (-not (Test-Path -LiteralPath \$iconPath)) {
try {
New-Item -ItemType Directory -Force -Path \$iconDir | Out-Null
@ -1244,9 +1282,11 @@ if (Test-Path -LiteralPath \$iconPath) {
(Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs')
)
\$created = @()
\$firstShortcut = \$false
foreach (\$dir in \$locations) {
if (-not \$dir -or -not (Test-Path \$dir)) { continue }
\$linkPath = Join-Path \$dir '$_css_lnk_name_ps'
if (-not (Test-Path -LiteralPath \$linkPath)) { \$firstShortcut = \$true }
\$shortcut = \$WshShell.CreateShortcut(\$linkPath)
\$shortcut.TargetPath = \$targetExe
\$shortcut.Arguments = '$_css_sc_args_ps'
@ -1255,27 +1295,43 @@ foreach (\$dir in \$locations) {
\$shortcut.Save()
\$created += \$linkPath
}
# Force Explorer to re-read EACH new shortcut's icon so it renders immediately
# instead of a stale/blank (generic) icon. The reliable, NON-disruptive fix
# (no explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM,
# SHCNF_PATHW, <lnk>) -- the global SHCNE_ASSOCCHANGED alone does not recover a
# stale item. Also clear the on-disk icon cache for heavier staleness.
try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {}
try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {}
\$iconChanged = \$false
if (\$hasIcon) {
if (-not \$preIconHash) {
\$iconChanged = \$true
} else {
try {
\$postIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash
\$iconChanged = (\$postIconHash -ne \$preIconHash)
} catch { \$iconChanged = \$true }
}
} elseif (\$preIconHash) {
\$iconChanged = \$true
}
# Per-item refresh always (cheap, non-disruptive) so the rewritten .lnk renders
# immediately instead of a stale/blank (generic) icon. The reliable fix (no
# explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, SHCNF_PATHW,
# <lnk>) -- the global SHCNE_ASSOCCHANGED alone does not recover a stale item.
try {
Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int e, uint f, string a, System.IntPtr b);' -ErrorAction SilentlyContinue
foreach (\$p in \$created) { try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, \$p, [System.IntPtr]::Zero) } catch {} }
[UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, \$null, [System.IntPtr]::Zero)
} catch {}
# Win11 Start Menu keeps its own tile-icon cache (preserve start2.bin).
try {
\$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState'
if (Test-Path -LiteralPath \$smeh) {
Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch {}
# Heavier on-disk icon-cache clear + StartMenuExperienceHost tile rebuild
# (preserve start2.bin) only on first install or a real icon change, so a no-op
# WSL reinstall does not run a dropper-like clear-cache + kill cluster each time.
if (\$created.Count -gt 0 -and (\$firstShortcut -or \$iconChanged)) {
try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {}
try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {}
try {
\$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState'
if (Test-Path -LiteralPath \$smeh) {
Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch {}
}
WSLPS1_EOF
# Convert WSL path to Windows path for powershell.exe
@ -1458,12 +1514,19 @@ fi
# ── Install uv ──
tauri_log "STEP" "Installing uv package manager"
UV_MIN_VERSION="0.7.22"
UV_MIN_VERSION="0.8.16"
# When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables).
: "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"
export UV_COMPILE_BYTECODE_TIMEOUT
# uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read
# timeout for large wheel downloads. ":=" preserves any user override.
: "${UV_HTTP_RETRIES:=5}"
export UV_HTTP_RETRIES
: "${UV_HTTP_TIMEOUT:=180}"
export UV_HTTP_TIMEOUT
version_ge() {
# returns 0 if $1 >= $2
_a=$1
@ -1929,6 +1992,45 @@ get_torch_index_url() {
else echo "$_base/cpu"; fi
}
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ ($1) -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu.
_torch_flavor_tag() {
case "$1" in
*+cu[0-9]*) printf '%s\n' "$1" | sed -n 's/.*+\(cu[0-9][0-9]*\).*/\1/p' ;;
*+rocm*) echo "rocm" ;;
*+cpu*) echo "cpu" ;;
"") echo "" ;;
*) echo "cpu" ;;
esac
}
# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* ->
# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops.
_expected_torch_flavor_tag() {
_u="${1%/}"
_leaf="${_u##*/}"
case "$_leaf" in
cu[0-9]*) echo "$_leaf" ;;
cpu) echo "cpu" ;;
rocm*|gfx*) echo "rocm" ;;
*) echo "" ;;
esac
}
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
_u="${1%/}"
_leaf="${_u##*/}"
case "$_leaf" in
cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;;
*) echo "no" ;;
esac
}
get_radeon_wheel_url() {
# Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing
# contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/,
@ -2422,28 +2524,28 @@ if [ "$_MIGRATED" = true ]; then
# PyPI metadata still declares torch as a hard dep), then install
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps
# to prevent transitive torch resolution.
run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
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.6.7" unsloth-zoo
"unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
# 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.
run_install_cmd "install pydantic (with deps for compatible core)" \
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
_NO_TORCH_RT="$(_find_no_torch_runtime)"
if [ -n "$_NO_TORCH_RT" ]; then
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
fi
else
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.7" unsloth-zoo
"unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
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..."
run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
--no-deps --reinstall-package unsloth-zoo \
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
fi
@ -2452,13 +2554,13 @@ if [ "$_MIGRATED" = true ]; then
# fresh reinstall.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--force-reinstall
@ -2584,7 +2686,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
else
@ -2596,30 +2698,30 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# filelock / sympy / networkx which are not in the
# Radeon listing.
if [ -n "$_tri_whl" ]; then
run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install triton + PyTorch" uv pip install --python "$_VENV_PY" \
--find-links "$_RADEON_BASE_URL" \
"$_tri_whl" "$_torch_whl" "$_tv_whl" "$_ta_whl"
else
run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
--find-links "$_RADEON_BASE_URL" \
"$_torch_whl" "$_tv_whl" "$_ta_whl"
fi
fi
else
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
@ -2628,7 +2730,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# which is only useful once torch is present for training.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
;;
esac
@ -2639,31 +2741,31 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
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.6.7" unsloth-zoo
"unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd "install pydantic (with deps for compatible core)" \
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
_NO_TORCH_RT="$(_find_no_torch_runtime)"
if [ -n "$_NO_TORCH_RT" ]; then
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
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..."
run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
--no-deps --reinstall-package unsloth-zoo \
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
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..."
run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
--no-deps --reinstall-package unsloth-zoo \
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
elif [ -n "${UNSLOTH_INSTALL_REF:-}" ] && [ "${UNSLOTH_INSTALL_REF}" != "main" ] && [ "$PACKAGE_NAME" = "unsloth" ]; then
@ -2675,7 +2777,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth @ git+https://github.com/unslothai/unsloth@${UNSLOTH_INSTALL_REF}" unsloth-zoo
else
run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth -- "$PACKAGE_NAME"
fi
# aarch64 + NVIDIA (DGX Spark / GB10 / N1X): unsloth's x86_64-oriented cuXXX
@ -2696,11 +2798,11 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
*/rocm*|*/gfx*)
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--force-reinstall
@ -2713,15 +2815,50 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --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..."
run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
--no-deps --reinstall-package unsloth-zoo \
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
else
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME"
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME"
fi
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
# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it
# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op.
if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
_expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")
# Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip.
if [ -n "$_expected_torch_tag" ] && [ "$_expected_torch_tag" != "cpu" ]; then
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
fi
# Safety net (incl. AMD/WSL): GPU build expected but still CPU -> warn loudly.
if [ "$_installed_torch_tag" = "cpu" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi

6
pyproject.toml
View file

@ -72,7 +72,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.6.5",
"unsloth_zoo>=2026.6.6",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -93,7 +93,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.6.5",
"unsloth_zoo>=2026.6.6",
"torchvision",
"unsloth[triton]",
]
@ -583,7 +583,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.6.5",
"unsloth_zoo>=2026.6.6",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

398
scripts/scan_npm_packages.py
View file

@ -412,7 +412,7 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = {
"@uipath/functions-tool": {"1.0.1"},
"@uipath/access-policy-sdk": {"0.3.1"},
"@uipath/platform-tool": {"1.0.1"},
# Mini Shai-Hulud May-12 wave: @mistralai/* (npm) — separate from PyPI mistralai
# Mini Shai-Hulud May-12 wave: @mistralai/* (npm), separate from PyPI mistralai
# (https://www.aikido.dev/blog/mini-shai-hulud-is-back-tanstack-compromised).
"@mistralai/mistralai": {"2.2.2", "2.2.3", "2.2.4"},
"@mistralai/mistralai-gcp": {"1.7.1", "1.7.2", "1.7.3"},
@ -916,6 +916,204 @@ def _evidence(
LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare")
# ─────────────────────────────────────────────────────────────────────
# Code-only scanning for JS/TS sources. Blank `//` and `/* */` comments
# before matching (the top FP source: scary strings in JSDoc/changelog
# comments), tracking string/template/regex context so a `//` inside
# "http://..." is not mistaken for a comment. Strings are NOT blanked
# (droppers hide payloads there). Fail open on lexer confusion: the raw
# text is still scanned. JS sibling of scan_packages.py::_strip_noncode.
# ─────────────────────────────────────────────────────────────────────
_JS_FAMILY_SUFFIXES = (".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx")
# Keywords after which a `/` begins a regex literal (not division).
_REGEX_PRECEDING_KEYWORDS = frozenset(
{
"return",
"typeof",
"instanceof",
"in",
"of",
"new",
"delete",
"void",
"throw",
"yield",
"await",
"do",
"else",
"case",
}
)
_IDENT_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$")
def _slash_is_regex(prev_tok: str) -> bool:
"""Disambiguate a lone ``/``: regex literal vs division operator.
Biased toward regex when ambiguous -- regex state never blanks, so a
wrong guess only costs FP reduction (or a fail-open), never a missed
detection.
"""
if prev_tok == "":
return True # start of file -> expression position
if prev_tok in _REGEX_PRECEDING_KEYWORDS:
return True
last = prev_tok[-1]
if last.isalnum() or last in "_$)]":
return False # previous token ends a value -> division
return True # operators, punctuation, `{`, `}` -> regex (safe bias)
def _strip_js_noncode(text: str) -> str:
"""Blank JS/TS comments, preserving byte geometry. Fail-open on confusion."""
if "//" not in text and "/*" not in text:
return text # nothing to strip
n = len(text)
out = list(text)
nl = ("\n", "\r")
def _blank(a: int, b: int) -> None:
for k in range(a, b):
if out[k] not in nl:
out[k] = " "
state = "code"
prev_tok = ""
tmpl_stack: list[str] = []
i = 0
try:
while i < n:
c = text[i]
nxt = text[i + 1] if i + 1 < n else ""
if state == "code":
if c == "/" and nxt == "/":
start = i
i += 2
while i < n and text[i] not in nl:
i += 1
_blank(start, i)
continue
if c == "/" and nxt == "*":
start = i
i += 2
closed = False
while i < n:
if text[i] == "*" and i + 1 < n and text[i + 1] == "/":
i += 2
closed = True
break
i += 1
if not closed:
return text # unterminated block comment
_blank(start, i)
continue
if c == "'":
state = "sq"
i += 1
continue
if c == '"':
state = "dq"
i += 1
continue
if c == "`":
state = "tmpl"
i += 1
continue
if c == "/":
if _slash_is_regex(prev_tok):
state = "regex"
i += 1
continue
prev_tok = "/"
i += 1
continue
if c.isspace():
i += 1
continue
if c in _IDENT_CHARS:
j = i
while j < n and text[j] in _IDENT_CHARS:
j += 1
prev_tok = text[i:j]
i = j
continue
if c == "}" and tmpl_stack:
state = tmpl_stack.pop()
i += 1
continue
prev_tok = c
i += 1
continue
elif state in ("sq", "dq"):
q = "'" if state == "sq" else '"'
if c == "\\":
i += 2
continue
if c == q:
state = "code"
prev_tok = "_v"
i += 1
continue
if c in nl:
return text # unterminated string literal
i += 1
continue
elif state == "tmpl":
if c == "\\":
i += 2
continue
if c == "`":
state = "code"
prev_tok = "_v"
i += 1
continue
if c == "$" and nxt == "{":
tmpl_stack.append("tmpl")
state = "code"
prev_tok = "{"
i += 2
continue
i += 1
continue
elif state == "regex":
if c == "\\":
i += 2
continue
if c == "[":
state = "regex_cc"
i += 1
continue
if c == "/":
state = "code"
prev_tok = "_v"
i += 1
continue
if c in nl:
return text # unterminated regex literal
i += 1
continue
elif state == "regex_cc":
if c == "\\":
i += 2
continue
if c == "]":
state = "regex"
i += 1
continue
if c in nl:
return text
i += 1
continue
else:
return text
if state != "code" or tmpl_stack:
return text # unterminated construct -> fail open
except Exception:
return text
return "".join(out)
def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
try:
@ -1042,6 +1240,14 @@ def _host_in_outbound_context(text: str, host: str) -> bool:
def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
# Code-only scanning for JS/TS sources: blank comments before matching so
# an IOC host / `eval(atob)` example / campaign marker quoted in a comment
# cannot manufacture a false positive. Assigned string literals (where real
# droppers hide base64 payloads) are preserved. Non-JS text (json/yaml/sh/
# py/html) is scanned as-is -- this lexer only understands JS comments.
if rel.lower().endswith(_JS_FAMILY_SUFFIXES):
text = _strip_js_noncode(text)
# IOC substrings (literal, case-sensitive).
for needle, (sev, why) in KNOWN_IOC_STRINGS.items():
if needle in text:
@ -1236,6 +1442,131 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N
pass
# ─────────────────────────────────────────────────────────────────────
# Baseline allowlist: triaged known-good HIGH/CRITICAL findings so the gate
# can enforce without red-failing on rare legitimate-library behavior.
# Matched on ``(normalized package, package-relative path, pattern)`` -- not
# evidence text -- so a version bump does not reopen a finding, but a *new*
# kind of finding in a listed file is a different pattern and still fails.
# Mirrors scan_packages.py. Regenerate with ``--write-baseline``.
# ─────────────────────────────────────────────────────────────────────
_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json")
# Bumped when the entry-key semantics change. v2 keys on the package-relative
# path; v1 stored only a basename, so a v1 entry could suppress a same-named file
# in a different directory. A pre-v2 baseline with entries is ignored (fail
# closed) rather than mis-applied.
_BASELINE_SCHEMA_VERSION = 2
def _norm_pkg_name(display: str) -> str:
"""``@scope/pkg@1.2.3`` / ``pkg@1.2.3`` -> name without the version.
The version is the LAST ``@``-separated field; a leading ``@`` (scope)
is preserved. Lower-cased (npm names are case-insensitive). Sentinels
like ``<root>`` / ``<lockfile>`` pass through unchanged.
"""
s = (display or "").strip()
at = s.rfind("@")
if at > 0: # >0 so a leading @scope is not treated as the version sep
s = s[:at]
return s.lower()
_NPM_TARBALL_ROOT = "package/"
def _relpath_in_package(filename: str) -> str:
"""Path within the published package, stable across version bumps. npm
tarballs root every file at ``package/``; strip it so the key is the real
source path (``dist/index.js``) and a new file with the same basename in a
different directory is not silently suppressed."""
f = (filename or "").replace("\\", "/")
return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f
def _finding_key(f: Finding) -> tuple[str, str, str]:
"""Stable allowlist key: normalized package, package-relative path, pattern."""
return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern)
def _load_baseline(path: str) -> set[tuple[str, str, str]]:
"""Load an allowlist JSON into a set of match keys. Missing file -> empty."""
try:
with open(path, "r", encoding = "utf-8") as fh:
data = json.load(fh)
except FileNotFoundError:
return set()
except (OSError, json.JSONDecodeError) as exc:
print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
return set()
entries = data.get("entries", [])
if entries and data.get("version") != _BASELINE_SCHEMA_VERSION:
print(
f" [WARN] baseline schema v{data.get('version')} predates package-relative "
f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.",
file = sys.stderr,
)
return set()
keys: set[tuple[str, str, str]] = set()
for e in entries:
try:
keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"]))
except (KeyError, TypeError):
continue
return keys
def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int:
"""Persist at-or-above-threshold findings as an allowlist for triage."""
entries = []
seen: set[tuple[str, str, str]] = set()
for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)):
if _SEVERITY_RANK[f.severity] > threshold_rank:
continue
key = _finding_key(f)
if key in seen:
continue
seen.add(key)
entries.append(
{
"package": _norm_pkg_name(f.package),
"file": _relpath_in_package(f.filename),
"pattern": f.pattern,
"severity": f.severity,
"evidence": (f.evidence or f.detail)[:240],
}
)
doc = {
"_comment": (
"scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL "
"finding manually judged benign. Matched on (package, "
"package-relative path, pattern); evidence/severity are for review "
"only. Regenerate with --write-baseline AFTER reviewing every line."
),
"version": _BASELINE_SCHEMA_VERSION,
"entries": entries,
}
with open(path, "w", encoding = "utf-8") as fh:
json.dump(doc, fh, indent = 2, sort_keys = False)
fh.write("\n")
print(f" Wrote {len(entries)} baseline entr(y/ies) to {path}")
return len(entries)
def _partition_baseline(
findings: list[Finding], baseline: set[tuple[str, str, str]]
) -> tuple[list[Finding], list[Finding]]:
"""Split findings into (active, suppressed) by allowlist membership."""
if not baseline:
return list(findings), []
active, suppressed = [], []
for f in findings:
(suppressed if _finding_key(f) in baseline else active).append(f)
return active, suppressed
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(
description = "Pre-install npm tarball content scanner.",
@ -1263,6 +1594,30 @@ def main(argv: list[str] | None = None) -> int:
"Medium and below print but exit 0."
),
)
parser.add_argument(
"--baseline",
metavar = "FILE",
default = None,
help = (
"Allowlist JSON of triaged known-good findings to suppress. "
"Defaults to scan_npm_packages_baseline.json next to this script "
"if present."
),
)
parser.add_argument(
"--no-baseline",
action = "store_true",
help = "Ignore the auto-discovered baseline allowlist.",
)
parser.add_argument(
"--write-baseline",
metavar = "FILE",
default = None,
help = (
"Write the current at/above-threshold findings to FILE as an "
"allowlist, then exit 0. Review every entry before committing it."
),
)
args = parser.parse_args(argv)
lockfile = Path(args.lockfile).resolve()
@ -1341,7 +1696,46 @@ def main(argv: list[str] | None = None) -> int:
"critical": CRITICAL,
}[args.fail_on]
threshold_rank = _SEVERITY_RANK[threshold]
blocking = [f for f in all_findings if _SEVERITY_RANK[f.severity] <= threshold_rank]
# --write-baseline: persist the full current at/above-threshold set as the
# new allowlist (ignoring any loaded baseline), then exit 0. A hard error
# means the scan was incomplete, so warn -- a baseline baked from a partial
# run would silently allow whatever failed to download.
if args.write_baseline:
if hard_errors:
print(
f" [WARN] {len(hard_errors)} hard error(s): baseline may be "
"incomplete (some packages did not scan).",
file = sys.stderr,
)
_write_baseline(args.write_baseline, all_findings, threshold_rank)
return 0
# Baseline allowlist: suppress triaged, known-good findings so the CI gate
# can be enforcing without red-failing on legitimate-library noise.
if args.no_baseline:
baseline_path = None
elif args.baseline:
baseline_path = args.baseline
elif os.path.isfile(_DEFAULT_BASELINE_PATH):
baseline_path = _DEFAULT_BASELINE_PATH
else:
baseline_path = None
baseline = _load_baseline(baseline_path) if baseline_path else set()
active, suppressed = _partition_baseline(all_findings, baseline)
if suppressed:
crit_s = sum(1 for f in suppressed if f.severity == CRITICAL)
high_s = sum(1 for f in suppressed if f.severity == HIGH)
print(
f"\n[scan-npm] {len(suppressed)} finding(s) suppressed by baseline "
f"{baseline_path} ({crit_s} CRITICAL, {high_s} HIGH).",
flush = True,
)
# Exit code: 1 on a hard error, or a NON-baselined finding at/above the
# threshold. This is the signal CI gates on once the baseline is clean.
blocking = [f for f in active if _SEVERITY_RANK[f.severity] <= threshold_rank]
if hard_errors or blocking:
if blocking:
print(

5
scripts/scan_npm_packages_baseline.json Normal file
View file

@ -0,0 +1,5 @@
{
"_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.",
"version": 2,
"entries": []
}

871
scripts/scan_packages.py
View file

@ -33,10 +33,24 @@ Examples:
python scan_packages.py --fix -r requirements.txt
python scan_packages.py --fix --max-search 20 -r requirements.txt
# Triage to a baseline once, then gate on anything NEW
python scan_packages.py -r requirements.txt --write-baseline scripts/scan_packages_baseline.json
python scan_packages.py -r requirements.txt # auto-loads the baseline, exits 0 if only baselined findings remain
False positives:
.py files are scanned code-only: comments and bare docstrings/doctests are
blanked before pattern matching (line numbers preserved), so prose, usage
examples and `>>>` doctests cannot trip a finding. Residual findings that
are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored
test fixture) are suppressed via a reviewed baseline allowlist, matched on
(package, basename(file), check). A NEW kind of finding in an already-listed
file is a different check and still fails. This mirrors the Hugging Face Hub
approach (ClamAV/picklescan: low-FP, signature/structural, surface status).
Exit codes:
0 -- no CRITICAL or HIGH findings
1 -- CRITICAL or HIGH findings detected
2 -- no packages specified
0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline)
1 -- non-baselined CRITICAL or HIGH findings detected
2 -- no packages specified, or scan incomplete (pip download failure)
"""
import argparse
@ -50,6 +64,8 @@ import subprocess
import sys
import tarfile
import tempfile
import tokenize
import urllib.parse
import urllib.request
import zipfile
from dataclasses import dataclass, field
@ -213,17 +229,26 @@ RE_ARCHIVE_STAGING = re.compile(
)
# Anti-analysis / sandbox evasion / debugger detection
# NB: deliberately does NOT include a bare ``platform.system() ... Linux/Windows
# /Darwin`` branch. Under re.DOTALL that matched across the whole file -- any
# cross-platform library (typer, packaging, pandas, pymupdf, ...) trips it -- so
# it had ~zero precision and only generated false positives. OS detection alone
# is not an anti-analysis signal; the debugger/VM/long-sleep signals below are.
RE_ANTI_ANALYSIS = re.compile(
r"\bptrace\b"
r"|\bsys\s*\.\s*gettrace\s*\("
r"|\bsys\s*\.\s*settrace\b"
r"|\bTracerPid\b"
r"|\b/proc/self/status\b"
# /proc/self/status is read to scrape TracerPid for anti-debug. A leading
# \b here is unsatisfiable (\b never holds between a non-word boundary and
# "/"), so the old pattern was dead; a lookbehind that only forbids a
# preceding word char or path separator lets `open("/proc/self/status")`
# and `cat /proc/self/status` match while avoiding mid-path partials.
r"|(?<![\w/])/proc/self/status\b"
r"|\bIsDebuggerPresent\b"
r"|\bvirtualbox\b.*\bhardware\b"
r"|\bvmware\b.*\bdetect\b"
r"|\btime\.sleep\s*\(\s*(?:[3-9]\d{2,}|[1-9]\d{3,})\s*\)" # long sleep (anti-sandbox)
r"|\bplatform\.\s*system\b.*\bif\b.*\b(?:Linux|Windows|Darwin)\b",
r"|\btime\.sleep\s*\(\s*(?:[3-9]\d{2,}|[1-9]\d{3,})\s*\)", # long sleep (anti-sandbox)
re.IGNORECASE | re.DOTALL,
)
@ -493,9 +518,159 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
return findings
# A STRING after one of these tokens (and before a NEWLINE) is a bare
# docstring/doctest/prose statement -- the dominant FP source -- so we blank it.
# A string after `=` or `(` is real code and is never blanked.
_LINE_START_TOKENS = frozenset({tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT})
def _is_fstring(tok_string: str) -> bool:
"""True if a STRING token is an f-string (3.10/3.11 emit one STRING token).
A bare f-string statement evaluates its expressions at import, so unlike an
inert docstring it must never be blanked.
"""
q = min((tok_string.find(c) for c in "'\"" if c in tok_string), default = -1)
return q > 0 and "f" in tok_string[:q].lower()
def _strip_noncode(content: str, blank_comments: bool = True) -> str:
"""Blank comments and bare docstrings so IOC patterns see code only.
Removed regions become spaces (newlines kept) so line numbers stay exact for
_extract_evidence. Fails open on tokenizer errors (the raw text is still
fully scanned, so a real detection is never lost). ``blank_comments=False``
keeps comments (only strings/docstrings blanked) to isolate the span that
exec() could actually run.
"""
try:
toks = list(tokenize.generate_tokens(io.StringIO(content).readline))
except (tokenize.TokenError, IndentationError, SyntaxError, ValueError):
return content
spans: list[tuple[int, int, int, int]] = [] # (srow, scol, erow, ecol)
prev_significant = tokenize.NEWLINE # start-of-file behaves like a new line
n = len(toks)
for i, tok in enumerate(toks):
ttype = tok.type
if ttype == tokenize.COMMENT:
if blank_comments:
spans.append((*tok.start, *tok.end))
continue # transparent; never advances prev_significant
if (
ttype == tokenize.STRING
and prev_significant in _LINE_START_TOKENS
and not _is_fstring(tok.string) # f-strings execute; never blank them
):
# Bare string only if it is the whole statement: next significant
# token must close the logical line.
j = i + 1
while j < n and toks[j].type in (tokenize.COMMENT, tokenize.NL):
j += 1
if j < n and toks[j].type == tokenize.NEWLINE:
spans.append((*tok.start, *tok.end))
prev_significant = ttype
continue
if ttype in (
tokenize.NL,
tokenize.NEWLINE,
tokenize.INDENT,
tokenize.DEDENT,
tokenize.ENCODING,
):
prev_significant = ttype
continue
prev_significant = ttype
if not spans:
return content
buf = content.splitlines(keepends = True)
for srow, scol, erow, ecol in spans:
for row in range(srow, erow + 1):
line = buf[row - 1]
if line.endswith("\n"):
body, nl = line[:-1], "\n"
elif line.endswith("\r"):
body, nl = line[:-1], "\r"
else:
body, nl = line, ""
start = scol if row == srow else 0
end = ecol if row == erow else len(body)
end = min(end, len(body))
if start < end:
body = body[:start] + (" " * (end - start)) + body[end:]
buf[row - 1] = body + nl
return "".join(buf)
# Payload carriers that are suspicious when hidden in a blanked region (a
# docstring/string) of a file that can dynamically execute strings.
_HIDDEN_PAYLOAD_PATTERNS = (
(RE_LARGE_BLOB, "large base64 blob"),
(RE_EMBEDDED_KEYS, "embedded key material"),
(RE_MAY12_IOC, "Shai-Hulud IOC string"),
(RE_OBFUSCATION, "marshal/compile/obfuscation"),
)
def _hidden_payload_findings(
original: str, stripped: str, filename: str, package: str
) -> list[Finding]:
"""Flag payloads that live only in the blanked (docstring/string) region of
a file that contains exec/eval. Such a string is invisible to code-only
scanning yet ``exec(__doc__)`` / ``exec(<str>)`` could still run it."""
if not RE_EXEC_EVAL.search(stripped):
return []
# Only docstrings/strings run via exec(__doc__)/exec(<str>); comments cannot.
# Isolate that span: keep comments as real code, take what string-blanking
# removed (length-preserved, so offsets stay exact for _extract_evidence).
code = _strip_noncode(original, blank_comments = False)
removed = "".join(o if o != s else " " for o, s in zip(original, code))
out = []
def _hidden(pat):
# Carrier present in a blanked region but NOT in real code. A carrier in
# real code is already caught by the normal check, so restricting to
# blanked-only avoids re-flagging legitimate in-code constants.
return bool(pat.search(removed)) and not pat.search(stripped)
for pat, label in _HIDDEN_PAYLOAD_PATTERNS:
if _hidden(pat):
out.append(
Finding(
HIGH,
package,
filename,
"exec/eval with payload hidden in a docstring/string",
f"{label}: {_extract_evidence(removed, pat)}",
)
)
# Fetch-then-run dropper: a network call AND an os/subprocess exec that both
# live in the blanked region. Search the removed span directly (not "absent
# from real code") so a benign visible network/subprocess call cannot mask
# the docstring payload.
if RE_NETWORK.search(removed) and RE_SUBPROCESS.search(removed):
out.append(
Finding(
HIGH,
package,
filename,
"exec/eval with hidden network+exec payload",
f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}",
)
)
return out
def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
"""Run all .py-specific checks."""
findings = []
# Code-only scanning: strip comments/docstrings up front so prose, doctests
# and usage examples cannot manufacture false positives. Aligns with the
# Hugging Face Hub model (ClamAV/picklescan: low-FP, signature/structural).
original = content
content = _strip_noncode(content)
findings = _hidden_payload_findings(original, content, filename, package)
basename = os.path.basename(filename)
is_setup = basename in ("setup.py", "setup.cfg")
is_init = basename == "__init__.py"
@ -937,7 +1112,13 @@ def _extract_evidence(
pattern: re.Pattern,
max_matches: int = 3,
) -> str:
"""Pull matching lines as evidence snippets."""
"""Pull matching lines as evidence snippets.
Falls back to a whole-content search when the pattern only matches across
line boundaries (several IOC regexes use ``re.DOTALL``). Without this an
anti-analysis / archive-staging finding could report empty evidence, making
the baseline entry impossible to review.
"""
lines = content.splitlines()
matches = []
for i, line in enumerate(lines, 1):
@ -948,7 +1129,17 @@ def _extract_evidence(
matches.append(f"L{i}: {snippet}")
if len(matches) >= max_matches:
break
return " | ".join(matches) if matches else ""
if matches:
return " | ".join(matches)
# Multiline (DOTALL) match: report the line where the match begins.
m = pattern.search(content)
if m:
line_no = content.count("\n", 0, m.start()) + 1
snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else ""
if len(snippet) > 160:
snippet = snippet[:160] + "..."
return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: <multiline match>"
return ""
# Non-Python checkers
@ -1390,6 +1581,394 @@ _PIP_DOWNLOAD_PIN_FLAGS = [
_RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]")
# sdist fallback. `--only-binary :all:` never builds an sdist (no setup.py
# exec), but a wheel-less project then can't be fetched at all and one such
# package fails the whole --with-deps resolve (exit 2) -- a coverage hole. So on
# resolve failure we drop to per-spec and fetch any sdist-only package's raw
# tarball from the PyPI JSON API for scan_archive() to read statically: no pip,
# no build, same no-exec guarantee. Transport failures are still exit 2; only
# "no wheel" is downgraded to a direct fetch.
# How many levels of indirect-dep recovery to chase (a wheel dep whose own child
# is sdist-only, and so on). Bounded with dedup so recovery always terminates.
_MAX_DEP_FOLLOWUP_DEPTH = 2
_SDIST_DOWNLOAD_TIMEOUT = 180
# Never fetch an archive larger than we would be willing to scan (iter_archive_files cap).
_MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES
# Direct sdist bytes only ever come from PyPI's own CDN; refuse anything else.
_TRUSTED_PYPI_HOSTS = frozenset({"files.pythonhosted.org", "pypi.org", "pypi.python.org"})
def _spec_pin_version(spec: str) -> str | None:
"""Return the ``==X.Y.Z`` pin from a spec, or None if unpinned."""
m = _RE_PYPI_SPEC_VERSION.search(spec)
return m.group(1) if m else None
def _pypi_json(name: str, version: str | None = None) -> dict | None:
"""Fetch PyPI metadata JSON (read-only HTTPS GET, no exec); None on error.
With ``version`` it fetches that release's document, whose ``requires_dist``
is accurate for the pin (the project-level doc describes only the latest)."""
url = "https://pypi.org/pypi/" + urllib.parse.quote(name, safe = "")
if version:
url += "/" + urllib.parse.quote(version, safe = "")
url += "/json"
try:
req = urllib.request.Request(url, headers = {"Accept": "application/json"})
with urllib.request.urlopen(req, timeout = 30) as resp:
if getattr(resp, "status", 200) != 200:
return None
data = resp.read(16 * 1024 * 1024) # metadata is small; cap regardless
return json.loads(data.decode("utf-8", errors = "replace"))
except Exception:
return None
def _release_files(meta: dict, version: str | None) -> list[dict]:
"""Files for a pinned version, else the latest release's. A pin that is
absent or empty returns [] (never the latest) so a yanked/bad pin fails
closed instead of a different artifact being scanned in its place."""
if version is not None:
return meta.get("releases", {}).get(version) or []
return meta.get("urls", []) or []
def _release_has_wheel(meta: dict, version: str | None) -> bool:
"""True if the (pinned or latest) release publishes any bdist_wheel."""
return any(f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version))
def _is_trusted_pypi_url(url: str) -> bool:
"""Only download sdist bytes from PyPI's own hosts, over HTTPS."""
try:
parsed = urllib.parse.urlparse(url)
except Exception:
return False
return parsed.scheme == "https" and parsed.hostname in _TRUSTED_PYPI_HOSTS
_MARKER_ENV_VARS = (
"sys_platform",
"platform_system",
"platform_machine",
"platform_release",
"platform_version",
"platform_python_implementation",
"os_name",
"python_version",
"python_full_version",
"implementation_name",
"implementation_version",
)
def _marker_holds_by_default(marker: str) -> bool:
"""Keep (scan) a dep unless its marker is purely ``extra``-gated. The scanner
runs on one OS/Python but a package may be installed on another, so a marker
that can be true on a different target (``sys_platform == 'win32'``,
``python_version == '3.13'``) is always kept; only a marker depending solely
on ``extra`` and false with no extra requested is dropped. Conservative: on
any uncertainty, keep (over-scan, never silently skip)."""
m = marker.strip()
if not m or "extra" not in m:
return True # no extra gate: installed by default on some target -> scan
if any(v in m for v in _MARKER_ENV_VARS):
return True # also platform/python gated: true on some target -> scan
# Pure extra marker: decide by evaluating with no extra requested.
try:
from packaging.markers import Marker, default_environment
env = default_environment()
env["extra"] = ""
return bool(Marker(m).evaluate(env))
except Exception:
# packaging missing/unparseable: drop only a pure positive extra-equality.
return re.fullmatch(r"\s*extra\s*==\s*['\"][^'\"]+['\"]\s*", m) is None
def _requires_dist_names(meta: dict) -> list[str]:
"""Transitive dep specs (name + version specifier) from metadata, to recover
a sdist-only package's tree. The specifier is kept so a pinned malicious
version is fetched, not latest. Drops deps whose marker cannot hold for a
default install."""
info = meta.get("info", {}) or {}
reqs = info.get("requires_dist") or []
specs: list[str] = []
for r in reqs:
if not isinstance(r, str):
continue
head = r
if ";" in r:
head, marker = r.split(";", 1)
if not _marker_holds_by_default(marker):
continue
if not _RE_NAME.match(head.strip()):
continue
# "torch (>=1.10)" / "torch >=1.10" -> "torch>=1.10" (pip-friendly).
specs.append(re.sub(r"\s+", "", head).replace("(", "").replace(")", ""))
return specs
def _requires_dist_for(
name: str,
version: str | None,
project_meta: dict,
errors: list[str] | None = None,
) -> list[str]:
"""Declared deps for the pinned version, read from that release's metadata
(its ``requires_dist`` can differ from latest). Unpinned uses the
project-level (latest) document. A pinned version whose own metadata cannot
be fetched returns [] (never latest's deps) and, when ``errors`` is given,
records an incomplete-scan error so a partial tree is not read as "no deps"."""
if not version:
return _requires_dist_names(project_meta)
vmeta = _pypi_json(name, version)
if vmeta is None:
msg = f"metadata fetch failed for pinned {name}=={version}; dependency scan incomplete"
if errors is None:
print(f" [WARN] {msg}", file = sys.stderr)
else:
errors.append(msg)
return []
return _requires_dist_names(vmeta)
def _download_sdist_direct(
name: str,
version: str | None,
dest: str,
*,
meta: dict | None = None,
) -> tuple[str | None, str | None]:
"""Fetch a project's sdist tarball directly from PyPI (no pip, no build).
Returns ``(filepath, error)``, one non-None. Suffix preserved for the archive
reader; bounded by ``_MAX_SDIST_BYTES`` and restricted to PyPI's CDN.
"""
if meta is None:
meta = _pypi_json(name)
if meta is None:
return None, f"PyPI metadata fetch failed for {name}"
picked: tuple[str, str] | None = None
for f in _release_files(meta, version):
if f.get("packagetype") == "sdist" and f.get("url") and f.get("filename"):
picked = (f["filename"], f["url"])
break
if picked is None:
return None, f"no sdist published for {name} (version={version or 'latest'})"
fname, url = picked
if not _is_trusted_pypi_url(url):
return None, f"refusing non-PyPI sdist URL for {name}: {url[:80]}"
# basename + sanitize keeps the path inside dest; the char class preserves
# the real `.tar.gz` / `.zip` suffix so the archive reader picks the format.
safe_fname = _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz"
out = os.path.join(dest, safe_fname)
try:
req = urllib.request.Request(url, headers = {"Accept": "application/octet-stream"})
with urllib.request.urlopen(req, timeout = _SDIST_DOWNLOAD_TIMEOUT) as resp:
if getattr(resp, "status", 200) != 200:
return None, f"sdist HTTP {getattr(resp, 'status', '?')} for {name}"
data = resp.read(_MAX_SDIST_BYTES + 1)
if len(data) > _MAX_SDIST_BYTES:
return None, f"sdist for {name} exceeds {_MAX_SDIST_BYTES} byte cap"
with open(out, "wb") as fh:
fh.write(data)
print(
f" [INFO] fetched sdist directly (no build) for {name}: {safe_fname}",
file = sys.stderr,
)
return out, None
except Exception as exc:
return None, f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}"
def _pip_download_with_deps(
specs: list[str],
dest: str,
env: dict,
*,
timeout: int = 600,
) -> tuple[int, str]:
"""One `pip download --with-deps --only-binary :all:` call. Returns (rc, stderr)."""
cmd = [
sys.executable,
"-m",
"pip",
"download",
*_PIP_DOWNLOAD_PIN_FLAGS,
"--dest",
dest,
] + list(specs)
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = timeout, env = env)
return proc.returncode, proc.stderr or ""
except subprocess.TimeoutExpired:
return 124, "pip download (with deps) timed out"
def _collect_flat_dir(dest: str, results: list[tuple[str, str]]) -> None:
"""Append every archive in a flat dest dir as (pkg_name, path)."""
for fname in sorted(os.listdir(dest)):
fpath = os.path.join(dest, fname)
if os.path.isfile(fpath):
pkg_name = fname.split("-")[0].replace("_", "-").lower()
results.append((pkg_name, fpath))
def _resolve_per_spec_with_deps(
specs: list[str], dest: str, env: dict, download_errors: list[str]
) -> None:
"""Fallback when the bulk --with-deps resolve fails: resolve each spec alone.
A still-failing spec is probed against PyPI: sdist-only -> direct fetch (deps
recovered one level); wheel-present but tree-unresolvable -> a --no-deps fetch
of just that package. Only a genuine fetch failure errors (caller exits 2);
unfetchable indirect deps are warned, since the named package is still scanned.
"""
sdist_dep_followups: list[str] = []
for spec in specs:
name = _extract_pkg_name(spec)
version = _spec_pin_version(spec)
cmd = [
sys.executable,
"-m",
"pip",
"download",
*_PIP_DOWNLOAD_PIN_FLAGS,
"--dest",
dest,
spec,
]
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env)
except subprocess.TimeoutExpired:
download_errors.append(f"per-spec --with-deps timed out for {spec}")
continue
if proc.returncode == 0:
continue # archives landed in dest; collected by the caller
meta = _pypi_json(name)
if meta is not None and not _release_has_wheel(meta, version):
fpath, serr = _download_sdist_direct(name, version, dest, meta = meta)
if fpath is None:
download_errors.append(serr or f"sdist fetch failed for {name}")
continue
sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors))
continue
# Has a wheel but the full transitive tree won't co-resolve
# (ResolutionImpossible) -- typically a package the requirement file
# installs with --no-deps by design (e.g. descript-audio-codec, whose
# own pins conflict). Fetch just the package itself with --no-deps so it
# is still scanned; its conflicting deps are out of scope here (the file
# excludes them on purpose). Only a genuine fetch failure is an error.
nd_cmd = [
sys.executable,
"-m",
"pip",
"download",
"--no-deps",
*_PIP_DOWNLOAD_PIN_FLAGS,
"--dest",
dest,
spec,
]
try:
nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env)
except subprocess.TimeoutExpired:
download_errors.append(f"per-spec --no-deps timed out for {spec}")
continue
if nd.returncode == 0:
print(
f" [INFO] {name}: full tree unresolvable; scanned the package "
f"alone (--no-deps), recovering deps individually.",
file = sys.stderr,
)
# The --with-deps failure may have been a sdist-only TRANSITIVE dep,
# which --no-deps skips. Recover the declared deps so that class is
# still scanned (each is fetched as a wheel or direct sdist below).
if meta is not None:
sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors))
continue
# --no-deps also failed: last-ditch sdist fetch at the pinned version.
if meta is not None:
fpath, _serr = _download_sdist_direct(name, version, dest, meta = meta)
if fpath is not None:
continue
download_errors.append(
f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}"
)
# Recover the transitive deps of sdist-only packages. A depth-bounded,
# deduped worklist so a wheel dep whose own child is sdist-only is itself
# fetched (--no-deps) and scanned -- not silently dropped -- and that child
# is then recovered in turn. `dep` carries the version specifier so a pinned
# version is fetched.
seen: set[str] = set()
worklist: list[tuple[str, int]] = [(d, 0) for d in sdist_dep_followups]
while worklist:
dep, depth = worklist.pop()
dep_name = _extract_pkg_name(dep)
key = _norm_pkg(dep_name)
if key in seen:
continue
seen.add(key)
dep_ver = _spec_pin_version(dep)
cmd = [
sys.executable,
"-m",
"pip",
"download",
*_PIP_DOWNLOAD_PIN_FLAGS,
"--dest",
dest,
dep,
]
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env)
except subprocess.TimeoutExpired:
print(f" [WARN] dep download timed out for {dep}", file = sys.stderr)
continue
if proc.returncode == 0:
continue
meta = _pypi_json(dep_name)
if meta is None:
print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
continue
if not _release_has_wheel(meta, dep_ver):
fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
if fpath is None:
print(f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr)
elif depth < _MAX_DEP_FOLLOWUP_DEPTH:
worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
continue
# Wheel published but its tree won't co-resolve (a sdist-only child).
# Fetch the dep alone so it is scanned, then chase its own declared deps.
nd_cmd = [
sys.executable,
"-m",
"pip",
"download",
"--no-deps",
*_PIP_DOWNLOAD_PIN_FLAGS,
"--dest",
dest,
dep,
]
try:
nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env)
except subprocess.TimeoutExpired:
print(f" [WARN] dep --no-deps timed out for {dep}", file = sys.stderr)
continue
if nd.returncode == 0:
if depth < _MAX_DEP_FOLLOWUP_DEPTH:
worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
continue
fpath, _serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
if fpath is None:
print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
elif depth < _MAX_DEP_FOLLOWUP_DEPTH:
worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
def download_packages(
specs: list[str],
dest: str,
@ -1403,49 +1982,36 @@ def download_packages(
summaries. A non-empty ``download_errors`` MUST make the caller exit
non-zero so a partial scan can't masquerade as "0 findings, all clean".
with_deps=True downloads the full transitive tree in one pip call (flat dir);
with_deps=False (default) downloads each spec individually with --no-deps.
with_deps=True downloads the full transitive tree (flat dir); a bulk resolve
failure (sdist-only package or version conflict) degrades to per-spec
resolution + direct sdist fetch rather than blanking the shard.
with_deps=False (default) downloads each spec individually with --no-deps,
also falling back to a direct sdist fetch when no wheel exists.
"""
results: list[tuple[str, str]] = []
download_errors: list[str] = []
env = _pip_download_env()
if with_deps:
# Single pip download for all specs + transitive deps. `--only-binary
# :all:` refuses sdists so we never execute setup.py for metadata.
os.makedirs(dest, exist_ok = True)
cmd = [
sys.executable,
"-m",
"pip",
"download",
*_PIP_DOWNLOAD_PIN_FLAGS,
"--dest",
dest,
] + specs
try:
proc = subprocess.run(
cmd,
capture_output = True,
text = True,
timeout = 600, # transitive resolution is slow
env = env,
# Fast path: resolve + download the whole transitive tree in one call.
# `--only-binary :all:` refuses sdists so we never build for metadata.
rc, stderr = _pip_download_with_deps(specs, dest, env)
if rc != 0:
# Atomic resolve failed -- a sdist-only package, or a cross-package
# version conflict (ResolutionImpossible). Degrade to per-spec
# resolution so one bad spec can't blank the shard, then direct-fetch
# any sdist-only holdouts (no build). Genuine failures still record an
# error so the caller exits 2.
print(
f" [INFO] bulk --with-deps resolve failed "
f"({stderr.strip()[:160]}); falling back to per-spec resolution "
f"for {len(specs)} spec(s).",
file = sys.stderr,
)
if proc.returncode != 0:
msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
except subprocess.TimeoutExpired:
msg = "pip download (with deps) timed out"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
# Collect every archive that landed in dest
for fname in sorted(os.listdir(dest)):
fpath = os.path.join(dest, fname)
if os.path.isfile(fpath):
pkg_name = fname.split("-")[0].replace("_", "-").lower()
results.append((pkg_name, fpath))
_resolve_per_spec_with_deps(specs, dest, env, download_errors)
# Collect everything that landed (bulk OR per-spec OR direct sdist).
_collect_flat_dir(dest, results)
else:
for spec in specs:
raw_name = _extract_pkg_name(spec)
@ -1465,22 +2031,25 @@ def download_packages(
spec,
]
try:
proc = subprocess.run(
cmd,
capture_output = True,
text = True,
timeout = 120,
env = env,
)
if proc.returncode != 0:
msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
continue
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 120, env = env)
except subprocess.TimeoutExpired:
msg = f"pip download timed out for {spec}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
download_errors.append(f"pip download timed out for {spec}")
continue
if proc.returncode != 0:
# No wheel? Direct-fetch the sdist (no build) before erroring.
name = _extract_pkg_name(spec)
version = _spec_pin_version(spec)
meta = _pypi_json(name)
if meta is not None and not _release_has_wheel(meta, version):
fpath, serr = _download_sdist_direct(name, version, pkg_dir, meta = meta)
if fpath is not None:
results.append((spec, fpath))
continue
download_errors.append(serr or f"sdist fetch failed for {name}")
continue
download_errors.append(
f"pip download failed for {spec}: {proc.stderr.strip()[:300]}"
)
continue
for fname in os.listdir(pkg_dir):
@ -1722,8 +2291,10 @@ def find_safe_version(
scan_dir = os.path.join(tmpdir, f"{name}_{ver}")
os.makedirs(scan_dir, exist_ok = True)
downloaded = download_packages([spec], scan_dir)
downloaded, download_errors = download_packages([spec], scan_dir)
if not downloaded:
for err in download_errors:
print(f" [WARN] {err}", file = sys.stderr)
continue
clean = True
@ -1856,9 +2427,12 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
# If no pinned version, download to find what pip resolves
dl_dir = os.path.join(tmpdir, f"resolve_{pkg_name}")
os.makedirs(dl_dir, exist_ok = True)
downloaded = download_packages([pkg_name], dl_dir)
downloaded, download_errors = download_packages([pkg_name], dl_dir)
if downloaded:
current_ver = get_downloaded_version(downloaded[0][1])
else:
for err in download_errors:
print(f" [WARN] {err}", file = sys.stderr)
shutil.rmtree(dl_dir, ignore_errors = True)
if not current_ver:
@ -1940,6 +2514,113 @@ def _find_requirements_files(root: str) -> list[str]:
return sorted(results)
# Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can
# enforce without drowning in legitimate-library noise. Matched on
# ``(package, basename(filename), check)`` -- not evidence text -- so a version
# bump does not reopen a finding, but a *new* kind of finding in a listed file
# is a different check and still fails. Regenerate with ``--write-baseline``.
_DEFAULT_BASELINE_PATH = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json"
)
def _norm_pkg(name: str) -> str:
"""PEP 503-style normalization so requests/Requests/req_uests collapse."""
return re.sub(r"[-_.]+", "-", (name or "").strip().lower())
# Leading "<name>-<version>/" archive root of an sdist member, which carries the
# version. Stripping it (but keeping the rest of the path) gives a key that is
# stable across version bumps yet still distinguishes same-named files.
_RE_SDIST_ROOT = re.compile(r"^[^/]+-\d[^/]*/")
def _relpath_in_package(filename: str) -> str:
"""Package-relative path: drop an sdist's version-carrying archive root.
Wheel members are already package-relative (``numba/cuda/utils.py``); sdist
members sit under ``numba-0.60.0/...``, so strip that one leading segment.
"""
return _RE_SDIST_ROOT.sub("", filename, count = 1)
def _finding_key(f: Finding) -> tuple[str, str, str]:
"""Stable allowlist key: normalized package, package-relative path, check.
The package-relative path (not just basename) keeps the key stable across
version bumps while still distinguishing same-named files like ``utils.py``.
"""
return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check)
def _load_baseline(path: str) -> set[tuple[str, str, str]]:
"""Load an allowlist JSON into a set of match keys. Missing file -> empty."""
try:
with open(path, "r", encoding = "utf-8") as fh:
data = json.load(fh)
except FileNotFoundError:
return set()
except (OSError, json.JSONDecodeError) as exc:
print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
return set()
keys: set[tuple[str, str, str]] = set()
for e in data.get("entries", []):
try:
keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"]))
except (KeyError, TypeError):
continue
return keys
def _write_baseline(path: str, findings: list[Finding]) -> None:
"""Persist CRITICAL/HIGH findings as an allowlist for human triage."""
entries = []
seen: set[tuple[str, str, str]] = set()
for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)):
if f.severity not in (CRITICAL, HIGH):
continue
key = _finding_key(f)
if key in seen:
continue
seen.add(key)
entries.append(
{
"package": f.package,
"file": _relpath_in_package(f.filename),
"check": f.check,
"severity": f.severity,
"evidence": f.evidence[:240],
}
)
doc = {
"_comment": (
"scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding "
"manually judged benign. Matched on (package, package-relative file, "
"check); evidence/severity are for review only. Regenerate with "
"--write-baseline AFTER reviewing every line."
),
"version": 1,
"entries": entries,
}
with open(path, "w", encoding = "utf-8") as fh:
json.dump(doc, fh, indent = 2, sort_keys = False)
fh.write("\n")
print(f" Wrote {len(entries)} baseline entr(y/ies) to {path}")
def _partition_baseline(
findings: list[Finding], baseline: set[tuple[str, str, str]]
) -> tuple[list[Finding], list[Finding]]:
"""Split findings into (active, suppressed) by allowlist membership."""
if not baseline:
return list(findings), []
active, suppressed = [], []
for f in findings:
(suppressed if _finding_key(f) in baseline else active).append(f)
return active, suppressed
# Main
@ -1986,6 +2667,30 @@ def main() -> int:
metavar = "N",
help = "Max older versions to scan when searching for safe version (default: 10)",
)
parser.add_argument(
"--baseline",
metavar = "FILE",
default = None,
help = (
"Allowlist JSON of triaged known-good findings to suppress. "
f"Defaults to {os.path.basename(_DEFAULT_BASELINE_PATH)} next to this "
"script if present."
),
)
parser.add_argument(
"--no-baseline",
action = "store_true",
help = "Ignore the auto-discovered baseline allowlist.",
)
parser.add_argument(
"--write-baseline",
metavar = "FILE",
default = None,
help = (
"Write the current CRITICAL/HIGH findings to FILE as an allowlist, "
"then exit 0. Review every entry before committing it."
),
)
args = parser.parse_args()
# --scan-dir: auto-discover requirements files
@ -2066,11 +2771,34 @@ def main() -> int:
finally:
shutil.rmtree(tmpdir, ignore_errors = True)
print_findings(all_findings)
# Baseline allowlist: suppress triaged, known-good findings so the CI gate
# can be enforcing without red-failing on legitimate-library noise.
if args.no_baseline:
baseline_path = None
elif args.baseline:
baseline_path = args.baseline
elif os.path.isfile(_DEFAULT_BASELINE_PATH):
baseline_path = _DEFAULT_BASELINE_PATH
else:
baseline_path = None
baseline = _load_baseline(baseline_path) if baseline_path else set()
# --fix mode: auto-search for safe versions
if args.fix and all_findings:
critical_pkgs = {f.package for f in all_findings if f.severity == CRITICAL}
active, suppressed = _partition_baseline(all_findings, baseline)
print_findings(active)
if suppressed:
crit_s = sum(1 for f in suppressed if f.severity == CRITICAL)
high_s = sum(1 for f in suppressed if f.severity == HIGH)
med_s = sum(1 for f in suppressed if f.severity == MEDIUM)
print(
f"\n {len(suppressed)} finding(s) suppressed by baseline "
f"{baseline_path} "
f"({crit_s} CRITICAL, {high_s} HIGH, {med_s} MEDIUM)."
)
# --fix mode: auto-search for safe versions (only real, non-baselined ones)
if args.fix and active:
critical_pkgs = {f.package for f in active if f.severity == CRITICAL}
if critical_pkgs:
print(
f"\n --fix: Searching for safe versions of {len(critical_pkgs)} CRITICAL package(s)..."
@ -2079,6 +2807,7 @@ def main() -> int:
# Surface pip-download failures BEFORE the exit code so a partial download
# can't masquerade as "0 findings, all clean" (silent-failure hardening 4).
# Also keeps us from writing a baseline from an incomplete scan.
if download_errors:
print(
f"\n {'=' * 72}\n"
@ -2095,8 +2824,16 @@ def main() -> int:
)
return 2
# Exit code: 1 if any CRITICAL or HIGH
if any(f.severity in (CRITICAL, HIGH) for f in all_findings):
# --write-baseline: persist the full current CRITICAL/HIGH set as the new
# allowlist (ignoring any loaded baseline), then exit 0. Only reached once
# the scan is known complete.
if args.write_baseline:
_write_baseline(args.write_baseline, all_findings)
return 0
# Exit code: 1 only if a NON-baselined CRITICAL or HIGH remains. This is the
# signal CI gates on once the baseline reaches a clean run.
if any(f.severity in (CRITICAL, HIGH) for f in active):
return 1
return 0

1329
scripts/scan_packages_baseline.json Normal file
View file

File diff suppressed because it is too large Load diff

9
studio/backend/assets/configs/inference_defaults.json
View file

@ -277,6 +277,13 @@
"min_p": 0.01,
"repetition_penalty": 1.0
},
"minimax-m2.7": {
"temperature": 1.0,
"top_p": 0.95,
"top_k": 40,
"min_p": 0.01,
"repetition_penalty": 1.0
},
"minimax-m2.5": {
"temperature": 1.0,
"top_p": 0.95,
@ -390,7 +397,7 @@
"deepseek-r1", "deepseek-v3", "deepseek-ocr",
"glm-5", "glm-4",
"nemotron",
"minimax-m2.5", "minimax",
"minimax-m2.7", "minimax-m2.5", "minimax",
"gpt-oss", "granite-4",
"kimi-k2", "kimi",
"lfm2", "smollm", "olmo", "falcon", "ernie", "seed", "grok", "mimo"

35
studio/backend/auth/storage.py
View file

@ -4,9 +4,11 @@
"""SQLite storage for auth data (user credentials + JWT secret)."""
import hashlib
import hmac
import os
import secrets
import sqlite3
import threading
from datetime import datetime, timezone
from typing import Optional, Tuple
@ -236,6 +238,29 @@ def _pbkdf2_desktop_secret(raw_secret: str) -> str:
return _pbkdf2_api_key(raw_secret)
# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round
# KDF runs once per key instead of on every authenticated request. Keyed by a
# salted HMAC of the key (not the key itself); revocation/expiry are still
# enforced by the SQLite read on every call, so a cache hit only skips the KDF.
# Only keys present in the DB are cached, so unknown-key spam can't grow it.
_api_key_hash_cache: dict[str, str] = {}
_API_KEY_HASH_CACHE_MAX = 4096
_api_key_hash_cache_lock = threading.Lock()
def _api_key_cache_id(raw_key: str) -> str:
"""Cache id for a raw key: salted HMAC-SHA256 (not the key itself)."""
return hmac.new(
_get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256
).hexdigest()
def _reset_api_key_hash_cache() -> None:
"""Drop memoized derivations (tests / salt change)."""
with _api_key_hash_cache_lock:
_api_key_hash_cache.clear()
def is_initialized() -> bool:
"""Check if auth is ready for login (at least one user exists in DB)."""
conn = get_connection()
@ -704,7 +729,9 @@ def validate_api_key(raw_key: str) -> Optional[str]:
Also updates ``last_used_at`` on success.
"""
key_hash = _pbkdf2_api_key(raw_key)
cache_id = _api_key_cache_id(raw_key)
cached_hash = _api_key_hash_cache.get(cache_id)
key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key)
conn = get_connection()
try:
cur = conn.execute(
@ -714,6 +741,12 @@ def validate_api_key(raw_key: str) -> Optional[str]:
row = cur.fetchone()
if row is None:
return None
# Real key: memoize so later requests skip the KDF. Bounded; clear on overflow.
if cached_hash is None:
with _api_key_hash_cache_lock:
if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX:
_api_key_hash_cache.clear()
_api_key_hash_cache[cache_id] = key_hash
if not row["is_active"]:
return None
if row["expires_at"] is not None:

11
studio/backend/cloudflare_tunnel.py
View file

@ -49,6 +49,16 @@ def _windows_hidden_kwargs() -> dict:
return {"creationflags": flags} if flags else {}
def _lifetime_kwargs() -> dict:
"""Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy +
best-effort so this module still loads standalone (storage_roots-style)."""
try:
from utils.process_lifetime import child_popen_kwargs
return child_popen_kwargs()
except Exception:
return {}
def _asset_name() -> Optional[Tuple[str, bool]]:
"""(release asset filename, is_tgz) for this OS/arch, or None if unsupported."""
system = platform.system().lower()
@ -233,6 +243,7 @@ class CloudflareTunnel:
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
**_lifetime_kwargs(),
)
self._proc = proc
threading.Thread(

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

@ -176,6 +176,9 @@ class JobManager:
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
self._mp_q = mp_q
self._proc = proc

19
studio/backend/core/export/export.py
View file

@ -145,13 +145,19 @@ class ExportBackend:
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
``hf_token`` authenticates the actual weight load for gated/private
checkpoints, matching the token the worker used for the security preflight
(otherwise a gated repo passes scanning then 401s at from_pretrained).
Returns:
Tuple of (success: bool, message: str)
"""
token = hf_token if hf_token and hf_token.strip() else None
try:
logger.info(f"Loading checkpoint: {checkpoint_path}")
@ -169,8 +175,10 @@ class ExportBackend:
model_id = base_model or checkpoint_path
self._audio_type = detect_audio_type(model_id)
self.is_vision = not self._audio_type and is_vision_model(model_id)
# Token the type-detection probes too, else a gated multimodal base
# 404s here and falls through to the text loader.
self._audio_type = detect_audio_type(model_id, hf_token = token)
self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token)
if self._audio_type == "csm":
from unsloth import FastModel
@ -184,6 +192,7 @@ class ExportBackend:
auto_model = CsmForConditionalGeneration,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "whisper":
@ -197,6 +206,7 @@ class ExportBackend:
load_in_4bit = False,
auto_model = WhisperForConditionalGeneration,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "snac":
@ -207,6 +217,7 @@ class ExportBackend:
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "bicodec":
@ -218,6 +229,7 @@ class ExportBackend:
dtype = None if _IS_MLX else torch.float32,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
)
elif self._audio_type == "dac":
@ -228,6 +240,7 @@ class ExportBackend:
max_seq_length = max_seq_length,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
)
elif self.is_vision:
@ -238,6 +251,7 @@ class ExportBackend:
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -249,6 +263,7 @@ class ExportBackend:
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
)
if _IS_MLX:

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

@ -63,6 +63,19 @@ class ExportOrchestrator:
self._run_start_seq: int = 0
# True while an export op runs; SSE ends the stream 1s after this flips False.
self._export_active: bool = False
# Set by cancel_export(); reset when a new load/export run starts. Lets the
# caller distinguish a user cancel from a genuine subprocess crash.
self._cancel_requested: bool = False
# Last finished operation, so a client whose blocking POST was cut off by a
# Cloudflare tunnel timeout (524 at ~100s, while the op runs for minutes) can
# poll /api/export/status and still learn the real outcome. Guarded by
# _op_lock. `_op_seq` is a monotonic counter the client uses as a baseline to
# tell "my op finished" (seq grew) from a stale previous result.
self._op_lock = threading.Lock()
self._op_seq: int = 0
self._active_op_kind: Optional[str] = None
self._last_op: Optional[Dict[str, Any]] = None
atexit.register(self._cleanup)
logger.info("ExportOrchestrator initialized (subprocess mode)")
@ -119,6 +132,72 @@ class ExportOrchestrator:
"""True while an export / load / cleanup command is running."""
return self._export_active
def was_cancelled(self) -> bool:
"""True if the in-flight (or most recent) run was cancelled by the user."""
return self._cancel_requested
def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None:
"""Snapshot the just-finished op so status pollers can recover its outcome.
Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE
``_export_active`` is cleared, so a status read that observes the op as
inactive is guaranteed to also see this matching result.
"""
with self._op_lock:
self._op_seq += 1
status = "cancelled" if self._cancel_requested else ("success" if success else "error")
self._last_op = {
"seq": self._op_seq,
"kind": self._active_op_kind,
"status": status,
"output_path": output_path if success else None,
"error": None if success else (message or None),
}
def get_last_op(self) -> Optional[Dict[str, Any]]:
"""Return the last finished op record (or None), for status recovery."""
with self._op_lock:
return dict(self._last_op) if self._last_op is not None else None
def get_active_op_kind(self) -> Optional[str]:
"""Return the kind of the currently running op (or None when idle)."""
return self._active_op_kind
def cancel_export(self) -> bool:
"""Terminate the in-flight export subprocess immediately.
An export op holds ``self._lock`` for its whole duration (blocked in
``_wait_response``), so we deliberately do NOT take the lock here -- we
kill the worker process directly, which unblocks that wait and makes the
in-flight op return a failure the caller surfaces as "cancelled".
Only the export subprocess is touched; training and inference run in
their own subprocesses and are left untouched.
Returns True if a live subprocess was terminated, False if none ran.
"""
self._cancel_requested = True
proc = self._proc
if proc is None or not proc.is_alive():
return False
logger.info(
"Export cancel requested: terminating export subprocess (pid=%s)",
proc.pid,
)
try:
proc.terminate()
proc.join(timeout = 5)
except Exception:
pass
if proc.is_alive():
logger.warning("Export subprocess survived terminate, killing")
try:
proc.kill()
proc.join(timeout = 3)
except Exception:
pass
return True
# ------------------------------------------------------------------
# Subprocess lifecycle
# ------------------------------------------------------------------
@ -147,6 +226,9 @@ class ExportOrchestrator:
daemon = True,
)
self._proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep)
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
@ -301,6 +383,7 @@ class ExportOrchestrator:
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
approved_remote_code_fingerprint: Optional[str] = None,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""Load a checkpoint for export.
@ -312,13 +395,17 @@ class ExportOrchestrator:
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
"trust_remote_code": trust_remote_code,
"approved_remote_code_fingerprint": approved_remote_code_fingerprint,
"hf_token": hf_token,
}
with self._lock:
# Fresh log buffer so the UI sees only this run's output.
self.clear_logs()
self._cancel_requested = False
self._active_op_kind = "load_checkpoint"
self._export_active = True
op_success, op_message = False, ""
try:
# Always kill any existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
@ -336,6 +423,7 @@ class ExportOrchestrator:
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
op_success, op_message = False, str(exc)
return False, str(exc)
if resp.get("success"):
@ -343,15 +431,19 @@ class ExportOrchestrator:
self.is_vision = resp.get("is_vision", False)
self.is_peft = resp.get("is_peft", False)
logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path)
return True, resp.get("message", "Loaded successfully")
op_success, op_message = True, resp.get("message", "Loaded successfully")
return True, op_message
else:
error = resp.get("message", "Failed to load checkpoint")
logger.error("Failed to load checkpoint: %s", error)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
op_success, op_message = False, error
return False, error
finally:
self._record_op_finished(op_success, op_message, None)
self._active_op_kind = None
self._export_active = False
def export_merged_model(
@ -453,7 +545,10 @@ class ExportOrchestrator:
)
self.clear_logs()
self._cancel_requested = False
self._active_op_kind = f"export_{export_type}"
self._export_active = True
op_success, op_message, op_output_path = False, "", None
try:
cmd = {"type": "export", "export_type": export_type, **params}
try:
@ -462,14 +557,16 @@ class ExportOrchestrator:
f"export_{export_type}_done",
timeout = 3600, # GGUF for 30B+ models can take 30+ min
)
return (
resp.get("success", False),
resp.get("message", ""),
resp.get("output_path"),
)
op_success = resp.get("success", False)
op_message = resp.get("message", "")
op_output_path = resp.get("output_path")
return op_success, op_message, op_output_path
except RuntimeError as exc:
op_success, op_message = False, str(exc)
return False, str(exc), None
finally:
self._record_op_finished(op_success, op_message, op_output_path)
self._active_op_kind = None
self._export_active = False
def cleanup_memory(self) -> bool:
@ -481,7 +578,9 @@ class ExportOrchestrator:
self.is_peft = False
return True
self._active_op_kind = "cleanup"
self._export_active = True
success = False
try:
try:
self._send_cmd({"type": "cleanup"})
@ -498,6 +597,8 @@ class ExportOrchestrator:
self.is_peft = False
return success
finally:
self._record_op_finished(success, "", None)
self._active_op_kind = None
self._export_active = False
def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]:

86
studio/backend/core/export/worker.py
View file

@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
# Auto-enable trust_remote_code for NemotronH/Nano models.
if not trust_remote_code:
from utils.security.trusted_org import is_trusted_org_repo
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_cp_lower = checkpoint_path.lower()
if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")
if (
any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/"))
# Genuine first-party Hub repo only (not a local/spoof name starting
# with "unsloth/"); authenticated so private repos resolve.
and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token"))
):
trust_remote_code = True
logger.info(
@ -199,6 +205,81 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path,
)
# Malware gate: a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) every
# load. Local checkpoints have no Hub scan and are skipped in the helper; a
# LoRA merges its base weights, so gate that repo too.
from utils.security import evaluate_file_security, security_load_subdirs
malware_targets = [checkpoint_path]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
_hf_token = cmd.get("hf_token")
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token)
)
if _fs.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
},
)
return
# Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless
# pinned-approved. A LoRA merges its base model, whose code runs, so gate it too.
if trust_remote_code:
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [checkpoint_path]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a local or remote adapter's base so its base repo is gated too.
base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
if base_model:
consent_targets.append(base_model)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = cmd.get("hf_token"),
trust_remote_code = True,
approved_fingerprint = cmd.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": (
f"Checkpoint '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review and "
f"approve it to proceed."
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
},
)
return
try:
_send_response(
resp_queue,
@ -214,6 +295,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
hf_token = cmd.get("hf_token"),
)
_send_response(

295
studio/backend/core/inference/api_monitor.py Normal file
View file

@ -0,0 +1,295 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Small in-memory monitor for OpenAI-compatible API traffic."""
from __future__ import annotations
import threading
import time
import uuid
from collections import deque
from dataclasses import dataclass
from typing import Any, Optional
_MAX_ENTRIES = 50
_MAX_PROMPT_CHARS = 12000
_MAX_REPLY_CHARS = 12000
_PREVIEW_CHARS = 360
def _trim(text: Optional[str], limit: int) -> str:
if not text:
return ""
if len(text) <= limit:
return text
# Guard against limit < 3 (slice would underflow).
if limit <= 3:
return "..."[:limit]
return text[: limit - 3] + "..."
@dataclass
class ApiMonitorEntry:
id: str
endpoint: str
method: str
model: str
prompt: str
status: str
started_at: float
updated_at: float
subject: Optional[str] = None
# Monotonic anchors so duration math survives wall-clock steps (NTP).
started_monotonic: float = 0.0
finished_monotonic: Optional[float] = None
reply: str = ""
finished_at: Optional[float] = None
context_length: Optional[int] = None
prompt_tokens: Optional[int] = None
completion_tokens: Optional[int] = None
total_tokens: Optional[int] = None
total_tokens_authoritative: bool = False
error: Optional[str] = None
def snapshot(self, *, include_details: bool = True) -> dict[str, Any]:
duration_ms = None
if self.finished_monotonic is not None:
duration_ms = max(
0,
int((self.finished_monotonic - self.started_monotonic) * 1000),
)
elif self.finished_at is not None:
duration_ms = max(0, int((self.finished_at - self.started_at) * 1000))
context_usage = None
if self.total_tokens is not None and self.context_length:
context_usage = min(1.0, max(0.0, self.total_tokens / self.context_length))
payload = {
"id": self.id,
"endpoint": self.endpoint,
"method": self.method,
"model": self.model,
"prompt_preview": _trim(self.prompt, _PREVIEW_CHARS),
"reply_preview": _trim(self.reply, _PREVIEW_CHARS),
"prompt_truncated": len(self.prompt) > _PREVIEW_CHARS,
"reply_truncated": len(self.reply) > _PREVIEW_CHARS,
"status": self.status,
"started_at": self.started_at,
"updated_at": self.updated_at,
"finished_at": self.finished_at,
"duration_ms": duration_ms,
"context_length": self.context_length,
"context_usage": context_usage,
"prompt_tokens": self.prompt_tokens,
"completion_tokens": self.completion_tokens,
"total_tokens": self.total_tokens,
"error": self.error,
}
if include_details:
payload["prompt"] = self.prompt
payload["reply"] = self.reply
return payload
class ApiMonitor:
def __init__(self, max_entries: int = _MAX_ENTRIES):
self._entries: deque[ApiMonitorEntry] = deque()
self._max_entries = max(0, max_entries)
self._lock = threading.Lock()
def start(
self,
*,
endpoint: str,
method: str,
model: str,
prompt: str,
context_length: Optional[int] = None,
subject: Optional[str] = None,
) -> str:
now = time.time()
entry = ApiMonitorEntry(
id = f"apireq_{uuid.uuid4().hex[:12]}",
endpoint = endpoint,
method = method,
model = model or "default",
prompt = _trim(prompt, _MAX_PROMPT_CHARS),
status = "running",
started_at = now,
updated_at = now,
subject = subject,
started_monotonic = time.monotonic(),
context_length = context_length,
)
with self._lock:
self._entries.appendleft(entry)
self._trim_terminal_locked()
return entry.id
def append_reply(self, entry_id: Optional[str], text: str) -> None:
if not entry_id or not text:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is None:
return
# Preview is capped: once the "..." marker is present the head is
# frozen, so skip the per-chunk re-concat (avoids O(n^2) on long
# generations). A reply that landed exactly on the cap has no marker
# yet, so let one more append record the truncation before freezing.
if len(entry.reply) >= _MAX_REPLY_CHARS:
if not entry.reply.endswith("..."):
entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS)
entry.updated_at = time.time()
return
entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS)
entry.updated_at = time.time()
def set_reply(self, entry_id: Optional[str], text: str) -> None:
if not entry_id:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is None:
return
entry.reply = _trim(text, _MAX_REPLY_CHARS)
entry.updated_at = time.time()
def set_usage(
self,
entry_id: Optional[str],
*,
prompt_tokens: Optional[int] = None,
completion_tokens: Optional[int] = None,
total_tokens: Optional[int] = None,
context_length: Optional[int] = None,
) -> None:
if not entry_id:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is None:
return
if prompt_tokens is not None:
entry.prompt_tokens = prompt_tokens
if completion_tokens is not None:
entry.completion_tokens = completion_tokens
if total_tokens is not None:
entry.total_tokens = total_tokens
entry.total_tokens_authoritative = True
elif not entry.total_tokens_authoritative and (
prompt_tokens is not None or completion_tokens is not None
):
# Derive only when no authoritative total has been set;
# a later partial chunk must not clobber a provider total.
entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0)
if context_length is not None:
entry.context_length = context_length
entry.updated_at = time.time()
def finish(
self,
entry_id: Optional[str],
status: str = "completed",
) -> None:
if not entry_id:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is None:
return
# Idempotent: second call (e.g. [DONE] after the finally block
# already ran) must not move finished_*.
if entry.finished_at is not None:
return
now = time.time()
entry.status = status
entry.updated_at = now
entry.finished_at = now
entry.finished_monotonic = time.monotonic()
self._entries.remove(entry)
self._entries.appendleft(entry)
self._trim_terminal_locked()
def fail(self, entry_id: Optional[str], error: str) -> None:
if not entry_id:
return
with self._lock:
entry = self._find_locked(entry_id)
if entry is None:
return
if entry.finished_at is not None:
# Already terminal; refresh error text only.
if error:
entry.error = _trim(error, 1000)
return
now = time.time()
entry.status = "error"
entry.error = _trim(error, 1000)
entry.updated_at = now
entry.finished_at = now
entry.finished_monotonic = time.monotonic()
self._entries.remove(entry)
self._entries.appendleft(entry)
self._trim_terminal_locked()
def snapshot(
self,
*,
include_details: bool = True,
subject: Optional[str] = None,
) -> list[dict[str, Any]]:
with self._lock:
return [
entry.snapshot(include_details = include_details)
for entry in self._entries
if subject is None or entry.subject == subject
]
def get(
self,
entry_id: str,
*,
subject: Optional[str] = None,
) -> Optional[dict[str, Any]]:
with self._lock:
entry = self._find_locked(entry_id)
if entry is None:
return None
if subject is not None and entry.subject != subject:
return None
return entry.snapshot(include_details = True)
def active_count(self, *, subject: Optional[str] = None) -> int:
with self._lock:
return sum(
1
for entry in self._entries
if entry.status == "running" and (subject is None or entry.subject == subject)
)
def clear(self) -> None:
with self._lock:
self._entries.clear()
def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]:
for entry in self._entries:
if entry.id == entry_id:
return entry
return None
def _trim_terminal_locked(self) -> None:
terminal_seen = 0
kept: deque[ApiMonitorEntry] = deque()
for entry in self._entries:
if entry.status == "running":
kept.append(entry)
continue
if terminal_seen < self._max_entries:
kept.append(entry)
terminal_seen += 1
self._entries = kept
api_monitor = ApiMonitor()

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

File diff suppressed because it is too large Load diff

56
studio/backend/core/inference/llama_http.py Normal file
View file

@ -0,0 +1,56 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared pooled httpx.AsyncClient for NON-streaming calls to the local llama-server.
Streaming generation must NOT use this. It relies on ``Connection: close`` and
``max_keepalive_connections=0`` so a client disconnect tears down the upstream
socket and stops GPU decode (PR #5749). This pooled client is only for short
request/response proxy calls (non-streaming completions, embeddings) where
reusing a connection removes per-request setup cost. Per-request ``timeout`` is
still passed at each call site.
"""
from __future__ import annotations
import asyncio
import weakref
import httpx
_LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32)
def _new_client() -> httpx.AsyncClient:
try:
return httpx.AsyncClient(limits = _LIMITS)
except Exception:
# Mirror external_provider: an unsupported env proxy scheme can raise.
return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
# One client per running event loop: an httpx client binds its transport to the
# loop it first runs on, so a single global instance breaks across a lifespan
# restart or a second test loop. Weak keys let a finished loop drop its client.
_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = (
weakref.WeakKeyDictionary()
)
def nonstreaming_client() -> httpx.AsyncClient:
loop = asyncio.get_running_loop()
client = _clients.get(loop)
if client is None or client.is_closed:
client = _new_client()
_clients[loop] = client
return client
async def aclose() -> None:
clients = list(_clients.values())
_clients.clear()
for client in clients:
try:
await client.aclose()
except Exception:
pass

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

@ -13,7 +13,8 @@ Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
from __future__ import annotations
from typing import Iterable, Optional
import os
from typing import Iterable, Mapping, Optional
# Each group = every alias (short + long) of one hard-denied flag.
# Extend the matching group when llama.cpp adds a new alias.
@ -124,7 +125,9 @@ def is_managed_flag(flag: str) -> bool:
# from inherited extras so they can't last-wins-override an Apply that
# re-sets the same field.
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"})
_CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"})
_CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"})
_CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS
_SPEC_FLAGS: frozenset[str] = frozenset(
{
"--spec-default",
@ -133,13 +136,22 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
"--spec-ngram-size",
"--draft-min",
"--draft-max",
# MTP path (llama.cpp #22673). --model-draft and aliases are
# Studio-managed since the separate-drafter support (Gemma 4): an
# inherited copy must not last-wins-override the auto-detected
# drafter. Explicit extras for the current load are never stripped.
# MTP path (llama.cpp #22673). The drafter selectors (local --model-draft
# and HF --spec-draft-hf aliases) are Studio-managed since the separate-
# drafter support (Gemma 4): an inherited copy must not last-wins-override
# the auto-detected drafter. Explicit extras for the current load are never
# stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld,
# --spec-draft-device) are deliberately NOT stripped: the VRAM budget reads
# them via the same parsers the child honors, so they stay consistent on
# inherit, and stripping them would silently move a CPU-offloaded drafter
# back onto the GPU.
"--model-draft",
"-md",
"--spec-draft-model",
"--spec-draft-hf",
"-hfd",
"-hfrd",
"--hf-repo-draft",
"--spec-draft-n-max",
"--spec-draft-n-min",
"--spec-draft-p-min",
@ -274,6 +286,20 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
return _last_flag_value(args, _CACHE_FLAGS)
def parse_cache_override_per_axis(
args: Optional[Iterable[str]],
) -> tuple[Optional[str], Optional[str]]:
"""Last-wins --cache-type-k / --cache-type-v values kept apart, as (k, v).
parse_cache_override collapses both axes to one last-wins value; this keeps
them separate so an asymmetric K/V can be budgeted by its heavier axis.
"""
return (
_last_flag_value(args, _CACHE_TYPE_K_FLAGS),
_last_flag_value(args, _CACHE_TYPE_V_FLAGS),
)
def resolve_cache_type_kv(
args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str]
) -> Optional[str]:
@ -309,6 +335,60 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral
return override.strip().lower() == "tensor"
def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool:
"""True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio
emits --split-mode only on its tensor branch, so a tensor env on the layer
path would run the child tensor-parallel unbudgeted; this flips the budget
to tensor. Only tensor is heavier, so other modes are ignored."""
raw = (os.environ if env is None else env).get("LLAMA_ARG_SPLIT_MODE")
return bool(raw) and raw.strip().lower() == "tensor"
def _effective_tensor_parallel(
extra_args: Optional[Iterable[str]],
tensor_parallel: bool,
env: Optional[Mapping[str, str]] = None,
) -> bool:
"""Tensor-parallel decision including the inherited LLAMA_ARG_SPLIT_MODE env.
resolve_tensor_parallel (extras + toggle), flipped on when extras set no split
mode but the child inherits a tensor split env. Shared by load_model (which
budgets and launches it) and the tensor-fallback wrapper (so an env-only
tensor crash still retries layer split)."""
resolved = resolve_tensor_parallel(extra_args, tensor_parallel)
if (
not resolved
and parse_split_mode_override(extra_args) is None
and _env_split_mode_is_tensor(env)
):
return True
return resolved
def _tensor_parallel_matches_loaded(
extra_args: Optional[Iterable[str]],
requested_tensor_parallel: bool,
loaded_tensor_parallel: bool,
env: Optional[Mapping[str, str]] = None,
) -> bool:
"""Whether a duplicate load request matches a loaded server's tensor state.
Env-only tensor mode is a launch hint load_model may downgrade to layer split
(capacity/buffer), scrubbing the child env. So only let an inherited tensor env
raise a match against a server that *actually* launched tensor; on a downgraded
(layer) server the env is ignored, and an identical request would downgrade the
same way -- avoiding an endless reload of a healthy server."""
requested = resolve_tensor_parallel(extra_args, requested_tensor_parallel)
if (
loaded_tensor_parallel
and not requested
and parse_split_mode_override(extra_args) is None
and _env_split_mode_is_tensor(env)
):
requested = True
return requested == loaded_tensor_parallel
_MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"})
_MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"})

118
studio/backend/core/inference/llama_stats.py Normal file
View file

@ -0,0 +1,118 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Translate llama-server's Prometheus /metrics into a periodic, vLLM-style
engine-stats log line (generation/prompt throughput, requests in flight).
llama-server already computes these (it needs `--metrics`); this lifts them
into Studio's structured log so the terminal shows serving health, not just
per-request access lines. Emitted only while there is activity.
"""
import os
import re
import threading
import time
import urllib.request
# Prometheus body lines: "llamacpp:<name>[{labels}] <value>" (skip "#" HELP/TYPE).
_METRIC_RE = re.compile(r"^llamacpp:(\w+)(?:\{[^}]*\})?\s+([0-9.eE+-]+)", re.MULTILINE)
_OFF = {"0", "false", "no", "off"}
class LlamaServerStatsLogger:
"""Daemon poller that logs vLLM-style engine stats from llama-server.
Keeps retrying through transient scrape failures; the backend stops it via
stop() on unload/reload, so a brief /metrics stall does not silence stats.
"""
def __init__(
self,
base_url,
logger,
interval_s = 10.0,
):
self._url = f"{base_url.rstrip('/')}/metrics"
self._log = logger
self._interval = max(1.0, float(interval_s))
self._stop = threading.Event()
self._thread = None
def start(self):
if self._thread is None:
self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True)
self._thread.start()
def stop(self):
self._stop.set()
def _scrape(self):
try:
with urllib.request.urlopen(self._url, timeout = 3) as r:
if r.status != 200:
return None
body = r.read().decode("utf-8", "replace")
except Exception:
return None
out = {}
for k, v in _METRIC_RE.findall(body):
try: # a malformed value must not kill the daemon thread
out[k] = float(v)
except ValueError:
continue
return out
def _run(self):
misses = 0
prev = None # (monotonic_t, tokens_predicted_total, prompt_tokens_total)
while not self._stop.wait(self._interval):
m = self._scrape()
if not m:
misses += 1
if misses == 3: # transient stall (load/GC); keep polling.
self._log.debug("engine_stats: /metrics scrape failing, still retrying")
continue # real shutdown is driven by stop() from _kill_process
misses = 0
# Generation tokens come from tokens_predicted_total (counter) and
# predicted_tokens_seconds (gauge); n_decode_total counts
# llama_decode() calls, not tokens, so it must not feed tok/s.
now = time.monotonic()
predicted = m.get("tokens_predicted_total", 0.0)
prompt = m.get("prompt_tokens_total", 0.0)
gen_delta = prompt_delta = 0.0
if prev is not None and now > prev[0]:
dt = now - prev[0]
gen_delta = max(0.0, (predicted - prev[1]) / dt)
prompt_delta = max(0.0, (prompt - prev[2]) / dt)
prev = (now, predicted, prompt)
# Prefer llama.cpp's own throughput gauges; fall back to the counter
# delta for binaries that expose only the counters.
gen_tps = m.get("predicted_tokens_seconds") or gen_delta
prompt_tps = m.get("prompt_tokens_seconds") or prompt_delta
running, waiting = (
int(m.get("requests_processing", 0)),
int(m.get("requests_deferred", 0)),
)
# Gate on real activity this tick so a stale gauge never logs at idle.
if running or waiting or gen_delta or prompt_delta:
self._log.info(
"engine_stats",
gen_tok_s = round(float(gen_tps), 1),
prompt_tok_s = round(float(prompt_tps), 1),
running = running,
waiting = waiting,
)
def maybe_start_stats_logger(base_url, logger):
"""Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it."""
if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF:
return None
try:
interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10"))
except ValueError:
interval = 10.0
sl = LlamaServerStatsLogger(base_url, logger, interval)
sl.start()
return sl

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

@ -18,7 +18,6 @@ import atexit
import base64
import os
import signal
import structlog
from loggers import get_logger
import multiprocessing as mp
import queue
@ -30,15 +29,15 @@ from pathlib import Path
from typing import Any, Generator, Optional, Tuple, Union
from utils.hardware import prepare_gpu_selection
# Re-exported from the shared helper so GGUF, training, and inference share one
# type; kept importable here for backwards compatibility.
from utils.hf_xet_fallback import DownloadStallError
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")
class DownloadStallError(RuntimeError):
"""Raised when the worker reports no download progress for too long."""
# Dispatcher timeout constants (seconds)
_DISPATCH_READ_TIMEOUT = 30.0
_DISPATCH_POLL_INTERVAL = 0.5
@ -61,7 +60,6 @@ class InferenceOrchestrator:
self._cmd_queue: Any = None
self._resp_queue: Any = None
self._cancel_event: Any = None # mp.Event — set to cancel generation
self._lock = threading.Lock()
self._gen_lock = threading.Lock() # Serializes generation
# Dispatcher state for compare mode (adapter-controlled requests):
@ -76,7 +74,6 @@ class InferenceOrchestrator:
self.active_model_name: Optional[str] = None
self.models: dict = {}
self.loading_models: set = set()
self.loaded_local_models: list = []
from core.inference.defaults import get_default_models
self._static_models = get_default_models()
@ -84,8 +81,6 @@ class InferenceOrchestrator:
self._top_hub_cache: Optional[list[str]] = None
self._top_models_ready = threading.Event()
self._current_transformers_major: Optional[str] = None # "4" or "5"
atexit.register(self._cleanup)
logger.info("InferenceOrchestrator initialized (subprocess mode)")
@ -177,6 +172,9 @@ class InferenceOrchestrator:
daemon = True,
)
self._proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep)
logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
def _cancel_generation(self) -> None:
@ -389,6 +387,108 @@ class InferenceOrchestrator:
return
logger.warning("Timed out waiting for gen_done after cancel")
# ------------------------------------------------------------------
# Generation command + token-stream helpers (shared by all paths)
# ------------------------------------------------------------------
def _build_generate_cmd(
self,
request_id: str,
image_b64: Optional[str],
*,
messages: list = None,
system_prompt: str = "",
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.0,
use_adapter = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> dict:
"""Build the 'generate' command shared by the locked and dispatched paths."""
cmd = {
"type": "generate",
"request_id": request_id,
"messages": messages or [],
"system_prompt": system_prompt,
"image_base64": image_b64,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
# Only forward template kwargs the caller set, for older worker compat.
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
cmd["enable_thinking"] = enable_thinking
if reasoning_effort is not None:
cmd["reasoning_effort"] = reasoning_effort
if preserve_thinking is not None:
cmd["preserve_thinking"] = preserve_thinking
return cmd
def _consume_token_stream(
self,
read_one,
drain_on_cancel,
*,
crash_context: str,
cancel_event = None,
stats_holder: Optional[dict] = None,
read_timeout: float = 30.0,
) -> Generator[str, None, None]:
"""Yield tokens from a response stream until gen_done/gen_error.
``read_one(timeout)`` returns the next response (or None on timeout) and
owns the queue choice — the shared resp_queue under _gen_lock, or a
per-request mailbox on the dispatcher path — so this loop stays agnostic
of which queue is read. On cancel, ``drain_on_cancel()`` consumes the
cancel ack from that same source so stale events don't leak into the
next request.
"""
while True:
resp = read_one(read_timeout)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield f"Error: {self._subprocess_crash_message(crash_context)}"
return
continue
rtype = resp.get("type", "")
if rtype == "status":
continue
# Subprocess-level error (no request_id); request-scoped failures
# arrive as gen_error below.
if rtype == "error" and not resp.get("request_id"):
yield f"Error: {resp.get('error', 'Unknown error')}"
return
if rtype == "token":
# Cancel from route (e.g. SSE connection closed).
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
drain_on_cancel()
return
yield resp.get("text", "")
elif rtype == "gen_done":
if stats_holder is not None:
stats_holder["stats"] = resp.get("stats")
return
elif rtype == "gen_error":
yield f"Error: {resp.get('error', 'Unknown error')}"
return
# ------------------------------------------------------------------
# Dispatcher — per-request mailbox routing for compare mode
# ------------------------------------------------------------------
@ -451,13 +551,12 @@ class InferenceOrchestrator:
continue
# No matching mailbox (a _gen_lock reader or orphaned). Can't
# un-get from mp.Queue, so just log.
if rtype not in ("status",):
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
rid,
rtype,
)
# un-get from mp.Queue, so just log. (status was handled above.)
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
rid,
rtype,
)
def _generate_dispatched(
self,
@ -502,30 +601,23 @@ class InferenceOrchestrator:
if image is not None:
image_b64 = self._pil_to_base64(image)
cmd = {
"type": "generate",
"request_id": request_id,
"messages": messages or [],
"system_prompt": system_prompt,
"image_base64": image_b64,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
cmd["enable_thinking"] = enable_thinking
if reasoning_effort is not None:
cmd["reasoning_effort"] = reasoning_effort
if preserve_thinking is not None:
cmd["preserve_thinking"] = preserve_thinking
cmd = self._build_generate_cmd(
request_id,
image_b64,
messages = messages,
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# Create mailbox BEFORE sending command
mailbox: queue.Queue = queue.Queue()
@ -540,36 +632,22 @@ class InferenceOrchestrator:
yield f"Error: {exc}"
return
# Read tokens from our private mailbox
def read_mailbox(timeout):
try:
return mailbox.get(timeout = timeout)
except queue.Empty:
return None
# Read tokens from our private mailbox (the dispatcher owns resp_queue).
try:
while True:
try:
resp = mailbox.get(timeout = _DISPATCH_READ_TIMEOUT)
except queue.Empty:
# Timeout — check subprocess health
if not self._ensure_subprocess_alive():
yield f"Error: {self._subprocess_crash_message('generation')}"
return
continue
rtype = resp.get("type", "")
if rtype == "token":
# Cancel from route (e.g. SSE connection closed)
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
self._drain_mailbox(mailbox, timeout = 5.0)
return
yield resp.get("text", "")
elif rtype == "gen_done":
if stats_holder is not None:
stats_holder["stats"] = resp.get("stats")
return
elif rtype == "gen_error":
yield f"Error: {resp.get('error', 'Unknown error')}"
return
yield from self._consume_token_stream(
read_mailbox,
lambda: self._drain_mailbox(mailbox, timeout = 5.0),
crash_context = "generation",
cancel_event = cancel_event,
stats_holder = stats_holder,
read_timeout = _DISPATCH_READ_TIMEOUT,
)
finally:
with self._mailbox_lock:
self._mailboxes.pop(request_id, None)
@ -636,6 +714,7 @@ class InferenceOrchestrator:
load_in_4bit: bool = True,
hf_token: Optional[str] = None,
trust_remote_code: bool = False,
approved_remote_code_fingerprint: Optional[str] = None,
gpu_ids: Optional[list[int]] = None,
) -> bool:
"""Load a model for inference.
@ -659,6 +738,7 @@ class InferenceOrchestrator:
"hf_token": hf_token or "",
"gguf_variant": getattr(config, "gguf_variant", None),
"trust_remote_code": trust_remote_code,
"approved_remote_code_fingerprint": approved_remote_code_fingerprint,
"gpu_ids": gpu_ids,
}
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
@ -716,7 +796,6 @@ class InferenceOrchestrator:
)
if resp.get("success"):
self._current_transformers_major = needed_major
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
self.models[self.active_model_name] = {
@ -738,7 +817,8 @@ class InferenceOrchestrator:
logger.info("Model '%s' loaded successfully in subprocess", model_name)
return True
else:
error = resp.get("error", "Failed to load model")
# Worker reports failures (consent gate included) under "message".
error = resp.get("message") or resp.get("error") or "Failed to load model"
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
@ -862,6 +942,7 @@ class InferenceOrchestrator:
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
**_unused,
@ -924,6 +1005,7 @@ class InferenceOrchestrator:
session_id = session_id,
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,
)
def generate_with_adapter_control(
@ -982,127 +1064,42 @@ class InferenceOrchestrator:
self._wait_dispatcher_idle()
# Serialize generation: two concurrent readers on resp_queue would
# consume and drop each other's token events.
# consume and drop each other's token events. Hold _gen_lock across the
# cmd build + send + whole stream so we stay the sole resp_queue reader.
with self._gen_lock:
yield from self._generate_locked(
request_id = str(uuid.uuid4())
image_b64 = self._pil_to_base64(image) if image is not None else None
cmd = self._build_generate_cmd(
request_id,
image_b64,
messages = messages,
system_prompt = system_prompt,
image = image,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
stats_holder = stats_holder,
)
def _generate_locked(
self,
messages: list = None,
system_prompt: str = "",
image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.0,
cancel_event = None,
use_adapter = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
) -> Generator[str, None, None]:
"""Actual generation logic — must be called under _gen_lock."""
request_id = str(uuid.uuid4())
# Convert PIL Image to base64 if needed
image_b64 = None
if image is not None:
image_b64 = self._pil_to_base64(image)
cmd = {
"type": "generate",
"request_id": request_id,
"messages": messages or [],
"system_prompt": system_prompt,
"image_base64": image_b64,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
# Only forward template kwargs the caller set, for older worker compat.
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
cmd["enable_thinking"] = enable_thinking
if reasoning_effort is not None:
cmd["reasoning_effort"] = reasoning_effort
if preserve_thinking is not None:
cmd["preserve_thinking"] = preserve_thinking
try:
self._send_cmd(cmd)
except RuntimeError as exc:
yield f"Error: {exc}"
return
# We are the only resp_queue reader (under _gen_lock).
while True:
resp = self._read_resp(timeout = 30.0)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield f"Error: {self._subprocess_crash_message('generation')}"
return
continue
rtype = resp.get("type", "")
# Status messages — skip
if rtype == "status":
continue
# Error without request_id = subprocess-level error
resp_rid = resp.get("request_id")
if rtype == "error" and not resp_rid:
yield f"Error: {resp.get('error', 'Unknown error')}"
try:
self._send_cmd(cmd)
except RuntimeError as exc:
yield f"Error: {exc}"
return
if rtype == "token":
# Cancel from route (e.g. SSE connection closed)
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
# Wait for the cancel ack so stale events don't leak into
# the next request.
self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
elif rtype == "gen_done":
if stats_holder is not None:
stats_holder["stats"] = resp.get("stats")
return
elif rtype == "gen_error":
yield f"Error: {resp.get('error', 'Unknown error')}"
return
yield from self._consume_token_stream(
self._read_resp,
lambda: self._drain_until_gen_done(timeout = 5.0),
crash_context = "generation",
cancel_event = cancel_event,
stats_holder = stats_holder,
)
def reset_generation_state(self):
"""Cancel any ongoing generation and reset state."""
@ -1138,8 +1135,6 @@ class InferenceOrchestrator:
if not self.active_model_name:
raise RuntimeError("No active model")
import uuid
request_id = str(uuid.uuid4())
cmd = {
@ -1252,8 +1247,6 @@ class InferenceOrchestrator:
return
with self._gen_lock:
import uuid
request_id = str(uuid.uuid4())
# numpy array -> list for mp.Queue serialization
@ -1282,38 +1275,12 @@ class InferenceOrchestrator:
yield f"Error: {exc}"
return
# Yield tokens — same pattern as _generate_locked
while True:
resp = self._read_resp(timeout = 30.0)
if resp is None:
if not self._ensure_subprocess_alive():
yield ("Error: " + self._subprocess_crash_message("audio input generation"))
return
continue
rtype = resp.get("type", "")
if rtype == "status":
continue
if rtype == "error" and not resp.get("request_id"):
yield f"Error: {resp.get('error', 'Unknown error')}"
return
if rtype == "token":
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
elif rtype == "gen_done":
return
elif rtype == "gen_error":
yield f"Error: {resp.get('error', 'Unknown error')}"
return
yield from self._consume_token_stream(
self._read_resp,
lambda: self._drain_until_gen_done(timeout = 5.0),
crash_context = "audio input generation",
cancel_event = cancel_event,
)
# ------------------------------------------------------------------
# Local helpers (no subprocess needed)

6
studio/backend/core/inference/safetensors_agentic.py
View file

@ -154,6 +154,7 @@ def run_safetensors_tool_loop(
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
@ -517,7 +518,9 @@ def run_safetensors_tool_loop(
else:
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
needs_confirm = bool(confirm_tool_calls)
# Bypass wins over the confirm gate at the loop level too, so a
# direct internal caller passing both flags never prompts.
needs_confirm = bool(confirm_tool_calls) and not bypass_permissions
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
start_event = decision.tool_start_event()
@ -575,6 +578,7 @@ def run_safetensors_tool_loop(
timeout = eff_timeout,
session_id = session_id,
rag_scope = rag_scope,
disable_sandbox = bypass_permissions,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", decision.tool_name, exc)

22
studio/backend/core/inference/tensor_fallback.py
View file

@ -13,7 +13,7 @@ import logging
from typing import Awaitable, Callable, Optional
from core.inference.llama_server_args import (
resolve_tensor_parallel,
_effective_tensor_parallel,
strip_split_mode_only,
)
@ -34,18 +34,20 @@ async def load_with_tensor_fallback(
True on success; it *raises* on a hard crash (llama-server aborts on some
archs / older builds), which is treated the same as a False return.
Tensor mode can be requested by the toggle or by a ``--split-mode tensor``
in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether
tensor mode is actually engaged, and it strips ``--split-mode`` from the
extras so the layer retry can't relaunch the same failing tensor load. A
non-tensor load keeps its original contract and propagates exceptions.
Tensor mode can be requested by the toggle, by a ``--split-mode tensor`` in
``extra_args`` (an allowed shadow flag), or by an inherited
``LLAMA_ARG_SPLIT_MODE=tensor`` env (load_model engages it the same way), so
the retry is keyed on whether tensor mode is actually engaged, and it forces
``--split-mode layer`` on the retry so neither leftover extras nor the
inherited tensor env can relaunch the same failing tensor load. A non-tensor
load keeps its original contract and propagates exceptions.
``cancelled()`` distinguishes a real tensor-start failure from a user
cancellation: ``attempt_load`` also returns False when the load was
cancelled, so without this the helper would restart a load the user just
cancelled.
"""
tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor)
tensor_requested = _effective_tensor_parallel(extra_args, requested_tensor)
try:
success = await attempt_load(requested_tensor, extra_args)
except Exception as exc:
@ -67,4 +69,8 @@ async def load_with_tensor_fallback(
"(this model may not support tensor parallelism)",
label,
)
return await attempt_load(False, strip_split_mode_only(extra_args))
# Force --split-mode layer (CLI wins over env) so neither leftover extras nor
# an inherited LLAMA_ARG_SPLIT_MODE=tensor can re-engage tensor and re-crash
# the retry; load_model and the child both honor the explicit layer override.
layer_extras = strip_split_mode_only(extra_args) or []
return await attempt_load(False, [*layer_extras, "--split-mode", "layer"])

337
studio/backend/core/inference/tools.py
View file

@ -316,6 +316,196 @@ def _build_safe_env(workdir: str) -> dict[str, str]:
return env
# Credential env vars dropped even in bypass mode so tool code cannot read the
# operator's keys. Over-strips on purpose (a benign var is harmless to lose).
_BYPASS_ENV_SECRET_NAMES = frozenset(
{
"HF_TOKEN",
"HF_HUB_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"HUGGINGFACE_TOKEN",
"HUGGINGFACEHUB_API_TOKEN",
"WANDB_API_KEY",
"GH_TOKEN",
"GITHUB_TOKEN",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"GEMINI_API_KEY",
"GOOGLE_API_KEY",
"GROQ_API_KEY",
"OPENROUTER_API_KEY",
"REPLICATE_API_TOKEN",
"COHERE_API_KEY",
"MISTRAL_API_KEY",
"NGC_API_KEY",
"KAGGLE_KEY",
"MYSQL_PWD", # exact name: markers use PASSWD, not PWD (PWD is the cwd var)
"LD_PRELOAD",
# Auth brokers / capability handles: not secrets by value, but they
# hand the child the operator's live agent (ssh/gpg), kube config, or
# docker daemon. Names are listed because there is no value signal to
# key off. URL config vars (HTTP_PROXY, PIP_INDEX_URL, DATABASE_URL,
# ...) are intentionally NOT name-listed: a benign proxy/index without
# credentials must keep working in bypass mode, while a credentialed
# value is dropped by _is_secret_env_value() regardless of its name.
"SSH_AUTH_SOCK",
"SSH_AGENT_PID",
"GPG_AGENT_INFO",
"GNUPGHOME",
"KUBECONFIG",
"DOCKER_HOST",
}
)
_BYPASS_ENV_SECRET_PREFIXES = ("AWS_", "AZURE_", "GOOGLE_", "GCP_", "GCLOUD_", "DYLD_")
_BYPASS_ENV_SECRET_MARKERS = (
"TOKEN",
"API_KEY",
"APIKEY",
"SECRET",
"PASSWORD",
"PASSWD",
"CREDENTIAL",
"PRIVATE_KEY",
"AUTH", # e.g. NPM_CONFIG__AUTH (npm _auth), REDISCLI_AUTH
# Azure App Service connection strings: SQLCONNSTR_/CUSTOMCONNSTR_/... and
# WEBSITE_CONTENTAZUREFILECONNECTIONSTRING carry DB/storage credentials.
"CONNSTR",
"CONNECTIONSTRING",
)
# Non-secret hardening flags that match a secret prefix/marker but must be KEPT
# so bypass mode does not silently undo an operator's opt-out. AWS_EC2_METADATA_
# DISABLED tells the AWS SDK/CLI not to pull instance-role creds from IMDS;
# dropping it would re-open that path for a bypassed tool.
_BYPASS_ENV_KEEP_NAMES = frozenset(
{
"AWS_EC2_METADATA_DISABLED",
"AWS_EC2_METADATA_V1_DISABLED",
}
)
# Matches a URL that embeds userinfo before the host, covering both
# "scheme://user:pass@host" and token-only "scheme://token@host" (and
# percent-encoded variants). The userinfo must precede the first '/', so an '@'
# in a path or query does not false-positive. Used to scrub credential-bearing
# URL values regardless of the variable's name.
_URL_USERINFO_RE = re.compile(r"://[^/\s@]+@")
# Connection-string credential fields (ADO.NET / Azure storage / Service Bus):
# "...;Password=...", "...;AccountKey=...", "...;SharedAccessKey=...". Catches
# credential-bearing values whose names dodge the name classifier. "accesskey"
# also covers Shared/Secret AccessKey via substring; the Name fields (e.g.
# SharedAccessKeyName=) do not match since "=" must follow the keyword.
_SECRET_VALUE_RE = re.compile(r"(?i)(?:password|pwd|accountkey|accesskey)\s*=\s*[^\s;]")
# Names that hold no secret value but point SDKs at the operator's real
# home/cache/config (cached tokens, cred files), defeating the HOME repoint.
# Startup always sets HF_HOME (-> $HF_HOME/token), so this is the live leak.
# Dropped in bypass mode so tools fall back to the empty repointed HOME.
_BYPASS_ENV_CRED_LOCATION_NAMES = frozenset(
{
# HF cache roots (token lives under $HF_HOME/token)
"HF_HOME",
"HF_HUB_CACHE",
"HUGGINGFACE_HUB_CACHE",
"HF_XET_CACHE",
"TRANSFORMERS_CACHE",
"HF_DATASETS_CACHE",
"HF_ASSETS_CACHE",
# XDG base dirs (resolved before $HOME)
"XDG_CONFIG_HOME",
"XDG_CACHE_HOME",
"XDG_DATA_HOME",
# explicit cred/config file pointers honoured before $HOME
"NETRC",
"PGPASSFILE",
"BOTO_CONFIG",
"PIP_CONFIG_FILE",
"CLOUDSDK_CONFIG",
"KAGGLE_CONFIG_DIR",
"DOCKER_CONFIG",
"WANDB_DIR",
"WANDB_CONFIG_DIR",
"WANDB_CACHE_DIR",
# package-manager / git / cloud config pointers to real cred files
"NPM_CONFIG_USERCONFIG",
"NPM_CONFIG_GLOBALCONFIG",
"YARN_RC_FILENAME",
"GIT_CONFIG_GLOBAL",
"GIT_CONFIG_SYSTEM",
"CARGO_HOME",
"RCLONE_CONFIG",
# auth-helper scripts that hand creds to git/ssh
"GIT_ASKPASS",
"SSH_ASKPASS",
# shell startup hook: bash -c sources $BASH_ENV (can re-export secrets)
"BASH_ENV",
# Windows: HOMEDRIVE+HOMEPATH compose a home that bypasses HOME
"HOMEDRIVE",
"HOMEPATH",
}
)
# Windows profile dirs SDKs read creds under; repointed (not dropped) since
# callers expect them present.
_BYPASS_ENV_WINDOWS_PROFILE_VARS = ("USERPROFILE", "APPDATA", "LOCALAPPDATA")
def _is_secret_env_name(name: str) -> bool:
"""True if an env var name looks like it carries a credential."""
upper = name.upper()
if upper in _BYPASS_ENV_KEEP_NAMES:
return False # non-secret hardening flag; keep it
if upper in _BYPASS_ENV_SECRET_NAMES:
return True
if any(upper.startswith(p) for p in _BYPASS_ENV_SECRET_PREFIXES):
return True
return any(marker in upper for marker in _BYPASS_ENV_SECRET_MARKERS)
def _is_cred_location_env_name(name: str) -> bool:
"""True for vars that point SDKs at the real home/cache/config (cached creds)."""
return name.upper() in _BYPASS_ENV_CRED_LOCATION_NAMES
def _is_secret_env_value(value: str) -> bool:
"""True if a value embeds credentials regardless of its name.
Catches URL userinfo (``scheme://user:token@host`` in DATABASE_URL /
PIP_INDEX_URL / HTTP_PROXY) and connection-string credential fields
(``...;Password=...`` / ``...;AccountKey=...``) whose names dodge the name
classifier.
"""
if not value:
return False
return _URL_USERINFO_RE.search(value) is not None or _SECRET_VALUE_RE.search(value) is not None
def _build_bypass_env(workdir: str) -> dict[str, str]:
"""Env for bypass exec: full host env (unrestricted) minus credential vars,
with HOME/TMPDIR repointed at the workdir so SDKs cannot read cached creds.
Note: stripping the child env is necessary but not sufficient on its own -
a same-UID child can still read the parent's environment via procfs, so
callers also harden the parent (see _harden_parent_against_proc_env_leak).
"""
env = {
k: v
for k, v in os.environ.items()
if not _is_secret_env_name(k)
and not _is_secret_env_value(v)
and not _is_cred_location_env_name(k)
}
env["HOME"] = workdir
env["TMPDIR"] = workdir
# Windows tempfile / SDKs honour TEMP/TMP, not TMPDIR; repoint all three so
# the bypassed tool writes under the per-session sandbox dir on every OS.
env["TEMP"] = workdir
env["TMP"] = workdir
# Windows SDKs read creds under the profile dirs, not $HOME; repoint set
# ones to the workdir (HOMEDRIVE/HOMEPATH are dropped above).
for var in _BYPASS_ENV_WINDOWS_PROFILE_VARS:
if var in os.environ:
env[var] = workdir
return env
def _sandbox_preexec():
"""Best-effort sandbox setup for sandboxed subprocesses (modules are
resolved at import time so the forked child runs no imports)."""
@ -377,6 +567,65 @@ def _sandbox_preexec():
pass
def _bypass_preexec():
"""Minimal pre-exec for bypass exec: os.setsid() only.
Required, not a restriction: _kill_process_tree does killpg(getpgid(child)),
so without a new session a timeout/cancel would kill the Studio server too.
"""
try:
os.setsid()
except OSError:
pass
# Hardening the Studio parent is done once (PR_SET_DUMPABLE is process-global
# and sticky); guarded so repeated bypass calls do not re-issue the prctl.
_parent_proc_hardened = False
def _harden_parent_against_proc_env_leak() -> bool:
"""Make the Studio process's /proc/<pid>/environ unreadable to its children.
Stripping the child env is not enough on Linux: a bypassed same-UID child
runs unsandboxed and can read /proc/<getppid()>/environ to recover the
tool-executing process's *unfiltered* secrets (HF_TOKEN, cloud keys, ...).
Clearing the dumpable flag (PR_SET_DUMPABLE=0) reparents this process's
/proc entries to root, so a same-UID child can no longer read its environ.
Returns True when the process is hardened or hardening is unnecessary (no
/proc leak off Linux), and False when it is needed but could not be applied
(e.g. prctl denied by a seccomp policy). Callers must fail closed - refuse
the unsandboxed exec - when this returns False, rather than running with the
parent environ still readable.
Scope: this closes the direct parent read (the demonstrated leak). It is a
mitigation, not a full boundary - a bypassed tool is unsandboxed by design,
so it can still walk /proc to a same-UID *ancestor* (e.g. the launching
shell) or read on-disk credentials by absolute path. Complete isolation
needs a separate uid / PID+mount namespace, which is out of scope here; the
UI already warns the mode is dangerous. Applied lazily on first bypass exec
so non-bypass operation is unchanged.
"""
global _parent_proc_hardened
if _parent_proc_hardened:
return True
if sys.platform != "linux":
return True # no /proc/<pid>/environ same-UID leak to close
if _libc is None:
return False # on Linux but cannot issue prctl -> cannot harden
try:
# prctl(PR_SET_DUMPABLE=4, SUID_DUMP_DISABLE=0). ctypes returns the
# syscall result (-1 on failure) and does NOT raise, so check it.
ret = _libc.prctl(4, 0, 0, 0, 0)
except (OSError, AttributeError):
return False
if ret != 0:
return False
_parent_proc_hardened = True
return True
def _get_shell_cmd(command: str) -> list[str]:
"""Return the platform-appropriate shell invocation for a command string."""
if sys.platform == "win32":
@ -524,10 +773,10 @@ RENDER_HTML_TOOL = {
"function": {
"name": "render_html",
"description": (
"Render a self-contained HTML/CSS/JavaScript artifact for the user. "
"Render a self-contained HTML/CSS/JavaScript canvas for the user. "
"Call this at most once per assistant response unless the user "
"explicitly asks for changes in that response. Future user requests "
"for new artifacts may call render_html once. Put the entire document "
"for new canvases may call render_html once. Put the entire document "
"in code, including any CSS in <style> tags and JavaScript in <script> tags."
),
"parameters": {
@ -539,7 +788,7 @@ RENDER_HTML_TOOL = {
},
"title": {
"type": "string",
"description": "Short display title for the artifact.",
"description": "Short display title for the canvas.",
},
},
"required": ["code"],
@ -705,14 +954,14 @@ def _render_html_result(arguments: dict) -> str:
if isinstance(title, str) and title.strip():
safe_title = title.strip()[:120]
return (
f"Rendered HTML artifact: {safe_title}. Do not call render_html "
f"Rendered HTML canvas: {safe_title}. Do not call render_html "
"again in this response unless the user asks for changes. For a later "
"user request for a new artifact, call render_html once."
"user request for a new canvas, call render_html once."
)
return (
"Rendered HTML artifact. Do not call render_html again in this response "
"Rendered HTML canvas. Do not call render_html again in this response "
"unless the user asks for changes. For a later user request for a new "
"artifact, call render_html once."
"canvas, call render_html once."
)
@ -723,6 +972,7 @@ def execute_tool(
timeout: int | None = _TIMEOUT_UNSET,
session_id: str | None = None,
rag_scope: dict | None = None,
disable_sandbox: bool = False,
) -> str:
"""Execute a tool by name with the given arguments; returns a string.
@ -730,6 +980,9 @@ def execute_tool(
``session_id``: optional ID for per-conversation sandbox isolation.
``rag_scope``: hidden per-request RAG context the model never sees; consumed
by ``search_knowledge_base``.
``disable_sandbox``: Bypass Permissions; run python/terminal without the
safety checks, blocklist, or resource caps (secrets still stripped). Only
affects local code tools; web_search / MCP are unchanged.
"""
logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}")
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
@ -765,9 +1018,21 @@ def execute_tool(
timeout = effective_timeout,
)
if name == "python":
return _python_exec(arguments.get("code", ""), cancel_event, effective_timeout, session_id)
return _python_exec(
arguments.get("code", ""),
cancel_event,
effective_timeout,
session_id,
disable_sandbox = disable_sandbox,
)
if name == "terminal":
return _bash_exec(arguments.get("command", ""), cancel_event, effective_timeout, session_id)
return _bash_exec(
arguments.get("command", ""),
cancel_event,
effective_timeout,
session_id,
disable_sandbox = disable_sandbox,
)
return f"Unknown tool: {name}"
@ -2242,15 +2507,28 @@ def _python_exec(
cancel_event = None,
timeout: int = _EXEC_TIMEOUT,
session_id: str | None = None,
disable_sandbox: bool = False,
) -> str:
"""Execute Python code in a subprocess sandbox."""
"""Execute Python code in a subprocess sandbox.
disable_sandbox (Bypass Permissions): skip the safety analysis and rlimit
pre-exec, and use the host env minus secrets.
"""
if not code or not code.strip():
return "No code provided."
# Validate imports and code safety
error = _check_code_safety(code)
if error:
return error
# Validate imports and code safety (skipped when the sandbox is disabled)
if not disable_sandbox:
error = _check_code_safety(code)
if error:
return error
elif not _harden_parent_against_proc_env_leak():
# Close the /proc/<parent>/environ secret-recovery path first; if it
# cannot be applied, fail closed rather than leak the parent environ.
return (
"Execution error: could not harden the Studio process against "
"/proc environment reads; refusing bypass execution."
)
tmp_path = None
workdir = _get_workdir(session_id)
@ -2270,7 +2548,7 @@ def _python_exec(
with os.fdopen(fd, "w") as f:
f.write(code)
safe_env = _build_safe_env(workdir)
safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
@ -2279,7 +2557,7 @@ def _python_exec(
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _sandbox_preexec
popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
@ -2346,19 +2624,32 @@ def _bash_exec(
cancel_event = None,
timeout: int = _EXEC_TIMEOUT,
session_id: str | None = None,
disable_sandbox: bool = False,
) -> str:
"""Execute a bash command in a subprocess sandbox."""
"""Execute a bash command in a subprocess sandbox.
disable_sandbox (Bypass Permissions): skip the command blocklist and rlimit
pre-exec, and use the host env minus secrets.
"""
if not command or not command.strip():
return "No command provided."
# Block dangerous commands (shlex + regex based)
blocked = _find_blocked_commands(command)
if blocked:
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
# Block dangerous commands (skipped when the sandbox is disabled)
if not disable_sandbox:
blocked = _find_blocked_commands(command)
if blocked:
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
elif not _harden_parent_against_proc_env_leak():
# Close the /proc/<parent>/environ secret-recovery path first; if it
# cannot be applied, fail closed rather than leak the parent environ.
return (
"Execution error: could not harden the Studio process against "
"/proc environment reads; refusing bypass execution."
)
try:
workdir = _get_workdir(session_id)
safe_env = _build_safe_env(workdir)
safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
@ -2367,7 +2658,7 @@ def _bash_exec(
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _sandbox_preexec
popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW

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

@ -13,12 +13,10 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.
from __future__ import annotations
import base64
import structlog
from loggers import get_logger
import os
import queue as _queue
import sys
import threading
import time
import traceback
from io import BytesIO
@ -28,13 +26,19 @@ from typing import Any
logger = get_logger(__name__)
from utils.hardware import apply_gpu_ids
# studio/backend root, prepended to sys.path so the spawned subprocess can
# import the utils/core packages.
_BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent)
def _ensure_backend_on_path() -> None:
if _BACKEND_PATH not in sys.path:
sys.path.insert(0, _BACKEND_PATH)
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on path for utils imports.
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
_ensure_backend_on_path()
from utils.transformers_version import activate_transformers_for_subprocess
@ -63,160 +67,97 @@ def _resize_image(img, max_size: int = 800):
def _send_response(resp_queue: Any, response: dict) -> None:
"""Send a response to the parent process."""
"""Send a response to the parent process; stamps ``ts`` if absent."""
response.setdefault("ts", time.time())
try:
resp_queue.put(response)
except (OSError, ValueError) as exc:
logger.error("Failed to send response: %s", exc)
def _clean_token(value: str | None) -> str | None:
"""Normalize an HF token: blank or whitespace-only becomes None."""
return value if value and value.strip() else None
def _build_model_config(config: dict):
"""Build a ModelConfig from the config dict."""
from utils.models import ModelConfig
model_name = config["model_name"]
hf_token = config.get("hf_token")
hf_token = hf_token if hf_token and hf_token.strip() else None
gguf_variant = config.get("gguf_variant")
mc = ModelConfig.from_identifier(
model_id = model_name,
hf_token = hf_token,
gguf_variant = gguf_variant,
hf_token = _clean_token(config.get("hf_token")),
gguf_variant = config.get("gguf_variant"),
)
if not mc:
raise ValueError(f"Invalid model identifier: {model_name}")
return mc
def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, bool] | None:
"""Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
With *model_names*, only those models' ``blobs/`` dirs are checked (faster);
accepts multiple names so LoRA loads can watch adapter + base repos at once.
*has_incomplete* is True when any ``*.incomplete`` files exist (download
active). None means state could not be determined, so callers skip stall logic.
def _needs_nemotron_trust(model_name: str, hf_token: str | None = None) -> bool:
"""Whether *model_name* is a NemotronH/Nano model that needs trust_remote_code.
NemotronH/Nano have config-parsing bugs that require it. Must NOT match
Llama-Nemotron (standard Llama arch), so also require the unsloth/ or nvidia/
namespace, and a genuine first-party Hub repo (not a local path or a spoof
name starting with "unsloth/"). The repo check is authenticated so private
first-party repos still resolve, and runs only after the cheap checks pass.
"""
mn = model_name.lower()
if not (
any(sub in mn for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (mn.startswith("unsloth/") or mn.startswith("nvidia/"))
):
return False
from utils.security.trusted_org import is_trusted_org_repo
return is_trusted_org_repo(model_name, hf_token = hf_token)
def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool:
"""Reconcile load_in_4bit with a LoRA adapter's recorded training method.
lora -> base is full precision (4bit off); qlora -> base is quantized (4bit
on); unknown method -> force off only when the base is not a -bnb-4bit repo.
A missing or unreadable adapter_config.json leaves the value unchanged.
"""
if not (mc.is_lora and mc.path):
return load_in_4bit
adapter_cfg_path = Path(mc.path) / "adapter_config.json"
if not adapter_cfg_path.exists():
return load_in_4bit
import json
try:
from huggingface_hub.constants import HF_HUB_CACHE
cache = Path(HF_HUB_CACHE)
if not cache.exists():
return (0, False)
total = 0
has_incomplete = False
blobs_dirs: list[Path] = []
if model_names:
from utils.paths import resolve_cached_repo_id_case
for name in model_names:
if not name:
continue
# Skip local filesystem paths -- HF IDs (org/model) never start
# with / . ~ or contain backslashes.
if name.startswith(("/", ".", "~")) or "\\" in name:
continue
name = resolve_cached_repo_id_case(name)
# HF cache dir format: models--org--name (slashes -> --).
cache_dir_name = "models--" + name.replace("/", "--")
blobs_dir = cache / cache_dir_name / "blobs"
if blobs_dir.exists():
blobs_dirs.append(blobs_dir)
else:
blobs_dirs = list(cache.glob("models--*/blobs"))
for bdir in blobs_dirs:
for f in bdir.iterdir():
try:
if f.is_file():
total += f.stat().st_size
if f.name.endswith(".incomplete"):
has_incomplete = True
except OSError:
pass
return (total, has_incomplete)
except Exception as e:
logger.debug("Failed to determine HF download state: %s", e)
return None
def _start_heartbeat(
resp_queue: Any,
interval: float = 30.0,
stall_timeout: float = 180.0,
xet_disabled: bool = False,
model_names: list[str] | None = None,
) -> threading.Event:
"""Start a daemon thread that sends periodic status heartbeats.
A stall is reported only when ``*.incomplete`` files are present (download
active) AND cache size hasn't changed for *stall_timeout* seconds. When the
download finishes the timer resets, so post-download init (quantization, GPU
weight load) isn't misclassified as a stall. Returns a stop event.
"""
stop = threading.Event()
transport = "https" if xet_disabled else "xet"
def _beat():
state = _get_hf_download_state(model_names)
last_size = state[0] if state is not None else 0
last_change = time.monotonic()
while not stop.wait(interval):
state = _get_hf_download_state(model_names)
now = time.monotonic()
# Skip stall logic if we cannot measure the cache.
if state is None:
_send_response(
resp_queue,
{
"type": "status",
"message": f"Loading model ({transport} transport)...",
"ts": time.time(),
},
)
continue
current_size, has_incomplete = state
if current_size != last_size:
last_size = current_size
last_change = now
# Only fire stall while .incomplete files confirm an active download;
# reset the timer otherwise so model init isn't counted as a stall.
if not has_incomplete:
last_change = now
elif now - last_change >= stall_timeout:
_send_response(
resp_queue,
{
"type": "stall",
"message": (
f"Download appears stalled ({transport} transport) "
f"-- no progress for {int(now - last_change)}s"
),
"ts": time.time(),
},
)
# fire once -- the orchestrator will kill us
return
_send_response(
resp_queue,
{
"type": "status",
"message": f"Loading model ({transport} transport)...",
"ts": time.time(),
},
with open(adapter_cfg_path) as f:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora" and load_in_4bit:
logger.info("adapter_config.json says lora — setting load_in_4bit=False")
return False
if training_method == "qlora" and not load_in_4bit:
logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
return True
if (
not training_method
and mc.base_model
and "-bnb-4bit" not in mc.base_model.lower()
and load_in_4bit
):
logger.info(
"No training method, base model has no -bnb-4bit — setting load_in_4bit=False"
)
t = threading.Thread(target = _beat, daemon = True)
t.start()
return stop
return False
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
return load_in_4bit
def _handle_load(backend, config: dict, resp_queue: Any) -> None:
@ -224,70 +165,89 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
try:
mc = _build_model_config(config)
hf_token = config.get("hf_token")
hf_token = hf_token if hf_token and hf_token.strip() else None
hf_token = _clean_token(config.get("hf_token"))
load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True))
# Auto-detect quantization for LoRA adapters.
load_in_4bit = config.get("load_in_4bit", True)
if mc.is_lora and mc.path:
import json
from pathlib import Path
adapter_cfg_path = Path(mc.path) / "adapter_config.json"
if adapter_cfg_path.exists():
try:
with open(adapter_cfg_path) as f:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora" and load_in_4bit:
logger.info("adapter_config.json says lora — setting load_in_4bit=False")
load_in_4bit = False
elif training_method == "qlora" and not load_in_4bit:
logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
load_in_4bit = True
elif not training_method:
if (
mc.base_model
and "-bnb-4bit" not in mc.base_model.lower()
and load_in_4bit
):
logger.info(
"No training method, base model has no -bnb-4bit — setting load_in_4bit=False"
)
load_in_4bit = False
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
# Auto-enable trust_remote_code only for NemotronH/Nano (config parsing
# bugs require it). Must NOT match Llama-Nemotron (standard Llama arch).
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code:
model_name = config["model_name"]
_mn_lower = model_name.lower()
if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")
):
trust_remote_code = True
logger.info(
"Auto-enabled trust_remote_code for Nemotron model: %s",
model_name,
if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
trust_remote_code = True
logger.info(
"Auto-enabled trust_remote_code for Nemotron model: %s", config["model_name"]
)
# Malware gate: a poisoned pickle deserializes during from_pretrained even
# with trust_remote_code False, so check HF's security scan (metadata-only)
# every load. For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security, security_load_subdirs
malware_targets = [config["model_name"]]
if mc.is_lora and getattr(mc, "base_model", None):
malware_targets.append(str(mc.base_model))
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
)
if _fs.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
},
)
return
# Heartbeat every 30s so the orchestrator knows we're alive during slow loads.
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
# Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH
# unless pinned-approved. For a LoRA, gate the base whose code runs.
if trust_remote_code:
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [config["model_name"]]
if mc.is_lora and getattr(mc, "base_model", None):
consent_targets.append(str(mc.base_model))
# Scan adapter + base as one unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = hf_token,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review "
f"and approve it to proceed."
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
},
)
return
# Heartbeat keeps the orchestrator's inactivity deadline alive during slow
# loads; a no-progress Xet download is reported as a stall so the parent
# can respawn over HTTP. Watch model + base repos (base is the LoRA
# download bottleneck).
from utils.hf_xet_fallback import start_watchdog
# Watch model + base repos (base download is the LoRA bottleneck).
watch_repos = [mc.identifier]
base = getattr(mc, "base_model", None)
if base and str(base) != mc.identifier:
watch_repos.append(str(base))
heartbeat_stop = _start_heartbeat(
resp_queue,
interval = 30.0,
xet_disabled = xet_disabled,
model_names = watch_repos,
heartbeat_stop = start_watchdog(
repo_ids = watch_repos,
on_stall = lambda msg: _send_response(resp_queue, {"type": "stall", "message": msg}),
on_heartbeat = lambda msg: _send_response(resp_queue, {"type": "status", "message": msg}),
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
try:
success = backend.load_model(
@ -302,7 +262,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
heartbeat_stop.set()
if success:
# Build model_info for the parent to mirror.
model_info = {
"identifier": mc.identifier,
"display_name": mc.display_name,
@ -315,13 +274,11 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {}
)
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
_context_length = _entry.get("context_length")
if _context_length is not None:
model_info["context_length"] = int(_context_length)
@ -329,12 +286,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
logger.warning("context_length forward failed: %s", _ctx_exc)
# Forward chat_template_info so the parent can classify capabilities.
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
_tpl_info = _entry.get("chat_template_info")
if isinstance(_tpl_info, dict):
model_info["chat_template_info"] = {
@ -352,7 +303,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"type": "loaded",
"success": True,
"model_info": model_info,
"ts": time.time(),
},
)
else:
@ -362,7 +312,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"type": "loaded",
"success": False,
"error": "Failed to load model",
"ts": time.time(),
},
)
@ -374,7 +323,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"success": False,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
@ -440,7 +388,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"type": "token",
"request_id": request_id,
"text": cumulative_text,
"ts": time.time(),
},
)
@ -451,7 +398,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"request_id": request_id,
# usage/timings from the MLX backend (None elsewhere).
"stats": getattr(backend, "last_generation_stats", None),
"ts": time.time(),
},
)
logger.info("Finished text generation for request_id=%s", request_id)
@ -465,7 +411,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
@ -494,7 +439,6 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
"request_id": request_id,
"wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
"sample_rate": sample_rate,
"ts": time.time(),
},
)
logger.info("Finished audio generation for request_id=%s", request_id)
@ -508,7 +452,6 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
@ -557,7 +500,6 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
"type": "token",
"request_id": request_id,
"text": text_chunk,
"ts": time.time(),
},
)
@ -566,7 +508,6 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
{
"type": "gen_done",
"request_id": request_id,
"ts": time.time(),
},
)
logger.info("Finished audio input generation for request_id=%s", request_id)
@ -580,7 +521,6 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
@ -599,7 +539,6 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
{
"type": "unloaded",
"model_name": model_name,
"ts": time.time(),
},
)
except Exception as exc:
@ -610,7 +549,6 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
"type": "unloaded",
"model_name": model_name,
"error": str(exc),
"ts": time.time(),
},
)
@ -647,25 +585,30 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
model_name = config["model_name"]
# ── 0. MLX fast-path — skip torch/transformers ──
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
_ensure_backend_on_path()
from utils.hardware import hardware as _hw
_hw.detect_hardware()
if _hw.DEVICE == _hw.DeviceType.MLX:
# Non-fatal: fall through with the installed version, but log the cause
# instead of swallowing it (issue #6103).
try:
_activate_transformers_version(model_name)
except Exception:
pass
except Exception as exc:
logger.warning(
"Failed to activate transformers version for '%s' (MLX inference); "
"inference may fail if this model requires a specific version. Error: %s",
model_name,
exc,
)
try:
from core.inference.mlx_inference import MLXInferenceBackend
backend = MLXInferenceBackend()
_send_response(
resp_queue,
{"type": "status", "message": "Loading model...", "ts": time.time()},
{"type": "status", "message": "Loading model..."},
)
_handle_load(backend, config, resp_queue)
except Exception as exc:
@ -675,7 +618,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
"type": "error",
"error": f"MLX inference init failed: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
@ -707,7 +649,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
elif cmd_type == "reset":
cancel_event.set()
backend.reset_generation_state()
_send_response(resp_queue, {"type": "reset_ack", "ts": time.time()})
_send_response(resp_queue, {"type": "reset_ack"})
elif cmd_type == "status":
_send_response(
resp_queue,
@ -719,7 +661,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
for k, v in backend.models.items()
},
"loading": list(backend.loading_models),
"ts": time.time(),
},
)
elif cmd_type == "shutdown":
@ -733,7 +674,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
"request_id": cmd.get("request_id"),
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
@ -748,7 +688,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
@ -772,13 +711,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
{
"type": "status",
"message": "Importing Unsloth...",
"ts": time.time(),
},
)
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
_ensure_backend_on_path()
from core.inference.inference import InferenceBackend
@ -793,7 +729,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
@ -807,7 +742,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
{
"type": "status",
"message": "Loading model...",
"ts": time.time(),
},
)
@ -820,7 +754,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
"type": "error",
"error": f"Failed to initialize inference backend: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
@ -851,7 +784,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "load":
# Unload the current model before loading the new one.
if backend.active_model_name:
backend.unload_model(backend.active_model_name)
_handle_load(backend, cmd, resp_queue)
@ -879,7 +811,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
resp_queue,
{
"type": "reset_ack",
"ts": time.time(),
},
)
@ -898,22 +829,20 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
for name, info in backend.models.items()
},
"loading": list(backend.loading_models),
"ts": time.time(),
},
)
elif cmd_type == "shutdown":
logger.info("Shutdown command received, exiting")
for model_name in list(backend.models.keys()):
for name in list(backend.models.keys()):
try:
backend.unload_model(model_name)
backend.unload_model(name)
except Exception:
pass
_send_response(
resp_queue,
{
"type": "shutdown_ack",
"ts": time.time(),
},
)
return
@ -925,7 +854,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
{
"type": "error",
"error": f"Unknown command type: {cmd_type}",
"ts": time.time(),
},
)
@ -937,6 +865,5 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
"type": "error",
"error": f"Command '{cmd_type}' failed: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)

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

@ -30,6 +30,7 @@ import numpy as np
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
from utils.process_lifetime import child_popen_kwargs
from . import config
@ -267,6 +268,7 @@ class LlamaServerBackend:
text = True,
env = env,
**windows_hidden_subprocess_kwargs(),
**child_popen_kwargs(),
)
self._process = proc
self._port = port

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

@ -44,6 +44,62 @@ from utils.hardware import (
# doesn't crash on RDNA2/RDNA3 with older ROCm wheels.
if hasattr(torch._dynamo.config, "recompile_limit"):
torch._dynamo.config.recompile_limit = 64
def _ensure_real_packages(*names: str) -> None:
"""Stop `import <name>` from binding to a namespace-package shadow.
A directory named like the package but missing __init__.py on sys.path (a
stray checkout, a partial clone, or a polluted PYTHONPATH) makes the path
finder return a namespace package, so `from unsloth import FastLanguageModel`
dies with "cannot import name ... (unknown location)". A normal
site-packages install always wins, so only source/editable installs are
exposed. Drop the offending entries, import the real packages, then restore
sys.path so other modules on those entries keep importing.
"""
import importlib
import importlib.util
bad: set = set()
shadowed: list = []
for name in names:
try:
spec = importlib.util.find_spec(name)
except (ImportError, ValueError, AttributeError):
spec = None
# a real package exposes its __init__ via spec.origin; a namespace
# shadow has origin None/"namespace" and only search locations
if spec is None or spec.origin not in (None, "namespace"):
continue
dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])}
if not dirs:
continue
shadowed.append(name)
for entry in sys.path:
pkg = os.path.join(entry or os.getcwd(), name)
if os.path.realpath(pkg) in dirs and not os.path.isfile(
os.path.join(pkg, "__init__.py")
):
bad.add(entry)
if not bad:
return
saved = list(sys.path)
sys.path[:] = [e for e in sys.path if e not in bad]
for name in shadowed:
for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]:
del sys.modules[cached]
try:
importlib.invalidate_caches()
# Import unsloth before unsloth_zoo (names are dependency-first):
# unsloth.__init__ runs ROCm/Windows bnb fixes before it imports zoo,
# so importing zoo first here would skip them. Repeat import is a no-op.
for name in reversed(names):
importlib.import_module(name)
finally:
sys.path[:] = saved
_ensure_real_packages("unsloth_zoo", "unsloth")
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
@ -2175,6 +2231,9 @@ class UnslothTrainer:
for dataset_file in file_paths:
if os.path.isabs(dataset_file):
file_path = dataset_file
elif os.path.exists(dataset_file):
# A path relative to the current working directory (CLI usage)
file_path = os.path.abspath(dataset_file)
else:
file_path = str(resolve_dataset_path(dataset_file))
@ -2722,6 +2781,74 @@ class UnslothTrainer:
logger.error(f"Failed to start training thread: {e}")
return False
def _chat_template_renders_empty(self) -> bool:
"""True when the chat template renders a sample to empty text (base-model signature)."""
try:
ds = getattr(self.trainer, "train_dataset", None)
if ds is None or len(ds) == 0:
return False
row = ds[0]
messages = row.get("messages") if isinstance(row, dict) else None
if not messages:
return False
tok = self.tokenizer
if not hasattr(tok, "apply_chat_template"):
return False
rendered = tok.apply_chat_template(
messages, tokenize = False, add_generation_prompt = False
)
return not (isinstance(rendered, str) and rendered.strip())
except Exception:
return False
def _preflight_first_batch(self) -> Optional[str]:
"""Validate the first real batch before train(). A base model whose chat
template renders empty yields empty float32 input_ids that crash the
embedding on step 1; catch it here. Returns None for a valid batch."""
try:
loader = self.trainer.get_train_dataloader()
batch = next(iter(loader))
except StopIteration:
return None
except Exception as e:
model = self.model_name or "this model"
return (
f"Cannot start training: failed to build the first training batch "
f"for '{model}': {e}"
)
try:
input_ids = batch["input_ids"] if "input_ids" in batch else None
except Exception:
input_ids = getattr(batch, "input_ids", None)
if input_ids is None:
return None # some collators omit input_ids
seq_len = input_ids.shape[-1] if input_ids.ndim > 0 else 0
if not (input_ids.is_floating_point() or input_ids.numel() == 0 or seq_len == 0):
return None
model = self.model_name or "this model"
if self._chat_template_renders_empty():
low = model.lower()
suffix = (
f" such as '{model}-Instruct'"
if not any(t in low for t in ("instruct", "chat", "-it", "_it"))
else ""
)
return (
f"Cannot start training: the chat template for '{model}' produced "
f"no text for your dataset, so the first batch had empty token IDs. "
f"'{model}' looks like a base (pretrained) model without a chat "
f"template suited to conversational fine-tuning. Use the "
f"instruction-tuned variant{suffix} or provide a chat template."
)
return (
f"Cannot start training: the first batch produced invalid token IDs "
f"(dtype={input_ids.dtype}, length={seq_len}). Check that your dataset "
f"columns are mapped correctly for '{model}'."
)
def _train_worker(self, dataset: Dataset, **training_args):
"""Worker function for training (runs in separate thread)"""
try:
@ -3415,6 +3542,13 @@ class UnslothTrainer:
training_args.get("max_steps", 0),
)
# ========== START TRAINING ==========
# Fail fast on an invalid first batch (empty/float input_ids) vs a step-1 crash.
preflight_error = self._preflight_first_batch()
if preflight_error:
logger.error(preflight_error)
self._update_progress(error = preflight_error, is_training = False)
return
self._update_progress(total_steps = total_steps, status_message = "Starting training...")
logger.info("Starting training...\n")
self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint"))

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

@ -217,17 +217,36 @@ class TrainingBackend:
self._db_config: Optional[dict] = None
self._db_started_at: Optional[str] = None
# Xet -> HTTP model-load fallback state (config kept for the respawn).
self._last_full_config: Optional[dict] = None
self._in_model_load: bool = False
self._xet_fallback_used: bool = False
self._needs_xet_respawn: bool = False
logger.info("TrainingBackend initialized (subprocess mode)")
# ------------------------------------------------------------------
# Public API (called by routes/training.py)
# ------------------------------------------------------------------
def start_training(self, job_id: str, **kwargs) -> bool:
def start_training(
self,
job_id: str,
*,
before_spawn = None,
**kwargs,
) -> bool:
"""Spawn a subprocess to run the full training pipeline.
All kwargs are serialized into a config dict and sent to the worker.
Returns True if the subprocess started successfully.
``before_spawn`` is an optional no-arg callable run after synchronous
validation (start guards, config build, explicit gpu_ids) passes but
before VRAM-dependent auto GPU-selection and the spawn -- used to free
VRAM (e.g. unload chat) without tearing it down on a refused start, while
still letting auto-selection place training against the freed memory.
Hook failures never block the start.
"""
with self._lock:
if self._proc is not None and self._proc.is_alive():
@ -309,34 +328,61 @@ class TrainingBackend:
"tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
"resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
"trust_remote_code": kwargs.get("trust_remote_code", False),
"approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"),
"gpu_ids": kwargs.get("gpu_ids"),
"s3_config": kwargs.get("s3_config"),
# Flipped to True only by the HTTP-fallback respawn after a stall.
"disable_xet": kwargs.get("disable_xet", False),
}
# Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request.
if config["training_type"] == "Full Finetuning":
config["load_in_4bit"] = False
# Spawn into locals so state is untouched on failure.
# Split GPU validation from placement around the VRAM hook:
# * Explicit gpu_ids are validated here (raises -> the route returns 400
# before any teardown) and their placement is VRAM-independent, so it
# stays correct after the hook frees memory.
# * Auto-selection ranks GPUs by *free* VRAM, so it is deferred until
# after the hook frees export/chat -- otherwise it could pin training
# onto a GPU the hook is about to clear (and onto a kept chat model).
from utils.hardware import hardware as _hw
gpu_ids = kwargs.get("gpu_ids")
gpu_selection_kwargs = dict(
model_name = config["model_name"],
hf_token = config["hf_token"] or None,
training_type = config["training_type"],
load_in_4bit = config["load_in_4bit"],
batch_size = config.get("batch_size", 4),
max_seq_length = config.get("max_seq_length", 2048),
lora_rank = config.get("lora_r", 16),
target_modules = config.get("target_modules"),
gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
optimizer = config.get("optim", "adamw_8bit"),
)
defer_auto_selection = False
if _hw.DEVICE == _hw.DeviceType.MLX:
config["resolved_gpu_ids"] = None
config["gpu_selection"] = None
elif gpu_ids:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(gpu_ids, **gpu_selection_kwargs)
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
else:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
kwargs.get("gpu_ids"),
model_name = config["model_name"],
hf_token = config["hf_token"] or None,
training_type = config["training_type"],
load_in_4bit = config["load_in_4bit"],
batch_size = config.get("batch_size", 4),
max_seq_length = config.get("max_seq_length", 2048),
lora_rank = config.get("lora_r", 16),
target_modules = config.get("target_modules"),
gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
optimizer = config.get("optim", "adamw_8bit"),
)
defer_auto_selection = True
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
if defer_auto_selection:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs)
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
@ -358,6 +404,9 @@ class TrainingBackend:
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
return False
@ -386,6 +435,11 @@ class TrainingBackend:
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
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._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Assign subprocess handles after state reset.
self._event_queue = event_queue
@ -446,6 +500,100 @@ class TrainingBackend:
output_dir,
)
def _handle_stall_event(self, event: dict) -> None:
"""A worker reported a no-progress download stall.
On the first model-load, terminate the worker so the pump loop respawns it
over HTTP. A later stall (already on HTTP, or outside model-load) surfaces
as an error instead.
"""
msg = event.get("message", "Download stalled")
with self._lock:
recover = self._in_model_load and not self._xet_fallback_used
proc = self._proc
if recover:
self._xet_fallback_used = True
self._needs_xet_respawn = True
self._progress.status_message = (
"Model download stalled on Xet; retrying over HTTP..."
)
else:
self._progress.error = self._progress.error or (
"Model download stalled even over HTTP -- check your network connection"
)
if recover:
logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg)
else:
logger.error("Training download stalled with no further fallback: %s", msg)
# Terminate either way so the pump loop proceeds (respawn or finalize).
if proc is not None and proc.is_alive():
proc.terminate()
def _respawn_worker_disable_xet(self) -> None:
"""Respawn the worker once with HF_HUB_DISABLE_XET=1 after a model-load
stall. Runs on the exiting pump thread, reaps the terminated worker, and
starts a fresh worker + pump. DB/progress run-state is preserved so the
history row is not duplicated; the new worker re-formats and loads over HTTP.
"""
config = self._last_full_config
if config is None:
logger.error("Cannot respawn training worker: no stored config")
return
with self._lock:
old_proc = self._proc
if old_proc is not None:
old_proc.join(timeout = 5.0)
if old_proc.is_alive():
old_proc.kill()
old_proc.join(timeout = 2.0)
config = {**config, "disable_xet": True}
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
try:
with 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,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
new_pump.start()
def is_training_active(self) -> bool:
"""Check if training is currently active."""
with self._lock:
@ -566,6 +714,14 @@ class TrainingBackend:
for e in self._drain_queue(self._event_queue):
self._handle_event(e)
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Runs on THIS exiting pump thread and starts a fresh pump (never joins
# the current thread); DB run-state is preserved.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
@ -597,6 +753,19 @@ class TrainingBackend:
db_action: Optional[str] = None
db_action_kwargs: dict = {}
# Model-load lifecycle + stall recovery (no DB metrics); handled first.
if etype == "model_load_started":
with self._lock:
self._in_model_load = True
return
if etype == "model_load_completed":
with self._lock:
self._in_model_load = False
return
if etype == "stall":
self._handle_stall_event(event)
return
with self._lock:
if etype == "progress":
self._progress.step = event.get("step", self._progress.step)

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

@ -1087,6 +1087,27 @@ def _activate_transformers_version(model_name: str) -> None:
activate_transformers_for_subprocess(model_name)
def _activate_transformers_version_or_warn(model_name: str) -> None:
"""Activate the required transformers version for the MLX fast-path.
Unlike the non-MLX path (which treats activation failure as fatal and
reports it via the event queue), the MLX path is intentionally non-fatal:
it falls through with whatever transformers version is installed. The
failure used to be swallowed by a bare ``except: pass``, leaving no trace
and only a confusing downstream crash. Log a warning instead so the cause
is visible, while keeping the fall-through behaviour.
"""
try:
_activate_transformers_version(model_name)
except Exception as exc:
logger.warning(
"Failed to activate transformers version for '%s' (MLX); "
"training may fail if this model requires a specific version. Error: %s",
model_name,
exc,
)
def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]:
if width <= 0 or height <= 0 or target <= 0:
return width, height
@ -1458,6 +1479,75 @@ def _run_mlx_training(event_queue, stop_queue, config):
model_random_state = random_seed if _model_seed is None else int(_model_seed)
_lora_seed = config.get("lora_random_state")
lora_random_state = random_seed if _lora_seed is None else int(_lora_seed)
# Malware gate (MLX): a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) first.
# For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security
malware_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
from utils.security import security_load_subdirs
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
)
if _fs.blocked:
_send(
"error",
error = _fs.reason,
error_kind = "malware_blocked",
security = _fs.response_payload(),
)
return
# Consent gate (MLX): the CUDA path gates in run_training_process, but MLX returns
# before that, so scan auto_map code here before FastMLXModel runs it. Block
# CRITICAL/HIGH unless pinned-approved; for a LoRA, gate the base whose code runs.
if config.get("trust_remote_code", False):
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
base_model = get_base_model_from_lora_identifier(
model_name, config.get("hf_token") or None
)
if base_model:
consent_targets.append(base_model)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = hf_token,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
_send(
"error",
error = (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review it and "
f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
),
error_kind = "remote_code_blocked",
remote_code = _rc.response_payload(),
)
return
model, tokenizer = FastMLXModel.from_pretrained(
model_name,
load_in_4bit = config.get("load_in_4bit", True),
@ -1992,6 +2082,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # before imports
# HTTP-fallback respawn: disable Xet before any huggingface_hub import (the
# var is read at import time). Mirrors core/inference/worker.py.
from utils.hf_xet_fallback import child_should_disable_xet
if child_should_disable_xet(config):
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
print(
"Xet transport disabled for this training worker (HF_HUB_DISABLE_XET=1).",
file = sys.stderr,
flush = True,
)
# Offline auto-detect: skip ~25s of HF retries per call when DNS is dead.
if "HF_HUB_OFFLINE" not in os.environ:
import socket as _socket
@ -2057,11 +2160,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
return
# Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.)
# before any transformers/mlx-lm imports in _run_mlx_training.
try:
_activate_transformers_version(model_name)
except Exception:
pass # Non-fatal: fall through with whatever version is installed
# Must happen before any transformers/mlx-lm imports in _run_mlx_training.
# Non-fatal: fall through with whatever version is installed, but log
# the failure instead of swallowing it (issue #6103).
_activate_transformers_version_or_warn(model_name)
try:
_run_mlx_training(event_queue, stop_queue, config)
except Exception as exc:
@ -2093,11 +2195,16 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# NemotronH needs trust_remote_code=True to work around config-parsing bugs.
# Other 5.x models are native and don't need it (it bypasses the compiler,
# disabling fused CE). Must NOT match Llama-Nemotron (standard Llama arch).
from utils.security.trusted_org import is_trusted_org_repo
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_lowered = model_name.lower()
if (
any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/"))
# Confirm a genuine first-party Hub repo (not a local/spoofed name starting
# with "unsloth/"); authenticated so private first-party repos resolve.
and is_trusted_org_repo(model_name, hf_token = config.get("hf_token") or None)
and not config.get("trust_remote_code", False)
):
config["trust_remote_code"] = True
@ -2106,6 +2213,81 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
model_name,
)
# 1a. Malware gate: a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) first.
# For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security
malware_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
from utils.security import security_load_subdirs
_ls_hf = config.get("hf_token") or None
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = _ls_hf, load_subdirs = security_load_subdirs(target, _ls_hf)
)
if _fs.blocked:
event_queue.put(
{
"type": "error",
"error": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
}
)
return
# 1a'. Consent gate: scan auto_map Python before it runs; refuse CRITICAL/HIGH
# unless pinned-approved.
if config.get("trust_remote_code", False):
from utils.security import evaluate_remote_code_consent_for_targets
# A LoRA adapter's base is where custom code runs, so gate it too.
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
base_model = get_base_model_from_lora_identifier(
model_name, config.get("hf_token") or None
)
if base_model:
consent_targets.append(base_model)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = config.get("hf_token") or None,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
event_queue.put(
{
"type": "error",
"error": (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review it and "
f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
}
)
return
# ── 1b. Install fast-path kernel libraries for the chosen model.
# 1) causal-conv1d ALWAYS runs eagerly via the substring path: some SSM
# modeling files lazy_load it without calling is_causal_conv1d_available.
@ -2780,18 +2962,33 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
cpt_trains_embeddings = False
# ── 4c. Load training model (uses VRAM — dataset already formatted) ──
# Watchdog lets the parent recover a stalled Xet download via respawn.
_send_status(event_queue, "Loading model...")
success = trainer.load_model(
model_name = model_name,
max_seq_length = config["max_seq_length"],
load_in_4bit = config["load_in_4bit"],
full_finetuning = not use_lora,
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),
is_dataset_audio = config.get("is_dataset_audio", False),
trust_remote_code = config.get("trust_remote_code", False),
gpu_ids = config.get("resolved_gpu_ids"),
from utils.hf_xet_fallback import start_watchdog
event_queue.put({"type": "model_load_started", "ts": time.time()})
_load_watchdog_stop = start_watchdog(
repo_ids = [model_name],
on_stall = lambda msg: event_queue.put(
{"type": "stall", "message": msg, "ts": time.time()}
),
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
try:
success = trainer.load_model(
model_name = model_name,
max_seq_length = config["max_seq_length"],
load_in_4bit = config["load_in_4bit"],
full_finetuning = not use_lora,
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),
is_dataset_audio = config.get("is_dataset_audio", False),
trust_remote_code = config.get("trust_remote_code", False),
gpu_ids = config.get("resolved_gpu_ids"),
)
finally:
_load_watchdog_stop.set()
event_queue.put({"type": "model_load_completed", "ts": time.time()})
if not success or trainer.should_stop:
if trainer.should_stop:
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
@ -3110,6 +3307,73 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
# Malware gate (embedding): a poisoned pickle deserializes on load even with
# trust_remote_code False, so check HF's security scan (metadata-only) first.
# For a LoRA, gate the base whose weights deserialize.
from utils.security import evaluate_file_security
malware_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
_base = get_base_model_from_lora_identifier(model_name, hf_token)
if _base:
malware_targets.append(_base)
except Exception as exc:
logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
from utils.security import security_load_subdirs
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
)
if _fs.blocked:
event_queue.put(
{
"type": "error",
"error": _fs.reason,
"error_kind": "malware_blocked",
"security": _fs.response_payload(),
"ts": time.time(),
}
)
return
# Consent gate (embedding): scan any auto_map code before it runs; block
# CRITICAL/HIGH unless pinned-approved. A no-op without auto_map.
if config.get("trust_remote_code", False):
from utils.security import evaluate_remote_code_consent_for_targets
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
_cbase = get_base_model_from_lora_identifier(model_name, hf_token)
if _cbase:
consent_targets.append(_cbase)
except Exception as exc:
logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
# Scan adapter + base as one combined unit, pinned by a single fingerprint.
_rc = evaluate_remote_code_consent_for_targets(
consent_targets,
hf_token = hf_token,
trust_remote_code = True,
approved_fingerprint = config.get("approved_remote_code_fingerprint"),
)
if _rc.blocked:
event_queue.put(
{
"type": "error",
"error": (
f"Model '{_rc.model_name}' ships custom code flagged as "
f"{_rc.max_severity} by the security scan. Review it and "
f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
),
"error_kind": "remote_code_blocked",
"remote_code": _rc.response_payload(),
"ts": time.time(),
}
)
return
model = FastSentenceTransformer.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,

8
studio/backend/hub/schemas/downloads.py
View file

@ -25,8 +25,8 @@ class DownloadModelRequest(BaseModel):
description = "Quantization label (e.g. 'Q4_K_M'). Required for GGUF repos.",
)
use_xet: bool = Field(
False,
description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.",
True,
description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
)
@ -125,8 +125,8 @@ class DownloadDatasetRequest(BaseModel):
repo_id: str = Field(..., description = "HuggingFace dataset repo ID")
use_xet: bool = Field(
False,
description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.",
True,
description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
)

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

@ -157,7 +157,8 @@ async def download_dataset_response(
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
key = _download_job_key(repo_id)
transport = download_lifecycle.resolve_transport(body.use_xet)
use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
transport = download_lifecycle.resolve_transport(use_xet)
claimed, claim_state = _registry.claim(
key,
@ -183,7 +184,7 @@ async def download_dataset_response(
spawn = lambda: download_lifecycle.spawn_worker(
["--repo-id", repo_id, "--dataset"],
hf_token,
use_xet = body.use_xet,
use_xet = use_xet,
),
hf_token = hf_token,
label = repo_id,

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

@ -3,6 +3,7 @@
from __future__ import annotations
import logging
import os
import signal
import subprocess
@ -20,11 +21,27 @@ from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import EXIT_CANCELLED
from hub.utils.state_dir import RepoType
logger = logging.getLogger(__name__)
def backend_dir() -> Path:
return Path(__file__).resolve().parent.parent.parent
def resolve_effective_use_xet(use_xet: bool) -> bool:
"""Downgrade an Xet request to HTTP when hf_xet is unavailable, so a defaulted
or explicit Xet request never hard-fails on installs without the Xet extra."""
if not use_xet:
return False
reason = download_registry.download_transport_unavailable_reason(
download_registry.TRANSPORT_XET
)
if reason is None:
return True
logger.warning("Xet transport unavailable, falling back to HTTP: %s", reason)
return False
def resolve_transport(use_xet: bool) -> str:
transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
unavailable_reason = download_registry.download_transport_unavailable_reason(transport)

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

@ -64,7 +64,7 @@ def _spawn_download_worker(
repo_id: str,
variant: Optional[str],
hf_token: Optional[str],
use_xet: bool = False,
use_xet: bool = True,
protected_blob_hashes: Optional[frozenset[str]] = None,
) -> subprocess.Popen:
args = ["--repo-id", repo_id]
@ -96,7 +96,8 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
detail = f"Invalid gguf_variant: {variant!r}",
)
key = _download_job_key(repo_id, variant)
transport = download_lifecycle.resolve_transport(body.use_xet)
use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
transport = download_lifecycle.resolve_transport(use_xet)
variant_blob_hashes = frozenset()
variant_progress_blob_hashes = frozenset()
completed_baseline_bytes = 0
@ -171,7 +172,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
repo_id,
variant,
hf_token,
use_xet = body.use_xet,
use_xet = use_xet,
protected_blob_hashes = protected_blob_hashes,
),
hf_token = hf_token,

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

@ -26,6 +26,7 @@ from hub.utils.hf_cache_state import (
from hub.utils.gguf import (
extract_quant_label,
iter_hf_cache_snapshots,
is_big_endian_gguf_path,
list_gguf_variants,
list_gguf_variants_from_hf_cache,
list_local_gguf_variants,
@ -482,7 +483,10 @@ async def get_gguf_variants_response(
by_filename[key] = max(by_filename.get(key, 0), size)
if _is_mmproj_filename(f.name) or _is_mtp_drafter_path(rel):
continue
q = extract_quant_label(rel).lower()
q = extract_quant_label(rel)
if is_big_endian_gguf_path(rel, q):
continue
q = q.lower()
by_quant[q] = by_quant.get(q, 0) + size
if by_filename:
cached_filenames_by_snapshot.append(by_filename)
@ -521,33 +525,51 @@ async def get_gguf_variants_response(
return False
def _any_mmproj_cached(filenames: frozenset[str]) -> bool:
return any(
if any(
by_filename.get(name.lower()) is not None
for by_filename in cached_filenames_by_snapshot
for name in filenames
):
return True
return any(
_is_mmproj_filename(name.rsplit("/", 1)[-1])
for by_filename in cached_filenames_by_snapshot
for name in by_filename
)
def _quant_bytes_present(quant: str, size_bytes: int) -> bool:
# Small rounding tolerance for symlinks vs real sizes.
if size_bytes <= 0:
return False
return any(
by_quant.get(quant, 0) >= size_bytes * 0.99
for by_quant in cached_quant_bytes_by_snapshot
)
def _is_fully_downloaded(variant) -> bool:
requirement = requirements_by_quant.get(variant.quant.lower())
if requirement is None:
if variant.size_bytes == 0:
return False
quant = variant.quant.lower()
# Allow small rounding tolerance (symlinks vs real sizes).
return any(
by_quant.get(quant, 0) >= variant.size_bytes * 0.99
for by_quant in cached_quant_bytes_by_snapshot
quant = variant.quant.lower()
requirement = requirements_by_quant.get(quant)
# Vision repos ship an mmproj adapter; any precision on disk suffices.
if (
requirement is not None
and _filenames_cached(
requirement.main_filenames,
requirement.main_size_bytes,
)
and (
not requirement.mmproj_filenames
or _any_mmproj_cached(requirement.mmproj_filenames)
)
if not _filenames_cached(
requirement.main_filenames,
requirement.main_size_bytes,
):
return True
# Byte fallback so a present quant isn't demoted by a filename mismatch;
# vision repos still need an mmproj cached (any precision).
if not _quant_bytes_present(quant, variant.size_bytes):
return False
# Vision repos ship an mmproj adapter per variant. Any mmproj
# precision on disk suffices (the loader picks whichever is present);
# requiring the API-preferred one would falsely demote variants.
if requirement.mmproj_filenames and not _any_mmproj_cached(
requirement.mmproj_filenames,
if (
requirement is not None
and requirement.mmproj_filenames
and not _any_mmproj_cached(requirement.mmproj_filenames)
):
return False
return True

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

@ -177,6 +177,17 @@ def _scan_models_dir(
return found
def _safe_is_dir(path: Path) -> bool:
"""``Path.is_dir()`` treating an unreadable path (``PermissionError`` /
``OSError`` on a restricted ``~/.cache/huggingface/hub``) as "not a
directory", so the inventory skips that source instead of 500ing the Hub page.
"""
try:
return path.is_dir()
except OSError:
return False
def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
blobs_dir = repo_dir / "blobs"
if not blobs_dir.is_dir():
@ -191,7 +202,7 @@ def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
if not cache_dir.exists() or not cache_dir.is_dir():
if not _safe_is_dir(cache_dir):
return []
discovered: List[tuple[Path, str, Optional[float]]] = []
@ -502,11 +513,11 @@ async def _collect_models_from_default_sources(
local_models = await _scan_source("models directory", _scan_models_dir, models_root)
local_models += await _scan_source("HF cache", _scan_hf_cache, hf_cache_dir)
if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve():
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)
if (
hf_default.is_dir()
_safe_is_dir(hf_default)
and hf_default.resolve() != hf_cache_dir.resolve()
and hf_default.resolve() != legacy_hf.resolve()
):
@ -621,9 +632,9 @@ async def list_local_models_response(models_dir: str = "./models") -> LocalModel
ollama_dirs = ollama_model_dirs()
allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir]
if legacy_hf.is_dir():
if _safe_is_dir(legacy_hf):
allowed_roots.append(legacy_hf)
if hf_default.is_dir():
if _safe_is_dir(hf_default):
allowed_roots.append(hf_default)
allowed_roots.extend([studio_root(), outputs_root()])

27
studio/backend/hub/tests/test_download_lifecycle.py Normal file
View file

@ -0,0 +1,27 @@
# 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 hub.services import download_lifecycle
def _set_xet_reason(monkeypatch, reason):
monkeypatch.setattr(
download_lifecycle.download_registry,
"download_transport_unavailable_reason",
lambda _transport: reason,
)
def test_resolve_effective_use_xet_keeps_http_when_not_requested(monkeypatch):
_set_xet_reason(monkeypatch, "should not be consulted")
assert download_lifecycle.resolve_effective_use_xet(False) is False
def test_resolve_effective_use_xet_keeps_xet_when_available(monkeypatch):
_set_xet_reason(monkeypatch, None)
assert download_lifecycle.resolve_effective_use_xet(True) is True
def test_resolve_effective_use_xet_downgrades_when_xet_unavailable(monkeypatch):
_set_xet_reason(monkeypatch, "Xet transport is unavailable because hf_xet is not installed.")
assert download_lifecycle.resolve_effective_use_xet(True) is False

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

@ -78,6 +78,28 @@ class TestExtractQuantToken:
assert labels == {"Q4_K_M", "Q8_0"}
def test_big_endian_detection_ignores_model_name_be_token():
assert gguf.is_big_endian_gguf_path("model-Q4_K_M-be.gguf", "Q4_K_M")
assert gguf.is_big_endian_gguf_path("model-Q4_K_M_be_infill.gguf", "Q4_K_M")
assert not gguf.is_big_endian_gguf_path("foo-be-Q4_K_M.gguf", "Q4_K_M")
assert not gguf.is_big_endian_gguf_path("Q4_K_M/foo-be.gguf", "Q4_K_M")
assert gguf.pick_best_gguf(["model-Q4_K_M-be.gguf", "model-Q4_K_M.gguf"]) == (
"model-Q4_K_M.gguf"
)
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)
variants, has_vision = gguf.list_local_gguf_variants(str(tmp_path))
assert has_vision is False
assert [(v.quant, v.filename, v.size_bytes) for v in variants] == [
("Q4_K_M", "model-Q4_K_M.gguf", 10)
]
@pytest.mark.parametrize("repo_id", ["bert-base-uncased", "owner/repo"])
def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id):
assert paths.is_valid_repo_id(repo_id)
@ -304,6 +326,22 @@ def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj():
)
def test_gguf_variant_requirements_skip_big_endian_sibling():
requirements = gguf_variants._build_gguf_variant_requirements(
[
_sibling("model-Q4_K_M-be.gguf", 100, "main-be"),
_sibling("model-Q4_K_M.gguf", 10, "main-le"),
]
)
req = requirements["q4_k_m"]
assert req.main_size_bytes == 10
assert req.main_hashes == frozenset({"main-le"})
assert req.main_filenames == frozenset({"model-Q4_K_M.gguf"})
assert req.target_filenames == ("model-Q4_K_M.gguf",)
def test_worker_gguf_variant_plan_matches_service_requirement(monkeypatch):
siblings = [
_sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),

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

@ -109,6 +109,33 @@ def is_gguf_filename(filename: str) -> bool:
return filename.lower().endswith(".gguf")
_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE)
def is_big_endian_gguf_path(path: str, quant: str = "") -> bool:
normalized = path.replace("\\", "/")
name = normalized.rsplit("/", 1)[-1]
stem = name.rsplit(".", 1)[0].lower()
quant_key = quant.strip().lower()
quant_index = stem.find(quant_key) if quant_key else -1
parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else ""
quant_in_parent_only = (
bool(parent)
and quant_index < 0
and (
(quant_key and quant_key in parent)
or (not quant_key and _GGUF_QUANT_RE.search(parent) is not None)
)
)
for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem):
if quant_index >= 0 and quant_index < match.start():
return True
tail = stem[match.end() :].lstrip("._-")
if not tail or _GGUF_QUANT_RE.search(tail) is None:
return not quant_in_parent_only
return False
# Cap recursive walks so a huge or system path cannot run unbounded.
_MAX_LOCAL_SCAN_ENTRIES = 100_000
@ -143,7 +170,10 @@ def pick_best_gguf(filenames: list[str]) -> Optional[str]:
gguf_files = [
name
for name in filenames
if is_gguf_filename(name) and not is_mmproj_filename(name) and not is_mtp_drafter_path(name)
if is_gguf_filename(name)
and not is_mmproj_filename(name)
and not is_mtp_drafter_path(name)
and not is_big_endian_gguf_path(name, extract_quant_label(name))
]
if not gguf_files:
return None
@ -323,6 +353,20 @@ def list_partial_gguf_variants_from_state(
return variants, has_vision
def resolve_local_gguf_path(repo_id: str, gguf_variant: Optional[str]) -> Optional[str]:
"""Absolute path to the (shard-1) GGUF file for ``repo_id`` + ``gguf_variant``
if it is already downloaded in the HF cache, else ``None``. Read-only — never
triggers a download. Lets callers read header metadata before a load."""
for snapshot in iter_hf_cache_snapshots(repo_id):
variants, _ = list_local_gguf_variants(str(snapshot))
for variant in variants:
if gguf_variant is None or variant.quant == gguf_variant:
candidate = snapshot / variant.filename
if candidate.is_file():
return str(candidate)
return None
def list_gguf_variants(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[list[GgufVariantInfo], bool, Optional[list]]:
@ -372,6 +416,8 @@ def list_gguf_variants(
has_vision = True
continue
quant = extract_quant_label(filename)
if is_big_endian_gguf_path(filename, quant):
continue
quant_totals[quant] = quant_totals.get(quant, 0) + int(getattr(sibling, "size", 0) or 0)
quant_first_file.setdefault(quant, filename)
@ -424,6 +470,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
if is_mtp_drafter_path(rel):
continue
quant = extract_quant_label(rel)
if is_big_endian_gguf_path(rel, quant):
continue
quant_totals[quant] = quant_totals.get(quant, 0) + size
quant_first_file.setdefault(quant, rel)

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

@ -9,6 +9,7 @@ from typing import Optional, Sequence
from hub.utils.download_manifest import ExpectedFile
from hub.utils.gguf import (
extract_quant_label,
is_big_endian_gguf_path,
is_gguf_filename,
is_mmproj_filename,
is_mtp_drafter_path,
@ -68,6 +69,7 @@ def is_main_gguf_variant_path(path: str, variant: str) -> bool:
is_gguf_filename(path)
and not is_mmproj_filename(path)
and not is_mtp_drafter_path(path)
and not is_big_endian_gguf_path(path, variant)
and extract_quant_label(path).lower() == variant.lower()
)
@ -140,6 +142,8 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]:
if is_mmproj_filename(name) or is_mtp_drafter_path(name):
continue
quant = extract_quant_label(name).lower()
if is_big_endian_gguf_path(name, quant):
continue
main.setdefault(quant, []).append(sibling)
plans: dict[str, GgufVariantPlan] = {}

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

@ -19,6 +19,16 @@ TRANSPORT_MARKER_NAME = ".transport"
INCOMPLETE_SUFFIX = ".incomplete"
def _safe_is_dir(path: Path) -> bool:
"""``Path.is_dir()`` returning False instead of raising when the path or a
parent is unreadable (e.g. a restricted ``~/.cache/huggingface/hub``), so
cache enumeration skips that root rather than 500ing."""
try:
return path.is_dir()
except OSError:
return False
def hf_cache_root(*, create: bool = False) -> Optional[Path]:
try:
from huggingface_hub import constants as hf_constants
@ -31,7 +41,7 @@ def hf_cache_root(*, create: bool = False) -> Optional[Path]:
except OSError:
return None
return root
return root if root.is_dir() else None
return root if _safe_is_dir(root) else None
def hf_cache_roots() -> list[Path]:
@ -41,7 +51,7 @@ def hf_cache_roots() -> list[Path]:
seen: set[str] = set()
def _add(path: Optional[Path]) -> None:
if path is None or not path.is_dir():
if path is None or not _safe_is_dir(path):
return
try:
key = str(path.resolve())

150
studio/backend/loggers/handlers.py
View file

@ -8,65 +8,137 @@ filter_sensitive_data (structlog processor for sanitization), and
get_logger (factory for structured loggers).
"""
import os
import re
import time
from typing import Callable
import structlog
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from utils.native_path_leases import redact_native_paths
logger = structlog.get_logger(__name__)
def _env_int(name: str, default: int) -> int:
try:
raw = (os.environ.get(name) or "").strip()
return int(raw) if raw else default
except ValueError:
return default
# Drop duplicate successful-GET access logs repeated within the window: the SPA
# fans one cache invalidation into many identical list fetches; only the first
# informs. Loading polls, mutations, and errors are unaffected. 0 = log all.
_ACCESS_LOG_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS", 300)
# Pure-liveness/UI polls whose access line carries no signal beyond "client still
# polling" (state changes are logged by their own modules). Collapsed to a longer
# heartbeat instead of one line per poll; first hit and any error still log. 0 = off.
_QUIET_POLL_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS", 10000)
_QUIET_POLL_PATHS = {
"/api/health",
"/api/auth/status",
"/api/inference/status",
"/api/inference/monitor",
}
_DEDUP_MAP_MAX = 4096
_NATIVE_PATH_LEASE_RE = re.compile(
r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
)
_EXCLUDED_PATHS = {
"/api/train/status",
"/api/train/metrics",
"/api/train/hardware",
"/api/system",
}
_EXCLUDED_SUFFIXES = (
".png",
".jpg",
".jpeg",
".svg",
".ico",
".woff",
".woff2",
".ttf",
)
class LoggingMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: Callable) -> Response:
start_time = time.time()
class LoggingMiddleware:
"""ASGI request logger that avoids BaseHTTPMiddleware streaming wrappers."""
def __init__(self, app: ASGIApp) -> None:
self.app = app
# (method, path, query, status_code) -> monotonic ts of the last EMITTED log.
self._last_log: dict[tuple[str, str, bytes, int], float] = {}
def _is_redundant_repeat(
self, method: str, path: str, query: bytes, status_code: int, now: float
) -> bool:
"""True if an identical GET/2xx log fired < window ago. The query string
is part of the identity, so distinct query-driven GETs are not collapsed.
Mutations and non-2xx are never deduped. Quiet-poll paths use a longer
heartbeat window. Stamps only on emit, so steady polls still log."""
if method != "GET" or not (200 <= status_code < 300):
return False
window_ms = _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS
if window_ms <= 0:
return False
key = (method, path, query, status_code)
last = self._last_log.get(key)
if last is not None and (now - last) * 1000.0 < window_ms:
return True
self._last_log[key] = now
if len(self._last_log) > _DEDUP_MAP_MAX:
cutoff = now - (max(_ACCESS_LOG_DEDUP_MS, _QUIET_POLL_DEDUP_MS) / 1000.0)
self._last_log = {k: v for k, v in self._last_log.items() if v >= cutoff}
return False
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
path = scope["path"]
excluded = (
path in _EXCLUDED_PATHS
or path.startswith("/assets/")
or path.endswith(_EXCLUDED_SUFFIXES)
)
start_time = time.perf_counter()
status_code = 500
async def send_wrapper(message: Message) -> None:
nonlocal status_code
if message["type"] == "http.response.start":
status_code = message["status"]
await send(message)
try:
response = await call_next(request)
process_time = (time.time() - start_time) * 1000
EXCLUDED_PATHS = {
"/api/train/status",
"/api/train/metrics",
"/api/train/hardware",
"/api/system",
}
is_excluded = (
request.url.path in EXCLUDED_PATHS
or request.url.path.startswith("/assets/")
or request.url.path.endswith(
(".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf")
)
)
if not is_excluded:
logger.info(
"request_completed",
method = request.method,
path = request.url.path,
status_code = response.status_code,
process_time_ms = round(process_time, 2),
)
return response
except Exception as e:
await self.app(scope, receive, send_wrapper)
except Exception as exc:
logger.error(
"request_failed",
path = request.url.path,
method = request.method,
error = str(e),
path = path,
method = scope["method"],
status_code = status_code,
error = str(exc),
process_time_ms = round((time.perf_counter() - start_time) * 1000, 2),
exc_info = True,
)
raise
else:
end_time = time.perf_counter()
if not excluded and not self._is_redundant_repeat(
scope["method"], path, scope.get("query_string", b""), status_code, end_time
):
logger.info(
"request_completed",
method = scope["method"],
path = path,
status_code = status_code,
process_time_ms = round((end_time - start_time) * 1000, 2),
)
def filter_sensitive_data(logger, method_name, event_dict):

124
studio/backend/main.py
View file

@ -15,6 +15,15 @@ from dataclasses import asdict
# Suppress C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
# Pin GPU index ordering to PCI bus id before any torch import creates a CUDA
# context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi
# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from
# nvidia-smi data can resolve to a different physical card via
# CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See
# utils/hardware/hardware.py for the full rationale; set here too so the entry
# process is covered before its heavy ML imports.
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
@ -250,7 +259,7 @@ if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
# warnings.filterwarnings("ignore", category=DeprecationWarning)
# warnings.filterwarnings("ignore", module="triton.*")
from fastapi import Depends, FastAPI, HTTPException, Request
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
@ -296,6 +305,7 @@ from utils.hardware import (
import utils.hardware.hardware as _hw_module
from utils.cache_cleanup import clear_unsloth_compiled_cache
from utils.lifespan_shutdown import run_lifespan_shutdown
from utils.native_path_leases import native_path_leases_supported
from utils.update_status import (
get_studio_install_source_status,
@ -463,9 +473,16 @@ async def lifespan(app: FastAPI):
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
yield
await asyncio.to_thread(terminate_hub_downloads)
_hw_module.DEVICE = None
clear_unsloth_compiled_cache()
from core.inference.llama_http import aclose as _close_llama_http
await _close_llama_http()
await run_lifespan_shutdown(
terminate_hub_downloads,
clear_unsloth_compiled_cache,
_hw_module,
)
app = FastAPI(
@ -488,8 +505,7 @@ app.add_middleware(LoggingMiddleware)
# img/media-src allow any https origin so HF model-card assets render (mirrors
# tauri.conf.json); scripts/frames/connect-src stay same-origin + HF.
from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402
from starlette.requests import Request as _StarletteRequest # noqa: E402
from starlette.datastructures import MutableHeaders # noqa: E402
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
@ -545,28 +561,51 @@ def _build_csp(script_nonce: "str | None" = None) -> str:
)
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
"""Set baseline security headers; splice per-response inline-script nonces into CSP."""
class SecurityHeadersMiddleware:
"""Set baseline security headers; splice per-response inline-script nonces into CSP.
async def dispatch(self, request: _StarletteRequest, call_next):
response = await call_next(request)
# Strip the internal nonce hand-off header so it never reaches the client
nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER)
if nonce is not None:
del response.headers[_CSP_SCRIPT_NONCE_HEADER]
response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
# Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and
# DENY would block serve_kernel_port_as_iframe regardless of CSP.
if not _IS_COLAB and request.url.path != _ARTIFACT_PREVIEW_FRAME_PATH:
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("Referrer-Policy", "no-referrer")
response.headers.setdefault(
"Permissions-Policy",
"camera=(), microphone=(self), geolocation=()",
)
response.headers["server"] = "unsloth-studio"
return response
Pure ASGI (not BaseHTTPMiddleware) so streaming responses are not wrapped in
an anyio stream. Header logic mirrors the prior version exactly via
MutableHeaders on the response-start message.
"""
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.app(scope, receive, send)
return
path = scope.get("path", "")
async def send_wrapper(message):
if message["type"] == "http.response.start":
# ASGI headers are an iterable; coerce to a list so MutableHeaders
# can mutate in place even if a server sends a tuple or omits it.
raw = message.setdefault("headers", [])
if not isinstance(raw, list):
raw = list(raw)
message["headers"] = raw
headers = MutableHeaders(raw = raw)
# Strip the internal nonce hand-off header so it never reaches the client
nonce = headers.get(_CSP_SCRIPT_NONCE_HEADER)
if nonce is not None:
del headers[_CSP_SCRIPT_NONCE_HEADER]
headers.setdefault("Content-Security-Policy", _build_csp(nonce))
# Omit X-Frame-Options in Colab: CSP frame-ancestors handles it, and
# DENY would block serve_kernel_port_as_iframe regardless of CSP.
if not _IS_COLAB and path != _ARTIFACT_PREVIEW_FRAME_PATH:
headers.setdefault("X-Frame-Options", "DENY")
headers.setdefault("X-Content-Type-Options", "nosniff")
headers.setdefault("Referrer-Policy", "no-referrer")
headers.setdefault(
"Permissions-Policy",
"camera=(), microphone=(self), geolocation=()",
)
headers["server"] = "unsloth-studio"
await send(message)
await self.app(scope, receive, send_wrapper)
app.add_middleware(SecurityHeadersMiddleware)
@ -873,6 +912,10 @@ async def health_check(request: Request):
"version": UNSLOTH_VERSION,
"studio_version": STUDIO_VERSION,
"device_type": device_type,
# API-screen fields (authed-only; they fingerprint how the host is exposed).
"cloudflare_url": getattr(request.app.state, "cloudflare_url", None),
"server_url": getattr(request.app.state, "server_url", None),
"secure": bool(getattr(request.app.state, "secure", False)),
}
@ -963,17 +1006,40 @@ async def get_gpu_visibility(current_subject: str = Depends(get_current_subject)
@app.get("/api/system/hardware")
async def get_hardware_info(current_subject: str = Depends(get_current_subject)):
def get_hardware_info(
include_details: bool = Query(False), current_subject: str = Depends(get_current_subject)
):
"""Return GPU name, total VRAM, and key ML package versions.
Gated behind auth alongside /api/system -- same fingerprinting concern.
/api/system/gpu-visibility is also auth-gated.
``include_details`` is for About/diagnostics. The default response stays
cheap for callers that only need the primary GPU summary, like training
method auto-selection. Sync def (not async): hardware/detail probes can
shell out, and FastAPI runs sync endpoints in a threadpool.
"""
from utils.hardware import get_gpu_summary, get_package_versions
return {
body = {
"gpu": get_gpu_summary(),
"versions": get_package_versions(),
}
if include_details:
from utils.llama_cpp_update import get_installed_llama_version
# All backend-visible GPUs (respects CUDA_VISIBLE_DEVICES), so multi-GPU
# hosts list every device -- get_gpu_summary alone reports only the primary.
# Sort by visible_ordinal: the nvidia-smi path returns rows in physical order,
# so under a reordering CUDA_VISIBLE_DEVICES (e.g. "5,3") labeling by array
# index would otherwise disagree with the GPU 0/1 the backend actually sees.
devices = get_backend_visible_gpu_info().get("devices", [])
body["gpus"] = [
{"name": d.get("name"), "vram_total_gb": d.get("memory_total_gb")}
for d in sorted(devices, key = lambda d: d.get("visible_ordinal", 0))
]
body["llama_cpp"] = get_installed_llama_version()
return body
# ============ Serve Frontend (Optional) ============

39
studio/backend/models/export.py
View file

@ -51,6 +51,14 @@ class LoadCheckpointRequest(BaseModel):
False,
description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
hf_token: Optional[str] = Field(
None,
description = "Hugging Face token used to scan/load gated checkpoints and their base models.",
)
class ExportStatusResponse(BaseModel):
@ -68,6 +76,37 @@ class ExportStatusResponse(BaseModel):
False,
description = "True if the loaded checkpoint is a PEFT (LoRA) model",
)
is_export_active: bool = Field(
False,
description = "True while a load / export / cleanup operation is running",
)
# Recovery fields: when a blocking export POST is cut off by a Cloudflare tunnel
# timeout (524 at ~100s), the client polls this endpoint to learn the real
# outcome of the operation that kept running on the backend.
active_op_kind: Optional[str] = Field(
None,
description = "Kind of the currently running op (load_checkpoint / export_* / cleanup)",
)
last_op_seq: int = Field(
0,
description = "Monotonic counter of finished ops; client baseline to detect 'my op finished'",
)
last_op_kind: Optional[str] = Field(
None,
description = "Kind of the most recently finished op",
)
last_op_status: Optional[str] = Field(
None,
description = "Outcome of the most recently finished op: success / error / cancelled",
)
last_op_output_path: Optional[str] = Field(
None,
description = "Output path of the most recently finished op, if it produced one",
)
last_op_error: Optional[str] = Field(
None,
description = "Error message of the most recently finished op, if it failed",
)
class ExportOperationResponse(BaseModel):

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

@ -42,6 +42,10 @@ class LoadRequest(BaseModel):
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
chat_template_override: Optional[str] = Field(
None,
description = "Custom Jinja2 chat template to use instead of the model's default",
@ -125,6 +129,16 @@ class ValidateModelRequest(BaseModel):
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
# Intended load settings so validate's coexistence check matches the follow-up
# /load; defaults preserve old behavior for callers that omit them.
max_seq_length: int = Field(0, ge = 0, le = 1048576)
load_in_4bit: bool = Field(True)
gpu_ids: Optional[List[int]] = Field(None)
include_context_length: bool = Field(
False,
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.",
)
class ValidateModelResponse(BaseModel):
@ -144,6 +158,16 @@ class ValidateModelResponse(BaseModel):
False,
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
requires_security_review: bool = Field(
False,
description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a "
"malicious pickle), so the load is hard-blocked pending review.",
)
context_length: Optional[int] = Field(
None,
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
class GenerateRequest(BaseModel):
@ -196,9 +220,15 @@ class LoadResponse(BaseModel):
False,
description = "Whether model supports thinking/reasoning mode (enable_thinking or reasoning_effort)",
)
reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = (
Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)",
)
)
reasoning_effort_levels: List[str] = Field(
default_factory = list,
description = "Discrete reasoning_effort levels the template offers when reasoning_style is 'enable_thinking_effort' (e.g. ['high', 'max']); empty otherwise",
)
reasoning_always_on: bool = Field(
False,
@ -308,9 +338,15 @@ class InferenceStatusResponse(BaseModel):
supports_reasoning: bool = Field(
False, description = "Whether the active model supports reasoning/thinking mode"
)
reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = (
Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)",
)
)
reasoning_effort_levels: List[str] = Field(
default_factory = list,
description = "Discrete reasoning_effort levels the template offers when reasoning_style is 'enable_thinking_effort' (e.g. ['high', 'max']); empty otherwise",
)
reasoning_always_on: bool = Field(
False, description = "Whether reasoning is always on (not toggleable)"
@ -581,6 +617,16 @@ class ChatMessage(BaseModel):
return self
class ThinkingConfig(BaseModel):
"""Anthropic-compatible thinking/reasoning configuration.
Use type='disabled' to turn off thinking, or type='enabled' to turn it on.
Only type is read; extra fields (e.g. budget_tokens) are ignored, since
Studio sets provider thinking budgets itself.
"""
type: Literal["disabled", "enabled"] = "disabled"
class ChatCompletionRequest(BaseModel):
"""OpenAI-compatible chat completion request.
@ -694,6 +740,11 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] When true, keep historical <think> blocks from past assistant turns in the prompt (Qwen3.6 templates). Independent of enable_thinking / reasoning_effort.",
)
thinking: Optional[ThinkingConfig] = Field(
None,
description = "[Anthropic-compatible] Thinking configuration. "
"Use {type: 'disabled'} to disable thinking, {type: 'enabled'} to enable.",
)
enable_tools: Optional[bool] = Field(
None,
description = "[x-unsloth] Enable tool calling for supported models",
@ -717,6 +768,10 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.",
)
bypass_permissions: Optional[bool] = Field(
False,
description = "[x-unsloth] Bypass Permissions: when true, skip the tool-call confirmation gate AND disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits). Secret env vars are still stripped. Takes precedence over confirm_tool_calls.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@ -952,6 +1007,20 @@ class ChatCompletionRequest(BaseModel):
msg.tool_call_id = picked
return self
@model_validator(mode = "after")
def _map_thinking_to_enable_thinking(self) -> "ChatCompletionRequest":
"""Map Anthropic-style ``thinking`` parameter to internal ``enable_thinking``.
``thinking: {type: 'enabled'}`` sets ``enable_thinking = True`` and
``thinking: {type: 'disabled'}`` sets ``enable_thinking = False``.
``enable_thinking`` takes precedence when both are provided so that
callers who already use the internal field are unaffected. Invalid
``thinking`` shapes are rejected at validation time (422).
"""
if self.thinking is not None and self.enable_thinking is None:
self.enable_thinking = self.thinking.type == "enabled"
return self
class ToolConfirmRequest(BaseModel):
session_id: Optional[str] = None
@ -1531,6 +1600,10 @@ class AnthropicMessagesRequest(BaseModel):
enabled_tools: Optional[list[str]] = None
session_id: Optional[str] = None
cancel_id: Optional[str] = None
bypass_permissions: Optional[bool] = Field(
False,
description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
)
model_config = {"extra": "allow"}
@model_validator(mode = "before")

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

@ -53,6 +53,24 @@ class CheckpointListResponse(BaseModel):
)
class ExportSizeResponse(BaseModel):
"""Model fp16/bf16-equivalent size; size fields are null when unknown."""
model: str = Field(..., description = "Model id or path the estimate was computed for")
fp16_bytes: Optional[int] = Field(
None,
description = "Estimated FP16/BF16-equivalent on-disk size in bytes, or null if unknown",
)
total_params: Optional[int] = Field(
None,
description = "Estimated total parameter count (fp16_bytes // 2), or null if unknown",
)
source: str = Field(
"unavailable",
description = "How the estimate was derived (e.g. safetensors, config, local, vllm, unavailable)",
)
class ModelDetails(BaseModel):
"""Model configuration and metadata; used for both list and detail views"""

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

@ -108,6 +108,10 @@ class TrainingStartRequest(BaseModel):
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
# Dataset parameters
hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier")

7
studio/backend/requirements/overrides.txt
View file

@ -1,2 +1,5 @@
# Torch AO overrides (installed with --force-reinstall --no-cache-dir)
torchao==0.14.0
# torchao is installed by studio/install_python_stack.py, which selects the
# version matching the torch release actually installed in the venv (torchao's
# C++ extensions are built against one exact torch version, so a fixed pin here
# would skip them on a newer torch). See _select_torchao_spec /
# _probe_installed_torch_version in that file.

83
studio/backend/routes/chat_history.py
View file

@ -18,9 +18,11 @@ from storage.studio_db import (
CorruptSettingsError,
clear_chat_history,
count_chat_threads,
count_forks_for_message,
delete_chat_threads,
delete_chat_project,
ensure_chat_project_workspace,
fork_chat_thread,
get_chat_project,
get_chat_thread,
get_chat_message,
@ -56,6 +58,8 @@ class ChatThread(BaseModel):
createdAt: int
openaiCodeExecContainerId: Optional[str] = None
anthropicCodeExecContainerId: Optional[str] = None
forkedFromThreadId: Optional[str] = None
forkedFromMessageId: Optional[str] = None
class ChatThreadPatch(BaseModel):
@ -518,6 +522,85 @@ async def put_settings(
) from exc
class ChatForkRequest(BaseModel):
messageId: str
newThreadId: str
createdAt: int
class ChatForkResponse(BaseModel):
thread: ChatThread
messages: list[ChatMessage]
containerSnapshotWarning: Optional[str] = None
class ChatForkCountResponse(BaseModel):
count: int
@router.post("/threads/{thread_id}/fork", response_model = ChatForkResponse)
async def fork_thread(
thread_id: str,
payload: ChatForkRequest,
current_subject: str = Depends(get_current_subject),
):
"""Fork a thread at `messageId` -- creates a new thread with
ancestor msgs [root..messageId] copied with fresh ids. Both
code-exec container ids reset on the fork. OpenAI snapshot is a
best-effort enhancement; failure surfaces as
`containerSnapshotWarning` and the fork still succeeds with a
clean sandbox.
"""
import uuid
source = get_chat_thread(thread_id)
if source is None:
raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found")
if get_chat_message(thread_id, payload.messageId) is None:
raise HTTPException(
status_code = 404,
detail = f"Message {payload.messageId} not found in thread {thread_id}",
)
base_title = source.get("title") or "New Chat"
new_title = f"fork · {base_title}"
forked = fork_chat_thread(
source_thread_id = thread_id,
branch_message_id = payload.messageId,
new_thread_id = payload.newThreadId,
new_title = new_title,
created_at = payload.createdAt,
id_factory = lambda: str(uuid.uuid4()),
)
if forked is None:
raise HTTPException(status_code = 500, detail = "Fork failed")
messages = list_chat_messages(payload.newThreadId)
# Best-effort OpenAI container snapshot. Stub: a follow-up patch can
# call /v1/containers list+download / create+upload here and patch
# the new openaiCodeExecContainerId. For v1 we always start clean
# and surface the same warning regardless of provider so the UI can
# show a consistent "sandbox starts fresh" toast.
warning: Optional[str] = None
if source.get("openaiCodeExecContainerId") or source.get("anthropicCodeExecContainerId"):
warning = "Sandbox starts fresh in fork; files from parent are not carried over."
return ChatForkResponse(
thread = ChatThread(**forked),
messages = [ChatMessage(**m) for m in messages],
containerSnapshotWarning = warning,
)
@router.get(
"/threads/{thread_id}/messages/{message_id}/forks",
response_model = ChatForkCountResponse,
)
async def get_fork_count(
thread_id: str,
message_id: str,
current_subject: str = Depends(get_current_subject),
):
return ChatForkCountResponse(count = count_forks_for_message(thread_id, message_id))
@router.get("/export", response_model = ChatExportResponse)
async def export_history(current_subject: str = Depends(get_current_subject)):
from datetime import datetime, timezone

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

@ -50,40 +50,14 @@ logger = get_logger(__name__)
async def load_checkpoint(
request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject)
):
"""Load a checkpoint into the export backend (ExportBackend.load_checkpoint)."""
"""Load a checkpoint into the export backend (ExportBackend.load_checkpoint).
Export runs in its own subprocess and is allowed to run in parallel with
training and inference. We deliberately do NOT stop training or unload the
chat model here -- if the GPU runs out of memory the load/export fails with
a clear error instead of tearing down the user's other running workloads.
"""
try:
# Free GPU memory: shut down running inference/training subprocesses
# before loading the export checkpoint (they'd compete for VRAM).
try:
from core.inference import get_inference_backend
inf = get_inference_backend()
if inf.active_model_name:
logger.info(
"Unloading inference model '%s' to free GPU memory for export",
inf.active_model_name,
)
inf._shutdown_subprocess()
inf.active_model_name = None
inf.models.clear()
except Exception as e:
logger.warning("Could not unload inference model: %s", e)
try:
from core.training import get_training_backend
trn = get_training_backend()
if trn.is_training_active():
logger.info("Stopping active training to free GPU memory for export")
trn.stop_training()
# Wait for the training subprocess to exit, else it may still hold GPU memory.
for _ in range(60): # up to 30s
if not trn.is_training_active():
break
await asyncio.sleep(0.5)
else:
logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
except Exception as e:
logger.warning("Could not stop training: %s", e)
backend = get_export_backend()
# Run in a worker thread (spawns and waits on a subprocess, can take
# minutes) so the event loop stays free to serve the live log SSE stream.
@ -93,6 +67,8 @@ async def load_checkpoint(
max_seq_length = request.max_seq_length,
load_in_4bit = request.load_in_4bit,
trust_remote_code = request.trust_remote_code,
approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
hf_token = request.hf_token,
)
if not success:
@ -136,15 +112,51 @@ async def cleanup_export_memory(current_subject: str = Depends(get_current_subje
)
@router.post("/cancel", response_model = ExportOperationResponse)
async def cancel_export(current_subject: str = Depends(get_current_subject)):
"""Cancel the in-flight export by terminating its worker subprocess.
Only the export subprocess is killed; training and inference run in their
own subprocesses and keep going.
"""
try:
backend = get_export_backend()
cancelled = await asyncio.to_thread(backend.cancel_export)
return ExportOperationResponse(
success = True,
message = "Export cancelled" if cancelled else "No active export to cancel",
)
except Exception as e:
logger.error(f"Error cancelling export: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to cancel export",
)
@router.get("/status", response_model = ExportStatusResponse)
async def get_export_status(current_subject: str = Depends(get_current_subject)):
"""Get export backend status (loaded checkpoint, model type, PEFT flag)."""
try:
backend = get_export_backend()
last_op = backend.get_last_op()
# Relativise the recovered output path the same way the per-op POST response
# does, so the success banner shows an identical path on either route.
last_op_output_path = None
if last_op and last_op.get("output_path"):
details = _export_details(last_op["output_path"])
last_op_output_path = (details or {}).get("output_path")
return ExportStatusResponse(
current_checkpoint = backend.current_checkpoint,
is_vision = bool(getattr(backend, "is_vision", False)),
is_peft = bool(getattr(backend, "is_peft", False)),
is_export_active = bool(backend.is_export_active()),
active_op_kind = backend.get_active_op_kind(),
last_op_seq = int(last_op["seq"]) if last_op else 0,
last_op_kind = last_op.get("kind") if last_op else None,
last_op_status = last_op.get("status") if last_op else None,
last_op_output_path = last_op_output_path,
last_op_error = last_op.get("error") if last_op else None,
)
except Exception as e:
logger.error(f"Error getting export status: {e}", exc_info = True)
@ -154,6 +166,59 @@ async def get_export_status(current_subject: str = Depends(get_current_subject))
)
@router.get("/logs")
async def get_export_logs(
since: Optional[int] = Query(
None,
description = "Return log entries with seq strictly greater than this cursor.",
),
current_subject: str = Depends(get_current_subject),
):
"""Tunnel-safe JSON fallback for the live export log stream.
The SSE endpoint (`/logs/stream`) is the low-latency path, but some reverse
proxies -- notably Cloudflare quick tunnels (`*.trycloudflare.com`) used by
`--secure` mode -- buffer `text/event-stream` responses and only flush when
the stream closes, so over the tunnel the browser sees nothing for the whole
export ("connecting..." with no logs). This endpoint returns the same
ring-buffer lines as a short, complete JSON response that no proxy buffers,
so the frontend can poll it and still show logs in near real time.
Shares the orchestrator's monotonic `seq` cursor with the SSE stream, so the
two transports can run together and the client de-dupes by seq.
"""
try:
backend = get_export_backend()
# No cursor on the first poll of a run: start from the run-start snapshot
# so the client gets every line since the run began (matches the SSE
# default), not the entire historical ring buffer.
if since is None:
cursor = backend.get_run_start_seq()
else:
cursor = max(0, int(since))
entries, new_cursor = backend.get_logs_since(cursor)
return {
"entries": [
{
"seq": int(entry.get("seq", 0)),
"stream": entry.get("stream", "stdout"),
"line": entry.get("line", ""),
"ts": entry.get("ts"),
}
for entry in entries
],
"cursor": new_cursor,
"active": bool(backend.is_export_active()),
}
except Exception as e:
logger.error(f"Error getting export logs: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to get export logs",
)
def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]:
"""Best-effort registration so absolute exports show up in local scans."""
try:

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

File diff suppressed because it is too large Load diff

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

@ -55,6 +55,9 @@ class LlamaUpdateStatusResponse(BaseModel):
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 Update would fetch, in bytes."
)
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)

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

@ -77,13 +77,17 @@ def _validate_url(url: str) -> str:
return trimmed
parsed = urlparse(trimmed)
if parsed.scheme not in ("http", "https"):
detail = (
"MCP server address must start with http:// or https:// "
"(for example https://example.com/mcp)."
)
# Host-scoped wording: self-hosted hosts can opt in via the env var.
if _looks_like_command(trimmed):
detail += " Running a local command is not enabled on this server."
detail = (
"Local commands aren't enabled on this server. To allow them, "
"set UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 and restart Studio, or use "
"an http:// or https:// URL instead."
)
else:
detail = (
"MCP server address must start with http:// or https:// "
"(for example https://example.com/mcp)."
)
raise HTTPException(status_code = 400, detail = detail)
if not parsed.netloc:
raise HTTPException(status_code = 400, detail = "url is missing a host")

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

@ -3,6 +3,7 @@
"""Model management API routes."""
import asyncio
import hashlib
import json
import os
@ -10,7 +11,7 @@ import shutil
import sys
import uuid
from pathlib import Path
from fastapi import APIRouter, Body, Depends, HTTPException, Query
from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
from typing import List, Optional
import structlog
from loggers import get_logger
@ -41,12 +42,21 @@ def _safe_is_dir(path) -> bool:
def _is_hidden_model(*values: str | None) -> bool:
"""True if any id/path is the RAG embedding model (EMBEDDING_MODEL or
EMBED_GGUF_REPO basename), so pickers hide it (GGUF and non-GGUF)."""
EMBED_GGUF_REPO basename) or the llama.cpp install validation probe
(ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
None are usable chat models; the probe can be cached as a side effect of
installing the prebuilt llama-server and otherwise sorts smallest, so it
would be auto-selected."""
from core.rag import config as rag_config
needles = (
rag_config.EMBEDDING_MODEL.split("/")[-1].lower(),
rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(),
# The validation probe's repo (matches the cached repo id) and its exact
# filename (matches the on-disk path). The filename carries the .gguf so
# it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``.
"ggml-org/models",
"stories260k.gguf",
)
return any(v and any(n in v.lower() for n in needles) for v in values)
@ -81,6 +91,7 @@ try:
from utils.models.model_config import (
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
is_audio_input_type,
)
from core.inference import get_inference_backend
@ -112,6 +123,7 @@ except ImportError:
from utils.models.model_config import (
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
is_audio_input_type,
)
from core.inference import get_inference_backend
@ -138,6 +150,7 @@ from models import (
from models.models import (
BrowseEntry,
BrowseFoldersResponse,
ExportSizeResponse,
GgufVariantDetail,
GgufVariantsResponse,
ModelType,
@ -1521,16 +1534,22 @@ async def get_model_config(
except Exception:
pass
# Fallback: try AutoConfig directly.
# Fallback: read raw config.json (declarative fields only) -- a selection-time
# metadata probe that must never execute a repo's auto_map Python.
if max_position_embeddings is None:
try:
from transformers import AutoConfig as _AutoConfig
from utils.transformers_version import _load_config_json
from types import SimpleNamespace
_trust = model_name.lower().startswith("unsloth/")
_ac = _AutoConfig.from_pretrained(
model_name, trust_remote_code = _trust, token = hf_token
)
max_position_embeddings = _get_max_position_embeddings(_ac)
_cfg = _load_config_json(model_name, hf_token = hf_token)
if _cfg is not None:
def _to_ns(d):
if isinstance(d, dict):
return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()})
return d
max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg))
except Exception:
pass
@ -1563,6 +1582,195 @@ async def get_model_config(
)
@router.post("/remote-code-scan")
async def scan_model_remote_code(
model_name: str = Body(..., embed = True),
hf_token: Optional[str] = Body(None, embed = True),
current_subject: str = Depends(get_current_subject),
):
"""Scan a model's ``auto_map`` custom code so the UI can show findings before
the user enables ``trust_remote_code``. Code-free: reads ``config.json`` and
statically scans the repo ``.py`` (never loads the model). Returns
``has_remote_code`` plus the severity-tagged findings + a pinning fingerprint.
POST (not GET) so the ``hf_token`` for gated repos travels in the body and
never lands in a URL, browser history, or access log.
"""
try:
from utils.security import preflight_remote_code_consent_for_targets
if not is_local_path(model_name):
model_name = resolve_cached_repo_id_case(model_name)
# Scan the adapter AND the base together (a LoRA runs both repos' code; a pickle
# can live in either), pinned by one combined fingerprint. Snapshot the primary's
# cache state BEFORE resolving the base: for a remote adapter that resolve
# downloads adapter_config.json, which would otherwise hide the adapter from
# cleanup on decline. On error treat as pre-existing so a decline never deletes it.
try:
_primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache(model_name)
except Exception:
_primary_preexisting = True
security_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so its code/weights are scanned too.
_base = get_base_model_from_lora_identifier(model_name, hf_token)
if _base:
security_targets.append(_base)
except Exception:
pass
security_targets = list(dict.fromkeys(security_targets))
# Record every repo OUR scan is first to pull into the cache (adapter, base, and
# external auto_map repos like owner/name--module.Class), so a decline purges
# exactly what was downloaded. Computed BEFORE the preflight downloads, against
# every cache the discard searches, so a repo the user already had is not deleted.
from utils.security.remote_code_scan import external_auto_map_repos
scan_created_repos: list = []
_seen_created: set = set()
def _mark_scan_created(repo: str, *, preexisting: Optional[bool] = None) -> None:
if not repo or repo in _seen_created:
return
_seen_created.add(repo)
try:
already = (
preexisting
if preexisting is not None
else (is_local_path(repo) or _repo_in_any_hf_cache(repo))
)
if not already:
scan_created_repos.append(repo)
except Exception:
pass
for _target in security_targets:
# Use the pre-base-resolution snapshot for the primary (see above).
_mark_scan_created(
_target, preexisting = _primary_preexisting if _target == model_name else None
)
for _ext in external_auto_map_repos(_target, hf_token):
_mark_scan_created(_ext)
decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token)
payload = decision.response_payload()
payload["requires_trust_remote_code"] = decision.has_remote_code
# created_by_scan = primary flag (older clients); scan_created_repos drives cleanup.
payload["created_by_scan"] = model_name in scan_created_repos
payload["scan_created_repos"] = scan_created_repos
# Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can
# hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map.
from utils.security import evaluate_file_security, security_load_subdirs
unsafe_files: list = []
security_blocked = False
for _target in security_targets:
_sec = evaluate_file_security(
_target, hf_token = hf_token, load_subdirs = security_load_subdirs(_target, hf_token)
)
security_blocked = security_blocked or _sec.blocked
unsafe_files.extend(_sec.unsafe_files)
payload["unsafe_files"] = unsafe_files
payload["security_blocked"] = security_blocked
if security_blocked:
# Non-approvable hard block: approvable False hides "Enable and continue", and
# requires_trust_remote_code forces the dialog open even with no custom code.
payload["approvable"] = False
payload["requires_trust_remote_code"] = True
payload["error_kind"] = "malware_blocked"
return payload
except Exception as e:
raise log_and_http_error(
e,
500,
"Failed to scan model remote code",
event = "models.remote_code_scan_failed",
log = logger,
)
@router.post("/discard-remote-code")
async def discard_remote_code_download(
model_name: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject)
):
"""Purge a repo the consent scan downloaded after the user DECLINED its custom
code, so untrusted code is not left on disk.
Safety: only ever deletes a metadata-only cache entry the scan created. It
refuses a local path (never touches user files), a currently-loaded model, and
any repo that has weight files cached (``*.safetensors`` / ``*.bin`` /
``*.gguf``) -- i.e. a model the user actually downloaded. The frontend only
calls this when the scan reported ``created_by_scan``.
"""
if is_local_path(model_name):
return {"deleted": False, "reason": "local"}
if not _is_valid_repo_id(model_name):
return {"deleted": False, "reason": "invalid"}
# Never delete a model that is loaded for inference.
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 = llama_backend.model_identifier.lower()
if loaded == model_name.lower() or loaded.startswith(model_name.lower()):
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()):
return {"deleted": False, "reason": "loaded"}
except Exception:
pass
_WEIGHTS = (
".safetensors",
".bin",
".pt",
".pth",
".h5",
".msgpack",
".gguf",
".onnx",
".ckpt",
)
try:
target_repo = None
hf_cache = None
for cache in _all_hf_cache_scans():
for repo_info in cache.repos:
if repo_info.repo_type != "model":
continue
if repo_info.repo_id.lower() == model_name.lower():
target_repo, hf_cache = repo_info, cache
break
if target_repo is not None:
break
if target_repo is None:
return {"deleted": False, "reason": "not_cached"}
# Hard guard: a repo with weights is a real model the user has -- leave it.
for rev in target_repo.revisions:
for f in rev.files:
if f.file_name.lower().endswith(_WEIGHTS):
return {"deleted": False, "reason": "has_weights"}
revision_hashes = [rev.commit_hash for rev in target_repo.revisions]
if not revision_hashes:
return {"deleted": False, "reason": "not_cached"}
hf_cache.delete_revisions(*revision_hashes).execute()
logger.info("Discarded declined remote-code download: %s", model_name)
return {"deleted": True}
except Exception as e:
logger.warning("Could not discard remote-code download for %s: %s", model_name, e)
return {"deleted": False, "reason": "error"}
@router.get("/loras")
async def scan_loras(
outputs_dir: str = Query(
@ -1979,7 +2187,11 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get
@router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)):
async def check_vision_model(
model_name: str,
hf_token: Optional[str] = Query(None),
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is a vision model.
@ -1987,7 +2199,8 @@ async def check_vision_model(model_name: str, current_subject: str = Depends(get
"""
try:
logger.info(f"Checking if vision model: {model_name}")
is_vision = is_vision_model(model_name)
# Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision).
is_vision = is_vision_model(model_name, hf_token = hf_token)
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
return VisionCheckResponse(
@ -2106,7 +2319,11 @@ async def get_gguf_variants(
size = f.stat().st_size
except OSError:
continue # broken symlink / unreadable: skip
q = _extract_quant_label(f.name).lower()
rel = f.relative_to(snap).as_posix()
q = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, q):
continue
q = q.lower()
by_quant[q] = by_quant.get(q, 0) + size
if by_quant:
cached_bytes_by_quant_per_snapshot.append(by_quant)
@ -2182,8 +2399,12 @@ async def get_gguf_download_progress(
for f in _iter_gguf_paths(entry):
if _is_mmproj_filename(f.name):
continue
fname = f.name.lower().replace("-", "").replace("_", "")
if not variant_lower or variant_lower in fname:
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:
@ -2329,6 +2550,51 @@ def _get_repo_size_cached(repo_id: str) -> int:
return 0
def _repo_in_any_hf_cache(model_name: str) -> bool:
"""Whether ``model_name`` already exists in ANY HF cache the discard searches
(active, legacy, default).
``created_by_scan`` must be True only when the scan itself first pulled the repo;
checking just the active cache (``get_cache_path``) would mark a repo the user
already had in a legacy/default cache as scan-created, so declining the consent
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,
)
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
# 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.
for cache in candidates:
try:
if (cache / dirname).exists():
return True
if cache.is_dir():
for entry in cache.iterdir():
if entry.name.lower() == dirname_lower and entry.is_dir():
return True
except Exception:
continue
return False
def _all_hf_cache_scans():
"""scan_cache_dir for the active, legacy, and default HF caches.
@ -2744,3 +3010,121 @@ async def list_checkpoints(
event = "models.list_checkpoints_failed",
log = logger,
)
# Successful estimates only, keyed by model id (token-independent, never stored).
# Failures are not cached so a transient offline/gated error can recover later.
_EXPORT_SIZE_CACHE: dict[str, tuple[int, int, str]] = {}
def _is_sizable_local_path(model: str) -> bool:
"""True only for local paths under a Studio data root.
Containment is decided lexically (no filesystem access) before the path is
touched, then the path is symlink-resolved and re-checked so a symlink
inside a root can't point the sizer outside it. A user-controlled path thus
can't trigger a scan of an arbitrary dir.
"""
from utils.paths import outputs_root, exports_root, studio_root
from utils.paths.storage_roots import cache_root
def _lexical(p: str) -> str:
# Lexical only (no filesystem read); normpath collapses '..'.
return os.path.normpath(os.path.abspath(os.path.expanduser(p)))
raw_roots = [studio_root(), outputs_root(), exports_root(), cache_root()]
roots = []
for root in raw_roots:
try:
roots.append(_lexical(str(root)))
except (OSError, RuntimeError, ValueError):
continue
try:
candidate = _lexical(model)
except (OSError, RuntimeError, ValueError):
return False
for root in roots:
if candidate == root or candidate.startswith(root + os.sep):
# Contained lexically; resolve symlinks and re-verify the real path
# is still under a root before touching the filesystem.
try:
real = os.path.realpath(candidate)
except (OSError, RuntimeError, ValueError):
return False
for raw in raw_roots:
try:
real_root = os.path.realpath(str(raw))
except (OSError, RuntimeError, ValueError):
continue
if real == real_root or real.startswith(real_root + os.sep):
return os.path.exists(real)
return False
return False
def _export_size_cached(
model: str, hf_token: Optional[str]
) -> tuple[Optional[int], Optional[int], str]:
"""Estimate a model's fp16/bf16-equivalent size in bytes (+ total params).
Memoizes successful results by model id; never raises (failures return
(None, None, "unavailable") and are not cached). Blocking I/O; call off-thread.
"""
cached = _EXPORT_SIZE_CACHE.get(model)
if cached is not None:
return cached
try:
from utils.hardware.hardware import (
_resolve_model_identifier_for_gpu_estimate,
estimate_fp16_model_size_bytes,
)
# A local LoRA adapter is sized via its base model, which the sizer
# reads from the adapter config; re-validate that resolved base so a
# crafted adapter can't redirect the local scan outside the roots.
if is_local_path(model):
base = _resolve_model_identifier_for_gpu_estimate(model, hf_token = hf_token)
if is_local_path(base) and not _is_sizable_local_path(base):
return None, None, "unavailable"
fp16_bytes, source = estimate_fp16_model_size_bytes(model, hf_token = hf_token)
if not fp16_bytes or fp16_bytes <= 0:
return None, None, source or "unavailable"
result = (int(fp16_bytes), int(fp16_bytes) // 2, source)
_EXPORT_SIZE_CACHE[model] = result
return result
except Exception as e: # a size hint must never break export
logger.warning("Could not estimate export size for '%s': %s", model, e)
return None, None, "unavailable"
@router.get("/export-size", response_model = ExportSizeResponse)
async def get_export_size(
model: str = Query(..., description = "Base model id or local model path to size"),
hf_token: Optional[str] = Header(None, alias = "X-HF-Token"),
current_subject: str = Depends(get_current_subject),
):
"""Estimate a model's fp16/bf16-equivalent size for the Export page.
Returns nulls with HTTP 200 when the size can't be determined. The HF token
(for gated repos) comes from the X-HF-Token header so it never hits URLs/logs.
"""
if is_local_path(model):
if not _is_sizable_local_path(model):
return ExportSizeResponse(
model = model, fp16_bytes = None, total_params = None, source = "unavailable"
)
resolved = model
else:
resolved = resolve_cached_repo_id_case(model)
# Blocking network/disk I/O: run off the event loop.
fp16_bytes, total_params, source = await asyncio.to_thread(
_export_size_cached, resolved, hf_token
)
return ExportSizeResponse(
model = resolved,
fp16_bytes = fp16_bytes,
total_params = total_params,
source = source,
)

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

@ -248,49 +248,99 @@ async def start_training(
"output_dir": resume_output_dir,
"resume_from_checkpoint": request.resume_from_checkpoint,
"trust_remote_code": request.trust_remote_code,
"approved_remote_code_fingerprint": request.approved_remote_code_fingerprint,
"gpu_ids": request.gpu_ids,
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
}
# Training page has no trust_remote_code toggle; as a safety net consult
# YAML model defaults directly so models that need it always get it.
# Training page has no trust_remote_code toggle, so honor the YAML default
# -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
# local path or a name merely starting with "unsloth/".
if not training_kwargs["trust_remote_code"]:
from utils.security.trusted_org import is_trusted_org_repo
model_defaults = load_model_defaults(request.model_name)
yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
if yaml_trust:
if yaml_trust and is_trusted_org_repo(
request.model_name, hf_token = request.hf_token or None
):
logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
training_kwargs["trust_remote_code"] = True
# Free GPU memory: shut down any running inference/export subprocesses
# before training (they'd compete for VRAM otherwise).
try:
from core.inference import get_inference_backend
inf_backend = get_inference_backend()
if inf_backend.active_model_name:
logger.info(
"Unloading inference model '%s' to free GPU memory for training",
inf_backend.active_model_name,
elif yaml_trust:
logger.warning(
"YAML sets trust_remote_code=True for %s but it is not a trusted "
"first-party repo; leaving disabled (user can opt in explicitly).",
request.model_name,
)
inf_backend._shutdown_subprocess()
inf_backend.active_model_name = None
inf_backend.models.clear()
except Exception as e:
logger.warning("Could not unload inference model: %s", e)
try:
from core.export import get_export_backend
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
logger.info("Shutting down export subprocess to free GPU memory for training")
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
exp_backend.is_peft = False
except Exception as e:
logger.warning("Could not shut down export subprocess: %s", e)
# Free VRAM for training: stop export, unload chat unless it can coexist.
# A before_spawn hook -> runs only after start_training's guards pass, so
# we never tear down chat/export VRAM for a start that is then refused.
def _free_vram_for_training() -> None:
try:
from core.export import get_export_backend
exp_backend = get_export_backend()
# Tear down the export subprocess whenever an export is in flight,
# not just once a checkpoint is loaded: during the load phase
# current_checkpoint is still unset while the worker is already
# allocating GPU memory, so gate on is_export_active() too.
if exp_backend.current_checkpoint or exp_backend.is_export_active():
logger.info("Shutting down export subprocess to free GPU memory for training")
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
exp_backend.is_peft = False
except Exception as e:
logger.warning("Could not shut down export subprocess: %s", e)
# start_training spawns a subprocess (non-blocking).
success = backend.start_training(job_id = job_id, **training_kwargs)
try:
from routes.training_vram import (
can_keep_chat_during_training,
free_chat_models_for_training,
summarize_resident_chat,
)
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,
)
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)
except Exception as e:
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
if not success:
progress_error = backend.trainer.training_progress.error

342
studio/backend/routes/training_vram.py Normal file
View file

@ -0,0 +1,342 @@
# 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.
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.
"""
from typing import Any, Dict, List, Optional, Tuple
from loggers import get_logger
logger = get_logger(__name__)
# keep iff usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB. Conservative:
# the probe sees only the chat model's current footprint, so reserve headroom for
# estimate error + KV-cache growth (KEEP_FLOOR_GB ~= 2 GB load buffer + 2 GB chat).
SAFETY_MARGIN = 1.15
KEEP_FLOOR_GB = 4.0
# Each extra GPU contributes less than its raw free memory (sharding overhead).
_MULTI_GPU_OVERHEAD = 0.85
def _free_vram_by_index(devices: List[Dict[str, Any]]) -> Dict[int, float]:
"""Map GPU index -> free VRAM (GB) from a get_visible_gpu_utilization() device list."""
free_by_index: Dict[int, float] = {}
for device in devices:
total_gb = device.get("vram_total_gb")
used_gb = device.get("vram_used_gb")
if total_gb is None or used_gb is None:
continue
free_by_index[device["index"]] = max(total_gb - used_gb, 0.0)
return free_by_index
def summarize_resident_chat() -> Dict[str, Any]:
"""Report which chat models hold GPU memory (resident even while loading). Never raises."""
hf_name: Optional[str] = None
gguf_name: Optional[str] = None
loading: bool = False
try:
from core.inference import get_inference_backend
inf = get_inference_backend()
# active_model_name is set only on success; a mid-load model sits in
# loading_models while already holding VRAM -> both count as resident.
if inf.active_model_name or inf.loading_models:
hf_name = inf.active_model_name or next(iter(inf.loading_models), None)
# Any in-flight load (incl. a replacement while the old model is still
# active) can't be sized -> flag it so the caller frees instead of keeps.
if inf.loading_models:
loading = True
except Exception as e:
logger.warning("Could not inspect inference backend: %s", e)
try:
from routes.inference import get_llama_cpp_backend
llama = get_llama_cpp_backend()
# is_active (not is_loaded): a mid-start server already allocates VRAM.
# A confirmed CPU-only server (_gpu_offload_active is False) holds no VRAM.
if llama.is_active and getattr(llama, "_gpu_offload_active", None) is not False:
gguf_name = llama.model_identifier or "gguf"
if not getattr(llama, "is_loaded", False): # still loading -> size unknown
loading = True
except Exception as e:
logger.warning("Could not inspect GGUF backend: %s", e)
return {
"hf": hf_name,
"gguf": gguf_name,
"loading": loading,
"any": bool(hf_name or gguf_name),
}
def can_keep_chat_during_training(
*,
model_name: str,
hf_token: Optional[str],
training_type: str,
load_in_4bit: bool,
batch_size: int,
max_seq_length: int,
lora_rank: int,
target_modules: Optional[List[str]],
gradient_checkpointing: str,
optimizer: str,
gpu_ids: Optional[List[int]],
) -> Tuple[bool, Dict[str, Any]]:
"""Decide if a resident chat model can coexist with training given free VRAM.
Reuses training's own estimator/selector so the decision matches later
placement. Default-deny: anything we can't size returns False (unload).
"""
try:
from utils.hardware import (
DeviceType,
auto_select_gpu_ids,
estimate_required_model_memory_gb,
get_device,
get_visible_gpu_utilization,
resolve_requested_gpu_ids,
)
if get_device() != DeviceType.CUDA:
return False, {"mode": "non_cuda", "reason": "non_cuda"}
# Full finetuning runs in 16-bit, so ignore the 4-bit request or we under-count.
effective_4bit = False if training_type == "Full Finetuning" else load_in_4bit
hf_token_arg = hf_token or None
est_kwargs = dict(
hf_token = hf_token_arg,
training_type = training_type,
load_in_4bit = effective_4bit,
batch_size = batch_size,
max_seq_length = max_seq_length,
lora_rank = lora_rank,
target_modules = target_modules,
gradient_checkpointing = gradient_checkpointing,
optimizer = optimizer,
)
if gpu_ids:
# Explicit GPUs: the selector does no VRAM math, so size it here.
try:
resolved = resolve_requested_gpu_ids(gpu_ids)
except ValueError:
# Invalid ids -> start_training will 400 first, so don't unload.
return True, {"mode": "explicit", "reason": "invalid_gpu_ids"}
required_gb, est_meta = estimate_required_model_memory_gb(model_name, **est_kwargs)
if required_gb is None:
return False, {"mode": "explicit", "reason": "estimate_unavailable"}
free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
# A requested GPU missing from the device list contributes 0.
free_vals = [free_by_index.get(i, 0.0) for i in resolved]
ranked = sorted(free_vals, reverse = True)
usable_gb = (
ranked[0] + sum(f * _MULTI_GPU_OVERHEAD for f in ranked[1:]) if ranked else 0.0
)
aggregate_fits = usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
# Activations don't shard: enforce a per-GPU floor so an uneven split
# (e.g. free [45, 10]) can't be kept into an OOM the aggregate misses.
per_gpu_fits = True
min_free_gb = min(free_vals) if free_vals else 0.0
if len(resolved) > 1:
min_per_gpu_gb = est_meta.get("vram_breakdown", {}).get(
f"min_per_gpu_{len(resolved)}"
)
if min_per_gpu_gb is not None:
per_gpu_fits = min_free_gb >= min_per_gpu_gb
keep = aggregate_fits and per_gpu_fits
return keep, {
"mode": "explicit",
"required_gb": required_gb,
"usable_gb": round(usable_gb, 3),
"min_free_gb": round(min_free_gb, 3),
}
# Auto: same call start_training makes later; reuse its sizing metadata.
_selected, meta = auto_select_gpu_ids(model_name, **est_kwargs)
mode = meta.get("selection_mode")
required_gb = meta.get("required_gb")
usable_gb = meta.get("usable_gb")
keep = (
mode == "auto"
and required_gb is not None
and usable_gb is not None
and usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
)
return keep, {
"mode": mode,
"required_gb": required_gb,
"usable_gb": usable_gb,
}
except Exception as e:
# Never let a sizing failure keep a chat model loaded into a training OOM.
logger.warning("Chat-coexistence probe failed; will unload: %s", e)
return False, {"reason": "probe_error", "error": str(e)}
def can_load_chat_during_training(
*,
model_name: str,
hf_token: Optional[str],
load_in_4bit: bool,
max_seq_length: int,
requested_gpu_ids: Optional[List[int]],
is_gguf: bool = False,
required_override_gb: Optional[float] = None,
) -> Tuple[bool, Dict[str, Any]]:
"""Decide if a NEW chat model can load without OOMing active training (inverse
of can_keep_chat_during_training: training is already resident, so size the
chat model against the free VRAM that remains). Sizes/places it the same way
the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an
even-share per-GPU floor for device_map="balanced"; GGUF sizes from
required_override_gb over the visible pool. `load_in_4bit` must be effective
(LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any
CUDA case it can't size, so a load never OOMs training."""
try:
from utils.hardware import (
DeviceType,
auto_select_gpu_ids,
estimate_required_model_memory_gb,
get_device,
get_visible_gpu_utilization,
resolve_requested_gpu_ids,
)
if get_device() != DeviceType.CUDA:
return True, {"mode": "non_cuda", "reason": "non_cuda"}
est_kwargs = dict(
hf_token = hf_token or None,
training_type = None, # inference sizing of the chat model itself
load_in_4bit = load_in_4bit,
max_seq_length = max_seq_length or 2048,
)
# HF auto: reuse the loader's selector; fits iff its pick clears the margin.
if not requested_gpu_ids and not is_gguf:
_selected, meta = auto_select_gpu_ids(model_name, **est_kwargs)
mode = meta.get("selection_mode")
required_gb = meta.get("required_gb")
usable_gb = meta.get("usable_gb")
needed_gb = (
round(required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB, 3)
if required_gb is not None
else None
)
fits = (
mode == "auto"
and required_gb is not None
and usable_gb is not None
and usable_gb >= needed_gb
)
return fits, {
"mode": mode,
"required_gb": required_gb,
"usable_gb": usable_gb,
"needed_gb": needed_gb,
}
# Explicit GPUs, or GGUF: size directly and check live free VRAM.
required_gb = required_override_gb
if required_gb is None:
required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs)
if required_gb is None:
mode = "explicit" if requested_gpu_ids else "gguf"
return False, {"mode": mode, "reason": "estimate_unavailable"}
free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
if requested_gpu_ids:
# Invalid ids -> load_model 400s first, so don't block; missing id = 0.
try:
resolved = resolve_requested_gpu_ids(requested_gpu_ids)
except ValueError:
return True, {"mode": "explicit", "reason": "invalid_gpu_ids"}
free_vals = [free_by_index.get(i, 0.0) for i in resolved]
mode = "explicit"
else:
# GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate.
free_vals = list(free_by_index.values())
mode = "gguf"
if not free_vals:
return False, {"mode": mode, "reason": "no_visible_gpus"}
ranked = sorted(free_vals, reverse = True)
usable_gb = ranked[0] + sum(f * _MULTI_GPU_OVERHEAD for f in ranked[1:])
needed_gb = required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
aggregate_fits = usable_gb >= needed_gb
# device_map="balanced" shards across GPUs: an even-share floor stops one
# near-full GPU hiding behind aggregate capacity. GGUF self-places, no floor.
min_free_gb = min(free_vals)
per_gpu_fits = True
if mode == "explicit" and len(free_vals) > 1:
per_gpu_fits = min_free_gb >= needed_gb / len(free_vals)
return aggregate_fits and per_gpu_fits, {
"mode": mode,
"required_gb": round(required_gb, 3),
"usable_gb": round(usable_gb, 3),
"needed_gb": round(needed_gb, 3),
"min_free_gb": round(min_free_gb, 3),
}
except Exception as e:
# Never let a sizing failure load a chat model into a training OOM.
logger.warning("Chat-load coexistence probe failed; will refuse: %s", e)
return False, {"reason": "probe_error", "error": str(e)}
def free_chat_models_for_training(reason: str) -> List[str]:
"""Unload every resident chat model (HF/MLX orchestrator + GGUF server) to free
VRAM for training. Each backend isolated. Returns labels of what was freed."""
freed: List[str] = []
try:
from core.inference import get_inference_backend
inf = get_inference_backend()
if inf.active_model_name or inf.loading_models:
name = inf.active_model_name or next(iter(inf.loading_models), None)
logger.info(
"Unloading inference model '%s' to free GPU memory for training (%s)",
name,
reason,
)
inf._shutdown_subprocess()
inf.active_model_name = None
inf.models.clear()
inf.loading_models.clear()
freed.append(f"hf:{name}")
except Exception as e:
logger.warning("Could not unload inference model: %s", e)
try:
from routes.inference import get_llama_cpp_backend
llama = get_llama_cpp_backend()
# CPU-only GGUF holds no VRAM, so killing it can't help (see summarize).
if llama.is_active and getattr(llama, "_gpu_offload_active", None) is not False:
name = llama.model_identifier or "gguf"
logger.info(
"Unloading GGUF chat model '%s' to free GPU memory for training (%s)",
name,
reason,
)
llama.unload_model()
freed.append(f"gguf:{name}")
except Exception as e:
logger.warning("Could not unload GGUF chat model: %s", e)
return freed

180
studio/backend/run.py
View file

@ -395,32 +395,76 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
pass
def _emit_startup_output(host: str, port: int, display_host: str) -> None:
"""Print the access banner plus any post-startup warnings.
def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> str:
"""One-line tool-policy summary for the plain-server startup banner, so a
network-reachable launch is never silent about code execution."""
if enable_tools is False:
return "Server-side tools are DISABLED (--disable-tools)."
state = (
"ENABLED (--enable-tools)"
if enable_tools
else "ENABLED by default (per-request setting honored)"
)
if secure:
return (
f"Server-side tools are {state}, reachable via the authenticated "
"Cloudflare HTTPS tunnel. Anyone with the API key can run code on "
"this machine. Do not share the API key. Pass --disable-tools to turn off."
)
from utils.host_policy import is_external_host
Extracted from ``_run`` so the banner/warning wiring is testable. The
``localhost``-to-::1 mismatch warning and the wildcard reachability
check are mutually exclusive (the mismatch helper returns None for any
non-127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the
trailing stop hint is emitted exactly once.
"""
if host in ("0.0.0.0", "::") or is_external_host(host):
return (
f"Server-side tools are {state} and this port is network-reachable. "
"Anyone who can reach it with the API key can run code on this "
"machine. Do not share the API key. Pass --disable-tools to turn off."
)
return f"Server-side tools are {state} for loopback. Pass --disable-tools to turn off."
def _emit_tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> None:
print(_tool_policy_notice(host, secure, enable_tools), flush = True)
def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None) -> None:
"""Secure-mode banner: only the Cloudflare link (loopback has no public raw URL)."""
print("")
print("🦥 Unsloth Studio is running (secure)")
print("─" * 52)
_print_cloudflare_line()
print(f" On this machine only: http://127.0.0.1:{port}/")
print("─" * 52)
_emit_tool_policy_notice("127.0.0.1", True, enable_tools)
print_studio_stop_hint()
def _emit_startup_output(
host: str,
port: int,
display_host: str,
secure: bool = False,
enable_tools: "Optional[bool]" = None,
) -> None:
"""Print the access banner, post-startup warnings, the tool-policy notice,
then a single stop hint. Extracted from ``_run`` so the wiring is testable."""
if secure:
_emit_secure_startup_output(port, enable_tools)
return
wildcard_bind = host in ("0.0.0.0", "::")
localhost_mismatch_url = _localhost_ipv6_mismatch_url(host, port)
# For wildcard binds, run the reachability check between the URL
# section and the stop hint so the stop hint stays last.
print_studio_access_banner(
port = port,
bind_host = host,
display_host = display_host,
include_stop_hint = not wildcard_bind and not localhost_mismatch_url,
include_stop_hint = False,
)
if localhost_mismatch_url:
_print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port)
print_studio_stop_hint()
elif wildcard_bind:
_verify_global_reachability(display_host, port)
_print_cloudflare_line()
print_studio_stop_hint()
_emit_tool_policy_notice(host, False, enable_tools)
print_studio_stop_hint()
def _print_cloudflare_line() -> None:
@ -623,6 +667,13 @@ def _graceful_shutdown(server = None):
except Exception as e:
logger.warning("Error stopping Cloudflare tunnel: %s", e)
# 7. Backstop sweep for any adopted child the steps above missed.
try:
from utils.process_lifetime import terminate_all
terminate_all()
except Exception as e:
logger.warning("Error in process-lifetime sweep: %s", e)
logger.info("All subprocesses cleaned up")
@ -811,6 +862,25 @@ def _setup_server_disk_logging():
return log_path
def _cloudflare_tunnel_should_start(
*, cloudflare: bool, host: str, secure: bool, api_only: bool, is_colab: bool
) -> bool:
"""Whether to start the Cloudflare tunnel. --secure tunnels a loopback bind too;
non-secure keeps the 0.0.0.0-only rule. Colab/api-only never tunnel."""
return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab
def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
"""Honor an explicit --enable-tools/--disable-tools; None leaves the policy
unset (tools default on, per-request enable_tools honored). Host is never
inspected here."""
if enable_tools is None:
return
from state.tool_policy import set_tool_policy
set_tool_policy(enable_tools)
def run_server(
host: str = "127.0.0.1",
port: int = 8888,
@ -819,6 +889,8 @@ def run_server(
api_only: bool = False,
llama_parallel_slots: int = 1,
cloudflare: bool = True,
secure: bool = False,
enable_tools: "Optional[bool]" = None,
):
"""
Start the FastAPI server.
@ -830,6 +902,8 @@ def run_server(
silent: Suppress startup messages
api_only: API server only, no frontend (for Tauri desktop app)
llama_parallel_slots: parallel slots for llama-server
enable_tools: explicit --enable-tools/--disable-tools policy; None leaves
the default (tools on, per-request enable_tools honored)
Note:
Signal handlers are NOT registered here so embedders (e.g. Colab) keep
@ -837,6 +911,24 @@ def run_server(
"""
global _server, _shutdown_event
# Reap every child if the parent dies abnormally (terminal close, Task
# Manager kill, SIGKILL); must run before any child can spawn.
from utils.process_lifetime import initialize_parent_lifetime
initialize_parent_lifetime()
# --secure exposes only the Cloudflare link: force a loopback bind so the raw
# port is never public (even with -H 0.0.0.0), and reject the contradictory combo.
if secure and not cloudflare:
raise SystemExit(
"A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link"
)
if secure:
host = "127.0.0.1"
# `unsloth studio run` installs its own resolved policy and passes None here.
_apply_cli_tool_policy(enable_tools)
# Windows cp1252 can't encode emoji; reconfigure stdout to UTF-8.
if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"):
try:
@ -975,6 +1067,13 @@ def run_server(
# backend, not whatever a proxy/tunnel exposed. For ephemeral binds (port==0)
# leave it unset so handlers fall back to the request scope / base_url.
app.state.server_port = port if port and port > 0 else None
# Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP.
if port and port > 0:
_direct_host = _resolve_external_ip() if host == "0.0.0.0" else host
app.state.server_url = f"http://{_direct_host}:{port}"
else:
app.state.server_url = None
app.state.secure = secure
app.state.llama_parallel_slots = llama_parallel_slots
# Expose a shutdown callable before the server accepts requests so
@ -1025,6 +1124,9 @@ def run_server(
import atexit
atexit.register(_remove_pid_file)
from utils.process_lifetime import terminate_all
atexit.register(terminate_all)
# Output port for Tauri (api-only), only after sockets bind and startup done.
if api_only:
@ -1036,7 +1138,13 @@ def run_server(
global _cloudflare_url
_cloudflare_url = None
app.state.cloudflare_url = None
_cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB
_cloudflare_enabled = _cloudflare_tunnel_should_start(
cloudflare = cloudflare,
host = host,
secure = secure,
api_only = api_only,
is_colab = _IS_COLAB,
)
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
@ -1049,8 +1157,19 @@ def run_server(
except Exception as e:
logger.debug("Cloudflare tunnel skipped: %s", e)
# --secure fails closed: no tunnel means no public link, so exit rather than
# silently fall back to a raw port.
if secure and not _cloudflare_url:
print(
"A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link",
file = sys.stderr,
flush = True,
)
_graceful_shutdown(_server)
sys.exit(1)
if not silent:
_emit_startup_output(host, port, display_host)
_emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools)
return app
@ -1094,6 +1213,31 @@ if __name__ == "__main__":
help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 "
"(default on; --no-cloudflare to disable)",
)
parser.add_argument(
"--secure",
action = argparse.BooleanOptionalAction,
default = False,
help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
"if the tunnel can't start. Without it, --not-secure also serves the raw "
"0.0.0.0 port, which is reachable from anywhere on the network",
)
# Tri-state tool policy: no flag -> None (tools on, per-request honored);
# --enable-tools/--disable-tools force on/off.
parser.add_argument(
"--enable-tools",
dest = "enable_tools",
action = "store_true",
default = None,
help = "Force server-side tools (web search, code execution) on for "
"every request. Default: on for every bind, per-request setting honored.",
)
parser.add_argument(
"--disable-tools",
dest = "enable_tools",
action = "store_false",
default = None,
help = "Force server-side tools off for every request.",
)
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct
# backend launches; `unsloth studio run` always passes its own value (4).
_PARALLEL_MIN = 1
@ -1113,6 +1257,10 @@ if __name__ == "__main__":
args = parser.parse_args()
if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
if args.secure and not args.cloudflare:
parser.error(
"--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare"
)
kwargs = dict(
host = args.host,
@ -1121,6 +1269,8 @@ if __name__ == "__main__":
api_only = args.api_only,
llama_parallel_slots = args.parallel,
cloudflare = args.cloudflare,
secure = args.secure,
enable_tools = args.enable_tools,
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)

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

@ -216,6 +216,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
created_at INTEGER NOT NULL,
openai_code_exec_container_id TEXT,
anthropic_code_exec_container_id TEXT,
forked_from_thread_id TEXT,
forked_from_message_id TEXT,
FOREIGN KEY(project_id) REFERENCES chat_projects(id) ON DELETE CASCADE
)
"""
@ -229,6 +231,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT")
if "anthropic_code_exec_container_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT")
if "forked_from_thread_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT")
if "forked_from_message_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
@ -956,6 +962,8 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict:
"createdAt": data["created_at"],
"openaiCodeExecContainerId": data.get("openai_code_exec_container_id"),
"anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"),
"forkedFromThreadId": data.get("forked_from_thread_id"),
"forkedFromMessageId": data.get("forked_from_message_id"),
}
@ -999,8 +1007,8 @@ def upsert_chat_thread(thread: dict) -> dict:
conn.execute(
"""
INSERT INTO chat_threads
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
model_type = excluded.model_type,
@ -1010,7 +1018,9 @@ def upsert_chat_thread(thread: dict) -> dict:
archived = excluded.archived,
created_at = excluded.created_at,
openai_code_exec_container_id = excluded.openai_code_exec_container_id,
anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id
anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id,
forked_from_thread_id = excluded.forked_from_thread_id,
forked_from_message_id = excluded.forked_from_message_id
""",
(
thread["id"],
@ -1023,6 +1033,8 @@ def upsert_chat_thread(thread: dict) -> dict:
int(thread["createdAt"]),
thread.get("openaiCodeExecContainerId"),
thread.get("anthropicCodeExecContainerId"),
thread.get("forkedFromThreadId"),
thread.get("forkedFromMessageId"),
),
)
conn.commit()
@ -1048,6 +1060,14 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]:
"anthropic_code_exec_container_id",
patch.get("anthropicCodeExecContainerId"),
),
"forkedFromThreadId": (
"forked_from_thread_id",
patch.get("forkedFromThreadId"),
),
"forkedFromMessageId": (
"forked_from_message_id",
patch.get("forkedFromMessageId"),
),
}
assignments = []
values = []
@ -1445,6 +1465,126 @@ def sync_chat_messages(
conn.close()
def fork_chat_thread(
source_thread_id: str,
branch_message_id: str,
new_thread_id: str,
new_title: str,
created_at: int,
id_factory,
) -> Optional[dict]:
"""Atomically clone thread + ancestor msgs `[root..branch_message_id]`
into a new thread. Returns the new thread dict (with messages copied)
or None if source missing.
Reset both code-exec container ids -- per-provider snapshot is handled
by the route layer (best-effort, OpenAI only).
`id_factory()` produces fresh message uuids; injected for testability.
"""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
src = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,)
).fetchone()
if src is None:
conn.rollback()
return None
# Verify branch msg belongs to source thread.
branch_row = conn.execute(
"SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?",
(source_thread_id, branch_message_id),
).fetchone()
if branch_row is None:
conn.rollback()
return None
# Walk ancestry from branch msg back to root via parent_id chain.
ancestry: list[sqlite3.Row] = []
cursor_row = branch_row
seen: set[str] = set()
while cursor_row is not None and cursor_row["id"] not in seen:
ancestry.append(cursor_row)
seen.add(cursor_row["id"])
parent = cursor_row["parent_id"]
if not parent:
break
cursor_row = conn.execute(
"SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?",
(source_thread_id, parent),
).fetchone()
ancestry.reverse() # root .. branch msg
# Map old msg id -> new msg id for parent_id rewriting.
id_map: dict[str, str] = {row["id"]: id_factory() for row in ancestry}
src_dict = dict(src)
conn.execute(
"""
INSERT INTO chat_threads
(id, title, model_type, model_id, pair_id, project_id, archived, created_at,
openai_code_exec_container_id, anthropic_code_exec_container_id,
forked_from_thread_id, forked_from_message_id)
VALUES (?, ?, ?, ?, ?, ?, 0, ?, NULL, NULL, ?, ?)
""",
(
new_thread_id,
new_title,
src_dict["model_type"],
src_dict.get("model_id") or "",
None, # pairId: forks always standalone (compare-mode disabled v1)
src_dict.get("project_id"),
int(created_at),
source_thread_id,
branch_message_id,
),
)
conn.executemany(
"""
INSERT INTO chat_messages
(id, thread_id, parent_id, role, content_json, attachments_json,
metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
[
(
id_map[row["id"]],
new_thread_id,
id_map.get(row["parent_id"]) if row["parent_id"] else None,
row["role"],
row["content_json"],
row["attachments_json"],
row["metadata_json"],
int(row["created_at"]),
)
for row in ancestry
],
)
conn.commit()
thread_row = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,)
).fetchone()
return _chat_thread_from_row(thread_row) if thread_row is not None else None
except Exception:
conn.rollback()
raise
finally:
conn.close()
def count_forks_for_message(thread_id: str, message_id: str) -> int:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT COUNT(*) FROM chat_threads
WHERE forked_from_thread_id = ? AND forked_from_message_id = ?
""",
(thread_id, message_id),
).fetchone()
return int(row[0]) if row is not None else 0
finally:
conn.close()
def list_chat_messages(thread_id: str) -> list[dict]:
conn = get_connection()
try:

60
studio/backend/tests/test_amd_apu_unified_memory.py
View file

@ -47,6 +47,30 @@ def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected):
assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected
def test_apu_guard_scopes_to_selected_gpu(monkeypatch):
# Mixed host: physical id 0 = discrete gfx1100, 1 = gfx1151 APU.
for _m in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
monkeypatch.delenv(_m, raising = False)
monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", ["gfx1100", "gfx1151"]))
# Selecting only the dGPU, or an empty selection, must not be unified-memory.
assert LlamaCppBackend._amd_apu_wants_unified_memory([0]) is False
assert LlamaCppBackend._amd_apu_wants_unified_memory([]) is False
# Selecting the APU, or no selection, does.
assert LlamaCppBackend._amd_apu_wants_unified_memory([1]) is True
assert LlamaCppBackend._amd_apu_wants_unified_memory() is True
def test_apu_guard_honors_hip_visible_devices_mask(monkeypatch):
# ROCm resolves ids via HIP first: the mask exposes only the APU as ordinal 0
# but physical id 1, so the selection [1] must still match.
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", ["gfx1151"]))
assert LlamaCppBackend._amd_apu_wants_unified_memory([1]) is True
assert LlamaCppBackend._amd_apu_wants_unified_memory([0]) is False
def test_cpu_no_cuda_returns_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False))
assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
@ -55,3 +79,39 @@ def test_cpu_no_cuda_returns_false(monkeypatch):
def test_missing_torch_returns_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", None)
assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
_GB = 1024**3
_MIB_PER_GB = 1024
# Module-level (not a class attr) so it stays a plain function, not a bound method.
_shortfall = LlamaCppBackend._apu_ram_shortfall_message
class TestApuRamShortfall:
"""On a unified-memory APU the weights load into system RAM, so a model
larger than available RAM (the field case: a 64.6 GB GGUF on a WSL VM capped
well below the ROCm-reported APU budget) must be refused before spawning,
not left to OOM-kill the Studio process."""
def test_field_case_wsl_cap_refuses(self):
# 64.6 GB weights, ~46 GB available (WSL VM): refuse with guidance.
msg = _shortfall(int(64.6 * _GB), 46 * _MIB_PER_GB)
assert msg is not None
assert "65 GB" in msg and "46 GB" in msg
assert ".wslconfig" in msg
def test_bare_metal_fits_allows(self):
# Same model, ~92 GB available (no WSL cap): allow.
assert _shortfall(int(64.6 * _GB), 92 * _MIB_PER_GB) is None
def test_unknown_available_never_refuses(self):
assert _shortfall(int(64.6 * _GB), None) is None
def test_boundary_at_headroom(self):
# 20 GB weights, headroom 2 GB. avail 23 GB -> fits; 21 GB -> refuse.
assert _shortfall(20 * _GB, 23 * _MIB_PER_GB) is None
assert _shortfall(20 * _GB, 21 * _MIB_PER_GB) is not None
def test_available_system_memory_is_int_or_none(self):
v = LlamaCppBackend._available_system_memory_mib()
assert v is None or (isinstance(v, int) and v > 0)

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

@ -33,6 +33,7 @@ from core.inference.anthropic_compat import (
AnthropicStreamEmitter,
AnthropicPassthroughEmitter,
)
from core.inference.api_monitor import ApiMonitor
from routes.inference import (
_build_tool_action_nudge,
_normalize_anthropic_openai_images,
@ -40,6 +41,7 @@ from routes.inference import (
_anthropic_requested_studio_tools,
_anthropic_passthrough_stream,
_anthropic_tool_non_streaming,
_monitor_anthropic_sse_line,
anthropic_messages,
)
from state.tool_policy import reset_tool_policy, set_tool_policy
@ -50,6 +52,46 @@ from io import BytesIO as _BytesIO
from types import SimpleNamespace
def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch):
import routes.inference as inf_mod
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monitor_id = monitor.start(
endpoint = "/v1/messages",
method = "POST",
model = "m",
prompt = "hi",
)
for payload in (
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "tool_use",
"id": "toolu_1",
"name": "lookup",
"input": {},
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": '{"query":"weather"}',
},
},
{"type": "content_block_stop", "index": 0},
):
_monitor_anthropic_sse_line(monitor_id, f"data: {json.dumps(payload)}")
entry = monitor.get(monitor_id)
assert entry is not None
assert entry["reply"] == 'Tool call: lookup\nInput: {"query":"weather"}'
# =====================================================================
# Tool nudge tests
# =====================================================================
@ -72,7 +114,7 @@ class TestToolActionNudge:
assert "Use code execution for math" in nudge
assert "render_html" not in nudge
def test_balanced_nudge_preserves_compact_web_tip_and_artifact_gate(self):
def test_balanced_nudge_preserves_compact_web_tip_and_canvas_gate(self):
nudge = _build_tool_action_nudge(
tools = [
{"type": "function", "function": {"name": "web_search"}},
@ -840,7 +882,7 @@ class TestAnthropicToolNonStreaming:
"type": "tool_end",
"tool_name": "render_html",
"tool_call_id": "call_0",
"result": "Rendered HTML artifact.",
"result": "Rendered HTML canvas.",
}
response = asyncio.run(_anthropic_tool_non_streaming(_run_gen, "msg_1", "m"))
@ -1385,7 +1427,7 @@ def _mock_backend(monkeypatch, **overrides):
def _gen_plain(**kwargs):
calls.append(("plain", kwargs))
yield {"type": "content", "text": "ok"}
yield "ok"
def _gen_tools(**kwargs):
calls.append(("tools", kwargs))
@ -1396,6 +1438,8 @@ def _mock_backend(monkeypatch, **overrides):
is_vision = False,
supports_tools = True,
model_identifier = "test-model",
context_length = 4096,
count_chat_tokens = lambda *args, **kwargs: 2,
generate_chat_completion = _gen_plain,
generate_chat_completion_with_tools = _gen_tools,
calls = calls,
@ -1426,6 +1470,112 @@ def _reset_policy():
class TestAnthropicMessagesToolRouting:
class _Request:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/messages")
method = "POST"
async def is_disconnected(self):
return False
@staticmethod
def _consume_response(response):
async def _consume():
chunks = []
async for chunk in response.body_iterator:
chunks.append(chunk)
return chunks
return _drive(_consume())
def test_plain_non_streaming_records_api_monitor_entry(self, monkeypatch):
import routes.inference as inf_mod
_mock_backend(monkeypatch, context_length = 2048)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
payload = _basic_payload()
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
assert response.status_code == 200
[entry] = monitor.snapshot()
assert entry["endpoint"] == "/v1/messages"
assert entry["status"] == "completed"
assert entry["model"] == "test-model"
assert entry["prompt_preview"] == "user: hi"
assert entry["reply_preview"] == "ok"
assert entry["context_length"] == 2048
assert monitor.active_count() == 0
def test_tool_use_non_streaming_records_api_monitor_reply(self, monkeypatch):
import routes.inference as inf_mod
def _gen_tools(**_kwargs):
yield {
"type": "tool_start",
"tool_call_id": "call_1",
"tool_name": "lookup",
"arguments": {"query": "weather"},
}
_mock_backend(
monkeypatch,
context_length = 2048,
generate_chat_completion_with_tools = _gen_tools,
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
payload = _basic_payload(
enable_tools = True,
tools = [{"type": "web_search_20250305", "name": "web_search"}],
)
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
assert response.status_code == 200
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
assert entry["reply_preview"] == 'Tool call: lookup({"query": "weather"})'
def test_plain_streaming_records_active_and_completed_monitor_entry(self, monkeypatch):
import routes.inference as inf_mod
_mock_backend(monkeypatch, context_length = 2048)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
payload = _basic_payload(stream = True)
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
assert monitor.active_count() == 1
self._consume_response(response)
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
assert entry["reply_preview"] == "ok"
assert entry["prompt_tokens"] == 2
assert entry["context_length"] == 2048
assert monitor.active_count() == 0
def test_plain_streaming_pre_response_cancel_finalizes_monitor(self, monkeypatch):
import routes.inference as inf_mod
async def _cancelled_before_response(*_args, **_kwargs):
raise asyncio.CancelledError()
_mock_backend(monkeypatch, context_length = 2048)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_anthropic_plain_stream", _cancelled_before_response)
payload = _basic_payload(stream = True)
with pytest.raises(asyncio.CancelledError):
_drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
[entry] = monitor.snapshot()
assert entry["status"] == "cancelled"
assert monitor.active_count() == 0
def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch):
_mock_backend(monkeypatch)
payload = _basic_payload(

198
studio/backend/tests/test_api_key_expiry.py Normal file
View file

@ -0,0 +1,198 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Expiry enforcement for API keys (tz-aware ``expires_at``) and JWT access
tokens (``exp`` claim). Both must surface as 401 on protected routes."""
from __future__ import annotations
import asyncio
import secrets
from datetime import datetime, timedelta, timezone
import pytest
from fastapi import HTTPException
from fastapi.security import HTTPAuthorizationCredentials
from auth import storage
from auth.authentication import create_access_token, get_current_subject
@pytest.fixture(autouse = True)
def isolated_auth_db(tmp_path, monkeypatch):
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password")
monkeypatch.setattr(storage, "_bootstrap_password", None)
monkeypatch.setattr(storage, "_api_key_pbkdf2_salt_cache", None)
storage._reset_api_key_hash_cache()
yield
storage._reset_api_key_hash_cache()
def seed_user():
storage.create_initial_user(
username = storage.DEFAULT_ADMIN_USERNAME,
password = "human-password-123",
jwt_secret = secrets.token_urlsafe(64),
)
def iso_from_now(**delta):
return (datetime.now(timezone.utc) + timedelta(**delta)).isoformat()
def make_key(expires_at):
raw, _row = storage.create_api_key(
username = storage.DEFAULT_ADMIN_USERNAME,
name = "test",
expires_at = expires_at,
)
return raw
def subject_of(token):
"""Run the real FastAPI auth dependency against a bearer token."""
credentials = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = token)
return asyncio.run(get_current_subject(credentials))
# --- validate_api_key (storage layer) ---------------------------------------
def test_unexpired_key_validates():
seed_user()
assert (
storage.validate_api_key(make_key(iso_from_now(days = 1))) == storage.DEFAULT_ADMIN_USERNAME
)
def test_never_expiring_key_validates():
seed_user()
assert storage.validate_api_key(make_key(None)) == storage.DEFAULT_ADMIN_USERNAME
def test_expired_key_rejected():
seed_user()
assert storage.validate_api_key(make_key(iso_from_now(seconds = -1))) is None
def test_key_expiring_far_in_past_rejected():
seed_user()
assert storage.validate_api_key(make_key(iso_from_now(days = -30))) is None
def test_revoked_key_rejected():
seed_user()
raw, row = storage.create_api_key(
username = storage.DEFAULT_ADMIN_USERNAME,
name = "doomed",
expires_at = iso_from_now(days = 1),
)
storage.revoke_api_key(storage.DEFAULT_ADMIN_USERNAME, int(row["id"]))
assert storage.validate_api_key(raw) is None
def test_unknown_key_rejected():
seed_user()
assert storage.validate_api_key(storage.API_KEY_PREFIX + secrets.token_hex(16)) is None
# --- get_current_subject (route dependency) ---------------------------------
def test_dependency_accepts_unexpired_key():
seed_user()
assert subject_of(make_key(iso_from_now(days = 1))) == storage.DEFAULT_ADMIN_USERNAME
def test_dependency_rejects_expired_key_as_401():
seed_user()
with pytest.raises(HTTPException) as exc:
subject_of(make_key(iso_from_now(seconds = -1)))
assert exc.value.status_code == 401
assert exc.value.detail == "Invalid or expired API key"
# --- JWT access-token expiry ------------------------------------------------
def test_dependency_accepts_unexpired_jwt():
seed_user()
token = create_access_token(storage.DEFAULT_ADMIN_USERNAME, timedelta(minutes = 5))
assert subject_of(token) == storage.DEFAULT_ADMIN_USERNAME
def test_dependency_rejects_expired_jwt_as_401():
seed_user()
token = create_access_token(storage.DEFAULT_ADMIN_USERNAME, timedelta(seconds = -1))
with pytest.raises(HTTPException) as exc:
subject_of(token)
assert exc.value.status_code == 401
assert exc.value.detail == "Invalid or expired token"
# --- derivation cache: speeds repeats without bypassing checks --------------
def test_cache_skips_pbkdf2_on_repeat(monkeypatch):
seed_user()
raw = make_key(iso_from_now(days = 1))
assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME # warms cache
calls = {"n": 0}
real = storage._pbkdf2_api_key
def counting(key):
calls["n"] += 1
return real(key)
monkeypatch.setattr(storage, "_pbkdf2_api_key", counting)
assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
assert calls["n"] == 0 # served from cache, KDF not re-run
def test_cache_does_not_bypass_revocation():
seed_user()
raw, row = storage.create_api_key(
username = storage.DEFAULT_ADMIN_USERNAME,
name = "revoke-after-cache",
expires_at = iso_from_now(days = 1),
)
assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME # cached
storage.revoke_api_key(storage.DEFAULT_ADMIN_USERNAME, int(row["id"]))
assert storage.validate_api_key(raw) is None # cache hit still re-checks is_active
def test_cache_does_not_bypass_expiry():
seed_user()
# Expires between the two calls: the first warms the cache, the second is still rejected.
near = (datetime.now(timezone.utc) + timedelta(milliseconds = 600)).isoformat()
raw = make_key(near)
assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
import time
time.sleep(0.8)
assert storage.validate_api_key(raw) is None
def test_unknown_key_not_cached():
seed_user()
bogus = storage.API_KEY_PREFIX + secrets.token_hex(16)
assert storage.validate_api_key(bogus) is None
cache_id = storage._api_key_cache_id(bogus)
assert cache_id not in storage._api_key_hash_cache # spam can't grow the cache
def test_create_api_key_route_stores_tz_aware_expiry():
from datetime import datetime as _dt
seed_user()
raw, row = storage.create_api_key(
username = storage.DEFAULT_ADMIN_USERNAME,
name = "route",
expires_at = iso_from_now(days = 30),
)
parsed = _dt.fromisoformat(row["expires_at"])
assert parsed.tzinfo is not None # tz-aware: comparison in validate_api_key won't raise
assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME

260
studio/backend/tests/test_api_monitor.py Normal file
View file

@ -0,0 +1,260 @@
# 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 core.inference.api_monitor import ApiMonitor, _trim
def test_api_monitor_tracks_reply_usage_and_context():
monitor = ApiMonitor(max_entries = 3)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "local-model",
prompt = "user: hello",
context_length = 100,
)
monitor.append_reply(entry_id, "hi")
monitor.append_reply(entry_id, " there")
monitor.set_usage(
entry_id,
prompt_tokens = 4,
completion_tokens = 6,
)
monitor.finish(entry_id)
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
assert entry["reply"] == "hi there"
assert entry["total_tokens"] == 10
assert entry["context_usage"] == 0.1
assert entry["duration_ms"] is not None
def test_api_monitor_summary_omits_full_prompt_and_reply():
monitor = ApiMonitor(max_entries = 3)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "local-model",
prompt = "p" * 500,
)
monitor.set_reply(entry_id, "r" * 500)
[summary] = monitor.snapshot(include_details = False)
assert "prompt" not in summary
assert "reply" not in summary
assert summary["prompt_preview"].endswith("...")
assert summary["reply_preview"].endswith("...")
assert summary["prompt_truncated"] is True
assert summary["reply_truncated"] is True
detail = monitor.get(entry_id)
assert detail is not None
assert detail["prompt"] == "p" * 500
assert detail["reply"] == "r" * 500
def test_api_monitor_filters_entries_by_subject():
monitor = ApiMonitor(max_entries = 3)
alice = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "alice prompt",
subject = "alice",
)
bob = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "bob prompt",
subject = "bob",
)
monitor.finish(bob)
alice_entries = monitor.snapshot(subject = "alice")
assert [entry["id"] for entry in alice_entries] == [alice]
assert monitor.get(bob, subject = "alice") is None
assert monitor.get(bob, subject = "bob")["id"] == bob
assert monitor.active_count(subject = "alice") == 1
assert monitor.active_count(subject = "bob") == 0
def test_api_monitor_keeps_bounded_recent_history():
monitor = ApiMonitor(max_entries = 2)
first = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "first",
)
second = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "second",
)
third = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "third",
)
monitor.finish(first)
monitor.finish(second)
monitor.finish(third)
entries = monitor.snapshot()
ids = [entry["id"] for entry in entries]
assert ids[0] == third
assert [entry["prompt"] for entry in entries] == ["third", "second"]
assert first not in ids
assert monitor.active_count() == 0
def test_api_monitor_keeps_running_entries_beyond_history_limit():
monitor = ApiMonitor(max_entries = 1)
running = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "running",
)
for prompt in ("done-1", "done-2", "done-3"):
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = prompt,
)
monitor.finish(entry_id)
entries = monitor.snapshot()
ids = [entry["id"] for entry in entries]
assert running in ids
assert monitor.active_count() == 1
monitor.finish(running)
[entry] = monitor.snapshot()
assert entry["id"] == running
assert entry["status"] == "completed"
assert monitor.active_count() == 0
def test_api_monitor_finish_is_idempotent():
monitor = ApiMonitor(max_entries = 2)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
)
monitor.finish(entry_id)
first = monitor.snapshot()[0]
monitor.finish(entry_id)
second = monitor.snapshot()[0]
assert first["finished_at"] == second["finished_at"]
assert first["duration_ms"] == second["duration_ms"]
def test_api_monitor_preserves_authoritative_total_tokens():
monitor = ApiMonitor(max_entries = 2)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
)
monitor.set_usage(
entry_id,
prompt_tokens = 10,
completion_tokens = 20,
total_tokens = 33,
)
# A later partial chunk omitting `total_tokens` must not clobber 33.
monitor.set_usage(entry_id, prompt_tokens = 11)
assert monitor.snapshot()[0]["total_tokens"] == 33
def test_api_monitor_recomputes_derived_total_tokens():
monitor = ApiMonitor(max_entries = 2)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
)
monitor.set_usage(entry_id, prompt_tokens = 10)
assert monitor.snapshot()[0]["total_tokens"] == 10
monitor.set_usage(entry_id, completion_tokens = 20)
entry = monitor.snapshot()[0]
assert entry["prompt_tokens"] == 10
assert entry["completion_tokens"] == 20
assert entry["total_tokens"] == 30
def test_api_monitor_duration_non_negative_under_clock_step(monkeypatch):
import core.inference.api_monitor as m
fake_now = [1000.0]
monkeypatch.setattr(m.time, "time", lambda: fake_now[0])
monitor = ApiMonitor(max_entries = 1)
entry_id = monitor.start(
endpoint = "/x",
method = "POST",
model = "m",
prompt = "hi",
)
fake_now[0] = 500.0
monitor.finish(entry_id)
assert monitor.snapshot()[0]["duration_ms"] >= 0
def test_api_monitor_trim_guards_tiny_limit():
assert _trim("abcdefgh", 2) == ".."
assert _trim("abcdefgh", 0) == ""
assert _trim("abcdefgh", 3) == "..."
assert _trim("abcdefgh", 4) == "a..."
assert _trim("abcdefgh", 100) == "abcdefgh"
def test_api_monitor_append_reply_caps_without_regrowing():
import core.inference.api_monitor as m
monitor = ApiMonitor(max_entries = 1)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "go",
)
monitor.append_reply(entry_id, "x" * (m._MAX_REPLY_CHARS + 500))
capped = monitor.snapshot()[0]["reply"]
assert len(capped) == m._MAX_REPLY_CHARS and capped.endswith("...")
# Chunks past the cap must not change or grow the stored preview.
monitor.append_reply(entry_id, "y" * 1000)
assert monitor.snapshot()[0]["reply"] == capped
def test_api_monitor_append_reply_exact_cap_then_more_marks_truncated():
import core.inference.api_monitor as m
monitor = ApiMonitor(max_entries = 1)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "go",
)
# A reply landing exactly on the cap has no "..." marker yet.
monitor.append_reply(entry_id, "x" * m._MAX_REPLY_CHARS)
assert not monitor.snapshot()[0]["reply"].endswith("...")
# One more chunk must record the truncation, not silently freeze.
monitor.append_reply(entry_id, "y")
reply = monitor.snapshot()[0]["reply"]
assert len(reply) == m._MAX_REPLY_CHARS and reply.endswith("...")

84
studio/backend/tests/test_api_perf_serialization.py Normal file
View file

@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""_model_json_response produces the same body as JSONResponse(model.model_dump())."""
import asyncio
import json
from typing import Optional
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import routes.inference as inference_route
from core.inference import llama_http
class _Usage(BaseModel):
prompt_tokens: int = 3
completion_tokens: int = 5
details: Optional[dict] = None
class _Choice(BaseModel):
index: int = 0
text: str = "hello"
logprobs: Optional[dict] = None
class _Resp(BaseModel):
id: str = "chatcmpl-abc"
object: str = "chat.completion"
created: int = 1700000000
model: str = "unsloth/SmolLM2-135M-Instruct-GGUF"
choices: list[_Choice] = [_Choice()]
usage: _Usage = _Usage()
system_fingerprint: Optional[str] = None
def _old_body(model) -> bytes:
# What the previous code emitted: dict -> Starlette json.dumps.
return JSONResponse(content = model.model_dump()).body
def test_body_matches_old_jsonresponse():
model = _Resp()
resp = inference_route._model_json_response(model)
# Same decoded JSON (key order is irrelevant once parsed), nulls preserved.
assert json.loads(resp.body) == json.loads(_old_body(model))
assert json.loads(resp.body)["system_fingerprint"] is None # null kept, not dropped
def test_media_type_and_status():
resp = inference_route._model_json_response(_Resp(), status_code = 200)
assert resp.media_type == "application/json"
assert resp.status_code == 200
err = inference_route._model_json_response(_Resp(), status_code = 503)
assert err.status_code == 503
def test_pooled_client_reused_within_loop_and_recreated_after_close():
async def _scenario():
a = llama_http.nonstreaming_client()
b = llama_http.nonstreaming_client()
assert a is b # reused within one loop
await llama_http.aclose()
assert a.is_closed
c = llama_http.nonstreaming_client() # must not return the closed client
assert c is not a and not c.is_closed
await llama_http.aclose()
asyncio.run(_scenario())
def test_pooled_client_is_per_event_loop():
clients = []
# Each asyncio.run uses a fresh loop; the pooled client must not leak across.
for _ in range(2):
async def _grab():
clients.append(llama_http.nonstreaming_client())
await llama_http.aclose()
asyncio.run(_grab())
assert clients[0] is not clients[1]

732
studio/backend/tests/test_bypass_permissions.py Normal file
View file

@ -0,0 +1,732 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for Bypass Permissions (skip confirmation + disable sandbox).
Covers the secret-name classifier, the two env builders, the
``disable_sandbox`` branch of ``_python_exec`` / ``_bash_exec`` (which env is
used, which pre-exec is used, and that safety checks / the blocklist are
skipped), the request-model default, the confirm-vs-bypass precedence rule the
route enforces, and that the agentic loop forwards ``disable_sandbox`` while
never gating under bypass.
Run with: ``PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_bypass_permissions.py -q``
"""
import os
import sys
import pytest
import core.inference.tools as tools
from core.inference.tools import (
_bash_exec,
_build_bypass_env,
_build_safe_env,
_is_cred_location_env_name,
_is_secret_env_name,
_is_secret_env_value,
_python_exec,
)
from core.inference.safetensors_agentic import run_safetensors_tool_loop
_POSIX_ONLY = pytest.mark.skipif(
sys.platform == "win32", reason = "preexec_fn / setsid are POSIX-only"
)
# ── secret-name classifier ──────────────────────────────────────────
@pytest.mark.parametrize(
"name",
[
"HF_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"WANDB_API_KEY",
"GH_TOKEN",
"GITHUB_TOKEN",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"AWS_SECRET_ACCESS_KEY",
"AWS_ACCESS_KEY_ID",
"AZURE_CLIENT_SECRET",
"GOOGLE_APPLICATION_CREDENTIALS",
"MY_DB_PASSWORD",
"x_api_key",
"SOME_PRIVATE_KEY",
"LD_PRELOAD",
],
)
def test_secret_names_are_flagged(name):
assert _is_secret_env_name(name) is True
@pytest.mark.parametrize(
"name", ["PATH", "HOME", "LANG", "TERM", "PWD", "SHELL", "HOSTVAR", "MY_VAR"]
)
def test_benign_names_are_not_flagged(name):
assert _is_secret_env_name(name) is False
# ── env builders ────────────────────────────────────────────────────
def test_bypass_env_keeps_benign_strips_secret_repoints_home(monkeypatch, tmp_path):
monkeypatch.setenv("HOSTVAR", "benign-123")
monkeypatch.setenv("HF_TOKEN", "secret-abc")
env = _build_bypass_env(str(tmp_path))
assert env.get("HOSTVAR") == "benign-123" # full host env inherited
assert "HF_TOKEN" not in env # ...minus secrets
assert env["HOME"] == str(tmp_path) # $HOME-based cred lookups defused
assert env["TMPDIR"] == str(tmp_path)
def test_safe_env_excludes_host_and_secret(monkeypatch, tmp_path):
monkeypatch.setenv("HOSTVAR", "benign-123")
monkeypatch.setenv("HF_TOKEN", "secret-abc")
env = _build_safe_env(str(tmp_path))
assert "HOSTVAR" not in env # whitelist build -> host vars never reach child
assert "HF_TOKEN" not in env
# ── Popen kwargs capture (no real execution) ────────────────────────
class _FakeProc:
returncode = 0
def communicate(self, timeout = None):
return ("FAKEOUT", None)
def poll(self):
return 0
def kill(self):
pass
@pytest.fixture
def captured_popen(monkeypatch):
cap = {}
def fake_popen(cmd, **kwargs):
cap["cmd"] = cmd
cap["kwargs"] = kwargs
return _FakeProc()
monkeypatch.setattr(tools.subprocess, "Popen", fake_popen)
return cap
@_POSIX_ONLY
def test_python_sandboxed_uses_sandbox_preexec_and_safe_env(captured_popen, monkeypatch):
monkeypatch.setenv("HF_TOKEN", "secret-abc")
_python_exec("print(1)", None, 5, "t", disable_sandbox = False)
assert captured_popen["kwargs"]["preexec_fn"] is tools._sandbox_preexec
assert "HF_TOKEN" not in captured_popen["kwargs"]["env"]
@_POSIX_ONLY
def test_python_bypass_uses_bypass_preexec_and_bypass_env(captured_popen, monkeypatch):
monkeypatch.setenv("HOSTVAR", "benign-xyz")
monkeypatch.setenv("HF_TOKEN", "secret-abc")
_python_exec("print(1)", None, 5, "t", disable_sandbox = True)
assert captured_popen["kwargs"]["preexec_fn"] is tools._bypass_preexec
env = captured_popen["kwargs"]["env"]
assert env.get("HOSTVAR") == "benign-xyz"
assert "HF_TOKEN" not in env
def test_bash_blocklist_enforced_when_sandboxed(captured_popen):
out = _bash_exec("rm -rf /", None, 5, "t", disable_sandbox = False)
assert "Blocked" in out
assert "cmd" not in captured_popen # never reached Popen
def test_bash_blocklist_skipped_when_bypassed(captured_popen):
out = _bash_exec("rm -rf /", None, 5, "t", disable_sandbox = True)
assert out == "FAKEOUT" # blocklist skipped -> reached (faked) execution
assert captured_popen["cmd"][0] in ("bash", "cmd")
@_POSIX_ONLY
def test_bash_bypass_uses_bypass_preexec(captured_popen):
_bash_exec("echo hi", None, 5, "t", disable_sandbox = True)
assert captured_popen["kwargs"]["preexec_fn"] is tools._bypass_preexec
# ── real end-to-end python execution under bypass ───────────────────
@_POSIX_ONLY
def test_python_bypass_real_exec_sees_host_env_but_not_secret(monkeypatch):
monkeypatch.setenv("HOSTVAR", "benign-xyz")
monkeypatch.setenv("HF_TOKEN", "secret-pqr")
code = (
"import os;"
"print('H=' + str(os.environ.get('HOSTVAR')),"
" 'T=' + str(os.environ.get('HF_TOKEN')))"
)
out = _python_exec(code, None, 30, "test-bypass", disable_sandbox = True)
assert "H=benign-xyz" in out # unrestricted: real host var visible
assert "T=None" in out # ...but the secret was stripped
assert "secret-pqr" not in out
# ── _bypass_preexec is setsid-only (no rlimits) ─────────────────────
@_POSIX_ONLY
def test_bypass_preexec_only_sets_session(monkeypatch):
calls = {"setsid": 0}
monkeypatch.setattr(
tools.os, "setsid", lambda: calls.__setitem__("setsid", calls["setsid"] + 1)
)
# _resource must not be touched by the bypass pre-exec.
if tools._resource is not None:
monkeypatch.setattr(
tools._resource,
"setrlimit",
lambda *a, **k: pytest.fail("bypass pre-exec must not set rlimits"),
)
tools._bypass_preexec()
assert calls["setsid"] == 1
# ── request model default ───────────────────────────────────────────
def test_request_model_bypass_default_false():
from models.inference import ChatCompletionRequest
assert ChatCompletionRequest.model_fields["bypass_permissions"].default is False
# ── confirm-vs-bypass precedence (mirrors the route rule) ───────────
@pytest.mark.parametrize(
"confirm,bypass,effective_confirm",
[
(False, False, False),
(True, False, True),
(False, True, False),
(True, True, False),
],
)
def test_confirm_precedence_rule(confirm, bypass, effective_confirm):
# The route computes: confirm_tool_calls = confirm and not bypass.
assert (bool(confirm) and not bool(bypass)) is effective_confirm
# ── agentic loop forwards disable_sandbox, never gates under bypass ──
_DEFAULT_TOOLS = [
{"type": "function", "function": {"name": "python"}},
{"type": "function", "function": {"name": "web_search"}},
]
def _tool_call(name, args_json):
return f'<tool_call>{{"name": "{name}", "arguments": {args_json}}}</tool_call>'
def _multi_turn(turns):
it = iter(turns)
def _gen(_messages):
try:
yield next(it)
except StopIteration:
return
return _gen
def test_loop_forwards_disable_sandbox_and_does_not_gate():
seen = []
def fake_exec(
name,
arguments,
*,
cancel_event = None,
timeout = None,
session_id = None,
rag_scope = None,
disable_sandbox = False,
):
seen.append(disable_sandbox)
return f"RAN[{name}]"
events = list(
run_safetensors_tool_loop(
single_turn = _multi_turn([_tool_call("python", '{"code": "x"}'), "done"]),
messages = [{"role": "user", "content": "hi"}],
tools = _DEFAULT_TOOLS,
execute_tool = fake_exec,
session_id = "s",
confirm_tool_calls = False, # route forces this off under bypass
bypass_permissions = True,
)
)
assert seen == [True] # disable_sandbox threaded through
starts = [e for e in events if e["type"] == "tool_start"]
assert starts and starts[0]["awaiting_confirmation"] is False
assert starts[0]["approval_id"] == ""
def test_loop_bypass_overrides_confirm_for_direct_callers():
# Even if a direct internal caller passes confirm_tool_calls=True, bypass
# must suppress the confirm gate at the loop level (not only at the route).
def fake_exec(
name,
arguments,
*,
cancel_event = None,
timeout = None,
session_id = None,
rag_scope = None,
disable_sandbox = False,
):
return f"RAN[{name}]"
events = list(
run_safetensors_tool_loop(
single_turn = _multi_turn([_tool_call("python", '{"code": "x"}'), "done"]),
messages = [{"role": "user", "content": "hi"}],
tools = _DEFAULT_TOOLS,
execute_tool = fake_exec,
session_id = "s",
confirm_tool_calls = True, # raw caller leaves this on...
bypass_permissions = True, # ...but bypass must still win
)
)
starts = [e for e in events if e["type"] == "tool_start"]
assert starts and starts[0]["awaiting_confirmation"] is False
assert starts[0]["approval_id"] == ""
def test_gguf_loop_confirm_gate_respects_bypass():
# The GGUF loop needs a live llama-server, so (per the other llama_cpp
# tests) assert via AST that its _needs_confirm gate applies the bypass
# precedence, mirroring the safetensors behavioral test above.
import ast
import inspect
import textwrap
llama_cpp = pytest.importorskip("core.inference.llama_cpp")
src = textwrap.dedent(
inspect.getsource(llama_cpp.LlamaCppBackend.generate_chat_completion_with_tools)
)
gates = [
node
for node in ast.walk(ast.parse(src))
if isinstance(node, ast.Assign)
and any(getattr(t, "id", None) == "needs_confirm" for t in node.targets)
]
assert gates, "could not find the needs_confirm gate in the GGUF loop"
names = {n.id for g in gates for n in ast.walk(g.value) if isinstance(n, ast.Name)}
assert "confirm_tool_calls" in names
assert "bypass_permissions" in names # bypass must suppress the GGUF gate
# ── broker / capability env vars are stripped (regression) ──────────
@pytest.mark.parametrize(
"name",
["SSH_AUTH_SOCK", "SSH_AGENT_PID", "GPG_AGENT_INFO", "GNUPGHOME", "KUBECONFIG"],
)
def test_broker_capability_names_are_flagged(name):
# Not secrets by value, but they hand the child the operator's live agent
# (ssh/gpg) or kube credentials, so bypass mode must drop them.
assert _is_secret_env_name(name) is True
# ── credential-bearing URL values stripped regardless of name ───────
@pytest.mark.parametrize(
"value",
[
"https://user:s3cr3t@feed.example.invalid/simple", # user:pass@
"https://ghp_deadbeef@github.com/org/private.git", # token-only@
"https://__token__@pypi.example.invalid/simple",
"https://ghp_1234:@npm.pkg.github.com/simple", # empty password
"postgres://dbuser:dbpass@db.example.invalid/app",
],
)
def test_url_userinfo_values_are_flagged(value):
assert _is_secret_env_value(value) is True
@pytest.mark.parametrize(
"value",
[
"https://example.invalid/simple", # no userinfo
"http://proxy.corp.example:8080", # benign proxy
"https://pypi.corp.example/simple", # benign internal index
"redis://localhost:6379/0", # no creds
"https://example.invalid/path?ref=a@b", # '@' only in query, not userinfo
],
)
def test_non_credential_url_values_are_not_flagged(value):
assert _is_secret_env_value(value) is False
def test_url_userinfo_value_is_stripped_even_with_benign_name(monkeypatch, tmp_path):
# NAME dodges the classifier, but the VALUE embeds userinfo -> must go.
monkeypatch.setenv("MY_FEED", "https://user:s3cr3t@feed.example.invalid/simple")
monkeypatch.setenv("REPO_URL", "https://ghp_deadbeef@github.com/org/private.git")
# A URL without credentials is harmless and should be kept.
monkeypatch.setenv("PLAIN_URL", "https://example.invalid/simple")
env = _build_bypass_env(str(tmp_path))
assert "MY_FEED" not in env
assert "REPO_URL" not in env
assert env.get("PLAIN_URL") == "https://example.invalid/simple"
def test_bypass_env_keeps_noncredential_proxy_and_index_urls(monkeypatch, tmp_path):
# Benign routing/config vars must survive bypass mode (proxy-only or
# internal-index networks); only credentialed values are dropped.
monkeypatch.setenv("HTTP_PROXY", "http://proxy.corp.example:8080")
monkeypatch.setenv("PIP_INDEX_URL", "https://pypi.corp.example/simple")
monkeypatch.setenv("PIP_EXTRA_INDEX_URL", "https://user:token@pypi.example.invalid/simple")
env = _build_bypass_env(str(tmp_path))
assert env["HTTP_PROXY"] == "http://proxy.corp.example:8080"
assert env["PIP_INDEX_URL"] == "https://pypi.corp.example/simple"
assert "PIP_EXTRA_INDEX_URL" not in env # this one carries credentials
# ── AWS IMDS-disable hardening flag is kept (regression) ────────────
def test_aws_imds_disable_flag_is_kept_but_creds_stripped(monkeypatch, tmp_path):
# AWS_EC2_METADATA_DISABLED is a non-secret opt-out: dropping it would let a
# bypassed boto/AWS-CLI call fall back to the instance role via IMDS even
# though the operator disabled that path. Keep it; drop the real creds.
monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "shhh")
assert _is_secret_env_name("AWS_EC2_METADATA_DISABLED") is False
assert _is_secret_env_name("AWS_ACCESS_KEY_ID") is True
env = _build_bypass_env(str(tmp_path))
assert env.get("AWS_EC2_METADATA_DISABLED") == "true"
assert "AWS_ACCESS_KEY_ID" not in env
assert "AWS_SECRET_ACCESS_KEY" not in env
# ── connection-string env vars are stripped (regression) ────────────
@pytest.mark.parametrize(
"name",
[
"SQLCONNSTR_DB", # Azure App Service injected connection strings
"MYSQLCONNSTR_DB",
"SQLAZURECONNSTR_DB",
"POSTGRESQLCONNSTR_DB",
"CUSTOMCONNSTR_CACHE",
"WEBSITE_CONTENTAZUREFILECONNECTIONSTRING",
],
)
def test_connection_string_names_are_flagged(name):
assert _is_secret_env_name(name) is True
@pytest.mark.parametrize(
"value",
[
"Server=tcp:db;Database=app;User ID=u;Password=p@ss;", # ADO.NET
"DefaultEndpointsProtocol=https;AccountName=x;AccountKey=abc123==;", # storage
"Endpoint=sb://x;SharedAccessKeyName=n;SharedAccessKey=zzz=", # Service Bus
],
)
def test_connection_string_values_are_flagged(value):
assert _is_secret_env_value(value) is True
@pytest.mark.parametrize(
"value",
[
"Server=tcp:db;Database=app;User ID=u;", # no password field
"Endpoint=sb://x;SharedAccessKeyName=n", # key NAME only, no secret
"AccountName=x;EndpointSuffix=core.windows.net", # no AccountKey
],
)
def test_connection_string_noncredential_values_are_not_flagged(value):
assert _is_secret_env_value(value) is False
def test_connection_string_value_stripped_even_with_benign_name(monkeypatch, tmp_path):
# NAME dodges the classifier, but the VALUE is a credentialed conn string.
monkeypatch.setenv("APP_DB", "Server=tcp:db;Database=app;User ID=u;Password=p@ss;")
monkeypatch.setenv("SQLCONNSTR_DB", "DefaultEndpointsProtocol=https;AccountKey=abc==")
env = _build_bypass_env(str(tmp_path))
assert "APP_DB" not in env # value-based catch
assert "SQLCONNSTR_DB" not in env # name-based catch
# ── temp dirs repointed on every platform (regression) ──────────────
def test_bypass_env_repoints_all_temp_vars(monkeypatch, tmp_path):
# Windows tempfile honours TEMP/TMP, not TMPDIR; all three must repoint.
monkeypatch.setenv("TEMP", "/host/tmp")
monkeypatch.setenv("TMP", "/host/tmp")
env = _build_bypass_env(str(tmp_path))
assert env["TMPDIR"] == str(tmp_path)
assert env["TEMP"] == str(tmp_path)
assert env["TMP"] == str(tmp_path)
# ── credential-location redirect vars are dropped (regression) ──────────
# Vars that point SDKs at the real home/cache/config (cached tokens), e.g.
# HF_HOME which startup always sets -> the live leak the HOME repoint missed.
@pytest.mark.parametrize(
"name",
[
"HF_HOME",
"HF_HUB_CACHE",
"HUGGINGFACE_HUB_CACHE",
"HF_XET_CACHE",
"TRANSFORMERS_CACHE",
"HF_DATASETS_CACHE",
"XDG_CONFIG_HOME",
"XDG_CACHE_HOME",
"XDG_DATA_HOME",
"NETRC",
"BOTO_CONFIG",
"PIP_CONFIG_FILE",
"CLOUDSDK_CONFIG",
"KAGGLE_CONFIG_DIR",
"DOCKER_CONFIG",
"WANDB_DIR",
"WANDB_CONFIG_DIR",
"NPM_CONFIG_USERCONFIG",
"NPM_CONFIG_GLOBALCONFIG",
"YARN_RC_FILENAME",
"GIT_CONFIG_GLOBAL",
"GIT_CONFIG_SYSTEM",
"CARGO_HOME",
"RCLONE_CONFIG",
"GIT_ASKPASS",
"SSH_ASKPASS",
"BASH_ENV",
"HOMEDRIVE",
"HOMEPATH",
],
)
def test_cred_location_names_are_flagged(name):
assert _is_cred_location_env_name(name) is True
@pytest.mark.parametrize("name", ["PATH", "HOME", "LANG", "PWD", "MY_VAR"])
def test_benign_names_not_flagged_as_cred_location(name):
assert _is_cred_location_env_name(name) is False
def test_bypass_env_drops_hf_home_so_cached_token_unreachable(monkeypatch, tmp_path):
# The live leak: startup sets HF_HOME at the real cache, whose $HF_HOME/token
# holds the operator's token. Repointing HOME does not stop huggingface_hub
# from reading $HF_HOME/token, so HF_HOME must be dropped in bypass mode.
real_cache = tmp_path / "real_hf_cache"
real_cache.mkdir()
(real_cache / "token").write_text("hf_cachedOperatorToken")
monkeypatch.setenv("HF_HOME", str(real_cache))
monkeypatch.setenv("HF_HUB_CACHE", str(real_cache / "hub"))
env = _build_bypass_env(str(tmp_path))
assert "HF_HOME" not in env # dropped -> HF falls back to $HOME/.cache (empty)
assert "HF_HUB_CACHE" not in env
def test_bypass_env_hf_token_resolves_outside_real_cache(monkeypatch, tmp_path):
# End-to-end: even when HF_HOME and XDG_CACHE_HOME both point at the real
# cache, the bypass env must make huggingface_hub resolve the token under the
# workdir (guards the XDG fallback chain, not just "HF_HOME absent").
pytest.importorskip("huggingface_hub")
import subprocess
real_cache = tmp_path / "real_hf"
real_cache.mkdir()
workdir = tmp_path / "sandbox"
workdir.mkdir()
monkeypatch.setenv("HF_HOME", str(real_cache))
monkeypatch.setenv("XDG_CACHE_HOME", str(real_cache))
monkeypatch.setenv("XDG_CONFIG_HOME", str(real_cache))
env = _build_bypass_env(str(workdir))
token_path = subprocess.run(
[
sys.executable,
"-c",
"import huggingface_hub.constants as c; print(c.HF_TOKEN_PATH)",
],
env = env,
capture_output = True,
text = True,
).stdout.strip()
assert str(real_cache) not in token_path # never the operator's cache
assert token_path.startswith(str(workdir)) # resolved under the sandbox
def test_bypass_env_drops_credential_config_path_vars(monkeypatch, tmp_path):
# NETRC / BOTO_CONFIG / PIP_CONFIG_FILE point clients at real credential
# files before $HOME, so they must not survive into the bypassed child.
monkeypatch.setenv("NETRC", "/home/op/.netrc")
monkeypatch.setenv("PGPASSFILE", "/home/op/.pgpass")
monkeypatch.setenv("BOTO_CONFIG", "/home/op/.boto")
monkeypatch.setenv("PIP_CONFIG_FILE", "/home/op/.pip/pip.conf")
env = _build_bypass_env(str(tmp_path))
assert "NETRC" not in env
assert "PGPASSFILE" not in env
assert "BOTO_CONFIG" not in env
assert "PIP_CONFIG_FILE" not in env
def test_bypass_env_strips_npm_auth_and_mysql_pwd(monkeypatch, tmp_path):
# NPM_CONFIG__AUTH (npm _auth, base64) and MYSQL_PWD dodge the URL-value
# check and the PASSWD marker, but must still be dropped.
monkeypatch.setenv("NPM_CONFIG__AUTH", "aGVsbG86c2VjcmV0")
monkeypatch.setenv("MYSQL_PWD", "db-password")
assert _is_secret_env_name("NPM_CONFIG__AUTH") is True
assert _is_secret_env_name("MYSQL_PWD") is True
env = _build_bypass_env(str(tmp_path))
assert "NPM_CONFIG__AUTH" not in env
assert "MYSQL_PWD" not in env
@_POSIX_ONLY
def test_bash_bypass_does_not_source_bash_env(monkeypatch, tmp_path):
# bash -c sources $BASH_ENV for non-interactive shells; an operator startup
# file could re-export stripped secrets, so a real bypass call must not see it.
startup = tmp_path / "startup.sh"
startup.write_text("export RECOVERED=leaked\n")
monkeypatch.setenv("BASH_ENV", str(startup))
out = _bash_exec("echo R=$RECOVERED", None, 30, "bash-env-test", disable_sandbox = True)
assert "R=leaked" not in out # BASH_ENV dropped -> startup not sourced
assert "R=" in out
def test_bypass_env_repoints_windows_profile_vars(monkeypatch, tmp_path):
# On Windows, SDKs read cached creds under USERPROFILE/APPDATA/LOCALAPPDATA,
# not $HOME. Set ones are repointed at the workdir; HOMEDRIVE/HOMEPATH drop.
monkeypatch.setenv("USERPROFILE", "/host/profile")
monkeypatch.setenv("APPDATA", "/host/profile/AppData/Roaming")
monkeypatch.setenv("LOCALAPPDATA", "/host/profile/AppData/Local")
monkeypatch.setenv("HOMEDRIVE", "C:")
monkeypatch.setenv("HOMEPATH", "\\Users\\op")
env = _build_bypass_env(str(tmp_path))
assert env["USERPROFILE"] == str(tmp_path)
assert env["APPDATA"] == str(tmp_path)
assert env["LOCALAPPDATA"] == str(tmp_path)
assert "HOMEDRIVE" not in env
assert "HOMEPATH" not in env
def test_bypass_env_does_not_add_unset_windows_profile_vars(monkeypatch, tmp_path):
# Only repoint Windows profile vars that were actually set (no pollution on
# Linux/macOS where they are absent).
monkeypatch.delenv("USERPROFILE", raising = False)
monkeypatch.delenv("APPDATA", raising = False)
monkeypatch.delenv("LOCALAPPDATA", raising = False)
env = _build_bypass_env(str(tmp_path))
assert "USERPROFILE" not in env
assert "APPDATA" not in env
assert "LOCALAPPDATA" not in env
# ── parent /proc env-leak hardening (regression) ────────────────────
@_POSIX_ONLY
def test_bypass_exec_hardens_parent_proc_env(monkeypatch, captured_popen):
# Stripping the child env is not enough: a same-UID child can read the
# parent's /proc environ. The exec paths must invoke the parent hardening
# when (and only when) the sandbox is disabled.
calls = {"n": 0}
def fake_harden():
calls["n"] += 1
return True
monkeypatch.setattr(tools, "_harden_parent_against_proc_env_leak", fake_harden)
_python_exec("print(1)", None, 5, "t", disable_sandbox = True)
_bash_exec("echo hi", None, 5, "t", disable_sandbox = True)
assert calls["n"] == 2
calls["n"] = 0
_python_exec("print(1)", None, 5, "t", disable_sandbox = False)
_bash_exec("echo hi", None, 5, "t", disable_sandbox = False)
assert calls["n"] == 0 # never hardened on the sandboxed path
def test_bypass_exec_fails_closed_when_hardening_fails(monkeypatch, captured_popen):
# If the parent cannot be hardened (e.g. prctl denied), the unsandboxed
# child must NOT run - otherwise the parent environ stays readable.
monkeypatch.setattr(tools, "_harden_parent_against_proc_env_leak", lambda: False)
out_py = _python_exec("print(1)", None, 5, "t", disable_sandbox = True)
out_sh = _bash_exec("echo hi", None, 5, "t", disable_sandbox = True)
assert "refusing bypass execution" in out_py
assert "refusing bypass execution" in out_sh
assert "cmd" not in captured_popen # never reached Popen
@_POSIX_ONLY
def test_proc_env_unreadable_after_hardening():
# Mechanism check: after hardening, a same-UID child can no longer read the
# parent process /proc environ. Restores the dumpable flag afterwards so the
# process-global state does not leak into later tests.
import subprocess
if tools._libc is None:
pytest.skip("no libc/prctl available")
pid = os.getpid()
probe = (
"try:\n"
f" open('/proc/{pid}/environ', 'rb').read()\n"
" print('READABLE')\n"
"except PermissionError:\n"
" print('DENIED')\n"
)
prev_dumpable = tools._libc.prctl(3, 0, 0, 0, 0) # PR_GET_DUMPABLE
prev_guard = tools._parent_proc_hardened
try:
# Establish a clean readable baseline: another test may have already
# cleared the dumpable flag on this process.
tools._libc.prctl(4, 1, 0, 0, 0) # PR_SET_DUMPABLE = 1
before = subprocess.run(
[sys.executable, "-c", probe], capture_output = True, text = True
).stdout.strip()
if before != "READABLE":
pytest.skip("/proc already restricted in this environment")
tools._parent_proc_hardened = False
assert tools._harden_parent_against_proc_env_leak() is True
after = subprocess.run(
[sys.executable, "-c", probe], capture_output = True, text = True
).stdout.strip()
assert after == "DENIED"
finally:
if prev_dumpable in (0, 1):
try:
tools._libc.prctl(4, prev_dumpable, 0, 0, 0)
except (OSError, AttributeError):
pass
tools._parent_proc_hardened = prev_guard
# ── Anthropic request model declares the field (regression) ─────────
def test_anthropic_request_model_bypass_default_false():
# Omitting the field on the Anthropic path must default to False rather than
# raising AttributeError (extra='allow' does not set absent attributes).
from models.inference import AnthropicMessagesRequest
assert AnthropicMessagesRequest.model_fields["bypass_permissions"].default is False
req = AnthropicMessagesRequest(model = "x", messages = [], max_tokens = 8)
assert bool(req.bypass_permissions) is False

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

@ -107,6 +107,48 @@ def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_
]
def test_is_hidden_model_hides_validation_probe_everywhere():
"""Every picker (model list, local, cached GGUF, cached models) gates on
_is_hidden_model, so hiding the probe here hides it in the search menu too.
Cover both forms callers pass: the reconstructed repo id and the on-disk
snapshot path."""
assert models_route._is_hidden_model("ggml-org/models")
assert models_route._is_hidden_model("ggml-org/models/tinyllamas/stories260K.gguf")
assert models_route._is_hidden_model(
None, "/hf/models--ggml-org--models/snapshots/abc/tinyllamas/stories260K.gguf"
)
assert not models_route._is_hidden_model("unsloth/gemma-3-270m-it-GGUF")
# The exact-filename needle must not hide a real repo that merely
# references stories260K in its name.
assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF")
def test_list_cached_gguf_hides_llama_validation_probe(monkeypatch, tmp_path):
"""The ggml-org/models / stories260K install validation probe can land in
the HF cache as a side effect of installing the prebuilt llama-server.
It is not a chat model (it sorts smallest and would be auto-selected), so
pickers must hide it while keeping real cached models."""
probe = _repo(
"ggml-org/models",
[_file("tinyllamas/stories260K.gguf", 1_000)],
tmp_path / "models--ggml-org--models",
)
real = _repo(
"unsloth/gemma-3-270m-it-GGUF",
[_file("gemma-3-270m-it-UD-Q4_K_XL.gguf", 200_000)],
tmp_path / "models--unsloth--gemma-3-270m-it-GGUF",
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [probe, real])]
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
repo_ids = [c["repo_id"] for c in result["cached"]]
assert "ggml-org/models" not in repo_ids
assert "unsloth/gemma-3-270m-it-GGUF" in repo_ids
def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tmp_path):
missing = _repo(
"Org/ReadmeOnly",
@ -503,6 +545,58 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
assert flags["F16"] is False
def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
siblings = [
SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100),
SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10),
]
monkeypatch.setattr(
"huggingface_hub.model_info",
lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
)
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.gguf").write_bytes(b"x" * 10)
result = asyncio.run(
models_route.get_gguf_variants(
repo_id = "org/repo", hf_token = None, current_subject = "test-user"
)
)
assert [(v.quant, v.filename, v.size_bytes, v.downloaded) for v in result.variants] == [
("Q4_K_M", "model-Q4_K_M.gguf", 10, True)
]
def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
variants = [
SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10),
]
monkeypatch.setattr(
models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False)
)
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"x" * 10)
result = asyncio.run(
models_route.get_gguf_variants(
repo_id = "org/repo", hf_token = None, current_subject = "test-user"
)
)
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)."""
@ -524,3 +618,45 @@ def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path):
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,
current_subject = "test-user",
)
)
assert result["downloaded_bytes"] == 0
assert result["progress"] == 0
def test_gguf_download_progress_counts_quant_subdir(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" / "Q4_K_M"
snap.mkdir(parents = True)
(snap / "foo.gguf").write_bytes(b"x" * 20_000)
result = asyncio.run(
models_route.get_gguf_download_progress(
repo_id = "org/repo",
variant = "Q4_K_M",
expected_bytes = 20_000,
current_subject = "test-user",
)
)
assert result["downloaded_bytes"] == 20_000
assert result["progress"] == 1.0

361
studio/backend/tests/test_capability_detection.py Normal file
View file

@ -0,0 +1,361 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Component A tests: capability detection must never execute model repo code.
Covers: load_model_config defaults trust_remote_code False; the _VISION_CHECK_SCRIPT
subprocess literal keeps remote code off; registry-backed vision/audio detection from
raw config.json (repo-code VLMs detected without execution; ForConditionalGeneration
false positives fixed); and the model-details / GPU probes never enable remote code.
"""
import json
from types import SimpleNamespace
from unittest.mock import patch, MagicMock
import pytest
from utils.models.model_config import (
load_model_config,
is_vision_model,
_is_vlm,
_raw_config_has_vision_config,
_vision_detection_cache,
_VISION_CHECK_SCRIPT,
_VLM_MODEL_TYPES,
_AUDIO_ONLY_MODEL_TYPES,
_VLM_CLASS_NAMES,
)
@pytest.fixture(autouse = True)
def _clear_vision_cache():
_vision_detection_cache.clear()
yield
_vision_detection_cache.clear()
def _write_model_dir(
tmp_path,
cfg,
with_evil_module = False,
):
"""Write a local model dir, optionally with an auto_map module that writes a sentinel
on import so accidental code execution during detection shows up on disk."""
(tmp_path / "config.json").write_text(json.dumps(cfg))
if with_evil_module:
sentinel = tmp_path / "PWNED_SENTINEL"
(tmp_path / "modeling_evil.py").write_text(
"import os\n"
f"open({str(sentinel)!r}, 'w').write('pwned')\n"
"class EvilConfig: pass\n"
"class EvilModel: pass\n"
)
return str(tmp_path)
# load_model_config default
class TestLoadModelConfigDefault:
@patch("transformers.AutoConfig.from_pretrained")
def test_default_off_with_token(self, fp):
load_model_config("org/m", token = "hf_x")
assert fp.call_args.kwargs["trust_remote_code"] is False
@patch("utils.models.model_config.without_hf_auth")
@patch("transformers.AutoConfig.from_pretrained")
def test_default_off_public(self, fp, no_auth):
from contextlib import nullcontext
no_auth.return_value = nullcontext()
load_model_config("org/m", use_auth = False)
assert fp.call_args.kwargs["trust_remote_code"] is False
@patch("transformers.AutoConfig.from_pretrained")
def test_default_off_cached_auth(self, fp):
load_model_config("org/m", use_auth = True)
assert fp.call_args.kwargs["trust_remote_code"] is False
@patch("transformers.AutoConfig.from_pretrained")
def test_explicit_true_forwarded(self, fp):
load_model_config("org/m", token = "t", trust_remote_code = True)
assert fp.call_args.kwargs["trust_remote_code"] is True
# subprocess script literal
def test_vision_check_script_disables_remote_code():
assert '"trust_remote_code": False' in _VISION_CHECK_SCRIPT
assert '"trust_remote_code": True' not in _VISION_CHECK_SCRIPT
# _is_vlm matrix (pure function, registry-backed)
def _cfg(**kw):
return SimpleNamespace(**kw)
class TestIsVlm:
def test_deepseek_ocr_vision_via_vision_config(self):
# auto_map repo-code model; vision-ness is declarative.
c = _cfg(
model_type = "deepseek_vl_v2",
architectures = ["DeepseekOCRForCausalLM"],
vision_config = {},
projector_config = {},
)
assert _is_vlm(c) is True
def test_kimi_vision_via_vision_config(self):
c = _cfg(
model_type = "kimi_k25",
architectures = ["KimiK25ForConditionalGeneration"],
vision_config = {},
)
assert _is_vlm(c) is True
def test_glm_flash_text_is_not_vision(self):
c = _cfg(model_type = "glm4_moe_lite", architectures = ["Glm4MoeLiteForCausalLM"])
assert _is_vlm(c) is False
def test_gemma4_vision_via_vision_config(self):
c = _cfg(
model_type = "gemma4_unified",
architectures = ["Gemma4UnifiedForConditionalGeneration"],
vision_config = {},
image_token_id = 1,
)
assert _is_vlm(c) is True
def test_t5_not_misclassified_as_vision(self):
# Regression: ForConditionalGeneration must NOT be a vision signal.
c = _cfg(model_type = "t5", architectures = ["T5ForConditionalGeneration"])
assert _is_vlm(c) is False
def test_bart_not_misclassified_as_vision(self):
c = _cfg(model_type = "bart", architectures = ["BartForConditionalGeneration"])
assert _is_vlm(c) is False
def test_whisper_audio_not_vision(self):
c = _cfg(model_type = "whisper", architectures = ["WhisperForConditionalGeneration"])
assert _is_vlm(c) is False
def test_csm_audio_not_vision(self):
c = _cfg(model_type = "csm", architectures = ["CsmForConditionalGeneration"])
assert _is_vlm(c) is False
def test_native_vlm_via_registry_model_type(self):
# llava is in the transformers vision registry.
assert "llava" in _VLM_MODEL_TYPES
c = _cfg(model_type = "llava", architectures = ["LlavaForConditionalGeneration"])
assert _is_vlm(c) is True
def test_native_vlm_via_registry_class_name(self):
# Class-name match works even if model_type were unknown.
cls = next(iter(_VLM_CLASS_NAMES))
c = _cfg(model_type = "something_unlisted", architectures = [cls])
assert _is_vlm(c) is True
def test_omni_audio_plus_vision_is_vision(self):
# An audio-registry model_type with an explicit vision sub-config is still vision.
audio_mt = next(iter(_AUDIO_ONLY_MODEL_TYPES - _VLM_MODEL_TYPES))
c = _cfg(model_type = audio_mt, architectures = ["X"], vision_config = {})
assert _is_vlm(c) is True
# _raw_config_has_vision_config (code-free reader, mocked HF download)
def _mock_raw_config(tmp_path, payload):
p = tmp_path / "config.json"
p.write_text(json.dumps(payload))
return p
class TestRawConfigVisionReader:
@pytest.mark.parametrize(
"payload,expected",
[
(
{
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"},
"vision_config": {},
"projector_config": {},
},
True,
),
(
{
"model_type": "kimi_k25",
"architectures": ["KimiK25ForConditionalGeneration"],
"auto_map": {"AutoConfig": "configuration_kimi_k25.KimiK25Config"},
"vision_config": {},
},
True,
),
({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False),
(
{
"model_type": "gemma4_unified",
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
"vision_config": {},
},
True,
),
({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False),
(
{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]},
False,
),
({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False),
],
)
def test_reader(self, tmp_path, payload, expected):
cfg_path = _mock_raw_config(tmp_path, payload)
with (
patch("utils.models.model_config.is_local_path", return_value = False),
patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)),
):
assert _raw_config_has_vision_config("org/model") is expected
def test_reader_never_executes_remote_code(self, tmp_path):
# Even with auto_map present, the reader only parses JSON: no AutoConfig touched.
cfg_path = _mock_raw_config(
tmp_path,
{
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"},
"vision_config": {},
},
)
with (
patch("utils.models.model_config.is_local_path", return_value = False),
patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)),
patch(
"transformers.AutoConfig.from_pretrained",
side_effect = AssertionError("AutoConfig must not be called"),
),
):
assert _raw_config_has_vision_config("org/deepseek-ocr") is True
# Probes: model-details + GPU estimate never execute remote code
def test_gpu_estimate_probe_is_code_free():
from utils.hardware import hardware
cfg = {
"model_type": "glm4_moe_lite",
"hidden_size": 4096,
"num_hidden_layers": 40,
"max_position_embeddings": 8192,
}
with (
patch("utils.transformers_version._load_config_json", return_value = cfg),
patch(
"transformers.AutoConfig.from_pretrained",
side_effect = AssertionError("AutoConfig must not be called"),
),
):
out = hardware._load_config_for_gpu_estimate("unsloth/GLM-4.7-Flash")
assert out.max_position_embeddings == 8192
assert out.hidden_size == 4096
def test_models_route_source_has_no_remote_code_probe():
# The metadata probe must never build a trust_remote_code=True loader; referencing
# the static consent scanner or the requires_trust_remote_code flag is fine.
import inspect
import routes.models as models_route
src = inspect.getsource(models_route)
assert "trust_remote_code = True" not in src
assert "trust_remote_code=True" not in src
# Adversarial end-to-end: is_vision_model + the two metadata probes never run auto_map.
def test_no_code_execution_on_detection(tmp_path):
# A malicious local auto_map -> modeling_evil must not execute through any probe.
cfg = {
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {
"AutoConfig": "modeling_evil.EvilConfig",
"AutoModel": "modeling_evil.EvilModel",
},
"vision_config": {"image_size": 1024},
"max_position_embeddings": 4096,
}
path = _write_model_dir(tmp_path, cfg, with_evil_module = True)
sentinel = tmp_path / "PWNED_SENTINEL"
from utils.hardware.hardware import _load_config_for_gpu_estimate
from utils.transformers_version import _load_config_json
result = is_vision_model(path)
ns = _load_config_for_gpu_estimate(path)
raw = _load_config_json(path)
assert not sentinel.exists(), "SECURITY FAILURE: auto_map code executed during detection"
assert result is True # detected as vision via raw vision_config, no exec
assert ns is not None and getattr(ns, "max_position_embeddings", None) == 4096
assert raw is not None and raw.get("model_type") == "deepseek_vl_v2"
@pytest.mark.parametrize(
"cfg, expected",
[
# repo-code VLMs (auto_map) detected via declarative vision_config
(
{
"model_type": "deepseek_vl_v2",
"architectures": ["DeepseekOCRForCausalLM"],
"auto_map": {"AutoConfig": "x.Y"},
"vision_config": {},
},
True,
),
(
{
"model_type": "kimi_k25",
"architectures": ["KimiK25ForConditionalGeneration"],
"auto_map": {"AutoConfig": "x.Y"},
"vision_config": {},
},
True,
),
# newer-native vision via vision_config
(
{
"model_type": "gemma4_unified",
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
"vision_config": {},
"image_token_id": 7,
},
True,
),
# text / seq2seq / audio that share the ForConditionalGeneration suffix
({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False),
({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False),
({"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, False),
({"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, False),
({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False),
# registry-native VLMs via model_type
({"model_type": "qwen2_vl", "architectures": ["Qwen2VLForConditionalGeneration"]}, True),
({"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, True),
],
)
def test_is_vision_model_end_to_end(tmp_path, cfg, expected):
path = _write_model_dir(tmp_path, cfg)
assert is_vision_model(path) is expected, f"{cfg['model_type']} expected vision={expected}"
def test_registry_derivation():
# Registry-derived sets are large and include the curated repo-code VLMs.
assert len(_VLM_MODEL_TYPES) >= 50, f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}"
assert (
len(_AUDIO_ONLY_MODEL_TYPES) >= 20
), f"_AUDIO_ONLY too small: {len(_AUDIO_ONLY_MODEL_TYPES)}"
for repo_vlm in ("deepseek_vl_v2", "kimi_k25", "phi3_v", "cogvlm2", "minicpmv"):
assert repo_vlm in _VLM_MODEL_TYPES, f"curated repo-code VLM {repo_vlm} missing"
for native in ("llava", "qwen2_vl"):
assert native in _VLM_MODEL_TYPES, f"registry-native VLM {native} missing"
for audio in ("whisper", "csm"):
assert audio in _AUDIO_ONLY_MODEL_TYPES, f"audio type {audio} missing"

174
studio/backend/tests/test_chat_history_routes.py
View file

@ -169,3 +169,177 @@ def test_record_import_ledger_rejects_oversize_payload():
chat_history.ChatImportLedgerRecordRequest(
threadIds = [f"id-{i}" for i in range(10_001)],
)
# ---------------------------------------------------------------------------
# /api/chat/threads/{id}/fork
# ---------------------------------------------------------------------------
def test_fork_thread_404_when_source_missing(monkeypatch):
monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: None)
with pytest.raises(HTTPException) as exc:
asyncio.run(
chat_history.fork_thread(
thread_id = "missing",
payload = chat_history.ChatForkRequest(
messageId = "m1",
newThreadId = "new",
createdAt = 1,
),
current_subject = "test-user",
)
)
assert exc.value.status_code == 404
def test_fork_thread_404_when_branch_message_missing(monkeypatch):
monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: {"id": _id, "title": "T"})
monkeypatch.setattr(chat_history, "get_chat_message", lambda _t, _m: None)
with pytest.raises(HTTPException) as exc:
asyncio.run(
chat_history.fork_thread(
thread_id = "src",
payload = chat_history.ChatForkRequest(
messageId = "missing",
newThreadId = "new",
createdAt = 1,
),
current_subject = "test-user",
)
)
assert exc.value.status_code == 404
def test_fork_thread_happy_path(monkeypatch):
source = {
"id": "src",
"title": "Original",
"modelType": "base",
"modelId": "m",
"pairId": None,
"archived": False,
"createdAt": 1,
"openaiCodeExecContainerId": None,
"anthropicCodeExecContainerId": None,
"forkedFromThreadId": None,
"forkedFromMessageId": None,
}
forked = {
**source,
"id": "new",
"title": "fork · Original",
"createdAt": 2,
"forkedFromThreadId": "src",
"forkedFromMessageId": "m1",
}
monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: source)
monkeypatch.setattr(
chat_history,
"get_chat_message",
lambda _t, _m: {
"id": _m,
"threadId": _t,
"role": "user",
"content": [],
"createdAt": 1,
},
)
monkeypatch.setattr(chat_history, "fork_chat_thread", lambda **_: forked)
monkeypatch.setattr(
chat_history,
"list_chat_messages",
lambda _id: [
{
"id": "n1",
"threadId": "new",
"parentId": None,
"role": "user",
"content": [],
"createdAt": 1,
}
],
)
response = asyncio.run(
chat_history.fork_thread(
thread_id = "src",
payload = chat_history.ChatForkRequest(
messageId = "m1",
newThreadId = "new",
createdAt = 2,
),
current_subject = "test-user",
)
)
assert response.thread.id == "new"
assert response.thread.title == "fork · Original"
assert response.thread.forkedFromThreadId == "src"
assert response.thread.forkedFromMessageId == "m1"
assert len(response.messages) == 1
assert response.containerSnapshotWarning is None
def test_fork_thread_warns_when_parent_had_container(monkeypatch):
source = {
"id": "src",
"title": "T",
"modelType": "base",
"modelId": "",
"pairId": None,
"archived": False,
"createdAt": 1,
"openaiCodeExecContainerId": "cnt_123",
"anthropicCodeExecContainerId": None,
"forkedFromThreadId": None,
"forkedFromMessageId": None,
}
monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: source)
monkeypatch.setattr(
chat_history,
"get_chat_message",
lambda _t, _m: {
"id": _m,
"threadId": _t,
"role": "user",
"content": [],
"createdAt": 1,
},
)
monkeypatch.setattr(
chat_history,
"fork_chat_thread",
lambda **_: {
**source,
"id": "new",
"title": "fork · T",
"forkedFromThreadId": "src",
"forkedFromMessageId": "m1",
"openaiCodeExecContainerId": None,
},
)
monkeypatch.setattr(chat_history, "list_chat_messages", lambda _id: [])
response = asyncio.run(
chat_history.fork_thread(
thread_id = "src",
payload = chat_history.ChatForkRequest(
messageId = "m1",
newThreadId = "new",
createdAt = 2,
),
current_subject = "test-user",
)
)
assert response.containerSnapshotWarning is not None
assert "fresh" in response.containerSnapshotWarning.lower()
def test_get_fork_count(monkeypatch):
monkeypatch.setattr(chat_history, "count_forks_for_message", lambda _t, _m: 3)
response = asyncio.run(
chat_history.get_fork_count(
thread_id = "t",
message_id = "m",
current_subject = "test-user",
)
)
assert response.count == 3

132
studio/backend/tests/test_chat_history_storage.py
View file

@ -378,3 +378,135 @@ def test_legacy_imports_ignores_empty(tmp_path, monkeypatch):
assert studio_db.upsert_chat_legacy_imports([]) == (0, 0)
assert studio_db.upsert_chat_legacy_imports(["", None]) == (0, 0) # type: ignore[list-item]
assert studio_db.list_chat_legacy_imports() == []
# ---------------------------------------------------------------------------
# fork_chat_thread
# ---------------------------------------------------------------------------
def _msg(mid: str, parent: str | None, t: int) -> dict:
return {
"id": mid,
"threadId": "src",
"parentId": parent,
"role": "user",
"content": [{"type": "text", "text": mid}],
"createdAt": t,
}
def test_fork_chat_thread_copies_ancestry_with_fresh_ids(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(
{**_thread("src"), "title": "Original", "openaiCodeExecContainerId": "cnt-x"}
)
# Linear chain: m1 -> m2 -> m3. Plus a sibling m4 off m2 (should NOT
# be copied since we fork at m3).
studio_db.sync_chat_messages(
"src",
[
_msg("m1", None, 1),
_msg("m2", "m1", 2),
_msg("m3", "m2", 3),
_msg("m4", "m2", 4), # sibling — must be excluded
],
)
counter = {"i": 0}
def id_factory():
counter["i"] += 1
return f"new-{counter['i']}"
forked = studio_db.fork_chat_thread(
source_thread_id = "src",
branch_message_id = "m3",
new_thread_id = "fork-1",
new_title = "fork · Original",
created_at = 99,
id_factory = id_factory,
)
assert forked is not None
assert forked["id"] == "fork-1"
assert forked["forkedFromThreadId"] == "src"
assert forked["forkedFromMessageId"] == "m3"
# Container ids reset on fork.
assert forked["openaiCodeExecContainerId"] is None
copied = studio_db.list_chat_messages("fork-1")
# 3 ancestors (m1, m2, m3); m4 excluded.
assert len(copied) == 3
# parent_id rewritten using new ids; root has parentId None.
assert copied[0]["parentId"] is None
assert copied[1]["parentId"] == copied[0]["id"]
assert copied[2]["parentId"] == copied[1]["id"]
# All new ids regenerated.
assert {m["id"] for m in copied}.isdisjoint({"m1", "m2", "m3"})
def test_fork_chat_thread_preserves_project_id(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_project(_project("project-1"))
studio_db.upsert_chat_thread({**_thread("src"), "projectId": "project-1"})
studio_db.upsert_chat_message(_msg("m1", None, 1))
forked = studio_db.fork_chat_thread(
source_thread_id = "src",
branch_message_id = "m1",
new_thread_id = "fork-1",
new_title = "fork · Original",
created_at = 99,
id_factory = lambda: "new-1",
)
assert forked is not None
assert forked["projectId"] == "project-1"
assert {thread["id"] for thread in studio_db.list_chat_threads(project_id = "project-1")} == {
"fork-1",
"src",
}
def test_fork_chat_thread_returns_none_for_missing_source(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
result = studio_db.fork_chat_thread(
source_thread_id = "nope",
branch_message_id = "m1",
new_thread_id = "fork",
new_title = "f",
created_at = 1,
id_factory = lambda: "x",
)
assert result is None
def test_count_forks_for_message(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread("src"))
studio_db.sync_chat_messages("src", [_msg("m1", None, 1)])
assert studio_db.count_forks_for_message("src", "m1") == 0
counter = {"i": 0}
def id_factory():
counter["i"] += 1
return f"id-{counter['i']}"
studio_db.fork_chat_thread(
source_thread_id = "src",
branch_message_id = "m1",
new_thread_id = "f1",
new_title = "f1",
created_at = 2,
id_factory = id_factory,
)
studio_db.fork_chat_thread(
source_thread_id = "src",
branch_message_id = "m1",
new_thread_id = "f2",
new_title = "f2",
created_at = 3,
id_factory = id_factory,
)
assert studio_db.count_forks_for_message("src", "m1") == 2

685
studio/backend/tests/test_chat_load_during_training.py Normal file
View file

@ -0,0 +1,685 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Loading a NEW chat model while training runs: can_load_chat_during_training
(VRAM fit check), _guard_chat_load_against_training and _effective_load_in_4bit
(409 + sizing wiring). The guard sizes the same effective load the backend will
perform (HF auto reuses the loader's selector, HF explicit applies a per-GPU
floor, GGUF sizes from on-disk weights, LoRA 4-bit->16-bit flips resolved first)
and leaves non-training/external loads untouched."""
import asyncio
import importlib.util
import sys
import types
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
from fastapi import HTTPException
from utils.hardware import DeviceType
import utils.hardware.hardware as _hw_module
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
# Load training_vram.py standalone (avoids the heavy routes/__init__.py); its
# lazy hardware imports still resolve against the patched utils.hardware names.
_spec = importlib.util.spec_from_file_location(
"training_vram_load_test", _BACKEND_ROOT / "routes" / "training_vram.py"
)
tv = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(tv)
class _GpuCacheResetMixin:
def tearDown(self):
_hw_module._physical_gpu_count = None
_hw_module._visible_gpu_count = None
def _devices(*free_specs):
"""Build a device list from (index, total, used) tuples."""
return [
{"index": i, "vram_total_gb": total, "vram_used_gb": used}
for (i, total, used) in free_specs
]
# ── can_load_chat_during_training: HF auto (reuses auto_select_gpu_ids) ───────
class TestCanLoadAutoHF(_GpuCacheResetMixin, unittest.TestCase):
def _run(self, *, selection_mode, required, usable):
meta = {"selection_mode": selection_mode, "required_gb": required, "usable_gb": usable}
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
patch("utils.hardware.auto_select_gpu_ids", return_value = ([0], meta)) as auto_mock,
):
ok, info = tv.can_load_chat_during_training(
model_name = "unsloth/Qwen3-1.7B",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
is_gguf = False,
)
return ok, info, auto_mock
def test_fits_with_margin(self):
# free 60 >= 8*1.15+4 = 13.2
ok, info, auto_mock = self._run(selection_mode = "auto", required = 8.0, usable = 60.0)
self.assertTrue(ok)
self.assertEqual(info["mode"], "auto")
self.assertAlmostEqual(info["needed_gb"], 13.2, places = 3)
auto_mock.assert_called_once() # mirrors the loader's own selection
def test_too_tight_refuses(self):
# free 10 < 8*1.15+4 = 13.2 -> refuse even though raw 10 > 8
ok, _, _ = self._run(selection_mode = "auto", required = 8.0, usable = 10.0)
self.assertFalse(ok)
def test_fallback_all_refuses(self):
# Selector couldn't confirm placement -> default-deny to protect training.
ok, info = self._run(selection_mode = "fallback_all", required = 8.0, usable = 999.0)[:2]
self.assertFalse(ok)
# ── can_load_chat_during_training: HF explicit (per-GPU floor) ────────────────
class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase):
def _run(
self,
*,
required,
devices,
gpu_ids,
resolved = None,
resolve_side_effect = None,
):
resolve_kwargs = (
{"side_effect": resolve_side_effect}
if resolve_side_effect
else {"return_value": resolved if resolved is not None else gpu_ids}
)
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
patch("utils.hardware.estimate_required_model_memory_gb", return_value = (required, {})),
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
patch("utils.hardware.resolve_requested_gpu_ids", **resolve_kwargs),
patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
):
ok, info = tv.can_load_chat_during_training(
model_name = "m",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = gpu_ids,
is_gguf = False,
)
return ok, info, auto_mock
def test_single_gpu_fits(self):
ok, info, auto_mock = self._run(required = 8.0, devices = _devices((0, 80, 20)), gpu_ids = [0])
self.assertTrue(ok)
self.assertEqual(info["mode"], "explicit")
auto_mock.assert_not_called() # explicit never calls the auto selector
def test_per_gpu_floor_blocks_uneven_split(self):
# free [45, 10]; aggregate 45 + 10*0.85 = 53.5 >= needed 27, but the 10 GB
# GPU is below the even-share floor 27/2 = 13.5 -> refuse (would OOM it).
ok, info, _ = self._run(
required = 20.0, devices = _devices((0, 80, 35), (1, 80, 70)), gpu_ids = [0, 1]
)
self.assertFalse(ok)
self.assertAlmostEqual(info["min_free_gb"], 10.0, places = 3)
def test_per_gpu_floor_passes_when_even(self):
# free [30, 30]; both clear the 13.5 even-share floor -> allow.
ok, _, _ = self._run(
required = 20.0, devices = _devices((0, 80, 50), (1, 80, 50)), gpu_ids = [0, 1]
)
self.assertTrue(ok)
def test_missing_gpu_counts_as_zero(self):
ok, _, _ = self._run(required = 5.0, devices = _devices((0, 80, 5)), gpu_ids = [3], resolved = [3])
self.assertFalse(ok)
def test_invalid_ids_does_not_block(self):
ok, info, _ = self._run(
required = 5.0,
devices = [],
gpu_ids = [99],
resolve_side_effect = ValueError("Invalid gpu_ids [99]"),
)
self.assertTrue(ok)
self.assertEqual(info["reason"], "invalid_gpu_ids")
# ── can_load_chat_during_training: GGUF (sized from on-disk weights) ──────────
class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
def _run(
self,
*,
devices,
required_override = None,
estimate = None,
):
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})),
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
):
ok, info = tv.can_load_chat_during_training(
model_name = "unsloth/gemma-GGUF",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
is_gguf = True,
required_override_gb = required_override,
)
return ok, info, auto_mock
def test_override_fits(self):
ok, info, auto_mock = self._run(devices = _devices((0, 80, 20)), required_override = 10.0)
self.assertTrue(ok)
self.assertEqual(info["mode"], "gguf")
auto_mock.assert_not_called() # GGUF never uses the HF auto selector
def test_no_per_gpu_floor_for_gguf(self):
# free [45, 10], override 20 -> needed 27, aggregate 53.5 >= 27. GGUF self-
# places, so the per-GPU floor that would block HF doesn't apply -> allow.
ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0)
self.assertTrue(ok)
def test_estimate_unavailable_refuses(self):
# No override and the estimator can't size it -> default-deny.
ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None)
self.assertFalse(ok)
self.assertEqual(info["reason"], "estimate_unavailable")
# ── can_load_chat_during_training: device-independent paths ──────────────────
class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase):
def test_non_cuda_allows(self):
with patch("utils.hardware.get_device", return_value = DeviceType.MLX):
ok, info = tv.can_load_chat_during_training(
model_name = "m",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
)
self.assertTrue(ok)
self.assertEqual(info["mode"], "non_cuda")
def test_no_visible_gpus_refuses(self):
# GGUF with an empty device list -> no candidate GPU -> default-deny.
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": []}),
patch("utils.hardware.auto_select_gpu_ids"),
):
ok, info = tv.can_load_chat_during_training(
model_name = "m",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
is_gguf = True,
required_override_gb = 8.0,
)
self.assertFalse(ok)
self.assertEqual(info["reason"], "no_visible_gpus")
def test_probe_exception_refuses(self):
with patch("utils.hardware.get_device", side_effect = RuntimeError("boom")):
ok, info = tv.can_load_chat_during_training(
model_name = "m",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
)
self.assertFalse(ok)
self.assertEqual(info["reason"], "probe_error")
# ── _guard_chat_load_against_training + _effective_load_in_4bit (route) ───────
def _load_inference_route():
spec = importlib.util.spec_from_file_location(
"inference_route_chatload_test", _BACKEND_ROOT / "routes" / "inference.py"
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _stub_guard_deps(
*,
training_active,
decision,
captured = None,
):
"""Inject the guard's two lazy imports (get_training_backend, can_load_chat_
during_training); `captured` records the can_load kwargs for assertions."""
core_training = types.ModuleType("core.training")
if isinstance(training_active, Exception):
def _raise():
raise training_active
core_training.get_training_backend = _raise
else:
core_training.get_training_backend = lambda: SimpleNamespace(
is_training_active = lambda: training_active
)
def _can_load(**kwargs):
if captured is not None:
captured.append(kwargs)
return decision
tv_stub = types.ModuleType("routes.training_vram")
tv_stub.can_load_chat_during_training = _can_load
return patch.dict(
sys.modules, {"core.training": core_training, "routes.training_vram": tv_stub}
)
class TestChatLoadGuardRoute(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.route = _load_inference_route()
def _guard(
self,
*,
config = None,
captured = None,
training_active,
decision,
):
config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
with _stub_guard_deps(
training_active = training_active, decision = decision, captured = captured
):
self.route._guard_chat_load_against_training(
config,
model_identifier = "unsloth/Qwen3-1.7B",
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
)
def test_noop_when_training_inactive(self):
self._guard(training_active = False, decision = (False, {})) # must not raise
def test_noop_when_training_state_unknown(self):
self._guard(training_active = RuntimeError("no backend"), decision = (False, {}))
def test_allows_when_fits(self):
self._guard(training_active = True, decision = (True, {"mode": "auto"}))
def test_refuses_with_headroom_number(self):
info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"}
with self.assertRaises(HTTPException) as exc:
self._guard(training_active = True, decision = (False, info))
self.assertEqual(exc.exception.status_code, 409)
self.assertIn("39 GB", exc.exception.detail) # reports needed_gb, not required_gb 30
self.assertNotIn("30 GB", exc.exception.detail)
self.assertIn("including safety headroom", exc.exception.detail)
self.assertNotIn("chat is disabled", exc.exception.detail.lower())
def test_refuses_generic_when_unsizable(self):
with self.assertRaises(HTTPException) as exc:
self._guard(training_active = True, decision = (False, {"reason": "estimate_unavailable"}))
self.assertEqual(exc.exception.status_code, 409)
self.assertIn("could not be verified", exc.exception.detail)
def test_gguf_config_passes_is_gguf_and_override(self):
captured = []
config = SimpleNamespace(is_gguf = True)
with patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {}),
)
self.assertEqual(captured[0]["is_gguf"], True)
self.assertEqual(captured[0]["required_override_gb"], 12.5)
class TestEffectiveLoadIn4bit(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.route = _load_inference_route()
def _write_adapter(self, tmpdir, payload):
import json
(Path(tmpdir) / "adapter_config.json").write_text(json.dumps(payload))
def test_non_lora_returns_request(self):
cfg = SimpleNamespace(is_lora = False, path = None, base_model = None)
self.assertTrue(self.route._effective_load_in_4bit(cfg, True))
def test_lora_method_flips_to_16bit(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
self._write_adapter(d, {"unsloth_training_method": "lora"})
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
# requested 4-bit, but a 'lora' adapter loads 16-bit
self.assertFalse(self.route._effective_load_in_4bit(cfg, True))
def test_qlora_method_keeps_4bit(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
self._write_adapter(d, {"unsloth_training_method": "qlora"})
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
self.assertTrue(self.route._effective_load_in_4bit(cfg, True))
def test_no_method_non_bnb_base_flips_to_16bit(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
self._write_adapter(d, {})
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "meta/Llama-3-8B")
self.assertFalse(self.route._effective_load_in_4bit(cfg, True))
def test_malformed_adapter_config_returns_request(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
(Path(d) / "adapter_config.json").write_text("[1, 2, 3]") # not a dict
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
self.assertTrue(self.route._effective_load_in_4bit(cfg, True)) # no crash
# ── validate_model integration (early refusal, real settings) ────────────────
class TestValidateRefusesDuringTraining(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.route = _load_inference_route()
def _validate(
self,
*,
training_active,
decision,
captured = None,
load_in_4bit = True,
):
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
model_path = "unsloth/Qwen3-1.7B", load_in_4bit = load_in_4bit, max_seq_length = 4096
)
cfg = SimpleNamespace(
identifier = "unsloth/Qwen3-1.7B",
display_name = "Qwen3-1.7B",
is_gguf = False,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
)
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
_stub_guard_deps(training_active = training_active, decision = decision, captured = captured),
):
return asyncio.run(self.route.validate_model(request, current_subject = "test-user"))
def test_ok_when_training_inactive(self):
resp = self._validate(training_active = False, decision = (False, {}))
self.assertTrue(resp.valid)
def test_refuses_when_wont_fit(self):
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0}
with self.assertRaises(HTTPException) as exc:
self._validate(training_active = True, decision = (False, info))
self.assertEqual(exc.exception.status_code, 409)
self.assertIn("training is running", exc.exception.detail)
def test_passes_real_load_settings_to_guard(self):
# validate must size with the request's settings, not hardcoded defaults.
captured = []
self._validate(
training_active = True, decision = (True, {}), captured = captured, load_in_4bit = False
)
self.assertEqual(captured[0]["load_in_4bit"], False)
self.assertEqual(captured[0]["max_seq_length"], 4096)
def test_rejects_gguf_with_gpu_ids_before_guard(self):
# /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0])
cfg = SimpleNamespace(
identifier = "x.gguf",
display_name = "x",
is_gguf = True,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
)
captured = []
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("x.gguf", "x.gguf", False),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
_stub_guard_deps(training_active = True, decision = (True, {}), captured = captured),
):
with self.assertRaises(HTTPException) as exc:
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(exc.exception.status_code, 400)
self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail)
self.assertEqual(captured, []) # guard never reached
# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────
class TestEstimateGgufRequiredGb(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.route = _load_inference_route()
def test_local_sums_split_shards(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
p = Path(d)
(p / "model-00001-of-00002.gguf").write_bytes(b"x" * 1000)
(p / "model-00002-of-00002.gguf").write_bytes(b"y" * 2000)
cfg = SimpleNamespace(
gguf_file = str(p / "model-00001-of-00002.gguf"),
gguf_mmproj_file = None,
gguf_mtp_file = None,
gguf_hf_repo = None,
gguf_variant = None,
)
gb = self.route._estimate_gguf_required_gb(cfg)
self.assertAlmostEqual(gb, 3000 / (1024**3), places = 9) # both shards
def test_remote_threads_token_and_adds_companions(self):
import utils.models.model_config as mc
cfg = SimpleNamespace(
gguf_file = None,
gguf_mmproj_file = None,
gguf_mtp_file = None,
gguf_hf_repo = "org/repo",
gguf_variant = "Q4_K_M",
)
variant = SimpleNamespace(quant = "Q4_K_M", size_bytes = 10 * 1024**3)
captured = {}
def fake_list(repo, hf_token = None):
captured["token"] = hf_token
return ([variant], True) # has_vision -> include mmproj
with (
patch.object(mc, "list_gguf_variants", fake_list),
patch.object(
self.route, "_remote_gguf_companion_bytes", return_value = 2 * 1024**3
) as comp,
):
gb = self.route._estimate_gguf_required_gb(cfg, hf_token = "tok")
self.assertEqual(captured["token"], "tok") # token threaded for gated repos
self.assertAlmostEqual(gb, 12.0, places = 6) # 10 GB variant + 2 GB companions
self.assertTrue(comp.call_args.kwargs["include_mmproj"])
def test_remote_unknown_variant_returns_none(self):
import utils.models.model_config as mc
cfg = SimpleNamespace(
gguf_file = None,
gguf_mmproj_file = None,
gguf_mtp_file = None,
gguf_hf_repo = "org/repo",
gguf_variant = "Q8_0",
)
with patch.object(
mc,
"list_gguf_variants",
return_value = ([SimpleNamespace(quant = "Q4_K_M", size_bytes = 1)], False),
):
self.assertIsNone(self.route._estimate_gguf_required_gb(cfg))
def test_local_adds_kv_cache(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
p = Path(d) / "model.gguf"
p.write_bytes(b"x" * 1000)
cfg = SimpleNamespace(
gguf_file = str(p),
gguf_mmproj_file = None,
gguf_mtp_file = None,
gguf_hf_repo = None,
gguf_variant = None,
)
with patch.object(self.route, "_estimate_gguf_kv_gb", return_value = 2.0):
gb = self.route._estimate_gguf_required_gb(cfg, max_seq_length = 8192)
self.assertAlmostEqual(gb, 1000 / (1024**3) + 2.0, places = 6) # weights + KV
def test_kv_helper_graceful_on_non_gguf(self):
import tempfile
with tempfile.TemporaryDirectory() as d:
p = Path(d) / "not-a.gguf"
p.write_bytes(b"not a gguf")
self.assertEqual(self.route._estimate_gguf_kv_gb(str(p), 4096), 0.0)
def test_kv_sizes_at_larger_of_max_seq_len_and_ctx_override(self):
# KV sized at the larger of max_seq_length and --ctx-size, else native.
seen = {}
class _FakeBackend:
_context_length = 2048
def _read_gguf_metadata(self, path):
pass
def _can_estimate_kv(self):
return True
def _estimate_kv_cache_bytes(
self,
ctx,
n_parallel = 1,
):
seen["ctx"] = ctx
seen["n_parallel"] = n_parallel
return ctx * n_parallel * (1024**2) # 1 MiB per ctx unit per slot
with patch.object(self.route, "LlamaCppBackend", _FakeBackend):
r = self.route
# --ctx-size override above max_seq_length -> override wins
self.assertAlmostEqual(
r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "131072"]), 128.0
)
self.assertEqual(seen["ctx"], 131072)
self.assertEqual(seen["n_parallel"], 1) # default single slot
# override below max_seq_length -> larger (max_seq_length) wins
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0)
self.assertEqual(seen["ctx"], 4096)
# no override, no max_seq_length -> native context fallback
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 0, None), 2.0)
self.assertEqual(seen["ctx"], 2048)
# malformed extras are ignored (fall back to max_seq_length)
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "oops"]), 4.0)
# --parallel slots scale the cache the same way the launcher does
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0)
self.assertEqual(seen["n_parallel"], 4)
# ── load_model integration: authoritative 409, and no unload before refusal ──
class TestLoadModelGuardIntegration(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.route = _load_inference_route()
def test_refusal_409_and_no_unload(self):
import contextlib
from unittest.mock import MagicMock
from models.inference import LoadRequest
inf = SimpleNamespace(active_model_name = None)
inf.unload_model = MagicMock()
inf._shutdown_subprocess = MagicMock()
llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None)
llama.unload_model = MagicMock()
cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None)
request = LoadRequest(model_path = "unsloth/Qwen3-1.7B")
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"}
with (
patch.object(self.route, "validate_extra_args", return_value = None),
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
),
patch.object(self.route, "resolve_effective_chat_template_override", return_value = None),
patch.object(self.route, "get_inference_backend", return_value = inf),
patch.object(self.route, "get_llama_cpp_backend", return_value = llama),
patch.object(self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext()),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
_stub_guard_deps(training_active = True, decision = (False, info)),
):
with self.assertRaises(HTTPException) as exc:
asyncio.run(
self.route.load_model(request, fastapi_request = MagicMock(), current_subject = "u")
)
self.assertEqual(exc.exception.status_code, 409)
# Guard runs before the unload step, so a refused load tears down nothing.
inf.unload_model.assert_not_called()
inf._shutdown_subprocess.assert_not_called()
llama.unload_model.assert_not_called()
if __name__ == "__main__":
unittest.main()

152
studio/backend/tests/test_compute_buffer.py Normal file
View file

@ -0,0 +1,152 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for ``_estimate_compute_buffer_bytes``: it scales with ``--parallel``,
tensor exceeds pipeline, and it is a safe upper bound on the allocations measured
on real hardware (Qwen3.6-27B-MTP: parallel 1/2/4/8 -> 36/492/1388/3220 MiB single
GPU, ~600 MiB/device tensor). No GPU, subprocess, or GGUF I/O."""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
sys.modules.setdefault("structlog", _structlog_stub)
# httpx -- only stub when the real library is missing. Unconditional stubbing
# shadows HTTPError/Response that huggingface_hub.errors imports at load time,
# silently breaking the transformers introspection tier in tests collected after
# this one (the stub leaks via sys.modules for the whole session).
try:
import httpx as _httpx_real # noqa: F401
except ImportError:
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"RequestError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
_httpx_stub.Response = type("Response", (), {})
_httpx_stub.Client = type(
"C",
(),
{
"__init__": lambda s, **kw: None,
"__enter__": lambda s: s,
"__exit__": lambda s, *a: None,
},
)
sys.modules["httpx"] = _httpx_stub
from core.inference.llama_cpp import LlamaCppBackend
MIB = 1024 * 1024
def _backend(vocab = 248320, embd = 5120):
"""Backend with just the dims the compute-buffer estimate reads."""
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._vocab_size = vocab
b._embedding_length = embd
return b
# Measured ground truth (MiB) the estimate must upper-bound.
_PIPELINE_MEASURED = {1: 36, 2: 492, 4: 1388, 8: 3220}
_TENSOR_MEASURED_PER_DEVICE = 600
class TestSafeUpperBound:
"""The estimate must be >= every measured allocation (never under-reserve)."""
@pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items()))
def test_pipeline_upper_bounds_measured(self, parallel, measured):
est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB
assert est >= measured, f"under-reserved at parallel={parallel}: {est:.0f} < {measured}"
@pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items()))
def test_pipeline_not_wildly_over(self, parallel, measured):
# Stay within ~2x of measured so we don't waste context (the point of
# replacing the flat reserve). parallel=1 is tiny in absolute terms.
est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB
assert est <= max(measured * 2.0, 128)
def test_tensor_upper_bounds_measured(self):
est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB
assert est >= _TENSOR_MEASURED_PER_DEVICE
def test_tensor_far_below_old_flat_reserve(self):
# The whole point: deterministic estimate << flat 5120 for this model.
est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB
assert est < LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
class TestScaling:
def test_grows_with_serving_slots(self):
b = _backend()
vals = [b._estimate_compute_buffer_bytes(n_parallel = p) for p in (1, 2, 4, 8)]
assert vals == sorted(vals) and vals[0] < vals[-1]
def test_parallel_1_is_small(self):
# Single-token decode: a few tens of MiB, not gigabytes.
est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1) / MIB
assert est < 128
def test_tensor_exceeds_pipeline_at_same_parallel(self):
b = _backend()
pipe = b._estimate_compute_buffer_bytes(n_parallel = 1)
tens = b._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True)
assert tens > pipe
def test_scales_with_vocab(self):
small = _backend(vocab = 32000)._estimate_compute_buffer_bytes(n_parallel = 4)
big = _backend(vocab = 256000)._estimate_compute_buffer_bytes(n_parallel = 4)
assert big > small
def test_scales_with_ubatch(self):
b = _backend()
lo = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 256)
hi = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 1024)
assert hi > lo
class TestFallback:
def test_zero_when_vocab_missing(self):
assert _backend(vocab = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0
def test_zero_when_embd_missing(self):
assert _backend(embd = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0
def test_zero_lets_tensor_plan_use_flat_fallback(self):
# When dims are missing, _plan_tensor_parallel must fall back to the flat
# reserve (defense-in-depth) rather than reserving 0 and OOMing.
b = _backend(vocab = None, embd = None)
b._n_layers = None # can't estimate KV -> floors ctx, still returns a plan
ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, 48000)], 8 * 1024**3, 8192)
assert gi == [0, 1] # both GPUs usable under the flat fallback
class TestParallel1Default:
"""At Studio's default --parallel 1 the buffer is negligible in pipeline."""
def test_default_n_parallel(self):
est = _backend()._estimate_compute_buffer_bytes() / MIB
assert est < 128

1414
studio/backend/tests/test_consent_gate.py Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -72,6 +72,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
fastapi.APIRouter = lambda: _Router()
fastapi.Body = lambda default = None, **_kwargs: default
fastapi.Depends = lambda dependency = None, **_kwargs: dependency
fastapi.Header = lambda default = None, **_kwargs: default
fastapi.HTTPException = _HTTPException
fastapi.Query = lambda default = None, **_kwargs: default
fastapi.Request = object
@ -156,6 +157,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
utils_model_config = types.ModuleType("utils.models.model_config")
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_audio_input_type = lambda *args, **kwargs: None
monkeypatch.setitem(
sys.modules,
@ -190,6 +192,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
for name in (
"BrowseEntry",
"BrowseFoldersResponse",
"ExportSizeResponse",
"GgufVariantDetail",
"GgufVariantsResponse",
"ScanFolderInfo",

244
studio/backend/tests/test_export_size_estimate.py Normal file
View file

@ -0,0 +1,244 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for GET /api/models/export-size (the Export page size estimate).
The endpoint must never raise and must degrade to nulls when size is unknown.
"""
import asyncio
import importlib.util
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
# Real Qwen3.6-35B-A3B: 35.95B params -> ~67 GiB bf16 (UI wrongly showed Q8 ~8.2 GB).
_QWEN35_PARAMS = 35_951_822_704
_QWEN35_FP16_BYTES = _QWEN35_PARAMS * 2
def _load_route_module(name: str, relative_path: str):
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
class TestExportSizeEndpoint(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.models_route = _load_route_module(
"models_route_module_for_export_size_test",
"routes/models.py",
)
def setUp(self):
self.models_route._EXPORT_SIZE_CACHE.clear()
def _call(self, model: str = "unsloth/Qwen3.6-35B-A3B"):
with (
patch.object(self.models_route, "is_local_path", return_value = False),
patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m),
):
return asyncio.run(
self.models_route.get_export_size(
model = model, hf_token = None, current_subject = "test-user"
)
)
def test_known_model_returns_bytes_and_params(self):
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (_QWEN35_FP16_BYTES, "safetensors"),
):
resp = self._call()
self.assertEqual(resp.fp16_bytes, _QWEN35_FP16_BYTES)
self.assertEqual(resp.total_params, _QWEN35_PARAMS)
self.assertEqual(resp.source, "safetensors")
self.assertEqual(resp.model, "unsloth/Qwen3.6-35B-A3B")
def test_moe_via_config_fallback(self):
# MoE sized via the sizer's config path -> source "config".
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (67 * (1024**3), "config"),
):
resp = self._call()
self.assertEqual(resp.fp16_bytes, 67 * (1024**3))
self.assertEqual(resp.total_params, (67 * (1024**3)) // 2)
self.assertEqual(resp.source, "config")
def test_unknown_size_returns_nulls_not_error(self):
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (None, "unavailable"),
):
resp = self._call()
self.assertIsNone(resp.fp16_bytes)
self.assertIsNone(resp.total_params)
self.assertEqual(resp.source, "unavailable")
def test_zero_size_treated_as_unknown(self):
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (0, "safetensors"),
):
resp = self._call()
self.assertIsNone(resp.fp16_bytes)
self.assertIsNone(resp.total_params)
def test_sizer_exception_is_swallowed(self):
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
side_effect = RuntimeError("boom"),
):
resp = self._call()
self.assertIsNone(resp.fp16_bytes)
self.assertEqual(resp.source, "unavailable")
def test_result_is_memoized_per_model(self):
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (_QWEN35_FP16_BYTES, "safetensors"),
) as mock_sizer:
first = self._call()
second = self._call()
self.assertEqual(first.fp16_bytes, second.fp16_bytes)
self.assertEqual(mock_sizer.call_count, 1)
def test_failures_are_not_cached(self):
# A transient failure must not poison the cache; a later call recovers.
with patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
side_effect = [(None, "unavailable"), (_QWEN35_FP16_BYTES, "safetensors")],
) as mock_sizer:
first = self._call()
second = self._call()
self.assertIsNone(first.fp16_bytes)
self.assertEqual(second.fp16_bytes, _QWEN35_FP16_BYTES)
self.assertEqual(mock_sizer.call_count, 2)
def test_token_is_forwarded_to_sizer(self):
with (
patch.object(self.models_route, "is_local_path", return_value = False),
patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m),
patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (_QWEN35_FP16_BYTES, "safetensors"),
) as mock_sizer,
):
asyncio.run(
self.models_route.get_export_size(
model = "unsloth/Private",
hf_token = "secret-token",
current_subject = "test-user",
)
)
self.assertEqual(mock_sizer.call_args.kwargs.get("hf_token"), "secret-token")
def test_arbitrary_local_path_is_not_scanned(self):
# Unsafe local paths must not be scanned -> unavailable.
with (
patch.object(self.models_route, "is_local_path", return_value = True),
patch.object(self.models_route, "_is_sizable_local_path", return_value = False),
patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer,
):
resp = asyncio.run(
self.models_route.get_export_size(
model = "/etc", hf_token = None, current_subject = "test-user"
)
)
self.assertIsNone(resp.fp16_bytes)
self.assertEqual(resp.source, "unavailable")
mock_sizer.assert_not_called()
def test_sizable_local_path_is_sized(self):
with (
patch.object(self.models_route, "is_local_path", return_value = True),
patch.object(self.models_route, "_is_sizable_local_path", return_value = True),
patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
side_effect = lambda m, **_kw: m,
),
patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (_QWEN35_FP16_BYTES, "local"),
),
):
resp = asyncio.run(
self.models_route.get_export_size(
model = "/root/.unsloth/studio/outputs/run",
hf_token = None,
current_subject = "test-user",
)
)
self.assertEqual(resp.fp16_bytes, _QWEN35_FP16_BYTES)
self.assertEqual(resp.source, "local")
def test_local_adapter_base_escaping_roots_is_rejected(self):
# A local adapter under a root whose resolved base points outside the
# roots (e.g. "/") must not be sized: the resolved base is re-validated.
adapter = "/root/.unsloth/studio/outputs/adapter"
with (
patch.object(self.models_route, "is_local_path", return_value = True),
patch.object(
self.models_route, "_is_sizable_local_path", side_effect = lambda p: p == adapter
),
patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
return_value = "/",
),
patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer,
):
resp = asyncio.run(
self.models_route.get_export_size(
model = adapter, hf_token = None, current_subject = "test-user"
)
)
self.assertIsNone(resp.fp16_bytes)
self.assertEqual(resp.source, "unavailable")
mock_sizer.assert_not_called()
def test_is_sizable_local_path_containment(self):
# Only paths under a trusted root are sizable; '..' can't escape.
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp) / "outputs"
inside = root / "run-1"
inside.mkdir(parents = True)
with (
patch("utils.paths.studio_root", return_value = root),
patch("utils.paths.outputs_root", return_value = root),
patch("utils.paths.exports_root", return_value = root),
patch("utils.paths.storage_roots.cache_root", return_value = root),
):
is_sizable = self.models_route._is_sizable_local_path
self.assertTrue(is_sizable(str(inside)))
self.assertTrue(is_sizable(str(root)))
self.assertFalse(is_sizable(str(root / "missing")))
self.assertFalse(is_sizable("/etc"))
self.assertFalse(is_sizable(str(root / ".." / "etc")))
# A symlink inside a root pointing outside it cannot escape.
escape = root / "escape"
os.symlink(tmp, escape)
self.assertFalse(is_sizable(str(escape)))
def test_local_weight_size_skips_nested_checkpoints(self):
# A run dir's intermediate checkpoint-*/global_step* snapshots must not
# be counted; only the model files at the root are summed.
from utils.hardware.hardware import _get_local_weight_size_bytes
with tempfile.TemporaryDirectory() as tmp:
run = Path(tmp)
(run / "model.safetensors").write_bytes(b"\0" * 1000)
for sub, size in (("checkpoint-60", 5000), ("global_step10", 7000)):
d = run / sub
d.mkdir()
(d / "model.safetensors").write_bytes(b"\0" * size)
self.assertEqual(_get_local_weight_size_bytes(str(run)), 1000)
if __name__ == "__main__":
unittest.main()

534
studio/backend/tests/test_file_security.py Normal file
View file

@ -0,0 +1,534 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the malware / unsafe-file gate (utils.security.file_security).
The gate reads HF's security scan (model_info securityStatus) metadata-only and never
downloads flagged files; only the Hub call is stubbed. Policy: block a non-"safe" level
(unknown levels fail closed), fail open when the scan is unavailable, skip local paths
only, no first-party exemption. The block is scoped to the load-path RCE vector (a
root-level code-executing file), so flagged safetensors and subdir pickles do not block.
"""
from types import SimpleNamespace
from unittest.mock import patch
import pytest
from utils.security import evaluate_file_security
def _patch_status(status):
"""Patch huggingface_hub.model_info to return one fixed security_repo_status."""
def _mi(*_args, **_kwargs):
return SimpleNamespace(security_repo_status = status)
return patch("huggingface_hub.model_info", side_effect = _mi)
def _patch_raises(exc = RuntimeError("offline")):
return patch("huggingface_hub.model_info", side_effect = exc)
def _patch_no_index():
"""Make the weight-index lookup find no index files (definitive: nothing sharded)."""
from huggingface_hub.utils import EntryNotFoundError
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
raise EntryNotFoundError(filename or "")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
def _patch_index(weight_map, index_filename = "pytorch_model.bin.index.json"):
"""Serve a root weight index mapping tensor names -> shard paths; others 404."""
import json
import tempfile
from pathlib import Path
from huggingface_hub.utils import EntryNotFoundError
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
if filename == index_filename:
p = Path(tempfile.mkdtemp()) / filename
p.write_text(json.dumps({"weight_map": weight_map}))
return str(p)
raise EntryNotFoundError(filename or "")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
def _patch_index_unreadable():
"""Make every index fetch fail transiently (inconclusive lookup -> fail closed)."""
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
raise RuntimeError("transient network error")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
def _patch_index_mixed(weight_map, readable_index, failing_index):
"""Serve one index cleanly while another fails transiently: the flagged shard is
listed only by the index we could not read, so a naive "read any index?" check breaks."""
import json
import tempfile
from pathlib import Path
from huggingface_hub.utils import EntryNotFoundError
def _dl(
repo_id = None,
filename = None,
token = None,
**kw,
):
if filename == readable_index:
p = Path(tempfile.mkdtemp()) / filename
p.write_text(json.dumps({"weight_map": weight_map}))
return str(p)
if filename == failing_index:
raise RuntimeError("transient network error")
raise EntryNotFoundError(filename or "")
return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
@pytest.mark.parametrize("level", ["unsafe", "suspicious", "malicious"])
def test_blocks_each_blocking_level(level):
status = {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}]}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": level}]
assert d.response_payload()["security_blocked"] is True
def test_ignores_safe_only():
status = {
"scansDone": True,
"filesWithIssues": [{"path": "model.safetensors", "level": "safe"}],
}
with _patch_status(status):
d = evaluate_file_security("good/repo")
assert d.blocked is False
assert d.unsafe_files == []
def test_blocks_unsafe_even_when_scans_not_done():
# scansDone is often False for clean repos; an already-flagged file must still block.
status = {"scansDone": False, "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}]}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
def test_fail_open_when_scan_unavailable():
# model_info returns no security_repo_status -> unknown -> allow.
with _patch_status(None):
d = evaluate_file_security("unknown/repo")
assert d.blocked is False
def test_fail_open_on_exception_offline():
with _patch_raises():
d = evaluate_file_security("offline/repo")
assert d.blocked is False
def test_fail_open_scans_done_no_issues():
with _patch_status({"scansDone": True, "filesWithIssues": []}):
d = evaluate_file_security("clean/repo")
assert d.blocked is False
def test_skips_local_path():
# A local path has no Hub scan; must not even call model_info.
with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")):
d = evaluate_file_security("/tmp/some/local/model")
assert d.blocked is False
assert "local" in d.reason
def test_remote_gguf_named_repo_is_still_scanned():
# Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a
# poisoned pickle smuggled into it is blocked.
status = {
"scansDone": True,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("evil/model.gguf")
assert d.blocked is True
assert d.unsafe_files
def test_skips_local_gguf_file():
# A local .gguf path is caught by is_local_path -- no Hub call.
with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")):
d = evaluate_file_security("/tmp/models/model.gguf")
assert d.blocked is False
assert "local" in d.reason
def test_no_first_party_exemption():
# A poisoned pickle in a first-party repo still blocks (compromised-repo defense).
status = {
"scansDone": True,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("unsloth/some-model")
assert d.blocked is True
def test_malformed_entries_are_ignored():
status = {
"scansDone": True,
"filesWithIssues": ["not-a-dict", {"path": "ok.pkl", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "ok.pkl", "level": "unsafe"}]
def test_response_payload_shape():
status = {"scansDone": True, "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}]}
with _patch_status(status):
payload = evaluate_file_security("evil/repo").response_payload()
assert set(payload) == {"unsafe_files", "security_blocked", "reason"}
assert payload["security_blocked"] is True
assert payload["unsafe_files"] == [{"path": "a.pkl", "level": "malicious"}]
# ── Load-path RCE scoping: block only files a load would actually deserialize ──
def test_flagged_safetensors_does_not_block():
# safetensors is tensor-only and cannot execute code, so a flag on one (often
# picklescan tripping on a sibling pickle) is not an RCE vector and must not block.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "model-00001-of-00004.safetensors", "level": "unsafe"}],
}
with _patch_status(status):
d = evaluate_file_security("nvidia/some-model")
assert d.blocked is False
assert d.unsafe_files == []
def test_flagged_subdirectory_pickle_does_not_block():
# from_pretrained reads only root weights; a flagged subdir pickle no root index
# references (e.g. a NeMo checkpoint) is never loaded, so it must not block.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "nemo/weights/common.pt", "level": "unsafe"},
{"path": "nemo/weights/__0_0.distcp", "level": "unsafe"},
],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("nvidia/some-model")
assert d.blocked is False
assert d.unsafe_files == []
def test_nemotron_h_shaped_status_loads():
# Real Nemotron-H-8B-Base-8K shape: flagged root safetensors + unreferenced nemo/
# pickles. None is a load-path vector, so it must load.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "nemo/weights/.metadata", "level": "unsafe"},
{"path": "nemo/weights/__0_0.distcp", "level": "unsafe"},
{"path": "nemo/weights/common.pt", "level": "unsafe"},
{"path": "model-00001-of-00004.safetensors", "level": "unsafe"},
{"path": "model-00002-of-00004.safetensors", "level": "unsafe"},
],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("nvidia/Nemotron-H-8B-Base-8K")
assert d.blocked is False
assert d.unsafe_files == []
def test_indexed_subdir_shard_blocks():
# A flagged subdir shard that a root index references IS deserialized, so it blocks.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"},
],
}
weight_map = {
"layer.0.weight": "shards/pytorch_model-00001-of-00002.bin",
"layer.1.weight": "shards/pytorch_model-00002-of-00002.bin",
}
with _patch_status(status), _patch_index(weight_map):
d = evaluate_file_security("evil/sharded")
assert d.blocked is True
assert d.unsafe_files == [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
]
def test_unindexed_subdir_pickle_does_not_block_when_index_present():
# An index exists but does not list the flagged subdir pickle -> not loaded -> no block.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "extras/notes.bin", "level": "unsafe"}],
}
weight_map = {"layer.0.weight": "pytorch_model-00001-of-00001.bin"}
with _patch_status(status), _patch_index(weight_map):
d = evaluate_file_security("org/has-index")
assert d.blocked is False
assert d.unsafe_files == []
def test_inconclusive_index_lookup_blocks_subdir_pickle():
# An unreadable index can't rule out that the flagged subdir pickle is a shard -> block.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "weights/model_part.bin", "level": "unsafe"}],
}
with _patch_status(status), _patch_index_unreadable():
d = evaluate_file_security("org/transient")
assert d.blocked is True
assert d.unsafe_files == [{"path": "weights/model_part.bin", "level": "unsafe"}]
def test_partial_index_read_with_transient_failure_blocks_subdir_pickle():
# The bin index (which would list the flagged shard) fails transiently; a partial
# path set is not definitive, so fail closed.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"},
],
}
# The readable index lists only benign shards; the flagged .bin is in the unread index.
safetensors_map = {"layer.0.weight": "model-00001-of-00001.safetensors"}
with (
_patch_status(status),
_patch_index_mixed(
safetensors_map,
readable_index = "model.safetensors.index.json",
failing_index = "pytorch_model.bin.index.json",
),
):
d = evaluate_file_security("evil/mixed-index")
assert d.blocked is True
assert d.unsafe_files == [
{"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
]
def test_root_pickle_alongside_safetensors_still_blocks():
# A real root pickle blocks even alongside a flagged safetensors; it is a load-path vector.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "model.safetensors", "level": "unsafe"},
{"path": "pytorch_model.bin", "level": "unsafe"},
],
}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": "unsafe"}]
def test_eicar_shaped_root_files_block():
# The canonical eicar repo ships its dangerous files at the ROOT, so it stays blocked.
status = {
"scansDone": True,
"filesWithIssues": [
{"path": "model_broken_X.pkl", "level": "unsafe"},
{"path": "danger.dat", "level": "unsafe"},
{"path": "eicar_test_file", "level": "unsafe"},
],
}
with _patch_status(status):
d = evaluate_file_security("mcpotato/42-eicar-street")
assert d.blocked is True
assert len(d.unsafe_files) == 3
def test_unknown_future_level_fails_closed():
# Hub schema drift: an unrecognized non-"safe" level (e.g. "infected") on a root pickle must block.
status = {"scansDone": True, "filesWithIssues": [{"path": "weights.bin", "level": "infected"}]}
with _patch_status(status):
d = evaluate_file_security("evil/repo")
assert d.blocked is True
assert d.unsafe_files == [{"path": "weights.bin", "level": "infected"}]
def test_pending_or_scanning_level_does_not_block():
# A not-yet-finished per-file scan state must not false-block.
for lvl in ("pending", "scanning", "queued", "unscanned", "error"):
status = {
"scansDone": False,
"filesWithIssues": [{"path": "pytorch_model.bin", "level": lvl}],
}
with _patch_status(status):
d = evaluate_file_security("some/repo")
assert d.blocked is False, lvl
# -- Subdir load roots: Spark-TTS / BiCodec load from_pretrained(<snapshot>/LLM) --
def test_flagged_pickle_under_load_subdir_blocks():
# A flagged pickle directly under a declared load subdir is a root-level artifact there.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",))
assert d.blocked is True
assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]
def test_flagged_pickle_under_subdir_without_load_root_does_not_block():
# Same file, but NOT declared a load root and not indexed -> not deserialized.
status = {
"scansDone": False,
"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}],
}
with _patch_status(status), _patch_no_index():
d = evaluate_file_security("org/not-a-load-root")
assert d.blocked is False
assert d.unsafe_files == []
def test_indexed_shard_under_load_subdir_blocks():
# An index inside the load subdir referencing a flagged shard makes it a vector.
status = {
"scansDone": False,
"filesWithIssues": [
{"path": "LLM/shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
],
}
weight_map = {
"layer.0.weight": "shards/pytorch_model-00001-of-00002.bin",
"layer.1.weight": "shards/pytorch_model-00002-of-00002.bin",
}
with (
_patch_status(status),
_patch_index(weight_map, index_filename = "LLM/pytorch_model.bin.index.json"),
):
d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",))
assert d.blocked is True
# -- Source files are the consent gate's domain, not a deserialization vector --
def test_flagged_root_python_helper_does_not_block():
# A root .py is never deserialized; repo code runs only via auto_map under the consent
# gate, so flagging it here would false-block.
status = {
"scansDone": True,
"filesWithIssues": [
{"path": "build_pickles.py", "level": "unsafe"},
{"path": "train.py", "level": "suspicious"},
],
}
with _patch_status(status):
d = evaluate_file_security("org/has-helper-scripts")
assert d.blocked is False
assert d.unsafe_files == []
def test_root_pickle_still_blocks_with_flagged_python_sibling():
# The .py exemption must not mask a genuine root pickle in the same repo.
status = {
"scansDone": True,
"filesWithIssues": [
{"path": "convert.py", "level": "unsafe"},
{"path": "pytorch_model.bin", "level": "unsafe"},
],
}
with _patch_status(status):
d = evaluate_file_security("evil/mixed")
assert d.blocked is True
# -- Alias resolution: scan the repo the loader actually fetches from --
def _patch_status_capture(status):
"""Like _patch_status, but records the repo id model_info was queried with."""
seen = {}
def _mi(repo, *_a, **_k):
seen["repo"] = repo
return SimpleNamespace(security_repo_status = status)
return patch("huggingface_hub.model_info", side_effect = _mi), seen
def test_spark_tts_llm_alias_scans_real_repo():
# "Spark-TTS-0.5B/LLM" loads as unsloth/Spark-TTS-0.5B with LLM as load root; scanning
# the literal alias 404s and fails open, missing a flagged LLM/ pickle.
status = {"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]}
cap, seen = _patch_status_capture(status)
with cap, patch("utils.paths.is_local_path", return_value = False), _patch_no_index():
d = evaluate_file_security("Spark-TTS-0.5B/LLM", load_subdirs = ())
assert seen["repo"] == "unsloth/Spark-TTS-0.5B" # scanned the real repo, not the alias
assert d.model_name == "unsloth/Spark-TTS-0.5B"
assert d.blocked is True
assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]
def test_non_llm_alias_is_not_rewritten():
# A normal repo id with one slash must be scanned as-is (no spurious rewrite).
status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}
cap, seen = _patch_status_capture(status)
with cap, patch("utils.paths.is_local_path", return_value = False):
d = evaluate_file_security("org/model")
assert seen["repo"] == "org/model"
assert d.model_name == "org/model"
def test_generic_slash_llm_repo_is_scanned_as_itself():
# A third-party repo merely ending in "/LLM" is not a bicodec alias, so it must be
# scanned as itself; rewriting to unsloth/<parent> would scan the wrong repo.
status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}
cap, seen = _patch_status_capture(status)
with cap, patch("utils.paths.is_local_path", return_value = False):
d = evaluate_file_security("evil/LLM")
assert seen["repo"] == "evil/LLM" # scanned the real repo, not unsloth/evil
assert d.model_name == "evil/LLM"
assert d.blocked is True
def test_security_load_subdirs_yaml_fallback(monkeypatch):
# Tokenizer detection failed, but a YAML default of audio_type=bicodec still yields LLM.
import utils.models.model_config as mc
from utils.security import security_load_subdirs
monkeypatch.setattr(mc, "detect_audio_type", lambda *_a, **_k: None)
monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"})
assert security_load_subdirs("unsloth/Spark-TTS-0.5B") == ("LLM",)
# A non-bicodec default contributes no subdir.
monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None})
assert security_load_subdirs("unsloth/Llama-3.2-1B") == ()

76
studio/backend/tests/test_gguf_metadata.py
View file

@ -13,6 +13,7 @@ from typing import Iterable, Mapping
from utils.models.gguf_metadata import (
is_mmproj_by_metadata,
pairing_score,
read_gguf_context_length,
read_gguf_general_metadata,
read_mmproj_audio_capability,
)
@ -21,6 +22,7 @@ from utils.models.gguf_metadata import (
_GGUF_MAGIC = 0x46554747
_VTYPE_STRING = 8
_VTYPE_UINT32 = 4
_VTYPE_UINT64 = 10
_VTYPE_ARRAY = 9
_VTYPE_BOOL = 7
@ -38,6 +40,10 @@ def _enc_kv_uint32(key: str, value: int) -> bytes:
return _enc_string(key) + struct.pack("<I", _VTYPE_UINT32) + struct.pack("<I", value)
def _enc_kv_uint64(key: str, value: int) -> bytes:
return _enc_string(key) + struct.pack("<I", _VTYPE_UINT64) + struct.pack("<Q", value)
def _enc_kv_bool(key: str, value: bool) -> bytes:
return _enc_string(key) + struct.pack("<I", _VTYPE_BOOL) + struct.pack("<B", 1 if value else 0)
@ -56,21 +62,29 @@ def _write_synthetic_gguf(
general_strings: Mapping[str, str],
*,
extra_uint32: Mapping[str, int] | None = None,
extra_uint64: Mapping[str, int] | None = None,
extra_string_arrays: Mapping[str, Iterable[str]] | None = None,
extra_bools: Mapping[str, bool] | None = None,
) -> Path:
"""Minimal GGUF: header + KV body, no tensors."""
extra_uint32 = extra_uint32 or {}
extra_uint64 = extra_uint64 or {}
extra_string_arrays = extra_string_arrays or {}
extra_bools = extra_bools or {}
kv_count = (
len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + len(extra_bools)
len(general_strings)
+ len(extra_uint32)
+ len(extra_uint64)
+ len(extra_string_arrays)
+ len(extra_bools)
)
body = b""
for k, v in general_strings.items():
body += _enc_kv_string(k, v)
for k, v in extra_uint32.items():
body += _enc_kv_uint32(k, v)
for k, v in extra_uint64.items():
body += _enc_kv_uint64(k, v)
for k, v in extra_string_arrays.items():
body += _enc_kv_string_array(k, v)
for k, v in extra_bools.items():
@ -100,6 +114,66 @@ def test_returns_none_for_non_gguf(tmp_path: Path):
assert read_gguf_general_metadata(str(p)) is None
def test_context_length_none_for_missing_file(tmp_path: Path):
assert read_gguf_context_length(str(tmp_path / "nope.gguf")) is None
def test_context_length_none_for_non_gguf(tmp_path: Path):
p = tmp_path / "garbage.gguf"
p.write_bytes(b"not a gguf file at all, just bytes")
assert read_gguf_context_length(str(p)) is None
def test_context_length_read_from_arch_namespaced_key(tmp_path: Path):
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.architecture": "llama"},
extra_uint32 = {"llama.context_length": 4096, "llama.block_count": 32},
)
assert read_gguf_context_length(str(p)) == 4096
def test_context_length_none_when_absent(tmp_path: Path):
# Architecture present but no <arch>.context_length key.
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.architecture": "llama"},
extra_uint32 = {"llama.block_count": 32},
)
assert read_gguf_context_length(str(p)) is None
def test_context_length_ignores_foreign_arch_key(tmp_path: Path):
# A context_length under a different arch namespace must not match.
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.architecture": "llama"},
extra_uint32 = {"qwen2.context_length": 8192},
)
assert read_gguf_context_length(str(p)) is None
def test_context_length_read_from_uint64(tmp_path: Path):
# Some models store context_length as a uint64 (vtype 10).
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.architecture": "qwen3"},
extra_uint64 = {"qwen3.context_length": 262144},
)
assert read_gguf_context_length(str(p)) == 262144
def test_context_length_zero_treated_as_absent(tmp_path: Path):
# A zero/garbage ceiling must read as None so the UI can't build a slider
# with max < min.
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.architecture": "llama"},
extra_uint32 = {"llama.context_length": 0},
)
assert read_gguf_context_length(str(p)) is None
def test_extracts_general_string_fields(tmp_path: Path):
p = _write_synthetic_gguf(
tmp_path / "model.gguf",

17
studio/backend/tests/test_gguf_routing.py
View file

@ -83,6 +83,23 @@ def test_detects_gguf_in_directory(tmp_path):
assert result.endswith("model-Q4_K_M.gguf")
def test_directory_auto_detect_ignores_big_endian_sibling(tmp_path):
be = tmp_path / "model-Q4_K_M-be.gguf"
be.write_bytes(b"x" * 100)
target = tmp_path / "model-Q4_K_M.gguf"
target.write_bytes(b"y" * 10)
result = detect_gguf_model(str(tmp_path))
assert result == str(target.resolve())
def test_direct_big_endian_file_is_not_detected(tmp_path):
gguf = tmp_path / "model-Q4_K_M-be.gguf"
gguf.write_bytes(b"")
assert detect_gguf_model(str(gguf)) is None
def test_directory_named_like_gguf_scans_inside(tmp_path):
"""A directory named *.gguf resolves the real .gguf inside, not itself."""
gguf_dir = tmp_path / "mymodel.gguf"

160
studio/backend/tests/test_gguf_xet_fallback_integration.py Normal file
View file

@ -0,0 +1,160 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Integration: GGUF Chat-Mode downloads route through the Xet->HTTP helper,
preserving cancellation and the best-effort companion contract. No GPU, no
network, no real subprocess (the helper is patched).
"""
from __future__ import annotations
import sys
import threading
import types as _types
from pathlib import Path
from unittest.mock import patch
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Heavy-dep stubbing; prefer the real structlog so a bare stub never leaks to
# later modules that log at import time.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
try:
import structlog # noqa: F401
except ImportError:
sys.modules["structlog"] = _types.ModuleType("structlog")
try:
import httpx # noqa: F401
except ImportError:
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"RequestError",
"HTTPStatusError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Response = type("Response", (), {})
_httpx_stub.Request = type("Request", (), {})
_httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None})
_httpx_stub.Client = type(
"Client",
(),
{
"__init__": lambda self, **k: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *a: None,
},
)
sys.modules.setdefault("httpx", _httpx_stub)
from huggingface_hub import constants as hf_constants
from core.inference.llama_cpp import LlamaCppBackend
from utils.hf_xet_fallback import DownloadStallError
REPO = "unsloth/vision-GGUF"
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
return tmp_path
def _build_cache(
root: Path,
repo_id: str,
files: dict[str, int],
sha: str = "a" * 40,
) -> Path:
repo_dir = root / f"models--{repo_id.replace('/', '--')}"
(repo_dir / "blobs").mkdir(parents = True, exist_ok = True)
snap = repo_dir / "snapshots" / sha
snap.mkdir(parents = True, exist_ok = True)
for rel, size in files.items():
(snap / rel).write_bytes(b"\0" * size)
return snap
def test_companion_routes_through_helper(hf_cache):
_build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1})
backend = LlamaCppBackend()
captured = {}
def fake_helper(
repo_id,
filename,
token = None,
**kwargs,
):
captured["filename"] = filename
captured["cancel_event"] = kwargs.get("cancel_event")
return f"/fake/{filename}"
with (
patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper),
):
out = backend._download_mmproj(hf_repo = REPO, hf_token = None)
assert out == "/fake/mmproj-vision-F16.gguf"
# _cancel_event must be threaded through so /unload can abort the download.
assert captured["cancel_event"] is backend._cancel_event
def test_companion_swallows_terminal_stall_to_none(hf_cache):
_build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1})
backend = LlamaCppBackend()
def stalling_helper(
repo_id,
filename,
token = None,
**kwargs,
):
raise DownloadStallError("both transports stalled")
with (
patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", stalling_helper),
):
out = backend._download_mmproj(hf_repo = REPO, hf_token = None)
assert out is None, "a companion download is best-effort; a terminal stall must not raise"
def test_companion_cancelled_skips_download(hf_cache):
_build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1})
backend = LlamaCppBackend()
backend._cancel_event.set()
called = {"n": 0}
def helper(
repo_id,
filename,
token = None,
**kwargs,
):
called["n"] += 1
return "/should-not-happen"
with (
patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", helper),
):
out = backend._download_mmproj(hf_repo = REPO, hf_token = None)
assert out is None
assert called["n"] == 0, "a cancelled load must not start a companion download"

8
studio/backend/tests/test_gpu_selection.py
View file

@ -754,8 +754,8 @@ class TestRouteErrors(unittest.TestCase):
with (
patch.object(training_route, "get_training_backend", return_value = DummyBackend()),
patch(
"core.inference.get_inference_backend",
return_value = SimpleNamespace(active_model_name = None),
"routes.training_vram.summarize_resident_chat",
return_value = {"any": False, "hf": None, "gguf": None},
),
patch(
"core.export.get_export_backend",
@ -794,8 +794,8 @@ class TestRouteErrors(unittest.TestCase):
with (
patch.object(training_route, "get_training_backend", return_value = DummyBackend()),
patch(
"core.inference.get_inference_backend",
return_value = SimpleNamespace(active_model_name = None),
"routes.training_vram.summarize_resident_chat",
return_value = {"any": False, "hf": None, "gguf": None},
),
patch(
"core.export.get_export_backend",

352
studio/backend/tests/test_hf_xet_fallback.py Normal file
View file

@ -0,0 +1,352 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP
transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on.
CPU-only, no network, no real subprocess (the per-attempt download seam is
monkeypatched).
"""
from __future__ import annotations
import subprocess
import sys
import threading
import time
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub heavy/unavailable deps before importing the module under test. Use the
# real structlog when present; a bare stub left in sys.modules would break later
# modules that log at import time.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
try:
import structlog # noqa: F401
except ImportError:
sys.modules["structlog"] = _types.ModuleType("structlog")
import huggingface_hub
from huggingface_hub import constants as hf_constants
import utils.hf_xet_fallback as xf
# --------------------------------------------------------------------------- #
# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total.
# --------------------------------------------------------------------------- #
REPO = "ztest/xet-watchdog"
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
return tmp_path
def _blobs_dir(root: Path, repo_id: str = REPO) -> Path:
d = root / f"models--{repo_id.replace('/', '--')}" / "blobs"
d.mkdir(parents = True, exist_ok = True)
return d
def _wait(
predicate,
timeout: float = 2.0,
step: float = 0.02,
) -> bool:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(step)
return predicate()
def test_constant_incomplete_fires_stall(hf_cache):
blobs = _blobs_dir(hf_cache)
(blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows
calls: list[str] = []
stop = xf.start_watchdog(
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
)
try:
assert _wait(
lambda: len(calls) >= 1, timeout = 3.0
), "watchdog never fired on a constant-size .incomplete"
finally:
stop.set()
assert "stalled" in calls[0].lower()
def test_growing_incomplete_never_stalls(hf_cache):
blobs = _blobs_dir(hf_cache)
part = blobs / "growing.incomplete"
part.write_bytes(b"\0" * 1024)
grow_stop = threading.Event()
def _grow():
size = 1024
while not grow_stop.wait(0.05):
size += 4096
part.write_bytes(b"\0" * size)
grower = threading.Thread(target = _grow, daemon = True)
grower.start()
calls: list[str] = []
stop = xf.start_watchdog(
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
)
try:
time.sleep(1.0) # well past stall_timeout, but bytes keep growing
assert calls == [], "watchdog fired despite continuous progress"
finally:
stop.set()
grow_stop.set()
def test_no_incomplete_never_stalls(hf_cache):
blobs = _blobs_dir(hf_cache)
(blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete
calls: list[str] = []
stop = xf.start_watchdog(
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
)
try:
time.sleep(0.8)
assert calls == [], "watchdog fired with no active .incomplete"
finally:
stop.set()
def test_stall_fires_at_most_once(hf_cache):
blobs = _blobs_dir(hf_cache)
(blobs / "frozen.incomplete").write_bytes(b"\0" * 2048)
calls: list[str] = []
stop = xf.start_watchdog(
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2
)
try:
assert _wait(lambda: len(calls) >= 1, timeout = 3.0)
time.sleep(0.6) # keep ticking; must not fire again
assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1"
finally:
stop.set()
def test_get_state_empty_cache(hf_cache):
assert xf.get_hf_download_state([REPO]) == (0, False)
def test_get_state_absent_cache_root(tmp_path, monkeypatch):
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache"))
assert xf.get_hf_download_state([REPO]) == (0, False)
def test_get_state_skips_local_paths(hf_cache):
# Filesystem paths are not HF repo IDs and must be ignored without error.
assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False)
def test_get_state_sparse_aware(hf_cache):
blobs = _blobs_dir(hf_cache)
sparse = blobs / "sparse.incomplete"
with open(sparse, "wb") as f:
f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks
st = sparse.stat()
if getattr(st, "st_blocks", 0) == 0:
pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable")
total, has_incomplete = xf.get_hf_download_state([REPO])
assert has_incomplete is True
assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks"
# --------------------------------------------------------------------------- #
# Transport policy: cached short-circuit, cancel, error propagation, and the
# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn.
# --------------------------------------------------------------------------- #
DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf"
@pytest.fixture(autouse = True)
def _no_real_cache_hit(monkeypatch):
"""Default: the cached probe misses; tests override it to force a hit."""
monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None)
class _FakeAttempt:
"""Records calls to the download seam and returns scripted results."""
def __init__(self, results):
self._results = list(results)
self.calls = []
def __call__(
self,
repo_id,
filename,
token,
*,
repo_type,
disable_xet,
cancel_event,
stall_timeout,
interval,
grace_period,
on_status,
):
self.calls.append(
_types.SimpleNamespace(
repo_id = repo_id,
filename = filename,
disable_xet = disable_xet,
repo_type = repo_type,
)
)
return self._results[len(self.calls) - 1]
def _install(monkeypatch, results):
fake = _FakeAttempt(results)
monkeypatch.setattr(xf, "_run_download_attempt", fake)
return fake
def test_cached_file_short_circuits(monkeypatch, tmp_path):
cached = tmp_path / "cached.gguf"
cached.write_bytes(b"\0" * 8)
monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached))
fake = _install(monkeypatch, []) # must not be called
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert out == str(cached)
assert fake.calls == [], "spawned a download for an already-cached file"
def test_cancel_before_start_raises_no_attempt(monkeypatch):
fake = _install(monkeypatch, [])
ev = threading.Event()
ev.set()
with pytest.raises(RuntimeError, match = "Cancelled"):
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev)
assert fake.calls == []
def test_nonstall_error_propagates_without_fallback(monkeypatch):
fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")])
with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"):
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback"
assert fake.calls[0].disable_xet is False
def test_immediate_success_uses_xet_only(monkeypatch):
prepared = []
monkeypatch.setattr(
"hub.utils.download_registry.prepare_cache_for_transport",
lambda *a, **k: prepared.append(a),
)
fake = _install(monkeypatch, [("ok", "/cache/model.gguf")])
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert out == "/cache/model.gguf"
assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False
assert prepared == [], "no cache prep should run when Xet succeeds first try"
def test_stall_then_http_fallback_succeeds(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)),
)
fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")])
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert out == "/cache/model.gguf"
assert len(fake.calls) == 2
assert fake.calls[0].disable_xet is False # Xet first
assert fake.calls[1].disable_xet is True # HTTP fallback
assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry"
def test_second_stall_raises_download_stall_error(monkeypatch):
monkeypatch.setattr(
"hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
)
fake = _install(monkeypatch, [("stall", None), ("stall", None)])
with pytest.raises(xf.DownloadStallError):
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert len(fake.calls) == 2
def test_cancelled_midattempt_raises_no_fallback(monkeypatch):
fake = _install(monkeypatch, [("cancelled", None)])
with pytest.raises(RuntimeError, match = "Cancelled"):
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert len(fake.calls) == 1
def test_per_file_independent_fallback(monkeypatch):
"""A stalled shard falls back; a sibling shard that succeeds does not."""
monkeypatch.setattr(
"hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
)
fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")])
assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a"
assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b"
assert [c.disable_xet for c in fake.calls] == [False, False, True]
# --------------------------------------------------------------------------- #
# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect
# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it).
# --------------------------------------------------------------------------- #
def _safe_path() -> str:
import os
return os.environ.get("PATH", "")
def test_disable_xet_constant_set_in_fresh_interpreter():
code = (
"from huggingface_hub import constants as c; "
"import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)"
)
proc = subprocess.run(
[sys.executable, "-c", code],
env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()},
capture_output = True,
text = True,
)
assert proc.returncode == 0, (
f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True "
f"(rc={proc.returncode}): {proc.stderr}"
)
def test_default_leaves_xet_enabled():
code = (
"from huggingface_hub import constants as c; "
"import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)"
)
proc = subprocess.run(
[sys.executable, "-c", code],
env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET
capture_output = True,
text = True,
)
assert proc.returncode == 0, (
f"without the env var, constants.HF_HUB_DISABLE_XET was not False "
f"(rc={proc.returncode}): {proc.stderr}"
)

8
studio/backend/tests/test_kv_cache_estimation.py
View file

@ -71,7 +71,7 @@ except ImportError:
)
sys.modules["httpx"] = _httpx_stub
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
# Helpers
@ -1484,8 +1484,8 @@ class TestServerFlags:
assert fitted < 32_768
def test_fit_mtp_engaged_returns_smaller_or_equal_context(self):
# MTP budget is 0.85 of available, non-MTP is 0.90; on a tight
# budget MTP must yield <= non-MTP.
# Flat MTP fallback budget is _CTX_FIT_VRAM_FRACTION - 0.05; non-MTP is
# the full fraction. On a tight budget MTP must yield <= non-MTP.
b = self._gqa_backend()
common = dict(
requested_ctx = 32_768,
@ -1518,7 +1518,7 @@ class TestServerFlags:
kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
assert kv_full > kv_default
# Budget = model + kv_default (rounded up) -- swa_full must not fit.
budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 1
budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1
fitted_default = b._fit_context_to_vram(
requested_ctx = ctx,
available_mib = int(budget_mib),

134
studio/backend/tests/test_lifespan_shutdown.py Normal file
View file

@ -0,0 +1,134 @@
# 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 run_lifespan_shutdown: a dead default executor (the
abrupt-shutdown teardown race) must not abort the remaining cleanup. The helper
is dependency-injected, so these need only structlog."""
import asyncio
import contextvars
import types
from utils.lifespan_shutdown import run_lifespan_shutdown
def _counter():
box = {"n": 0}
def _fn():
box["n"] += 1
return box, _fn
def test_run_lifespan_shutdown_survives_dead_default_executor():
term_box, terminate = _counter()
clear_box, clear = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
async def _drive():
loop = asyncio.get_running_loop()
# Kill the default executor to mimic the teardown race.
await asyncio.to_thread(lambda: None)
loop._default_executor.shutdown(wait = True)
await run_lifespan_shutdown(terminate, clear, hw)
asyncio.run(_drive())
assert term_box["n"] == 1, "terminate must run via inline fallback"
assert clear_box["n"] == 1, "clear must still run after the to_thread failure"
assert hw.DEVICE is None
def test_run_lifespan_shutdown_survives_shutdown_default_executor():
"""Production path: loop.shutdown_default_executor() makes run_in_executor raise
'Executor shutdown has been called'; the helper must still recover inline."""
term_box, terminate = _counter()
clear_box, clear = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
async def _drive():
await asyncio.get_running_loop().shutdown_default_executor()
await run_lifespan_shutdown(terminate, clear, hw)
asyncio.run(_drive())
assert term_box["n"] == 1, "terminate must run via inline fallback"
assert clear_box["n"] == 1
assert hw.DEVICE is None
def test_run_lifespan_shutdown_normal_path():
"""Healthy executor: each step runs exactly once."""
term_box, terminate = _counter()
clear_box, clear = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
asyncio.run(run_lifespan_shutdown(terminate, clear, hw))
assert term_box["n"] == 1
assert clear_box["n"] == 1
assert hw.DEVICE is None
def test_run_lifespan_shutdown_swallows_terminate_errors():
"""A terminate failure must not block later cleanup."""
clear_box, clear = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
def _boom():
raise ValueError("boom")
asyncio.run(run_lifespan_shutdown(_boom, clear, hw))
assert clear_box["n"] == 1, "later cleanup must run even when terminate raises"
assert hw.DEVICE is None
def test_run_lifespan_shutdown_swallows_clear_errors():
"""A clear failure must not raise out of shutdown."""
term_box, terminate = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
def _boom():
raise ValueError("boom")
asyncio.run(run_lifespan_shutdown(terminate, _boom, hw))
assert term_box["n"] == 1
assert hw.DEVICE is None
def test_run_lifespan_shutdown_does_not_retry_body_runtime_error():
"""A body-side RuntimeError (healthy executor) must run terminate once, not retry inline."""
term_box, _ = _counter()
clear_box, clear = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
def _boom():
term_box["n"] += 1
raise RuntimeError("body failed")
asyncio.run(run_lifespan_shutdown(_boom, clear, hw))
assert term_box["n"] == 1, "body RuntimeError must not be retried inline"
assert clear_box["n"] == 1, "later cleanup must still run"
assert hw.DEVICE is None
def test_run_lifespan_shutdown_preserves_contextvars():
"""terminate runs in a copy of the caller's context (parity with asyncio.to_thread)."""
cv = contextvars.ContextVar("unsloth_test_cv")
cv.set("bound-value")
seen = []
clear_box, clear = _counter()
hw = types.SimpleNamespace(DEVICE = "cuda:0")
def terminate():
seen.append(cv.get("UNSET"))
asyncio.run(run_lifespan_shutdown(terminate, clear, hw))
assert seen == ["bound-value"], "terminate must run with the caller's contextvars"
assert clear_box["n"] == 1
assert hw.DEVICE is None

151
studio/backend/tests/test_llama_cpp_context_fit.py
View file

@ -67,7 +67,11 @@ _httpx_stub.Client = type(
)
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_cpp import (
_CTX_FIT_VRAM_FRACTION,
LlamaCppBackend,
classify_gpu_offload_lines,
)
from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
@ -171,7 +175,7 @@ def _drive(
)
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * 0.90:
if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION:
best_cap = max(best_cap, capped)
if best_cap > 0:
max_available_ctx = best_cap
@ -557,3 +561,146 @@ class TestClassifyGpuOffload:
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
def test_offloaded_zero_count_returns_false(self):
# Authoritative count overrides any GPU-looking buffer line.
inst = self._backend(
[
"load_tensors: offloaded 0/33 layers to GPU",
"load_tensors: CUDA0 model buffer size = 21000.0 MiB",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
def test_offloaded_draft_then_main_returns_true(self):
# A small draft model (0/2) does not mask the main model (33/33).
inst = self._backend(
[
"load_tensors: offloaded 0/2 layers to GPU",
"load_tensors: offloaded 33/33 layers to GPU",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
def test_main_on_cpu_with_draft_on_gpu_returns_false(self):
# MTP: the small drafter fits on GPU (1/1) but the main model is on CPU
# (0/33). Decide on the largest model, so the warning still fires.
inst = self._backend(
[
"load_tensors: offloaded 0/33 layers to GPU",
"load_tensors: offloaded 1/1 layers to GPU",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
def test_main_on_gpu_with_draft_on_cpu_returns_true(self):
# Reverse: main model on GPU (33/33), drafter on CPU (0/1) -> no warning.
inst = self._backend(
[
"load_tensors: offloaded 33/33 layers to GPU",
"load_tensors: offloaded 0/1 layers to GPU",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
def test_cuda_host_buffer_excluded_returns_false(self):
# CUDA_Host is CPU-pinned memory, not a model offload.
inst = self._backend(
[
"load_tensors: CUDA_Host model buffer size = 500.0 MiB",
"load_tensors: CPU model buffer size = 21000.0 MiB",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
def test_device_info_gpu_row_alone_is_inconclusive(self):
# device_info lists available devices, not where the model loaded, so a
# GPU row alone is not proof of offload.
inst = self._backend(
[
"print_info: device_info:",
" - CUDA0 : 24564 MiB free",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is None
def test_cpu_buffers_with_gpu_device_row_returns_false(self):
# Definite CPU-only buffers must win over a GPU device-inventory row.
inst = self._backend(
[
"load_tensors: CPU model buffer size = 21000.0 MiB",
"print_info: device_info:",
" - CUDA0 : 24564 MiB free",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
def test_device_info_cpu_only_returns_false(self):
inst = self._backend(
[
"print_info: device_info:",
" - CPU : 64000 MiB free",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
def test_system_info_cuda_before_device_info_does_not_count(self):
# A compiled-in backend named in system_info is not proof of offload;
# only the device_info table (here CPU only) decides.
inst = self._backend(
[
"system_info: CUDA : ARCHS = 890 | n_threads = 8",
"print_info: device_info:",
" - CPU : 64000 MiB free",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
@pytest.mark.parametrize(
"marker",
["CUDA0", "ROCm0", "HIP0", "Metal", "Vulkan0", "OpenCL0", "SYCL0", "MUSA0", "CANN0"],
)
def test_all_gpu_buffer_markers_return_true(self, marker):
assert (
classify_gpu_offload_lines([f"load_tensors: {marker} model buffer size = 8000.0 MiB"])
is True
)
def test_module_level_no_signal_returns_none(self):
assert classify_gpu_offload_lines(["INFO starting server"]) is None
def test_select_gpus_ranks_by_usable_not_raw_free():
# 80 GB card (30 GB free -> 25.9 GB usable) vs 32 GB card (29 GB free -> 27.4
# GB usable). A 27 GB model fits the 32 GB card alone; raw-free ranking would
# try the 80 GB card first and split across both. Usable ranking picks [1].
gpus = [(0, 30000), (1, 29000)]
totals = {0: 81920, 1: 32607}
model = int(27000 * 1024 * 1024)
idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals)
assert idxs == [1] and use_fit is False
def test_select_gpus_reserves_per_device_overhead():
# Two 16 GB cards, ~15181 MiB usable each at 0.95 -> 30362 MiB pooled. A 30000
# MiB model fits the pool with no per-device overhead, but a layer split also
# pays ~1 GiB/extra-GPU; that pushes the 2-GPU need to 31024 MiB > pool, so a
# pin would OOM -> must fall back to --fit. Single-GPU fits add no overhead
# (Finding F1, the explicit/file-size multi-GPU pin gap).
gpus = [(0, 16000), (1, 16000)]
totals = {0: 16384, 1: 16384}
gib = 1024 * 1024 * 1024
model = int(30000 * 1024 * 1024)
idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals)
assert idxs == [0, 1] and use_fit is False # fits 2 GPUs without overhead
idxs2, use_fit2 = LlamaCppBackend._select_gpus(
model, gpus, total_by_idx = totals, per_device_overhead_bytes = gib
)
assert idxs2 is None and use_fit2 is True # overhead tips it past the pool
# A single-GPU fit is unchanged by the overhead (k=1 adds nothing).
small = int(15000 * 1024 * 1024)
a, _ = LlamaCppBackend._select_gpus(small, gpus, total_by_idx = totals)
b, _ = LlamaCppBackend._select_gpus(
small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib
)
assert a == [0] and b == [0]

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

@ -520,3 +520,114 @@ def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path)
assert info["latest_tag"] == "b9596-mix-aaa"
assert info["behind"] is True
assert info["stale"] is True
# update_download_size_bytes (banner download-size lookup).
def _patch_assets(monkeypatch, mapping):
"""Stub latest_release_assets with a per-repo {asset_name: size} lookup."""
monkeypatch.setattr(
fr,
"latest_release_assets",
lambda repo, *, force_refresh = False: mapping.get(repo),
)
def test_update_size_unsloth_prebuilt_exact_match(monkeypatch):
# The unsloth fork's own bundle (app-<tag>-<platform>): the want= exact match
# on app-<latest>-<suffix> wins.
marker = {
"asset": "app-b9190-linux-x64-cuda13-newer.tar.gz",
"published_repo": "unslothai/llama.cpp",
}
_patch_assets(
monkeypatch,
{
"unslothai/llama.cpp": {
"app-b9300-linux-x64-cuda13-newer.tar.gz": 123_456_789,
"app-b9300-windows-x64-cuda13-newer.zip": 999,
}
},
)
assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 123_456_789
def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch):
# macOS bundles use the upstream-style llama-<tag>-bin-macos-*, matched via the
# endswith fallback in the publish repo.
marker = {
"asset": "llama-b9190-bin-macos-arm64.tar.gz",
"published_repo": "unslothai/llama.cpp",
}
_patch_assets(
monkeypatch,
{"unslothai/llama.cpp": {"llama-b9300-bin-macos-arm64.tar.gz": 55_000_000}},
)
assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 55_000_000
def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch):
# #6338 P2: ggml-org ubuntu-* prebuilt lives in binary_repo, not the fork
# publish repo. The size must still resolve.
marker = {
"asset": "llama-b9190-bin-ubuntu-x64.tar.gz",
"published_repo": "unslothai/llama.cpp",
"binary_repo": "ggml-org/llama.cpp",
}
_patch_assets(
monkeypatch,
{
"unslothai/llama.cpp": {"app-b9300-linux-x64-cuda13-newer.tar.gz": 1},
"ggml-org/llama.cpp": {
"llama-b9673-bin-ubuntu-x64.tar.gz": 42_000_000,
"llama-b9673-bin-ubuntu-vulkan-x64.tar.gz": 7,
},
},
)
assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 42_000_000
def test_update_size_upstream_windows_uses_binary_repo(monkeypatch):
# Regression (#6338 P2): the Windows upstream CPU prebuilt uses a win-* token.
marker = {
"asset": "llama-b9190-bin-win-cpu-x64.zip",
"published_repo": "unslothai/llama.cpp",
"binary_repo": "ggml-org/llama.cpp",
}
_patch_assets(
monkeypatch,
{"ggml-org/llama.cpp": {"llama-b9673-bin-win-cpu-x64.zip": 33_000_000}},
)
assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 33_000_000
def test_update_size_no_matching_asset_fails_open(monkeypatch):
# A ROCm version drift (installed 6.4 vs latest 7.2) leaves no suffix match;
# the helper fails open to None rather than guessing a wrong artifact.
marker = {
"asset": "llama-b9190-bin-ubuntu-rocm-6.4-x64.tar.gz",
"published_repo": "unslothai/llama.cpp",
"binary_repo": "ggml-org/llama.cpp",
}
_patch_assets(
monkeypatch,
{"ggml-org/llama.cpp": {"llama-b9673-bin-ubuntu-rocm-7.2-x64.tar.gz": 9}},
)
assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") is None
def test_update_size_missing_inputs_fail_open(monkeypatch):
_patch_assets(
monkeypatch,
{"unslothai/llama.cpp": {"app-b9300-linux-x64-cpu.tar.gz": 5}},
)
# No marker, no latest tag, or no asset string -> None (never raise).
assert fr.update_download_size_bytes(None, "b9300", "unslothai/llama.cpp") is None
assert (
fr.update_download_size_bytes(
{"asset": "app-b9190-linux-x64-cpu.tar.gz"}, None, "unslothai/llama.cpp"
)
is None
)
assert fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") is None

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

@ -74,7 +74,7 @@ _httpx_stub.Client = type(
)
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
# Helpers
@ -140,7 +140,7 @@ def _compute_max_available_ctx(
)
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * 0.90:
if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION:
best_cap = max(best_cap, capped)
if best_cap > 0:
max_available_ctx = best_cap

150
studio/backend/tests/test_llama_cpp_mmproj_fallback.py
View file

@ -40,6 +40,9 @@ from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
_detect = LlamaCppBackend._is_projector_incompatibility
_strip = LlamaCppBackend._strip_mmproj_args
_signal_crash = LlamaCppBackend._is_signal_crash
_flash_off = LlamaCppBackend._with_flash_attn_off
_nonproj = LlamaCppBackend._output_has_nonprojector_diagnostic
# Real abort captured loading gemma-4 on a 3-day-old prebuilt (build b9496).
_GEMMA4_OLD_LLAMACPP_OUT = (
@ -103,6 +106,20 @@ class TestProjectorIncompatibilityDetector:
assert _detect(out) is False
class TestSignalCrashDetector:
"""_is_signal_crash flags a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS
or a Windows 0xC0000000+ fault); not a clean exit, hung (None), or an
external kill (SIGKILL/SIGTERM/SIGINT) that an OOM/unload would cause."""
@pytest.mark.parametrize("rc", [-11, -6, -4, -7, -8, 0xC0000005, 0xC000001D])
def test_program_faults_are_hard_crashes(self, rc):
assert _signal_crash(rc) is True
@pytest.mark.parametrize("rc", [0, 1, 2, 137, None, -9, -15, -2])
def test_clean_hung_or_external_kill_is_not(self, rc):
assert _signal_crash(rc) is False
# A realistic vision launch argv (mirrors the live "Starting llama-server"
# command), projector pair at the end.
_VISION_CMD = [
@ -170,6 +187,99 @@ class TestStripMmprojArgs:
assert cmd[-1] == "/p/mm.gguf" # input untouched
class TestFlashAttnOff:
"""_with_flash_attn_off is the least-destructive recovery rung: flip
'--flash-attn on' to 'off' (keeps vision + MTP), or None when there is
nothing to disable."""
def test_flips_on_to_off_keeping_vision_and_mtp(self):
out = _flash_off(_VISION_CMD)
assert out is not None
# FA disabled, every other capability (mmproj, MTP, ctx) preserved.
i = out.index("--flash-attn")
assert out[i + 1] == "off"
assert "--mmproj" in out and "--spec-default" in out
assert len(out) == len(_VISION_CMD)
def test_none_when_already_off(self):
assert _flash_off(["llama-server", "--flash-attn", "off", "-c", "4096"]) is None
def test_none_when_no_flash_attn(self):
assert _flash_off(["llama-server", "-m", "/m.gguf", "-c", "4096"]) is None
def test_returns_new_list_input_untouched(self):
cmd = ["llama-server", "--flash-attn", "on"]
out = _flash_off(cmd)
assert out == ["llama-server", "--flash-attn", "off"]
assert cmd[-1] == "on" # input not mutated
def test_flips_equals_form(self):
out = _flash_off(["llama-server", "--flash-attn=on", "-c", "4096"])
assert out == ["llama-server", "--flash-attn=off", "-c", "4096"]
def test_flips_fa_alias_and_auto(self):
assert _flash_off(["llama-server", "-fa", "auto"]) == ["llama-server", "-fa", "off"]
assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"]
def test_flips_every_occurrence_last_wins(self):
# extra_args can re-enable FA after Studio's flag; llama.cpp is last-wins,
# so one leftover 'on' would re-crash the retry. Every enable must flip.
cmd = ["llama-server", "--flash-attn", "on", "--mmproj", "/p", "--flash-attn", "on"]
out = _flash_off(cmd)
assert out is not None
assert "on" not in out
assert out.count("off") == 2
def test_none_when_equals_off(self):
assert _flash_off(["llama-server", "--flash-attn=off"]) is None
def test_none_when_user_off_wins_last(self):
# User appended 'off' after Studio's 'on'; effective (last-wins) is off,
# so there is nothing to retry.
assert _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn", "off"]) is None
def test_neutralizes_trailing_bare_flag(self):
# A bare --flash-attn reads as on under last-wins; it must be neutralized
# too, else the retry re-enables FA and re-crashes.
out = _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn"])
assert out == ["llama-server", "--flash-attn", "off", "--flash-attn=off"]
assert "on" not in out
def test_bare_flag_only(self):
assert _flash_off(["llama-server", "--flash-attn"]) == ["llama-server", "--flash-attn=off"]
assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"]
class TestNonProjectorDiagnostic:
"""_output_has_nonprojector_diagnostic gates the signal-only text-only retry:
a hard crash that already names OOM / a bad arch / a TP limit must surface
that error, not be silently downgraded to a non-vision session."""
@pytest.mark.parametrize(
"out",
[
_OOM_OUT,
_BAD_ARCH_OUT,
"ggml_backend_cuda_buffer_type_alloc: failed to allocate buffer",
"split_mode_tensor not implemented for this architecture",
],
)
def test_known_nonprojector_causes_match(self, out):
assert _nonproj(out) is True
@pytest.mark.parametrize(
"out",
[
"", # bare crash: no marker -> still eligible for the text-only retry
_GEMMA4_OLD_LLAMACPP_OUT, # a real projector abort must NOT be suppressed
_HEALTHY_VISION_OUT,
_PORT_OUT,
],
)
def test_bare_or_projector_output_does_not_match(self, out):
assert _nonproj(out) is False
class TestRetryContract:
"""The two helpers compose into the load_model retry decision."""
@ -185,3 +295,43 @@ class TestRetryContract:
# An OOM with --mmproj present must NOT be treated as a projector
# problem: load_model errors out instead of dropping vision.
assert _detect(_OOM_OUT) is False
def test_bare_segfault_with_mmproj_yields_text_only_retry(self):
# Field report: a -11 SIGSEGV on --mmproj has no projector line; the
# signal path fires only when no other diagnostic explains the crash.
out = "" # a SIGSEGV produced no projector-format line
assert _detect(out) is False
should_retry = _detect(out) or (_signal_crash(-11) and not _nonproj(out))
assert should_retry is True
retry_cmd = _strip(_VISION_CMD)
assert "--mmproj" not in retry_cmd and "-m" in retry_cmd
def test_signal_crash_with_oom_output_keeps_the_real_error(self):
# A hard fault that already printed an OOM must surface it, not silently
# drop --mmproj and tell the user to update llama.cpp.
assert _signal_crash(-6) is True
should_retry = _detect(_OOM_OUT) or (_signal_crash(-6) and not _nonproj(_OOM_OUT))
assert should_retry is False
def test_signal_crash_with_bad_arch_does_not_drop_vision(self):
should_retry = _detect(_BAD_ARCH_OUT) or (_signal_crash(-6) and not _nonproj(_BAD_ARCH_OUT))
assert should_retry is False
def test_clean_nonzero_exit_with_mmproj_does_not_retry(self):
# Clean non-zero exit (bad path, port bind) is not a hard crash; stay message-based.
assert (_detect(_MISSING_OUT) or _signal_crash(1)) is False
def test_signal_crash_tries_flash_attn_off_before_dropping_vision(self):
# Ladder order: a hard fault retries FA-off FIRST (keeps vision + MTP);
# only if THAT is None/fails do we fall back to stripping --mmproj.
assert _signal_crash(-11) is True
fa_retry = _flash_off(_VISION_CMD)
assert fa_retry is not None
assert "--mmproj" in fa_retry # vision preserved at this rung
# Last resort still available and strictly more destructive.
text_only = _strip(fa_retry)
assert "--mmproj" not in text_only
def test_external_kill_skips_flash_attn_retry(self):
# SIGKILL (-9, OOM killer) is not a program fault: no FA-off retry.
assert _signal_crash(-9) is False

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

@ -9,6 +9,8 @@ the _already_in_target_state mirror that prevents needless reloads.
from __future__ import annotations
import inspect
import os
import struct
import sys
import types as _types
@ -52,9 +54,12 @@ import pytest
from core.inference.llama_cpp import (
LlamaCppBackend,
_GPU_OFFLOAD_OVERRIDE_FLAGS,
_THREAD_OVERRIDE_FLAGS,
_backfill_usage_from_timings,
_build_ngram_mod_flags,
_canonicalize_spec_mode,
_extra_args_set_any_flag,
_extra_args_set_spec_type,
_is_mtp_model_name,
)
@ -315,6 +320,46 @@ def test_extra_args_set_spec_type_passes_on_non_spec_type_args(extra_args):
assert _extra_args_set_spec_type(extra_args) is False
@pytest.mark.parametrize(
"extra_args",
[
["-ngl", "12"],
["--gpu-layers", "12"],
["--n-gpu-layers=12"],
["-fit", "off"],
["--fit=off"],
],
)
def test_extra_args_detect_gpu_offload_overrides(extra_args):
assert _extra_args_set_any_flag(extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS) is True
@pytest.mark.parametrize("extra_args", [["-t", "8"], ["--threads=8"]])
def test_extra_args_detect_thread_overrides(extra_args):
assert _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS) is True
def test_windows_full_offload_flags_use_current_llama_server_args():
src = inspect.getsource(LlamaCppBackend.load_model)
stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens"
assert '"--cache-ram"' in src
assert '"--ctx-checkpoints"' in src
assert '"--no-cache-prompt"' in src
assert stale_checkpoint_flag not in src
def test_load_model_sets_threads_once():
src = inspect.getsource(LlamaCppBackend.load_model)
assert src.count('cmd.extend(["--threads", str(') == 1
def test_llama_cpp_annotations_stay_python39_safe():
src = inspect.getsource(LlamaCppBackend.generate_chat_completion)
helper_src = inspect.getsource(_extra_args_set_any_flag)
assert "Generator[str | dict" not in src
assert "set[str] | frozenset[str]" not in helper_src
def test_already_in_target_state_user_spec_type_override_matches_clean_backend():
# User --spec-type none suppressed auto-MTP; repeat /load must not re-promote.
backend = _mtp_backend(
@ -522,6 +567,38 @@ def test_probe_server_capabilities_detects_draft_mtp(tmp_path):
assert caps["supports_mtp"] is True
@_NEEDS_BASH
def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch):
fake = _make_fake_llama_server(
tmp_path / "llama-server",
"--spec-type none,mtp,ngram-simple\n",
)
captured = {}
monkeypatch.setattr(
"core.inference.llama_cpp.child_env_without_native_path_secret",
lambda: {"LD_LIBRARY_PATH": "/already-there"},
)
def fake_run(cmd, **kwargs):
captured["cmd"] = cmd
captured["env"] = kwargs.get("env")
return _types.SimpleNamespace(stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "")
monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["found"] is True
assert caps["supports_mtp"] is True
assert captured["cmd"] == [str(fake), "--help"]
assert captured["env"] is not None
ld_dirs = captured["env"]["LD_LIBRARY_PATH"].split(os.pathsep)
assert str(fake.parent) in ld_dirs
assert "/already-there" in ld_dirs
@_NEEDS_BASH
def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
# Renamed upstream: draft-mtp -> mtp.
@ -554,6 +631,9 @@ def test_probe_server_capabilities_handles_missing_binary():
caps = LlamaCppBackend.probe_server_capabilities("/no/such/llama-server")
assert caps["found"] is False
assert caps["supports_mtp"] is False
assert caps["supports_cache_ram"] is False
assert caps["supports_ctx_checkpoints"] is False
assert caps["supports_no_cache_prompt"] is False
# ngram-mod flag flavor detection (new vs legacy llama-server).
@ -588,6 +668,12 @@ _LEGACY_HELP = """\
--spec-type none,ngram-mod,ngram-simple comma-separated list of types of speculative decoding to use
"""
_CACHE_FLAGS_HELP = """\
--cache-ram N store prompt cache in RAM (default: 0)
--ctx-checkpoints N number of context checkpoints (default: 0)
--no-cache-prompt do not reuse prompt cache
"""
@_NEEDS_BASH
def test_probe_detects_post_rename_ngram_mod_flavor(tmp_path):
@ -634,6 +720,26 @@ def test_probe_no_ngram_mod_on_minimal_binary(tmp_path):
assert caps["supports_ngram_mod"] is False
@_NEEDS_BASH
def test_probe_detects_windows_cache_flags(tmp_path):
fake = _make_fake_llama_server(tmp_path / "llama-server", _CACHE_FLAGS_HELP)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["supports_cache_ram"] is True
assert caps["supports_ctx_checkpoints"] is True
assert caps["supports_no_cache_prompt"] is True
@_NEEDS_BASH
def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path):
fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n")
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["supports_cache_ram"] is False
assert caps["supports_ctx_checkpoints"] is False
assert caps["supports_no_cache_prompt"] is False
def test_build_ngram_mod_flags_new():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"})
assert flags == [
@ -1162,9 +1268,10 @@ def test_build_speculative_flags_user_draft_n_max_override(monkeypatch):
assert backend.spec_draft_n_max == 5
def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
# Outdated llama-server with no MTP support: forced MTP must degrade
# to spec-off (warned) rather than emit a bad --spec-type.
def test_build_speculative_flags_mtp_token_missing_emits_spec_default(monkeypatch):
# Outdated llama-server with no MTP support: forced MTP must degrade (warned)
# and emit --spec-default so an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI
# wins over env) can't make the child attempt MTP the gate budgeted off.
backend = _resolver_backend(monkeypatch, mtp_token = None)
flags = backend._build_speculative_flags(
speculative_type = "mtp",
@ -1176,10 +1283,11 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
binary = "/fake/llama-server",
)
assert "--spec-type" not in flags
# _speculative_type stays None (resolved emission was none); the user's
# choice is still reflected in _requested_spec_mode.
assert "--spec-default" in flags
# Degraded to non-speculative; the user's choice is still reflected.
assert backend.speculative_type == "default"
assert backend.requested_spec_mode == "mtp"
assert backend.speculative_type is None
assert backend.spec_fallback_reason == "binary_no_mtp"
def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch):

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

@ -139,3 +139,35 @@ class TestOllamaAndFallback:
def test_empty_output_is_safe(self):
msg = _classify("", None, None)
assert "llama-server failed to start" in msg
class TestOsKillReturncode:
"""SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named,
actionable message; SIGTERM (-15) is also unload/cancel/supervisor stop, so it
stays neutral; a recognized output still wins; a hard fault (-11) keeps the
generic fallback."""
def test_sigkill_with_no_output_names_oom(self):
msg = _classify("", "/models/big-bf16.gguf", "local/big", -9)
assert "signal 9" in msg
assert "out of memory" in msg.lower()
assert ".wslconfig" in msg
assert "GGUF file is valid" not in msg
def test_sigterm_is_neutral_not_oom(self):
msg = _classify("", "/models/big-bf16.gguf", "local/big", -15)
assert "signal 15" in msg
assert "terminated" in msg.lower()
assert "out of memory" not in msg.lower()
def test_specific_output_wins_over_os_kill_code(self):
msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image", -9)
assert "diffusion" in msg.lower()
assert "out of memory" not in msg.lower()
def test_signal_crash_code_keeps_generic_message(self):
# -11 is handled by the retry ladder; if it reaches here with no output
# it gets the generic fallback, not the OOM message.
msg = _classify("", "/models/x.gguf", "local/x", -11)
assert "GGUF file is valid" in msg
assert "out of memory" not in msg.lower()

20
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -106,7 +106,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
tool_call_id = "call_render_late"
first_stream = [
_sse({"content": "Here is the artifact.\n\n"}),
_sse({"content": "Here is the canvas.\n\n"}),
_sse(
{
"tool_calls": [
@ -140,7 +140,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "Rendered HTML artifact: Simple Red Square."
return "Rendered HTML canvas: Simple Red Square."
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
@ -168,7 +168,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
)
content_events = [e for e in events if e.get("type") == "content"]
assert content_events[0]["text"] == "Here is the artifact.\n\n"
assert content_events[0]["text"] == "Here is the canvas.\n\n"
first_content_index = next(
i for i, event in enumerate(events) if event.get("type") == "content"
@ -195,7 +195,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
# plus the structured tool call, preserving OpenAI-compatible ordering.
assert len(payloads) == 2
assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"]
assert assistant_messages[-1]["content"] == "Here is the artifact.\n\n"
assert assistant_messages[-1]["content"] == "Here is the canvas.\n\n"
assert assistant_messages[-1]["tool_calls"][0]["id"] == tool_call_id
assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html"
@ -257,7 +257,7 @@ def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch):
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "Rendered HTML artifact: First."
return "Rendered HTML canvas: First."
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
@ -346,7 +346,7 @@ def test_render_html_success_drops_tool_schema_before_final_pass(monkeypatch):
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
def fake_execute_tool(name, arguments, **_kwargs):
return "Rendered HTML artifact: Done."
return "Rendered HTML canvas: Done."
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
@ -783,7 +783,7 @@ def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(m
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "Rendered HTML artifact: One."
return "Rendered HTML canvas: One."
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
@ -869,7 +869,7 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch):
The post-tool model pass can say it will use render_html again without
emitting a tool call. That should be accepted as a final model mistake,
not turned into repeated internal re-prompts after the artifact already
not turned into repeated internal re-prompts after the canvas already
exists.
"""
@ -907,7 +907,7 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch):
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "Rendered HTML artifact: First."
return "Rendered HTML canvas: First."
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
@ -1146,7 +1146,7 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "Rendered HTML artifact: Forced."
return "Rendered HTML canvas: Forced."
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)

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

@ -308,6 +308,33 @@ def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path):
assert probes == {"resolve": 0, "version": 0}
def test_installed_version_skips_probe_while_job_runs(monkeypatch, tmp_path):
# Markerless build: get_installed_llama_version falls back to exec'ing
# `llama-server --version`. While the updater swaps the tree that exec can
# fail the installer's os.replace on Windows, so the About-panel probe must
# be skipped (return None) exactly like get_update_status's source probe.
binary = tmp_path / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # markerless: no UNSLOTH_PREBUILT_INFO.json
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
probed = {"n": 0}
def _count_version(b):
probed["n"] += 1
return 9585
monkeypatch.setattr(upd, "_installed_build_number", _count_version)
with upd._job_lock:
upd._job["state"] = upd._JOB_RUNNING
assert upd.get_installed_llama_version() is None
assert probed["n"] == 0 # never exec'd the binary mid-swap
upd._reset_job_for_tests() # back to idle -> probe runs
assert upd.get_installed_llama_version() == "b9585"
assert probed["n"] == 1
def test_status_update_available(monkeypatch, tmp_path):
binary = _write_install(tmp_path, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
@ -889,3 +916,49 @@ def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path):
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"
def test_status_update_available_includes_size(monkeypatch, tmp_path):
# Marker (prebuilt) update path attaches the download size of the asset the
# banner would fetch.
binary = _write_install(tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
monkeypatch.setattr(
freshness,
"latest_release_assets",
lambda repo, *, force_refresh = False: {
"app-b9518-linux-x64-cuda13-newer.tar.gz": 88_000_000
},
)
st = upd.get_update_status(force_refresh = True)
assert st["update_available"] is True
assert st["update_size_bytes"] == 88_000_000
def test_status_source_build_includes_update_size(monkeypatch, tmp_path):
# #6338 P3: a source build offered a prebuilt must carry the asset size too.
binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub") # no marker -> source build
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(
monkeypatch,
repo = "unslothai/llama.cpp",
release_tag = "b9585",
asset = "app-b9585-linux-x64-cpu.tar.gz",
)
monkeypatch.setattr(upd, "_installed_build_number", lambda b: None)
monkeypatch.setattr(
upd,
"latest_release_assets",
lambda repo, *, force_refresh = False: (
{"app-b9585-linux-x64-cpu.tar.gz": 77_000_000}
if repo == "unslothai/llama.cpp"
else None
),
)
st = upd.get_update_status()
assert st["source_build"] is True
assert st["update_available"] is True
assert st["update_size_bytes"] == 77_000_000

294
studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
View file

@ -281,6 +281,7 @@ def test_kill_process_records_timestamp_on_actual_kill():
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = None
backend._healthy = False
backend._stats_logger = None # _kill_process stops it in finally
backend._stdout_thread = None
backend._llama_log_fh = None
backend._last_kill_monotonic = 0.0
@ -308,6 +309,26 @@ def test_kill_process_records_timestamp_on_actual_kill():
assert before <= backend._last_kill_monotonic <= after
def test_kill_process_tolerates_partially_constructed_backend():
# Teardown must not AttributeError on a __new__-built backend that never ran
# __init__: _stats_logger / _stdout_thread / _llama_log_fh are left unset.
backend = LlamaCppBackend.__new__(LlamaCppBackend)
class _FakeProcess:
def terminate(self):
pass
def wait(self, timeout = None):
return 0
def kill(self):
pass
backend._process = _FakeProcess()
backend._kill_process()
assert backend._process is None
def test_helper_is_static_method_callable_off_class():
"""Pin the @staticmethod binding so call sites can invoke off the class."""
ctx, _state = _patch_probe([[]])
@ -315,3 +336,276 @@ def test_helper_is_static_method_callable_off_class():
LlamaCppBackend._wait_for_vram_settle(
**_kw(max_wait = 0.1, interval = 0.05),
)
# ---------------------------------------------------------------------------
# Startup orphan-reaper arms the settle clock (the "wrong card after restart"
# root cause: reaped VRAM frees lazily, so the first load must wait).
# ---------------------------------------------------------------------------
def test_kill_orphaned_servers_returns_count():
"""The reaper reports how many owned orphans it killed, so __init__ can
arm the settle wait. Only Studio-owned llama-server procs count."""
import os
mypid = os.getpid()
fake_path = "/tmp/unsloth-test-llama/llama-server"
killed: list[int] = []
class _FakeProc:
def __init__(self, pid, name, exe):
self.info = {"pid": pid, "name": name, "exe": exe}
def kill(self):
killed.append(self.info["pid"])
owned = _FakeProc(mypid + 1, "llama-server", fake_path) # exact-path match
foreign = _FakeProc(mypid + 2, "llama-server", "/usr/bin/llama-server")
unrelated = _FakeProc(mypid + 3, "python3", "/usr/bin/python3")
fake_psutil = _types.ModuleType("psutil")
fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {})
fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {})
fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {})
fake_psutil.process_iter = lambda attrs = None: [owned, foreign, unrelated]
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
):
n = LlamaCppBackend._kill_orphaned_servers()
assert n == 1, "only the Studio-owned orphan should be counted"
assert killed == [mypid + 1]
# No owned orphans -> zero, so __init__ leaves the cold-start sentinel.
fake_psutil.process_iter = lambda attrs = None: [foreign, unrelated]
killed.clear()
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
):
assert LlamaCppBackend._kill_orphaned_servers() == 0
assert killed == []
def test_startup_reaper_arms_settle_timestamp():
"""__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an
orphan (so the first load_model waits for VRAM to settle), and leaves the
0.0 cold-start sentinel when nothing was reaped."""
with patch.object(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 1)):
before = time.monotonic()
backend = LlamaCppBackend()
after = time.monotonic()
assert (
before <= backend._last_kill_monotonic <= after
), "a positive reap count must arm the settle clock"
with patch.object(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)):
backend_cold = LlamaCppBackend()
assert (
backend_cold._last_kill_monotonic == 0.0
), "no reap must leave the cold-start sentinel so the wait is skipped"
# ---------------------------------------------------------------------------
# Cross-session backstop: a server PID recorded at spawn is reaped on the next
# startup even when parent-death cleanup did not run (macOS, a best-effort
# PR_SET_PDEATHSIG / Job Object failure, or a pre-existing orphan), but ONLY when
# it is a true orphan (its parent is gone), it still is a llama-server, and its
# start-time identity matches. A live server (parent still running) is spared so a
# helper backend built in-process can never kill the active chat server.
# ---------------------------------------------------------------------------
class _FakeKillProc:
def terminate(self):
pass
def wait(self, timeout = None):
return 0
def kill(self):
pass
def poll(self):
return 0
def test_kill_process_clears_pidfile(tmp_path):
"""A real kill removes the recorded pidfile so a clean eject leaves no orphan marker."""
pidfile = tmp_path / "llama-server.pid"
pidfile.write_text("12345")
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = _FakeKillProc()
backend._healthy = False
backend._stdout_thread = None
backend._llama_log_fh = None
backend._last_kill_monotonic = 0.0
backend._stats_logger = None
with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)):
backend._kill_process()
assert not pidfile.exists()
def test_reap_recorded_pid_kills_recorded_server(tmp_path):
"""An orphaned recorded PID (parent gone) is killed and the pidfile cleared
when it is still a llama-server."""
import subprocess
proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
pidfile = tmp_path / "llama-server.pid"
pidfile.write_text(str(proc.pid))
try:
with (
patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
patch.object(
LlamaCppBackend,
"_pid_is_llama_server",
staticmethod(lambda pid: pid == proc.pid),
),
):
n = LlamaCppBackend._reap_recorded_pid()
assert n == 1
assert not pidfile.exists()
proc.wait(timeout = 5)
assert proc.poll() is not None
finally:
if proc.poll() is None:
proc.kill()
proc.wait(timeout = 5)
def test_record_then_reap_round_trip_identity_matches(tmp_path):
"""Full round trip: _record_server_pid writes pid:starttime, and an orphaned
reap whose recorded identity still matches DOES kill it."""
import subprocess
proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
pidfile = tmp_path / "llama-server.pid"
try:
with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)):
LlamaCppBackend._record_server_pid(proc.pid)
assert ":" in pidfile.read_text(), "a start-time identity must be recorded"
with (
patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
):
n = LlamaCppBackend._reap_recorded_pid()
assert n == 1, "a matching identity on a true orphan must be reaped"
proc.wait(timeout = 5)
assert proc.poll() is not None
assert not pidfile.exists()
finally:
if proc.poll() is None:
proc.kill()
proc.wait(timeout = 5)
def test_reap_recorded_pid_spares_live_server(tmp_path):
"""A recorded server whose parent is still alive (the running Studio) is NEVER
reaped, and its pidfile is kept. This is the finding-3 guard: a helper backend
constructed in-process must not kill the active chat server. Uses the REAL
_pid_parent_is_alive (the child's parent is this live test process)."""
import subprocess
proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
pidfile = tmp_path / "llama-server.pid"
pidfile.write_text(str(proc.pid))
try:
with (
patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
# Force the name check True so ONLY the parent-alive guard can spare it.
patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
):
n = LlamaCppBackend._reap_recorded_pid()
assert n == 0, "a live server with a running parent must not be reaped"
assert proc.poll() is None, "the live server must still be running"
assert pidfile.exists(), "the record is kept so a later orphan reap still works"
finally:
proc.kill()
proc.wait(timeout = 5)
def test_reap_recorded_pid_skips_pid_reuse(tmp_path):
"""A recorded PID recycled to a non-llama-server must NOT be killed (only the
stale pidfile is cleaned), so the user's vllm/games are never touched."""
import subprocess
proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
pidfile = tmp_path / "llama-server.pid"
pidfile.write_text(str(proc.pid))
try:
with (
patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: False)),
):
n = LlamaCppBackend._reap_recorded_pid()
assert n == 0
assert proc.poll() is None, "an unrelated reused PID must not be killed"
assert not pidfile.exists(), "stale pidfile is cleaned up"
finally:
proc.kill()
proc.wait(timeout = 5)
def test_reap_recorded_pid_skips_identity_mismatch(tmp_path):
"""An orphaned PID whose recorded start-time identity no longer matches has been
recycled; it must NOT be killed even if it now looks like a llama-server."""
import subprocess
proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
pidfile = tmp_path / "llama-server.pid"
pidfile.write_text(f"{proc.pid}:0.0") # stale identity that cannot match
try:
with (
patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
):
n = LlamaCppBackend._reap_recorded_pid()
assert n == 0, "a PID whose start-time identity changed must not be killed"
assert proc.poll() is None, "the recycled process must survive"
assert not pidfile.exists(), "stale pidfile is cleaned up"
finally:
proc.kill()
proc.wait(timeout = 5)
def test_reap_recorded_pid_windows_sigkill_fallback(tmp_path, monkeypatch):
"""On Windows signal.SIGKILL is undefined; the reaper must fall back to SIGTERM
(os.kill -> TerminateProcess) instead of crashing and leaving the orphan."""
import os as _os
import signal as _signal
monkeypatch.delattr(_signal, "SIGKILL", raising = False)
captured = {}
def _fake_kill(pid, sig):
captured["pid"] = pid
captured["sig"] = sig # recorded; do not actually signal anything
pidfile = tmp_path / "llama-server.pid"
pidfile.write_text("424242")
with (
patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
patch.object(_os, "kill", _fake_kill),
):
n = LlamaCppBackend._reap_recorded_pid()
assert n == 1
assert (
captured.get("sig") == _signal.SIGTERM
), "must fall back to SIGTERM when SIGKILL is absent"
assert not pidfile.exists()
def test_reap_recorded_pid_no_pidfile(tmp_path):
"""No pidfile -> nothing reaped, no error."""
pidfile = tmp_path / "llama-server.pid" # never created
with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)):
assert LlamaCppBackend._reap_recorded_pid() == 0

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

@ -78,6 +78,27 @@ def test_status_response_exposes_source_build():
rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1})
def test_status_response_exposes_update_size_bytes():
payload = {
"supported": True,
"update_available": True,
"stale": False,
"installed_tag": "b9493",
"latest_tag": "b9518",
"published_repo": "unslothai/llama.cpp",
"installed_at_utc": None,
"age_days": None,
"source_build": False,
"update_size_bytes": 123_456_789,
"job": {"state": "idle"},
}
model = rl.LlamaUpdateStatusResponse(**payload)
assert model.model_dump()["update_size_bytes"] == 123_456_789
# Omitted -> defaults to None (the offline / no-matching-asset case).
without = {k: v for k, v in payload.items() if k != "update_size_bytes"}
assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None
def test_status_handler_runs_off_event_loop(monkeypatch):
seen = {}

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

@ -24,6 +24,7 @@ _lsa = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_lsa)
is_managed_flag = _lsa.is_managed_flag
parse_cache_override = _lsa.parse_cache_override
parse_cache_override_per_axis = _lsa.parse_cache_override_per_axis
parse_ctx_override = _lsa.parse_ctx_override
parse_split_mode_override = _lsa.parse_split_mode_override
resolve_cache_type_kv = _lsa.resolve_cache_type_kv
@ -467,6 +468,25 @@ def test_parse_cache_override_rejects_malformed_values(args):
parse_cache_override(args)
@pytest.mark.parametrize(
"args, expected",
[
(["--cache-type-k", "f32", "--cache-type-v", "f16"], ("f32", "f16")),
(["-ctk", "q8_0", "-ctv", "q4_0"], ("q8_0", "q4_0")),
(["--cache-type-k=f32"], ("f32", None)),
(["--cache-type-v", "f16"], (None, "f16")),
(["-c", "4096"], (None, None)),
(None, (None, None)),
# Last-wins is kept per axis.
(["-ctk", "f16", "-ctk", "f32"], ("f32", None)),
],
)
def test_parse_cache_override_per_axis(args, expected):
# Unlike parse_cache_override (collapses both axes to one last-wins value),
# this keeps K and V apart so an asymmetric cache can be budgeted per axis.
assert parse_cache_override_per_axis(args) == expected
def test_resolve_cache_type_kv_uses_override_when_present():
assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0"
@ -649,6 +669,53 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec():
assert out == ["--top-k", "20"]
@pytest.mark.parametrize(
"selector",
[
["--spec-draft-hf", "org/repo"],
["-hfd", "org/repo"],
["-hfrd", "org/repo"],
["--hf-repo-draft", "org/repo"],
["--spec-draft-hf=org/repo"],
],
)
def test_strip_shadowing_flags_drops_hf_drafter_selectors_with_spec(selector):
# HF drafter selectors must reset on inherit like local --model-draft, or a
# stale inherited HF drafter last-wins over Studio's re-derived spec choice.
out = strip_shadowing_flags(
selector + ["--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = True,
strip_template = False,
)
assert out == ["--top-k", "20"]
def test_strip_shadowing_flags_keeps_draft_tuning_with_spec():
# Per-drafter tuning knobs are deliberately preserved: the VRAM budget reads
# them via the same parsers the child honors (so they stay consistent on
# inherit), and stripping --spec-draft-ngl would move a CPU drafter to GPU.
keep = [
"--spec-draft-type-k",
"q4_0",
"--spec-draft-type-v",
"q4_0",
"--spec-draft-ngl",
"0",
"--spec-draft-device",
"cpu",
]
out = strip_shadowing_flags(
list(keep),
strip_context = False,
strip_cache = False,
strip_spec = True,
strip_template = False,
)
assert out == keep
def test_strip_shadowing_flags_keeps_split_mode_when_not_requested():
# No tensor_parallel field supplied on the Apply -> an inherited
# --split-mode survives (mirrors the chat-template keep behavior).

128
studio/backend/tests/test_llama_stats.py Normal file
View file

@ -0,0 +1,128 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the llama-server /metrics -> engine_stats translator: generation
throughput comes from generated-token metrics (not llama_decode() calls), and
the unexposed kv_cache_usage_ratio is never fabricated into the log line."""
from core.inference.llama_stats import LlamaServerStatsLogger
class _Capture:
def __init__(self):
self.events = []
def info(self, event, **kw):
self.events.append((event, dict(kw)))
def debug(self, *a, **k):
pass
def _drive(snaps):
"""Run _run() synchronously over `snaps`, then stop deterministically."""
cap = _Capture()
lg = LlamaServerStatsLogger("http://127.0.0.1:0", cap)
lg._interval = 0.001 # bypass the 1s floor for a fast, synchronous run
state = {"i": 0}
def fake_scrape():
i = state["i"]
state["i"] += 1
if i >= len(snaps):
lg.stop()
return None
return snaps[i]
lg._scrape = fake_scrape
lg._run()
return [kw for ev, kw in cap.events if ev == "engine_stats"]
def test_gen_tok_s_uses_token_metrics_not_decode_calls():
# tokens_predicted_total jumps 95 while n_decode_total only moves 9; the
# gauge reports 95 tok/s. Decode-call rate (9) must not be reported.
snaps = [
{
"tokens_predicted_total": 0.0,
"prompt_tokens_total": 0.0,
"n_decode_total": 0.0,
"predicted_tokens_seconds": 95.0,
"prompt_tokens_seconds": 30.0,
"requests_processing": 1.0,
},
{
"tokens_predicted_total": 95.0,
"prompt_tokens_total": 30.0,
"n_decode_total": 9.0,
"predicted_tokens_seconds": 95.0,
"prompt_tokens_seconds": 30.0,
"requests_processing": 1.0,
},
]
stats = _drive(snaps)
assert stats, "expected engine_stats while a request is processing"
assert all(s["gen_tok_s"] == 95.0 for s in stats)
assert all(s["prompt_tok_s"] == 30.0 for s in stats)
def test_kv_cache_pct_not_emitted_when_metric_absent():
# llama.cpp does not expose kv_cache_usage_ratio, so it must not appear.
snaps = [
{
"tokens_predicted_total": 0.0,
"prompt_tokens_total": 0.0,
"predicted_tokens_seconds": 10.0,
"requests_processing": 1.0,
},
{
"tokens_predicted_total": 10.0,
"prompt_tokens_total": 5.0,
"predicted_tokens_seconds": 10.0,
"requests_processing": 1.0,
},
]
stats = _drive(snaps)
assert stats
assert all("kv_cache_pct" not in s for s in stats)
def test_scrape_parses_labelled_and_bare_metrics(monkeypatch):
# Prometheus samples may carry labels; both labelled and bare lines parse.
import core.inference.llama_stats as ls
body = (
'llamacpp:tokens_predicted_total{model="m"} 20\n'
'llamacpp:prompt_tokens_total{model="m"} 5\n'
"llamacpp:requests_processing 1\n"
"# HELP llamacpp:ignored ignored\n"
)
class _Resp:
status = 200
def read(self):
return body.encode()
def __enter__(self):
return self
def __exit__(self, *a):
return False
monkeypatch.setattr(ls.urllib.request, "urlopen", lambda *a, **k: _Resp())
m = ls.LlamaServerStatsLogger("http://127.0.0.1:0", _Capture())._scrape()
assert m["tokens_predicted_total"] == 20.0
assert m["prompt_tokens_total"] == 5.0
assert m["requests_processing"] == 1.0
def test_counter_delta_fallback_without_gauges():
# Older binaries expose only the counters; throughput falls back to deltas.
snaps = [
{"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0},
{"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0},
]
stats = _drive(snaps)
# running=1 keeps it emitting; gen_tok_s falls back to the (here zero) delta.
assert stats and all(s["gen_tok_s"] >= 0.0 for s in stats)

244
studio/backend/tests/test_logging_middleware.py Normal file
View file

@ -0,0 +1,244 @@
# 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 pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from starlette.staticfiles import StaticFiles
from loggers import handlers as hmod
from loggers.handlers import LoggingMiddleware
class _LogCapture:
def __init__(self):
self.events = []
def info(self, event, **kw):
self.events.append(("info", event, kw))
def error(self, event, **kw):
self.events.append(("error", event, kw))
@pytest.fixture
def logs(monkeypatch):
capture = _LogCapture()
monkeypatch.setattr(hmod, "logger", capture)
return capture
def _http_scope(path, method = "GET"):
return {"type": "http", "path": path, "method": method}
async def _noop_receive():
return {"type": "http.disconnect"}
def _run(coro):
return asyncio.run(coro)
def test_success_logs_status_and_forwards_chunks(logs):
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 206, "headers": []})
await send({"type": "http.response.body", "body": b"a", "more_body": True})
await send({"type": "http.response.body", "body": b"", "more_body": False})
seen = []
async def send(message):
seen.append(message)
_run(LoggingMiddleware(app)(_http_scope("/api/health"), _noop_receive, send))
assert [m["type"] for m in seen] == [
"http.response.start",
"http.response.body",
"http.response.body",
]
assert logs.events[0][1] == "request_completed"
assert logs.events[0][2]["status_code"] == 206
def test_excluded_asset_success_skips_log(logs):
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
async def send(message):
pass
for path in ("/assets/index.css", "/huggingface.svg", "/font.woff2"):
_run(LoggingMiddleware(app)(_http_scope(path), _noop_receive, send))
assert logs.events == []
def test_exception_logs_real_status_and_reraises(logs):
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 418, "headers": []})
raise RuntimeError("stream failed")
async def send(message):
pass
with pytest.raises(RuntimeError, match = "stream failed"):
_run(LoggingMiddleware(app)(_http_scope("/api/health"), _noop_receive, send))
assert logs.events[0][1] == "request_failed"
assert logs.events[0][2]["status_code"] == 418
assert logs.events[0][2]["error"] == "stream failed"
assert "process_time_ms" in logs.events[0][2]
def test_cancelled_error_propagates_without_error_log(logs):
async def app(scope, receive, send):
raise asyncio.CancelledError()
async def send(message):
pass
with pytest.raises(asyncio.CancelledError):
_run(LoggingMiddleware(app)(_http_scope("/api/health"), _noop_receive, send))
assert logs.events == []
def test_non_http_scope_passes_through(logs):
seen = []
async def app(scope, receive, send):
seen.append(scope["type"])
async def send(message):
pass
_run(LoggingMiddleware(app)({"type": "websocket", "path": "/ws"}, _noop_receive, send))
assert seen == ["websocket"]
assert logs.events == []
def test_duplicate_get_within_window_deduped(logs, monkeypatch):
monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000)
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
async def send(message):
pass
mw = LoggingMiddleware(app)
for _ in range(3):
_run(mw(_http_scope("/api/chat/projects"), _noop_receive, send))
# Only the first of the identical GET/200 burst is logged.
assert len(logs.events) == 1
assert logs.events[0][1] == "request_completed"
def test_mutations_and_errors_are_never_deduped(logs, monkeypatch):
monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000)
async def post_ok(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
async def get_404(scope, receive, send):
await send({"type": "http.response.start", "status": 404, "headers": []})
await send({"type": "http.response.body", "body": b""})
async def send(message):
pass
mw = LoggingMiddleware(post_ok)
for _ in range(2):
_run(mw(_http_scope("/api/chat/threads", method = "POST"), _noop_receive, send))
mw_404 = LoggingMiddleware(get_404)
for _ in range(2):
_run(mw_404(_http_scope("/api/models"), _noop_receive, send))
# 2 mutations + 2 errors all logged (dedup only touches GET/2xx).
assert len(logs.events) == 4
def test_quiet_poll_paths_use_longer_heartbeat_window(logs, monkeypatch):
# Burst dedup off, quiet-poll heartbeat on: only liveness paths collapse.
monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 0)
monkeypatch.setattr(hmod, "_QUIET_POLL_DEDUP_MS", 1000)
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
async def send(message):
pass
mw = LoggingMiddleware(app)
for _ in range(3):
_run(mw(_http_scope("/api/inference/monitor"), _noop_receive, send)) # quiet
for _ in range(3):
_run(mw(_http_scope("/api/chat/projects"), _noop_receive, send)) # normal
paths = [e[2]["path"] for e in logs.events]
assert paths.count("/api/inference/monitor") == 1 # collapsed to one heartbeat
assert paths.count("/api/chat/projects") == 3 # base dedup off -> all logged
def test_distinct_query_strings_are_not_deduped(logs, monkeypatch):
monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000)
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
async def send(message):
pass
def scope(query):
return {
"type": "http",
"path": "/api/models/browse-folders",
"method": "GET",
"query_string": query,
}
mw = LoggingMiddleware(app)
_run(mw(scope(b"path=/tmp/a"), _noop_receive, send))
_run(mw(scope(b"path=/tmp/b"), _noop_receive, send)) # distinct query -> logs
_run(mw(scope(b"path=/tmp/a"), _noop_receive, send)) # repeat of first -> deduped
# Two distinct query strings log; the immediate repeat of the first does not.
assert len(logs.events) == 2
def test_fastapi_static_asset_success_skips_log(tmp_path, logs):
assets_dir = tmp_path / "assets"
assets_dir.mkdir()
(assets_dir / "app.css").write_text("body { color: black; }", encoding = "utf-8")
app = FastAPI()
app.add_middleware(LoggingMiddleware)
@app.get("/api/health")
async def health():
return {"ok": True}
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
client = TestClient(app)
response = client.get("/api/health")
assert response.status_code == 200
assert logs.events[0][1] == "request_completed"
assert logs.events[0][2]["path"] == "/api/health"
log_count = len(logs.events)
response = client.get("/assets/app.css")
assert response.status_code == 200
assert response.text == "body { color: black; }"
assert len(logs.events) == log_count

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