diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml index bb7dcbf8e4..45ce231743 100644 --- a/.github/workflows/cross-platform-parity-ci.yml +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Runs installer parity and autostart opt-out tests on Windows and macOS. +# Runs installer parity and autostart opt-out tests across all three platforms. # -# Why: that test is the guard that install.sh and install.ps1 stay in -# sync, but today it only runs on ubuntu-latest (auto-discovered by -# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both -# installer scripts, and on Windows Path.read_text() defaults to the -# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already -# contains a U+274C) raises UnicodeDecodeError there even though Linux and -# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in -# #6166; this job keeps that from silently regressing by exercising the -# test on the platforms it claims parity for. Pure pytest, no GPU, -# sub-second, so the matrix is cheap. +# Why: the parity test guards that install.sh and install.ps1 stay in sync. +# It originally ran only on ubuntu-latest through studio-backend-ci.yml. +# On Windows, Path.read_text() defaults to the cp1252 locale encoding, so a +# non-cp1252 byte in install.sh raises UnicodeDecodeError even though Linux +# and macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this matrix keeps that from silently regressing. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. Linux also runs the POSIX rollback test +# under dash, matching the supported curl-to-sh installer path. name: Cross-platform parity @@ -23,6 +21,8 @@ on: - 'install.ps1' - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' + - 'tests/sh/test_install_rollback_lifecycle.sh' + - 'tests/studio/test_install_rollback_lifecycle.ps1' - '.github/workflows/cross-platform-parity-ci.yml' push: branches: [main] @@ -31,6 +31,8 @@ on: - 'install.ps1' - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' + - 'tests/sh/test_install_rollback_lifecycle.sh' + - 'tests/studio/test_install_rollback_lifecycle.ps1' - '.github/workflows/cross-platform-parity-ci.yml' workflow_dispatch: @@ -47,7 +49,7 @@ jobs: strategy: fail-fast: false matrix: - os: [windows-latest, macos-latest] + os: [ubuntu-latest, windows-latest, macos-latest] runs-on: ${{ matrix.os }} timeout-minutes: 10 steps: @@ -67,3 +69,10 @@ jobs: tests/python/test_cross_platform_parity.py tests/test_installer_skip_autostart.py -q + - name: PowerShell rollback lifecycle tests + if: runner.os == 'Windows' + shell: pwsh + run: pwsh -NoProfile -File tests/studio/test_install_rollback_lifecycle.ps1 + - name: POSIX rollback lifecycle tests + if: runner.os == 'Linux' + run: sh tests/sh/test_install_rollback_lifecycle.sh diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index b8f587b63e..3968f2e80a 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -193,6 +193,7 @@ jobs: --ignore=tests/sh \ --ignore=tests/studio/test_hardware_dispatch_matrix.py \ --ignore=tests/studio/test_is_mlx_dispatch_gate.py \ + --ignore=tests/studio/test_xpu_spoof_pipeline.py \ --ignore=tests/vllm_compat \ --ignore=tests/version_compat \ -m 'not server and not e2e' \ @@ -205,14 +206,15 @@ jobs: env: PYTHONPATH: ${{ github.workspace }}/studio UNSLOTH_COMPILE_DISABLE: '1' - # These two files mutate hardware.py module globals at runtime - # via the spoof fixtures, which leaks state into any other test - # that imports hardware. Run them in their own pytest invocation - # so the leak does not cross file boundaries. + # These files mutate hardware.py module globals at runtime via the + # spoof fixtures (CUDA/ROCm/XPU/MLX/CPU), which leaks state into any + # other test that imports hardware. Run them in their own pytest + # invocation so the leak does not cross file boundaries. run: | python -m pytest -q --tb=short \ tests/studio/test_hardware_dispatch_matrix.py \ - tests/studio/test_is_mlx_dispatch_gate.py + tests/studio/test_is_mlx_dispatch_gate.py \ + tests/studio/test_xpu_spoof_pipeline.py - name: Shell installer tests # Subset that does not depend on a writable / pristine install.sh @@ -228,6 +230,7 @@ jobs: tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_resolve_cuda_archs.sh \ + tests/sh/test_staged_validation_enabled.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ tests/sh/test_torch_flavor.sh \ diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 30280c281e..97eb07b2d8 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -231,12 +231,69 @@ jobs: mkdir -p logs/playwright_extra python tests/studio/playwright_extra_ui.py + - name: UI font size scaling regression (Playwright) + env: + BASE_URL: http://127.0.0.1:18894 + STUDIO_PW: ${{ env.STUDIO_EXTRA_NEW_PW }} + PW_ART_DIR: logs/playwright_fontscale + run: | + mkdir -p logs/playwright_fontscale + python tests/studio/playwright_ui_font_scale.py + - name: Stop second Unsloth if: always() run: | kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true sleep 2 + # Model-picker per-model-config regression (PR #7207 re-land of #6647). + # Fourth Unsloth on its own port; loads the tiny GGUF and drives the + # picker's run-settings surface: Context Length persists across a reload, + # Reset clears the stored override (never pins it), and the infra models + # (RAG embedder + llama.cpp probe) stay hidden from the picker. + - name: Reset auth + boot Unsloth for model-config tests (port 18898) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \ + > logs/studio_modelcfg.log 2>&1 & + echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18898 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then + jq -e '.status == "healthy"' /tmp/health4.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health4.json + + - name: Pass bootstrap pw for model-config test + run: | + NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$NEW" + echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive model-picker per-model-config with Playwright + env: + BASE_URL: http://127.0.0.1:18898 + STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }} + PW_ART_DIR: logs/playwright_modelcfg + STUDIO_UI_STRICT: '1' + GGUF_REPO: ${{ env.GGUF_REPO }} + GGUF_VARIANT: ${{ env.GGUF_VARIANT }} + STUDIO_MODEL_HINT: gemma-3-270m + run: | + mkdir -p logs/playwright_modelcfg + python tests/studio/playwright_model_config.py + + - name: Stop fourth Unsloth + if: always() + run: | + kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true + sleep 2 + # IME + multilingual paste regression (issue #5318 / PR #5327). # Third Unsloth on its own port so a hang here cannot poison the # earlier UI tests. No GGUF -- the bug surface is the composer. @@ -297,12 +354,15 @@ jobs: path: | logs/studio.log logs/studio_extra.log + logs/studio_modelcfg.log logs/studio_ime.log logs/install.log logs/server-logs/ logs/playwright logs/playwright-permissions-* logs/playwright_extra + logs/playwright_fontscale + logs/playwright_modelcfg logs/playwright_ime logs/studio-permissions-*.log retention-days: 7 diff --git a/README.md b/README.md index 6aa8f4f4c3..514454f985 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,13 @@ Replace `claude` with any supported agent: | OpenCode | `unsloth start opencode` | | Pi Coding Agent | `unsloth start pi` | +Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local +subagent: + +```bash +unsloth start claude --as-subagent --model unsloth/model-GGUF:quant +``` + ## ๐Ÿ“ฅ Install Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements. diff --git a/install.ps1 b/install.ps1 index f64021c629..0933edcdaf 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1295,13 +1295,82 @@ exit 0 $suffix++ $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix" } - Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop $script:StudioVenvRollbackDir = $candidate $script:StudioVenvRollbackTarget = $ExistingDir $script:StudioVenvRollbackActive = $true + # Publish the rollback state before the atomic rename so interruption + # cannot land after Move-Item but before cleanup knows where the old venv went. + try { + Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop + } catch { + # A collision or ordinary rename failure leaves the original in place. + # Keep state active only when the rename happened before interruption. + if (Test-Path -LiteralPath $ExistingDir) { + $script:StudioVenvRollbackActive = $false + $script:StudioVenvRollbackDir = $null + } + throw + } substep "previous environment preserved for rollback" } + function Remove-StudioVenvTreeWithRetry { + param( + [Parameter(Mandatory = $true)][string]$Path, + [Parameter(Mandatory = $true)][string]$Label + ) + $lastError = $null + for ($attempt = 1; $attempt -le 3; $attempt++) { + try { + Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop + } catch { + $lastError = $_.Exception.Message + } + if (-not (Test-Path -LiteralPath $Path)) { return $true } + if ($attempt -lt 3) { Start-Sleep -Milliseconds (250 * $attempt) } + } + Write-Host "[WARN] Could not remove $Label at $Path" -ForegroundColor Yellow + if ($lastError) { Write-Host " $lastError" -ForegroundColor Yellow } + return $false + } + + function Test-StudioVenvRollbackMustBePreserved { + param([Parameter(Mandatory = $true)][System.IO.FileSystemInfo]$Rollback) + # Preserve anything outside the installer's timestamp.PID[.suffix] format. + if ($Rollback.Name -notmatch '^unsloth_studio\.rollback\.[0-9]{14}\.([0-9]+)(?:\.[0-9]+)?$') { + return $true + } + $ownerPid = 0 + if (-not [int]::TryParse($Matches[1], [ref]$ownerPid)) { return $true } + if ($ownerPid -eq $PID) { return $true } + return $null -ne (Get-Process -Id $ownerPid -ErrorAction SilentlyContinue) + } + + function Remove-StaleStudioVenvRollbacks { + try { + $rollbacks = @( + Get-ChildItem -LiteralPath $StudioHome -Directory -Force -ErrorAction Stop | + Where-Object { $_.Name -like 'unsloth_studio.rollback.*' } + ) + } catch { + Write-Host "[WARN] Could not inspect stale environment rollbacks in $StudioHome" -ForegroundColor Yellow + Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow + return + } + foreach ($rollback in $rollbacks) { + if (($rollback.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -ne 0) { + Write-Host "[WARN] Refusing to remove rollback reparse point $($rollback.FullName)" -ForegroundColor Yellow + continue + } + # A concurrent installer may have moved its live venv aside. The PID + # in the generated name keeps this run from deleting its rescue copy. + if (Test-StudioVenvRollbackMustBePreserved -Rollback $rollback) { continue } + if (Remove-StudioVenvTreeWithRetry -Path $rollback.FullName -Label "stale environment rollback") { + substep "removed stale environment rollback $($rollback.Name)" + } + } + } + function Restore-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir @@ -1313,7 +1382,9 @@ exit 0 substep "restoring previous environment after failed install..." "Yellow" try { if (Test-Path -LiteralPath $target) { - Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue + if (-not (Remove-StudioVenvTreeWithRetry -Path $target -Label "incomplete environment")) { + throw "Could not remove incomplete environment at $target" + } } Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop substep "restored previous environment" @@ -1328,11 +1399,13 @@ exit 0 function Complete-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir - if ($backup -and (Test-Path -LiteralPath $backup)) { - Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue - } + # The replacement is committed. Disable restoration before deleting the + # backup so interruption cannot restore a partially deleted environment. $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null + if ($backup -and (Test-Path -LiteralPath $backup)) { + Remove-StudioVenvTreeWithRetry -Path $backup -Label "environment rollback" | Out-Null + } } # Raw torch.__version__ from $PythonExe's venv (last non-empty stdout line), or $null. @@ -1366,6 +1439,8 @@ exit 0 } catch { return $null } } + $studioVenvReplacementCommitted = $false + try { if (Test-Path -LiteralPath $VenvPython) { # env-mode: $StudioHome is a user-chosen workspace, so refuse to nuke an existing venv lacking Unsloth sentinels (-PathType Leaf rejects a directory at the sentinel path; accept the in-VENV ownership marker so partial-install retries aren't blocked). if ( @@ -1711,7 +1786,7 @@ exit 0 $nameArchTable = @( @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080) @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060) - @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) @@ -1880,16 +1955,18 @@ exit 0 # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with _strip_index_url_credentials (install.sh / py / setup.ps1). function Remove-IndexUrlCredentials { param([string]$Url) - $sep = $Url.IndexOf('://') + # Ordinal, not culture-aware: on non-English locales (e.g. th-TH) linguistic + # IndexOf treats "://" as ignorable, mis-locates it, and crashes Substring (issue #7279). + $sep = $Url.IndexOf('://', [System.StringComparison]::Ordinal) if ($sep -lt 0) { return $Url } $scheme = $Url.Substring(0, $sep) $rest = $Url.Substring($sep + 3) # Drop query / fragment (may hold auth tokens). $q = $rest.IndexOfAny([char[]]('?', '#')) if ($q -ge 0) { $rest = $rest.Substring(0, $q) } - $slash = $rest.IndexOf('/') + $slash = $rest.IndexOf('/', [System.StringComparison]::Ordinal) $authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest } - $at = $authority.LastIndexOf('@') + $at = $authority.LastIndexOf('@', [System.StringComparison]::Ordinal) $host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority } if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" } return "${scheme}://${host_}" @@ -2006,6 +2083,10 @@ exit 0 "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" + "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000) + "gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all" + "gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all" + "gfx1030" = "gfx103X-all" "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 } # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) hit a null-pointer bug in torch._C._grouped_mm on torch <2.11.0 (TheRock #5284 / #3284); force torch>=2.11.0 so pip skips the broken 2.10.0 wheels on the AMD index. The <2.12.0 ceiling (matches install_python_stack.py) blocks an unvalidated future 2.12.0+rocm wheel; bump both when 2.12.x is confirmed on gfx120X / Strix. @@ -2139,7 +2220,7 @@ exit 0 substep "upgrading unsloth in migrated environment..." if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo --no-deps, then runtime deps (typer, safetensors, transformers, etc.) --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins the matching pydantic-core (no-torch-runtime.txt below is --no-deps); all transitive deps are torch-free. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2151,7 +2232,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2258,7 +2339,7 @@ exit 0 substep "installing unsloth (this may take a few minutes)..." if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo --no-deps, then runtime deps (typer, safetensors, transformers, etc.) --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2273,11 +2354,11 @@ exit 0 # Freeze the installed torch trio so this with-deps resolve can't downgrade the pinned +cuXXX/+rocm build (twin of install.sh's _build_unsloth_torch_overrides). $script:TorchOverridesFile = New-UnslothTorchOverridesFile -PythonExe $VenvPython if ($script:TorchOverridesFile) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth --overrides $script:TorchOverridesFile "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth --overrides $script:TorchOverridesFile "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } Remove-Item -LiteralPath $script:TorchOverridesFile -Force -ErrorAction SilentlyContinue $script:TorchOverridesFile = $null } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } } } else { # Freeze the installed torch trio (see above) so the with-deps unsloth resolve can't strip the +cuXXX/+rocm suffix. @@ -2314,7 +2395,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --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) @@ -2340,6 +2421,13 @@ exit 0 } } + $installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim() + if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) { + step $PackageName "$installedPackageVersion installed" + } else { + substep "[WARN] installed $PackageName version could not be determined" "Yellow" + } + # โ”€โ”€ Enforce the installed torch flavor matches the detected GPU build. PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv keeps a stale torch+cpu against a CUDA index and setup.ps1 loops on "cpu != required cuXXX". Reinstall the right triplet when a GPU build is expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (a PEP 503 index uv resolves via --default-index). --no-torch / CPU-only hosts are no-ops. โ”€โ”€ if (-not $SkipTorch) { $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl @@ -2601,6 +2689,13 @@ exit 0 } Refresh-SessionPath # sync current session with registry Complete-StudioVenvRollback + $studioVenvReplacementCommitted = $true + Remove-StaleStudioVenvRollbacks + } finally { + if (-not $studioVenvReplacementCommitted) { + Restore-StudioVenvRollback + } + } # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy User PATH entry (Machine > User > current $env:Path) would win. if ($StudioRedirectMode -eq 'env' -and (Test-Path -LiteralPath $ShimExe)) { diff --git a/install.sh b/install.sh index 2c96ed360a..8808138386 100755 --- a/install.sh +++ b/install.sh @@ -444,14 +444,20 @@ _start_studio_venv_replacement() { _stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time") _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$" _suffix=0 - while [ -e "$_candidate" ]; do + while [ -e "$_candidate" ] || [ -L "$_candidate" ]; do _suffix=$((_suffix + 1)) _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix" done - mv "$_existing_dir" "$_candidate" _VENV_ROLLBACK_DIR="$_candidate" _VENV_ROLLBACK_TARGET="$_existing_dir" _VENV_ROLLBACK_ACTIVE=true + # Publish the rollback state before the atomic rename so a signal cannot + # land after mv but before the exit handlers know where the old venv went. + if ! mv "$_existing_dir" "$_candidate"; then + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + return 1 + fi substep "previous environment preserved for rollback" } @@ -472,13 +478,68 @@ _restore_studio_venv_replacement() { fi } -_commit_studio_venv_replacement() { - [ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0 - if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then - rm -rf "$_VENV_ROLLBACK_DIR" || true +_studio_venv_rollback_must_be_preserved() { + _rollback_name=${1##*/} + _rollback_metadata=${_rollback_name#unsloth_studio.rollback.} + _rollback_stamp=${_rollback_metadata%%.*} + _rollback_process=${_rollback_metadata#*.} + # Preserve anything outside the installer's timestamp.PID[.suffix] format. + [ "$_rollback_process" != "$_rollback_metadata" ] || return 0 + case "$_rollback_stamp" in + time) ;; + ''|*[!0-9]*) return 0 ;; + *) [ "${#_rollback_stamp}" -eq 14 ] || return 0 ;; + esac + _rollback_pid=${_rollback_process%%.*} + case "$_rollback_pid" in + ''|*[!0-9]*) return 0 ;; + esac + _rollback_suffix=${_rollback_process#*.} + if [ "$_rollback_suffix" != "$_rollback_process" ]; then + case "$_rollback_suffix" in ''|*[!0-9]*) return 0 ;; esac fi - _VENV_ROLLBACK_ACTIVE=false - _VENV_ROLLBACK_DIR="" + kill -0 "$_rollback_pid" 2>/dev/null +} + +_prune_stale_studio_venv_rollbacks() { + for _stale_rollback in "$STUDIO_HOME"/unsloth_studio.rollback.*; do + [ -d "$_stale_rollback" ] || continue + if [ -L "$_stale_rollback" ]; then + echo "โš ๏ธ Refusing to remove rollback symlink $_stale_rollback" >&2 + continue + fi + # A concurrent installer may have moved its live venv aside. The PID in + # the generated name keeps this successful run from deleting its rescue copy. + _studio_venv_rollback_must_be_preserved "$_stale_rollback" && continue + if rm -rf "$_stale_rollback"; then + substep "removed stale environment rollback ${_stale_rollback##*/}" + else + echo "โš ๏ธ Could not remove stale environment rollback $_stale_rollback" >&2 + fi + done +} + +_commit_studio_venv_replacement() { + if [ "$_VENV_ROLLBACK_ACTIVE" = true ]; then + _rollback_to_remove="$_VENV_ROLLBACK_DIR" + # The new environment is already committed. Clear the restore state + # before deletion so an interrupt cannot replace it with a half-deleted backup. + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + if [ -n "$_rollback_to_remove" ] && [ -d "$_rollback_to_remove" ]; then + if ! rm -rf "$_rollback_to_remove"; then + echo "โš ๏ธ Could not remove environment rollback $_rollback_to_remove" >&2 + fi + fi + fi + # Only prune older orphaned copies after the replacement has succeeded, so + # an interrupted install never discards the last known-good environment. + _prune_stale_studio_venv_rollbacks +} + +_cleanup_install_temporaries() { + [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true } _on_install_exit() { @@ -486,14 +547,28 @@ _on_install_exit() { if [ "$_status" -ne 0 ]; then _restore_studio_venv_replacement fi - [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true - [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true + _cleanup_install_temporaries exit "$_status" } -# Empty so an inherited value never reaches the trap's rm. + +_on_install_signal() { + _signal_status="$1" + # EXIT is disabled to avoid a second cleanup pass. Ignore further termination + # signals until the old environment is back in place. + trap - EXIT + trap '' HUP INT TERM + _restore_studio_venv_replacement + _cleanup_install_temporaries + exit "$_signal_status" +} +# Empty so an inherited value never reaches the trap's rm; only temp paths this +# script creates below (spaced-path dir, torch-trio overrides) are removed. _UV_OVERRIDE_TMPDIR="" _UNSLOTH_TORCH_OVERRIDES="" trap _on_install_exit EXIT +trap '_on_install_signal 129' HUP +trap '_on_install_signal 130' INT +trap '_on_install_signal 143' TERM # โ”€โ”€ Helper: download a URL to a file (supports curl and wget) โ”€โ”€ download() { @@ -519,6 +594,36 @@ _is_pkg_installed() { esac } +# โ”€โ”€ Helper: human-readable apt distro label for the sudo package prompt (#6207) โ”€โ”€ +# Reads /etc/os-release so the Accept? prompt can say which distro we detected and +# that packages come from that distro's official apt repos (not a tarball). +_apt_distro_description() { + # Plain ( ... ) subshell โ€” not $() โ€” so case/;; stays bash-3.2-safe on macOS. + # Bash 3.2 misparses case arms inside command substitution and errors on `;;`. + ( + if [ ! -r /etc/os-release ]; then + printf 'a debian-like system' + exit 0 + fi + # shellcheck disable=SC1091 + . /etc/os-release 2>/dev/null || true + if [ -n "${NAME:-}" ] && [ -n "${VERSION_ID:-}" ]; then + _ad_label="$NAME $VERSION_ID" + elif [ -n "${PRETTY_NAME:-}" ]; then + _ad_label="$PRETTY_NAME" + elif [ -n "${NAME:-}" ]; then + _ad_label="$NAME" + else + printf 'a debian-like system' + exit 0 + fi + case " ${ID:-} ${ID_LIKE:-} " in + *" debian "*|*" ubuntu "*) _ad_label="${_ad_label} (debian-like)" ;; + esac + printf '%s' "$_ad_label" + ) +} + # โ”€โ”€ Helper: install packages via apt, escalating to sudo only if needed โ”€โ”€ # Usage: _smart_apt_install pkg1 pkg2 pkg3 ... _smart_apt_install() { @@ -549,11 +654,14 @@ _smart_apt_install() { # Step 3: Escalate -- need elevated permissions for remaining packages if command -v sudo >/dev/null 2>&1; then + _ad_desc="$(_apt_distro_description)" echo "" echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo " WARNING: We require sudo elevated permissions to install:" echo " $_STILL_MISSING" - echo " If you accept, we'll run sudo now, and it'll prompt your password." + echo " Detected ${_ad_desc}." + echo " If you accept, we'll run sudo apt-get to install these packages" + echo " from your distro's official repositories (not a third-party tarball)." echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo "" printf " Accept? [Y/n] " @@ -1541,8 +1649,8 @@ _maybe_reroute_strixhalo_to_2404() { [ -e /dev/dxg ] || return 0 # A usable NVIDIA GPU means the CUDA path works here, so don't reroute for AMD. if _has_usable_nvidia_gpu; then return 0; fi - # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. - if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \ && ! _wsl_amd_gpu_name >/dev/null 2>&1; then return 0 fi @@ -1949,17 +2057,152 @@ _has_amd_rocm_gpu() { amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then return 0 elif [ -e /dev/kfd ] && \ - awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \ - gpu && amd { found=1 } END{ exit !found }' \ + awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \ /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then - # vendor_id 4098 = 0x1002 (AMD); the NVIDIA open kernel module can register KFD nodes too. + # vendor_id 4098 = 0x1002 (AMD) marks a GPU node: the KFD CPU node + # reports vendor_id 0, so any 4098 node is an AMD GPU. NVIDIA's open + # kernel module (driver 560+) registers KFD nodes as vendor_id 4318 + # (0x10DE), so this never false-positives on NVIDIA-only hosts. + # The prior check also required a gpu_id line, but gpu_id is a SIBLING + # sysfs file, not a line in properties -- it never matched, so the + # fallback silently missed every ROCm-less AMD host (issue: fresh + # Arch/CachyOS boxes reporting "no GPU detected"). return 0 fi return 1 } -# โ”€โ”€ Detect GPU and choose PyTorch index URL (mirrors Get-TorchIndexUrl) โ”€โ”€ -# CPU-only machines get the cpu index, avoiding the --torch-backend=auto solver dead-end. +# Returns 0 if an AMD display GPU is on the PCI bus even when ROCm can't use it +# (e.g. a Strix Halo iGPU with no /dev/kfd). Only sharpens the "no GPU detected" +# hint. vendor 0x1002 = AMD/ATI; class 0x03* = display controller. +_amd_gpu_present_via_pci() { + [ -d /sys/bus/pci/devices ] || return 1 + for _pci_vendor in /sys/bus/pci/devices/*/vendor; do + [ -r "$_pci_vendor" ] || continue + read -r _v < "$_pci_vendor" 2>/dev/null || continue + [ "$_v" = "0x1002" ] || continue + _cls="${_pci_vendor%vendor}class" + [ -r "$_cls" ] || continue + read -r _c < "$_cls" 2>/dev/null || continue + case "$_c" in 0x03*) return 0 ;; esac + done + return 1 +} + +# Map a gfx arch to the AMD pip index family (mirrors install.ps1 $archFamilyMap). +_amd_arch_index_family_for_gfx() { + case "$1" in + gfx1201|gfx1200) echo gfx120X-all ;; + gfx1151) echo gfx1151 ;; + gfx1150) echo gfx1150 ;; + gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;; + gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;; + gfx90a) echo gfx90a ;; + gfx908) echo gfx908 ;; + *) return 1 ;; + esac +} + +# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable). +_infer_amd_gfx_arch_from_gpu_name() { + case "$1" in + *"9070 XT"*|*9080*) echo gfx1201 ;; + *9070*|*9060*) echo gfx1200 ;; + *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;; + *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;; + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;; + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;; + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;; + *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;; + *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;; + *"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) echo gfx1034 ;; + *) return 1 ;; + esac +} + +# Best-effort gfx inference when ROCm tools can't see the GPU (unslothai#7301). +# Mirrors install.ps1 arch resolution on Windows ($HasROCm false, $ROCmGfxArch set). +_infer_linux_amd_gfx_arch() { + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then + printf '%s\n' "$(printf '%s' "$UNSLOTH_ROCM_GFX_ARCH" | tr '[:upper:]' '[:lower:]')" + return 0 + fi + # On WSL /proc/cpuinfo and lspci still report the host APU, but without the + # ROCDXG bridge (librocdxg over /dev/dxg) the AMD wheels can't reach the GPU; + # keep the CPU fallback there unless that runtime is present (the explicit + # override above still wins). Mirrors install_python_stack.py. + _gpu_evidence="" + if [ -e /dev/dxg ] || grep -qi microsoft /proc/version 2>/dev/null; then + for _d in /opt/rocm/lib /opt/rocm/lib64 /opt/rocm-*/lib /opt/rocm-*/lib64; do + { [ -e "$_d/librocdxg.so" ] || [ -e "$_d/librocdxg.so.1" ]; } && _rocdxg=1 && break + done + [ -n "${_rocdxg:-}" ] || return 1 + # WSL enumerates no PCI display device; /dev/dxg + librocdxg IS the + # GPU evidence there. + _gpu_evidence=1 + elif _amd_gpu_present_via_pci; then + _gpu_evidence=1 + fi + # /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received + # no AMD GPU, so the CPU-model text alone is not GPU evidence: require an + # AMD display device (PCI vendor 0x1002, class 0x03*) before trusting it. + # The lspci fallback below needs no gate; an AMD display line IS evidence. + if [ -n "$_gpu_evidence" ] && grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null; then + echo gfx1151 + return 0 + fi + if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then + echo gfx1150 + return 0 + fi + if command -v lspci >/dev/null 2>&1; then + # A non-AMD controller can enumerate first (Intel/ASPEED before an AMD + # dGPU), so scan every display-class line and take the first AMD one + # that maps. The vendor guard is case-SENSITIVE (a -i "ATI" would match + # "CorporATIon" on every Intel/NVIDIA line); whole-line matching also + # survives the 0000: PCI domain prefix. Mirrors install_python_stack.py. + _amd_disp=$(lspci -nn 2>/dev/null | grep -E 'VGA compatible controller|3D controller|Display controller' | grep -E 'AMD|ATI' || true) + while IFS= read -r _ln; do + [ -n "$_ln" ] || continue + if _gfx=$(_infer_amd_gfx_arch_from_gpu_name "$_ln"); then + echo "$_gfx" + return 0 + fi + done </dev/null 2>&1; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_pg" ] && command -v amd-smi >/dev/null 2>&1; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + if [ -z "$_pg" ]; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + fi + printf '%s\n' "$_pg" +} + +# โ”€โ”€ Detect GPU and choose PyTorch index URL โ”€โ”€ +# Mirrors Get-TorchIndexUrl in install.ps1. +# On CPU-only machines this returns the cpu index, avoiding the solver +# dead-end where --torch-backend=auto resolves to unsloth==2024.8. get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" @@ -2000,6 +2243,29 @@ get_torch_index_url() { if ! _has_amd_rocm_gpu; then echo "$_base/cpu"; return fi + # A generic rocm index is only safe when the gfx arch is readable: the + # Strix reroute (gfx1150/1151 -> arch-specific index) learns gfx from + # rocminfo/amd-smi, so if those are missing OR do not enumerate the GPU, an + # unknown-arch box might be Strix and would get the broken _grouped_mm + # wheels. Probe via the shared helper (override first, then rocminfo/amd-smi + # with visibility masks cleared); if the arch is unreadable, never guess a + # rocm index. A KFD-only host whose arch is still inferable from hardware + # IDs (PCI/cpuinfo/lspci) returns the cpu index and lets the runtime-less + # reroute below upgrade it to AMD per-arch wheels -- the reroute gate uses + # this same probe, so the handoff can't misfire. Only when inference fails + # too is CPU final, with the actionable warning. + _amd_gfx_probe=$(_probe_amd_gfx_arch) + if [ -z "$_amd_gfx_probe" ]; then + if _amd_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null) && \ + [ -n "$_amd_inferred_gfx" ] && \ + _amd_arch_index_family_for_gfx "$_amd_inferred_gfx" >/dev/null 2>&1; then + echo "[WARN] AMD GPU detected but rocminfo/amd-smi can't read its gfx arch -- inferring $_amd_inferred_gfx from hardware IDs." >&2 + echo "$_base/cpu"; return + fi + echo "[WARN] AMD GPU detected but its gfx arch can't be read (rocminfo/amd-smi missing or not enumerating the GPU) -- installing CPU-only PyTorch." >&2 + echo "[WARN] For GPU PyTorch, install or repair rocminfo/amd-smi (e.g. sudo pacman -S rocm-hip-sdk) and re-run this installer." >&2 + echo "$_base/cpu"; return + fi # AMD GPU confirmed -- detect ROCm version _rocm_tag="" _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ @@ -2016,7 +2282,11 @@ get_torch_index_url() { { command -v rpm >/dev/null 2>&1 && \ ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \ [ -n "$ver" ] && \ - printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null + printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null || _rocm_tag="" + # ^ || guard: when EVERY version source is missing (e.g. rocminfo present + # but rocm-core not installed, so dpkg-query/rpm exit 1), the whole || + # chain fails and set -e would kill the installer BEFORE the actionable + # no-version WARN below -- exactly the fresh-install case it exists for. # Validate _rocm_tag: must match "rocmX.Y" with major >= 1 case "$_rocm_tag" in rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1) @@ -2049,10 +2319,27 @@ get_torch_index_url() { esac return fi - # ROCm version unreadable from any source: warn rather than silently install CPU PyTorch. - echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2 - echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2 - echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + # AMD GPU confirmed (rocminfo/amd-smi or the KFD topology fallback) but + # no ROCm/HIP install was found to read the version from (amd-smi, + # /opt/rocm/.info/version, hipconfig, dpkg, rpm). This is the common + # fresh-install case: the GPU is real, but with no ROCm userspace the + # correct PyTorch build can't be selected. Warn with an actionable fix + # rather than silently installing CPU PyTorch. + # A user-set UNSLOTH_ROCM_GFX_ARCH seeded the probe above, so rocminfo/ + # amd-smi may still be unable to see the GPU; when the named arch maps to + # a wheel family, the runtime-less reroute (gated on the override) will + # install the AMD per-arch wheels -- a CPU-only warning here would be + # false for that path. Defer like the inferable-arch branch does. + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] && \ + _amd_arch_index_family_for_gfx "$_amd_gfx_probe" >/dev/null 2>&1; then + echo "[WARN] AMD GPU detected with no readable ROCm version, but UNSLOTH_ROCM_GFX_ARCH=$_amd_gfx_probe is set -- routing to AMD per-arch wheels." >&2 + echo "$_base/cpu"; return + fi + echo "[WARN] AMD GPU detected, but no ROCm/HIP install was found to select the matching GPU PyTorch build -- falling back to CPU-only PyTorch." >&2 + echo "[WARN] Install the ROCm/HIP SDK, then re-run this installer:" >&2 + echo "[WARN] Arch / CachyOS : sudo pacman -S rocm-hip-sdk" >&2 + echo "[WARN] other distros : https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "[WARN] Minimum required for version detection: amd-smi, hipconfig, /opt/rocm/.info/version, or the rocm-core package." >&2 echo "$_base/cpu"; return fi # Parse CUDA version from nvidia-smi (POSIX-safe): accept both "CUDA Version:" and the newer "CUDA UMD Version:". Bounded, C locale. @@ -2387,8 +2674,9 @@ _maybe_bootstrap_rocm_wsl() { fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). [ -e /dev/dxg ] || return 0 - # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also ask the Windows host. - if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also + # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \ && ! _wsl_amd_gpu_name >/dev/null 2>&1; then return 0 fi @@ -2463,7 +2751,84 @@ fi TORCH_INDEX_URL=$(get_torch_index_url) -# Export UNSLOTH_TORCH_BACKEND ("cuda"/"rocm"/"cpu") for downstream scripts; classify on the FINAL lowercased leaf only so a custom mirror path containing "rocm"/"gfx" can't mislabel a cu*/cpu index. +# Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo +# in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's +# per-arch wheels like install.ps1 does on Windows (unslothai#7301). +# Gated on the runtime probes NOT naming a gfx: either no AMD GPU is detected at +# all (_has_amd_rocm_gpu false), or the GPU is visible only through the +# env-independent KFD topology while rocminfo/amd-smi can't read its arch +# (KFD-only host, unslothai#7314 -- before the KFD detection fix these hosts +# reached this reroute via the false branch, so the empty-probe condition +# preserves that routing). A */cpu index chosen WITH a readable gfx +# (unsupported/unreadable ROCm version, after its own warning) is a deliberate +# fallback -- rerouting it would contradict that decision, and stays excluded +# because the shared probe returns its gfx. An explicit UNSLOTH_ROCM_GFX_ARCH +# override stays authoritative either way. +if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \ + ! _has_usable_nvidia_gpu && \ + { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu || \ + [ -z "$(_probe_amd_gfx_arch)" ]; } && \ + case "$(uname -s)" in Linux) true ;; *) false ;; esac && \ + case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then + # ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other + # arches, so an inferred/overridden gfx must not reroute arm64 to AMD wheels. + case "$TORCH_INDEX_URL" in + */cpu) + _linux_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null || true) + if [ -n "$_linux_inferred_gfx" ]; then + _amd_family=$(_amd_arch_index_family_for_gfx "$_linux_inferred_gfx") || _amd_family="" + if [ -n "$_amd_family" ]; then + _amd_mirror="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}" + while [ "${_amd_mirror%/}" != "$_amd_mirror" ]; do + _amd_mirror="${_amd_mirror%/}" + done + TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/" + # Hand the inferred arch to setup.sh (llama.cpp): it re-probes + # ROCm on its own, and on these runtime-less hosts its probes + # find nothing, so without this it classifies the box as + # non-ROCm and installs the CPU prebuilt while torch just got + # AMD per-arch wheels. setup.sh and install_llama_prebuilt.py + # both honor UNSLOTH_ROCM_GFX_ARCH, so exporting it is the + # whole handoff (a user-set override re-exports unchanged). + export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx" + case "$_linux_inferred_gfx" in + gfx1201|gfx1200|gfx1151|gfx1150) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + esac + echo "" >&2 + # KFD-only hosts reach this reroute with /dev/kfd present + # (that's what detected them), so don't claim it's missing. + if _has_amd_rocm_gpu; then + echo " [WARN] AMD GPU visible via the kernel driver (KFD) but rocminfo/amd-smi can't read its gfx arch; using $_linux_inferred_gfx." >&2 + else + echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2 + fi + echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2 + echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2 + echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2 + echo " [WARN] Tip: set UNSLOTH_ROCM_GFX_ARCH=$_linux_inferred_gfx to skip inference next time." >&2 + echo "" >&2 + fi + fi + ;; + esac +fi + +# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that +# downstream scripts (setup.sh -> install_python_stack.py) know what was +# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts. +# Classify on the FINAL path segment only: a custom UNSLOTH_PYTORCH_MIRROR +# whose base path happens to contain "rocm" or "gfx" must not mislabel a +# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix +# overrides in gfxNNNN/, so the trailing slash is stripped first). +# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD +# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror +# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so +# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in +# lockstep with the shared _torch_index_url_leaf extractor). _torch_index_leaf="${TORCH_INDEX_URL%%\?*}" _torch_index_leaf="${_torch_index_leaf%%#*}" # Strip ALL trailing slashes: .../cu128// must yield cu128, not empty. @@ -2508,20 +2873,64 @@ case "$TORCH_INDEX_URL" in fi ;; esac -# โ”€โ”€ Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo โ”€โ”€โ”€โ”€โ”€โ”€โ”€ -# gfx1151/gfx1150 segfault in torch._grouped_mm on ROCm 7.1. -case "$TORCH_INDEX_URL" in - */rocm7.1|*/rocm7.1.*) - # Index gfx tokens by HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + selected dGPU box isn't rerouted. - _gfx_all="" - if command -v rocminfo >/dev/null 2>&1; then - _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') +# 0 when a rocmX.Y index leaf ($1, the final path segment) is older than floor +# $2.$3 (int compare, so rocm7.2 < rocm7.13). Non-rocm leaves (gfx*, cu*, cpu) and +# non-numeric versions return 1. Leaf-based (like $_torch_index_leaf) so a mirror +# base holding its own rocm token compares the family leaf, not the base path. +_rocm_leaf_below() { + case "$1" in rocm[0-9]*.[0-9]*) : ;; *) return 1 ;; esac + _rb=${1#rocm}; _maj=${_rb%%.*}; _min=${_rb#*.}; _min=${_min%%.*} + case "$_maj$_min" in *[!0-9]*) return 1 ;; esac + if [ "$_maj" -lt "$2" ]; then return 0; fi + if [ "$_maj" -eq "$2" ] && [ "$_min" -lt "$3" ]; then return 0; fi + return 1 +} +# โ”€โ”€ Strix Halo / Strix Point: route to the AMD arch-specific index โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# gfx1151/gfx1150 need torch 2.11+rocm7.13 from repo.amd.com/rocm/whl/gfx/, +# which carries AMD's real fixes (the rocm7.1 _grouped_mm segfault, moe_utils.py:167, +# and later Strix kernel bugs). Every generic pytorch.org index below rocm7.13 lacks +# them (and the Radeon repo can be offline, unslothai#7264), so reroute a detected +# Strix GPU whenever the picked index is older than the arch build -- covers today's +# rocm6.0-7.2 and any future 7.x < 7.13; rocm7.13+ already has the fixes, so leave it. +case "$_torch_index_leaf" in + rocm[0-9]*) + # Collect every gfx token in rocminfo / amd-smi enumeration order + # (skip duplicates), then index by HIP_VISIBLE_DEVICES / + # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box + # where the user selected the dGPU does NOT get rerouted to the + # Strix per-gfx index. + # || true on each probe: no gfx match makes grep exit 1, which under + # set -euo pipefail would abort the installer before the next fallback + # runs (now that the case matches every rocm* index, not just rocm7.1). + # A user-supplied UNSLOTH_ROCM_GFX_ARCH overrides probing (mirrors setup.sh + # and the display block), so a Strix override still reaches the arch index. + _gfx_all=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]') + if [ -z "$_gfx_all" ] && command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) fi if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then - _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') - # Mirror the PowerShell `amd-smi static --asic` probe as a fallback. + _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + # PowerShell paths also probe `amd-smi static --asic`; mirror it + # so a host with hipinfo-less amd-smi reports the gfx target. if [ -z "$_gfx_all" ]; then - _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + fi + # get_torch_index_url reads the arch with ROCR/HIP masks cleared, so a + # mask hiding every agent (e.g. ROCR_VISIBLE_DEVICES=-1) still lands + # here on a generic rocm index; re-probe unmasked or a masked-out Strix + # box keeps the broken generic wheels. Partial masks never get here + # (they enumerate at least one agent above) and keep their selection. + # ${VAR+x} (not :-): a SET-but-empty mask also hides every agent and + # must trigger the re-probe too. + if [ -z "$_gfx_all" ] && [ -n "${ROCR_VISIBLE_DEVICES+x}${HIP_VISIBLE_DEVICES+x}" ]; then + if command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_gfx_all" ] && \ + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) fi fi _runtime_gfx="" @@ -2546,13 +2955,14 @@ case "$TORCH_INDEX_URL" in case "$_runtime_gfx" in gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; esac - if [ -n "$_strix_gfx" ]; then + # Skip rocm7.13+ generic indexes: they already ship the fixes, so the + # arch build (rocm7.13) would be a downgrade rather than a rescue. + if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then echo "" >&2 - echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2 - echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2 - echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2 - echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2 - echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo " [WARN] $_strix_gfx (Strix) detected -- routing to the AMD arch-specific index" >&2 + echo " [WARN] torch 2.11+rocm7.13 has AMD's real gfx1150/gfx1151 fixes (the ROCm 7.1" >&2 + echo " [WARN] _grouped_mm segfault, moe_utils.py:167, and later Strix kernel bugs)," >&2 + echo " [WARN] and is more reliable than the rocm7.2 index or an offline Radeon repo." >&2 echo "" >&2 # AMD's arch-specific index has the real _grouped_mm fix (torch 2.11.0+rocm7.13.0); UNSLOTH_AMD_ROCM_MIRROR overrides for air-gapped installs. _amd_strix_base="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}" @@ -2629,7 +3039,7 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then case "$_gpu_disp_mkt" in *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 - *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) @@ -2664,6 +3074,17 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then # Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only. step "gpu" "Apple Silicon (Metal, unified memory)" +elif _has_amd_rocm_gpu; then + if [ "$_torch_index_pinned" = true ]; then + # An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY pin skipped all probing; + # do not claim ROCm is unusable when a CPU/other index was requested. + step "gpu" "AMD GPU (torch index pinned: $_torch_index_leaf)" "$C_WARN" + else + # AMD GPU visible to the kernel but the torch index stayed CPU: no usable + # ROCm userspace to pick a wheel. "none" would repeat the false diagnosis + # this installer used to give. + step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)" "$C_WARN" + fi else step "gpu" "none (CPU-only)" "$C_WARN" fi @@ -2672,9 +3093,20 @@ fi case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then - substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" - if [ "$OS" = "wsl" ]; then - # WSL + no GPU: often an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet. + if [ "$_torch_index_pinned" = true ]; then + # An explicit CPU pin is a request, not a detection failure: + # skip the SDK guidance (ROCm may be perfectly healthy here). + substep "CPU-only PyTorch (index pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY)." + elif _has_amd_rocm_gpu; then + substep "AMD GPU detected, but no usable ROCm/HIP install -- installing CPU-only PyTorch." "$C_WARN" + substep "Install the ROCm/HIP SDK and re-run this installer for GPU PyTorch." "$C_WARN" + else + substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" + fi + if [ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]; then + # WSL + no GPU detected (detection above found nothing). Common + # cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet -- + # /dev/dxg present (graphics) but no ROCm runtime. _wsl_ubu_ver="" [ -r /etc/os-release ] && _wsl_ubu_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}") if [ -e /dev/dxg ]; then @@ -2698,6 +3130,13 @@ case "$TORCH_INDEX_URL" in substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself." else substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" + # Only when ROCm truly can't see the GPU: a detected-but-too-old + # ROCm (rocminfo works, wheels need 6.0+) has its own guidance. + if ! _has_amd_rocm_gpu && _amd_gpu_present_via_pci; then + substep "An AMD GPU is on the PCI bus but ROCm cannot see it (no /dev/kfd," "$C_WARN" + substep " rocminfo, or amd-smi). Install the ROCm kernel stack so /dev/kfd exists;" + substep " Strix Halo (gfx1151/gfx1150) needs a recent kernel (6.11+) and ROCm 7.x." + fi fi substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):" substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" @@ -2747,7 +3186,7 @@ if [ "$_MIGRATED" = true ]; then # No-torch: --no-deps installs (PyPI metadata still hard-deps torch), then torch-free runtime deps --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.7.4" "unsloth-zoo>=2026.7.4" + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" # Resolve pydantic WITH deps so its pydantic-core matches (all torch-free). run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2761,7 +3200,7 @@ if [ "$_MIGRATED" = true ]; then run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" ${_MLX_LM_EXCLUDE_ARG:-} [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" _UNSLOTH_TORCH_OVERRIDES="" fi @@ -2949,7 +3388,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # No-torch: install unsloth + unsloth-zoo --no-deps, then runtime deps --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2968,7 +3407,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ - --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" + --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.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..." @@ -2995,7 +3434,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --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..." @@ -3007,6 +3446,15 @@ else fi fi +_installed_package_version=$("$_VENV_PY" -c \ + 'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' \ + "$PACKAGE_NAME" 2>/dev/null || true) +if [ -n "$_installed_package_version" ]; then + step "$PACKAGE_NAME" "$_installed_package_version installed" +else + substep "[WARN] installed $PACKAGE_NAME version could not be determined" "$C_WARN" +fi + # โ”€โ”€ Enforce the installed torch flavor matches the detected GPU build โ”€โ”€ # PEP 440 ignores the +cpu/+cuXXX/+rocm local label, so uv keeps a stale torch==X+cpu against a GPU index; reinstall the right triplet, else warn loudly. if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then diff --git a/pyproject.toml b/pyproject.toml index 071258eb8f..0f57ecf4df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "typer", + "typer>=0.12.0", "rich", "pydantic", "pyyaml", @@ -42,7 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"} include-package-data = true [tool.setuptools.package-data] -unsloth_cli = ["codex_fallback_prompt.md"] +unsloth_cli = ["codex_fallback_prompt.md", "pi_subagent.ts"] studio = [ "*.sh", "*.ps1", @@ -74,7 +74,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.4", + "unsloth_zoo>=2026.7.6", "wheel>=0.42.0", "packaging", "numpy", @@ -93,9 +93,20 @@ huggingfacenotorch = [ "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", ] +# torchcodec backend for Gemma audio / datasets>=4 (#7225). +# Pick the audio-torch* pin matching your torch minor (see TORCH_TORCHCODEC). +audio-torch210 = [ + "torchcodec>=0.10.0,<0.11.0 ; python_version >= '3.10'", +] +audio-torch290 = [ + "torchcodec>=0.8.0,<0.10.0 ; python_version >= '3.10'", +] +audio-torch280 = [ + "torchcodec>=0.6.0,<0.8.0 ; python_version >= '3.9'", +] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.4", + "unsloth_zoo>=2026.7.6", "torchvision", "unsloth[triton]", ] @@ -532,16 +543,19 @@ cu126-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] kaggle = [ "unsloth[huggingface]", @@ -580,7 +594,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.4", + "unsloth_zoo>=2026.7.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", @@ -831,16 +845,19 @@ cu126-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] flashattentiontorch260abiFALSEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", @@ -1125,7 +1142,8 @@ intelgputorch210 = [ "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=1c4b44b36a557f7381e3076fb8843366742238648441d607c8d049c6da0f8886 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] intel-gpu-torch210 = [ - "unsloth[intelgputorch210]" + "unsloth[intelgputorch210]", + "unsloth[audio-torch210]", ] intelgputorch2110 = [ "unsloth_zoo[intelgpu]", @@ -1279,6 +1297,7 @@ rocm72-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] rocm711-torch2100 = [ "unsloth[amd]", @@ -1297,6 +1316,7 @@ rocm711-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] [project.urls] diff --git a/scripts/build_whisper_cpp.sh b/scripts/build_whisper_cpp.sh new file mode 100755 index 0000000000..9f7e4d4ef3 --- /dev/null +++ b/scripts/build_whisper_cpp.sh @@ -0,0 +1,71 @@ +#!/bin/sh +# Build whisper.cpp's whisper-server for Studio's GGUF dictation engine. +# +# Installs into the managed Studio home so the backend's binary discovery +# (core/inference/stt_ggml_sidecar.py::find_whisper_server_binary) picks it up: +# /whisper.cpp/build/bin/whisper-server (custom home) +# ~/.unsloth/whisper.cpp/build/bin/whisper-server (default) +# +# Usage: +# ./scripts/build_whisper_cpp.sh # build the pinned tag +# WHISPER_CPP_TAG=v1.9.0 ./scripts/build_whisper_cpp.sh +# +# Requires: git, cmake, a C/C++ toolchain (the same prerequisites as a +# llama.cpp source build). GPU backends are auto-detected by whisper.cpp's +# CMake (Metal on macOS; set GGML_CUDA=1 to force a CUDA build on Linux). + +set -eu + +WHISPER_CPP_SOURCE="${WHISPER_CPP_SOURCE:-https://github.com/ggml-org/whisper.cpp}" +WHISPER_CPP_TAG="${WHISPER_CPP_TAG:-v1.9.1}" + +STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}" +CUSTOM_STUDIO_HOME=false +if [ -n "$STUDIO_HOME" ]; then + CUSTOM_STUDIO_HOME=true + INSTALL_DIR="$STUDIO_HOME/whisper.cpp" +else + INSTALL_DIR="$HOME/.unsloth/whisper.cpp" +fi + +command -v git >/dev/null 2>&1 || { echo "ERROR: git is required" >&2; exit 1; } +command -v cmake >/dev/null 2>&1 || { echo "ERROR: cmake is required" >&2; exit 1; } + +# Same policy as studio/setup.sh's _assert_studio_owned_or_absent: never delete +# a directory under a custom Studio home unless Studio itself created it (the +# marker file below). Protects a user-managed whisper.cpp/src from rm -rf. +STUDIO_OWNED_MARKER=".unsloth-studio-owned" +if [ "$CUSTOM_STUDIO_HOME" = true ] && [ -e "$INSTALL_DIR" ] && \ + [ ! -f "$INSTALL_DIR/$STUDIO_OWNED_MARKER" ]; then + echo "ERROR: $INSTALL_DIR already exists and is not marked as an Unsloth-owned whisper.cpp build tree." >&2 + echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2 + exit 1 +fi + +echo "==> Building whisper.cpp ($WHISPER_CPP_TAG) into $INSTALL_DIR" +mkdir -p "$INSTALL_DIR" +: > "$INSTALL_DIR/$STUDIO_OWNED_MARKER" + +if [ ! -d "$INSTALL_DIR/src/.git" ]; then + rm -rf "$INSTALL_DIR/src" + git clone --depth 1 --branch "$WHISPER_CPP_TAG" "$WHISPER_CPP_SOURCE" "$INSTALL_DIR/src" +else + git -C "$INSTALL_DIR/src" fetch --depth 1 origin "$WHISPER_CPP_TAG" + git -C "$INSTALL_DIR/src" checkout FETCH_HEAD +fi + +CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF" +if [ "${GGML_CUDA:-0}" = "1" ]; then + CMAKE_FLAGS="$CMAKE_FLAGS -DGGML_CUDA=ON" +fi + +# shellcheck disable=SC2086 +cmake -S "$INSTALL_DIR/src" -B "$INSTALL_DIR/src/build" $CMAKE_FLAGS +NCPU="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)" +cmake --build "$INSTALL_DIR/src/build" --config Release --target whisper-server -j"$NCPU" + +mkdir -p "$INSTALL_DIR/build/bin" +cp "$INSTALL_DIR/src/build/bin/whisper-server" "$INSTALL_DIR/build/bin/whisper-server" + +echo "==> Installed $INSTALL_DIR/build/bin/whisper-server" +"$INSTALL_DIR/build/bin/whisper-server" --help >/dev/null 2>&1 && echo "==> Binary runs OK" diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index c1be7a63a4..7bcee47c66 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -95,8 +95,8 @@ COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-i # Source: pytorch/torchcodec compatibility matrix on its README. TORCH_TORCHCODEC: dict[str, set[str]] = { "2.10": {"0.10"}, - "2.9": {"0.7", "0.8", "0.9"}, - "2.8": {"0.6"}, + "2.9": {"0.8", "0.9"}, + "2.8": {"0.6", "0.7"}, "2.7": {"0.3", "0.4", "0.5"}, "2.6": {"0.2", "0.3"}, "2.5": {"0.1", "0.2"}, diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 1f7bc8dcc0..1c21f8da86 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "_comment": "scan_packages.py allowlist (reviewed). Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", "version": 1, "entries": [ { @@ -303,8 +303,8 @@ "file": "openai/_base_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b", - "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14" + "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6", + "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66" }, { "package": "openai", @@ -319,8 +319,8 @@ "file": "openai/auth/_workload.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:", - "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567" + "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()", + "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0" }, { "package": "openai", @@ -343,8 +343,8 @@ "file": "openai/resources/beta/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", - "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd", + "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f" }, { "package": "openai", @@ -359,16 +359,16 @@ "file": "openai/resources/realtime/realtime.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05", - "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89" + "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5", + "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650" }, { "package": "openai", "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", - "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" + "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac", + "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f" }, { "package": "openai", @@ -1545,6 +1545,78 @@ "severity": "HIGH", "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_mlx_save_export_regressions.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:ab5c587f9ec31a0cc10ee55698ab133a417148d9d3f371bbc81b1e13fa119c13", + "evidence_hash": "93a11159147aad94f353ec4d2e0b8486b256abef88cd96d741813222cd32b138" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_vision_collator_audio.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:2efe23ffbe2b91b8403aec9b700736919b59e5ca770f8e1f5501651b44b7d398", + "evidence_hash": "d416b79dd17b24214f3f7653ac01354507d7bf0fc464dee30a4a4b8998f063ba" + }, + { + "package": "openai", + "file": "openai/_base_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6", + "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66" + }, + { + "package": "openai", + "file": "openai/auth/_workload.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()", + "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0" + }, + { + "package": "openai", + "file": "openai/resources/beta/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd", + "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f" + }, + { + "package": "openai", + "file": "openai/resources/realtime/realtime.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5", + "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650" + }, + { + "package": "openai", + "file": "openai/resources/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac", + "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_gemma4_forced_float32_ple_dtype.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L277: compile(rewritten + _GEMMA4_PLE_CAST_HELPER, \"\", \"exec\") | L440: compile(on, \"\", \"exec\") | L468: compile(generated, \"\", \"exec\")\nExec: L19: exec(_GEMMA4_PLE_CAST_HELPER, namespace)", + "evidence_hash": "a85e24d8e7c431563cbd83b70f91a3b971abde0f37083d68e70984147960cc70" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_vision_collator_audio.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:022f81dd21acfc6a35a058de96132834c218404a9e37b3d09a7768a8c8f6c728", + "evidence_hash": "2d1e75446af120d9133a42aa8af426a839d3434d9dc109cc1d6c1b22ca1ddb75" } ] } diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 44282b2255..612d739806 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -1,134 +1,145 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "view-in-github", - "colab_type": "text" - }, - "source": [ - "\"Open" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6b87de59" + }, + "source": [ + "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", + "
\n", + "\n", + "\n", + " Join Discord if you need help + โญ Star us on Github โญ\n", + "
\n", + "\n", + "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", + "\n", + "### Unsloth Studio\n", + "\n", + "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", + "\n", + "\n", + "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", + "\n", + "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) โ€ข [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) โ€ข [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) โ€ข [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) โ€ข [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" + ], + "id": "6b87de59" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e4206349" + }, + "source": [ + "

" + ], + "id": "e4206349" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "27da2957" + }, + "source": [ + "### Setup: Clone repo and run setup" + ], + "id": "27da2957" + }, + { + "cell_type": "code", + "metadata": { + "id": "27e68f91" + }, + "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local", + "execution_count": null, + "outputs": [], + "id": "27e68f91" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3e1771a9" + }, + "source": [ + "### Start Unsloth Studio" + ], + "id": "3e1771a9" + }, + { + "cell_type": "code", + "metadata": { + "id": "277e431e" + }, + "source": [ + "import sys\n", + "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", + "from colab import start\n", + "\n", + "# On Colab, start() auto-opens a Cloudflare link and prints admin login credentials.\n", + "# Use the Cloudflare link above the ready card to open Studio (in-cell iframes often stay blank).\n", + "start()\n", + "\n", + "# To skip the Cloudflare tunnel and try the in-notebook proxy iframe only:\n", + "# start(cloudflare=False)" + ], + "execution_count": null, + "outputs": [], + "id": "277e431e" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f2b0c6a1" + }, + "source": [ + "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", + "\n", + "Some other resources:\n", + "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", + "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", + "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", + "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", + "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", + "\n", + "
\n", + " \n", + " \n", + " \n", + "\n", + " Join Discord if you need help + โญ๏ธ Star us on Github โญ๏ธ\n", + "\n", + " This notebook is licensed AGPL-3.0\n", + "
" + ], + "id": "f2b0c6a1" + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } }, - { - "cell_type": "markdown", - "id": "6b87de59", - "metadata": { - "id": "6b87de59" - }, - "source": [ - "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", - "
\n", - "\n", - "\n", - " Join Discord if you need help + โญ Star us on Github โญ\n", - "
\n", - "\n", - "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", - "\n", - "### Unsloth Studio\n", - "\n", - "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", - "\n", - "\n", - "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", - "\n", - "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) โ€ข [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) โ€ข [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) โ€ข [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) โ€ข [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" - ] - }, - { - "cell_type": "markdown", - "id": "e4206349", - "metadata": { - "id": "e4206349" - }, - "source": [ - "

" - ] - }, - { - "cell_type": "markdown", - "id": "27da2957", - "metadata": { - "id": "27da2957" - }, - "source": [ - "### Setup: Clone repo and run setup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "27e68f91", - "metadata": { - "id": "27e68f91" - }, - "outputs": [], - "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local" - }, - { - "cell_type": "markdown", - "id": "3e1771a9", - "metadata": { - "id": "3e1771a9" - }, - "source": [ - "### Start Unsloth Studio" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "277e431e", - "metadata": { - "id": "277e431e" - }, - "outputs": [], - "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)" - }, - { - "cell_type": "markdown", - "id": "f2b0c6a1", - "metadata": { - "id": "f2b0c6a1" - }, - "source": [ - "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", - "\n", - "Some other resources:\n", - "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", - "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", - "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", - "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", - "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", - "\n", - "
\n", - " \n", - " \n", - " \n", - "\n", - " Join Discord if you need help + โญ๏ธ Star us on Github โญ๏ธ\n", - "\n", - " This notebook is licensed AGPL-3.0\n", - "
" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "gpuType": "T4", - "provenance": [], - "include_colab_link": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat": 4, + "nbformat_minor": 5 } \ No newline at end of file diff --git a/studio/backend/assets/configs/full_finetune.yaml b/studio/backend/assets/configs/full_finetune.yaml index e398515f61..98c45dd851 100644 --- a/studio/backend/assets/configs/full_finetune.yaml +++ b/studio/backend/assets/configs/full_finetune.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: false use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/lora_text.yaml b/studio/backend/assets/configs/lora_text.yaml index 9cb6b8c700..6c6a4d8839 100644 --- a/studio/backend/assets/configs/lora_text.yaml +++ b/studio/backend/assets/configs/lora_text.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: false use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml index 841e8ba166..e569031a31 100644 --- a/studio/backend/assets/configs/model_defaults/default.yaml +++ b/studio/backend/assets/configs/model_defaults/default.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml index f7b49c75b7..7ac1c83e04 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml index be7da0f624..4cab9e9f96 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml @@ -30,6 +30,7 @@ lora: - "query" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml index d9e49bc0d5..c1f1c2a344 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml @@ -30,6 +30,7 @@ lora: - "value" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml index c3422d399f..7828feae81 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml index 529a56a527..5a4028f15b 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml @@ -29,6 +29,7 @@ lora: - "Wqkv" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml index 734115ec41..7645d11c98 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml index 1032449e8c..b746235f1f 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml index c8e5f35841..4964fea276 100644 --- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml index 251409c29d..e5f3344356 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml index 89b1d7f938..71c61f383a 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml index e3292b5972..3fe29cd800 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml index 98fe497912..cd4e3e0c4d 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml index bda5471643..97aa10e861 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml index 18392568bd..a1b1640fa2 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml index 434ac41b46..dbf60f04d4 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml index 5f0a7b26ce..54c7dd6cd4 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml index dd5ae51ab0..119440a585 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml index e53e163a04..d08e5e9547 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml index ebe344e382..a266d7a39b 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml index fb89a07133..970cac3259 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml index 4a089992ac..5bba4ccdc0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml index ae7524b7c6..ac5c6eca22 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml index 10c1abd8a5..68c2d35644 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml index fb5c1d9dea..175f9c0f17 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml index 189e5dc6b2..4f3834e7c0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml index aa51440b6a..d6d97f7e44 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml index e2d67bcb0b..4f1f54a4e6 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml index aa436117a1..127700b53b 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml index 3f2cb84a94..2412b3accf 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml @@ -37,6 +37,7 @@ lora: - "shared_mlp.output_linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml index ab756fe764..81b59c4323 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml @@ -37,6 +37,7 @@ lora: - "shared_mlp.output_linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml index 1a7a91e56f..6110d84a6c 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml index 7c7bb8dc3e..3c7fc7f238 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml index f73b0c09b6..2b0977e435 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml index ffefb29e24..1742c04a06 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml index cd986a6da1..f33726b0dd 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml index 55dd3144c6..79b30bd758 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml index 8c9cb07fb9..4ee9a5a8ed 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml index 32441c5674..da20663688 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml index 6bba9c9633..30e4440afb 100644 --- a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml +++ b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml @@ -30,6 +30,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml index f9833ce705..9bb0a93e63 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml index 0ba857cd40..ded3607a14 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml index 3476f2dd6d..2ac72f1c88 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml index eda04d21f9..a087ced1f3 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml index bcd0d20c8c..c9811f4f06 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml index 34a033e32f..e3659d9fb0 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml index 98105eaf38..ee17efc54d 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 72b5b018e1..ef836b9b55 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml @@ -33,6 +33,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index d20751b0c7..c80fad35a8 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml @@ -38,6 +38,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 8a80282a2a..034b5bd131 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml @@ -37,6 +37,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml index a973c2d4e4..d1a226be79 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml @@ -35,6 +35,7 @@ lora: - "out_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml index b0feafbd6e..1b8df5ced9 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml index 2c44c91eab..cecab7f083 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml @@ -37,6 +37,7 @@ lora: - "out_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml index e1fbc08e4d..730be338cf 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml index 2abdfd8ac3..a70ac0bd49 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 5a3c4abb48..90ead037f6 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml @@ -38,6 +38,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml index a6ce27620f..a97c557c31 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index 050774a8cd..6855ed6a35 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -33,6 +33,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml index c574714d78..1933fed2ba 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml index e803c842b3..fda4e64158 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml index 4de3d9437d..c3910e3e5b 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml index bb75b3ce52..765ffee938 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml @@ -36,6 +36,7 @@ lora: - "gate_up_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml index c305d328c2..39b30e9cee 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml index 6cee3d0949..f97e525798 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml index 20ba81df2c..e19b94ede2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml index 9930786c24..982f54b32f 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml index 775c7ce08f..5242128004 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml index 856db0c1b3..3559b636c6 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml index 5900392547..3bc6d69afc 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml index bd54b1d015..604b86dacd 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml index 9feb6dcaae..daed4ebccb 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml index a40eace253..05eef89b88 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml index c130771c32..b4580e6d71 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml index 2fb3a95c30..2eceb7d0de 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml @@ -36,6 +36,7 @@ lora: - "gate_up_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml index 152f4ae06a..032091880c 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml index 94fe000708..e0e7f4ee3d 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml index 3c325485d2..bb463849ed 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml index 5b47c3bdd2..23e2b89dd0 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/vision_lora.yaml b/studio/backend/assets/configs/vision_lora.yaml index 063a970316..a06f971523 100644 --- a/studio/backend/assets/configs/vision_lora.yaml +++ b/studio/backend/assets/configs/vision_lora.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: true use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py index e855f4078b..925404f47d 100644 --- a/studio/backend/auth/terminal_prompt.py +++ b/studio/backend/auth/terminal_prompt.py @@ -236,6 +236,10 @@ def prompt_for_password_change( out.write(f"Password must be at least {min_length} characters; try again.\n") out.flush() continue + if any(ch.isspace() for ch in new_password): + out.write("Password cannot contain spaces; try again.\n") + out.flush() + continue if is_current_password(new_password): out.write( "New password must differ from the current bootstrap password; try again.\n" diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index b1ddc74c32..78fce0c70a 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -20,6 +20,7 @@ import shutil import subprocess import sys import threading +import time from pathlib import Path from typing import Optional, Tuple @@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl _READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection _DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download +# A registered edge connection does not mean the hostname resolves yet, so the +# URL is fetched once before it is advertised. +_PUBLIC_PROBE_PATH = "/api/health" +_PUBLIC_PROBE_MARKER = "Unsloth UI Backend" +# One deadline for DNS propagation + the health probe, bounding the startup stall. +_PUBLIC_PROBE_TIMEOUT = 45.0 +_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0 +_PUBLIC_PROBE_RETRY_DELAY = 1.0 + +# Wait for the hostname via DoH first: an early OS lookup negative-caches the +# NXDOMAIN for up to 30 min. +_DNS_POLL_DELAY = 2.0 +# Retry transient DoH failures, but give up fast when DoH is blocked outright. +_DNS_MAX_DOH_ERRORS = 3 +_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A" + def _windows_hidden_kwargs() -> dict: """Suppress a child console window on Windows; no-op elsewhere.""" @@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]: return None +def _wait_for_dns(host: str, deadline: float) -> None: + import json + import urllib.request + + errors = 0 + while True: + answered = False + try: + req = urllib.request.Request( + _DOH_URL.format(host = host), + headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"}, + ) + with urllib.request.urlopen(req, timeout = 5) as response: + answered = bool(json.loads(response.read(65536)).get("Answer")) + errors = 0 + except Exception: + errors += 1 + if errors >= _DNS_MAX_DOH_ERRORS: + return + if answered: + return + remaining = deadline - time.monotonic() + if remaining <= 0: + return + time.sleep(min(_DNS_POLL_DELAY, remaining)) + + +def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool: + import json + import urllib.request + from urllib.parse import urlsplit + + deadline = time.monotonic() + timeout + host = urlsplit(url).hostname + if host: + _wait_for_dns(host, deadline) + + probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}" + while True: + try: + req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response: + body = response.read(4096) + if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER: + return True + except Exception: + pass + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining)) + + class CloudflareTunnel: """A cloudflared quick tunnel to http://localhost:. Best-effort throughout. @@ -322,11 +392,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ """Start a quick tunnel and return its public URL once it is actually serving, or None (best-effort). - Waits for cloudflared to both mint the URL and register an edge connection - before returning, so the caller never advertises a URL that yields Cloudflare - error 1033 (HTTP 530). If a URL is minted but no connection registers within - the window (e.g. quic is blocked on this network), retries once forcing the - http2 protocol. On any failure the tunnel is stopped and None is returned. + Waits for cloudflared to both mint the URL and register an edge connection, + then fetches /api/health over the public URL, so the caller never advertises + a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host. + If a URL is minted but no connection registers within the window (e.g. quic + is blocked on this network), retries once forcing the http2 protocol. On any + failure the tunnel is stopped and None is returned. """ global _active_tunnel, _shutdown_requested binary = ensure_cloudflared() @@ -349,9 +420,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ prior, _active_tunnel = _active_tunnel, tunnel if prior is not None: prior.stop() + registered = False try: tunnel.start() url = tunnel.wait_for_ready(timeout) + registered = url is not None + if url and not verify_public_url(url): + url = None except Exception: url = None if url: @@ -371,6 +446,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ # http2 will not help, so do not burn another window on it. if not saw_url: return None + # probe failure after registering is DNS propagation; http2 would not help + if registered: + return None return None diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 1762469bcf..baa18a2fec 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Colab helpers for Unsloth Studio. Uses Colab's built-in proxy. -""" +"""Colab helpers for Unsloth Studio. Uses Colab's built-in proxy.""" from pathlib import Path import sys @@ -22,11 +20,9 @@ logger = get_logger(__name__) def get_colab_url(port: int = 8888) -> str: - """ - Get the Colab proxy URL for a port. + """Get the Colab proxy URL for a port. - Retries up to 3 times, validating the result is a real HTTPS Colab URL. - Falls back to http://localhost:{port} only when all attempts fail. + Retries 3x validating a real HTTPS Colab URL; falls back to localhost on failure. """ import time as _time @@ -55,28 +51,244 @@ def get_colab_url(port: int = 8888) -> str: return fallback -def show_link(port: int = 8888, *, _url: "str | None" = None): - """Display a styled clickable link to the UI. - - *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip. - """ - from IPython.display import display, HTML - - url = _url if _url is not None else get_colab_url(port) - - # Truncated display URL; try/except so an odd URL shape still renders the link. +def _short_colab_url(url: str, port: int) -> str: + """Truncated display form of a Colab proxy URL; falls back to the full URL.""" try: port_prefix = f"{port}-" idx = url.index(port_prefix) next_dash = url.index("-", idx + len(port_prefix)) - short_url = url[: next_dash + 1] + "..." + return url[: next_dash + 1] + "..." except (ValueError, IndexError): - short_url = url + return url - # Plain-text line so the URL shows even if HTML display fails. - logger.info(f"๐ŸŒ Unsloth Studio URL: {url}") - html = f""" +def _is_colab_proxy_url(url: str, port: int) -> bool: + """True when *url* looks like a real Colab kernel proxy, not a localhost fallback.""" + return bool(url and isinstance(url, str) and url.startswith("https://") and str(port) in url) + + +def _is_colab_runtime() -> bool: + """True on a hosted Colab notebook kernel. + + Reuses the backend's main Colab detector (``/content`` + Colab env / ``google.colab``) + instead of a single env var, which is not always present on hosted runtimes. + """ + try: + from main import _IS_COLAB + return bool(_IS_COLAB) + except Exception: + return False + + +def _colab_login_credentials_path() -> Path: + from auth.storage import DB_PATH + return DB_PATH.parent / ".colab_notebook_login" + + +def _store_colab_login_credentials(username: str, password: str) -> None: + """Persist Colab admin credentials for notebook re-runs after interrupt.""" + path = _colab_login_credentials_path() + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(f"{username}\n{password}\n") + try: + import os + os.chmod(path, 0o600) + except OSError: + pass + except OSError as e: + logger.info(f"Could not persist Colab login credentials ({e}).") + + +def _load_colab_login_credentials() -> "tuple[str, str] | None": + """Return stored Colab admin credentials from a previous ``start()`` run, if any.""" + path = _colab_login_credentials_path() + try: + if not path.is_file(): + return None + lines = path.read_text().splitlines() + if len(lines) >= 2 and lines[0] and lines[1]: + return lines[0], lines[1] + except OSError as e: + logger.info(f"Could not load Colab login credentials ({e}).") + return None + + +def _clear_colab_login_credentials() -> None: + """Drop the cached Colab credentials once they no longer authenticate.""" + path = _colab_login_credentials_path() + try: + path.unlink(missing_ok = True) + except OSError as e: + logger.info(f"Could not clear Colab login credentials ({e}).") + + +def _colab_credentials_still_valid(username: str, password: str) -> bool: + """True when *password* still matches the stored admin hash. + + Guards against redisplaying a cached first-run password after the user has + changed the admin password through the app, which would print credentials + that no longer authenticate to the current Cloudflare tunnel. + """ + try: + from auth.storage import get_user_and_secret + from auth.hashing import verify_password + except Exception as e: + logger.info(f"Could not load auth to validate cached Colab credentials ({e}).") + return False + try: + row = get_user_and_secret(username) + if not row: + return False + salt, pwd_hash = row[0], row[1] + return bool(verify_password(password, salt, pwd_hash)) + except Exception as e: + logger.info(f"Could not validate cached Colab credentials ({e}).") + return False + + +def _colab_wants_cloudflare(cloudflare: "bool | None") -> bool: + """Resolve whether to open a Cloudflare tunnel. + + ``None`` auto-enables on real Colab (the in-cell proxy embed is often blank); + pass ``False`` to opt out. + """ + if cloudflare is not None: + return cloudflare + return _is_colab_runtime() + + +def _finalize_colab_admin_password() -> "tuple[str, str] | None": + """Clear the bootstrap-password gate on Colab so Cloudflare tunnels can start. + + Returns ``(username, password)`` for display in the notebook. On first run the + random admin password is finalized; on later runs (e.g. after interrupt) the + stored credentials are re-displayed so the Cloudflare link stays usable. + Anyone who can read this cell already controls the runtime. + """ + if not _is_colab_runtime(): + return None + try: + from auth.storage import ( + DEFAULT_ADMIN_USERNAME, + ensure_default_admin, + generate_bootstrap_password, + get_bootstrap_password, + requires_password_change, + update_password, + ) + except Exception as e: + logger.warning( + f"Could not load auth for Colab setup ({e}); Cloudflare link may be blocked." + ) + return None + + try: + ensure_default_admin() + username = DEFAULT_ADMIN_USERNAME + if not requires_password_change(username): + creds = _load_colab_login_credentials() + if creds is not None and _colab_credentials_still_valid(username, creds[1]): + return creds + # The admin password was changed through the app after the first run, + # so the cached copy is stale; drop it instead of printing dead credentials. + _clear_colab_login_credentials() + return None + password = get_bootstrap_password() or generate_bootstrap_password() + if not update_password(username, password): + logger.warning( + "Could not finalize Colab admin password; Cloudflare link may be blocked." + ) + return None + _store_colab_login_credentials(username, password) + return username, password + except Exception as e: + logger.warning( + f"Could not finalize Colab admin password ({e}); Cloudflare link may be blocked." + ) + return None + + +def _colab_login_html(username: str, password: str) -> str: + """Notebook card with Colab admin credentials (shown once after auto-finalize).""" + return f""" +
+

+ Unsloth Studio Login (Colab) +

+

+ Log in to Studio with the Cloudflare link above using these credentials. This cell + is visible only in your notebook session. +

+

+ Username: {username}
+ Password: {password} +

+
+ """ + + +def _show_colab_login_credentials(username: str, password: str) -> None: + """Display Colab admin credentials in the notebook output.""" + from IPython.display import HTML, display + + logger.info(f"๐Ÿ” Unsloth Studio login โ€” user: {username}") + display(HTML(_colab_login_html(username, password))) + + +def _ready_card_html( + url: str, + port: int, + *, + has_cloudflare_link: bool = False, + cloudflare_requested: bool = False, +) -> str: + """Branded ready card for the in-notebook Studio view. + + Colab ``*.prod.colab.dev`` proxy URLs are session-scoped and 404 when opened as a + top-level tab or on another device, so never ``window.open`` them. On real Colab the + Cloudflare link is the supported entry point because in-cell proxy embeds often stay blank. + """ + short_url = _short_colab_url(url, port) + if _is_colab_runtime() or _is_colab_proxy_url(url, port): + if has_cloudflare_link: + embed_note = ( + "Open Studio with the Cloudflare link above. In-cell proxy previews on " + "current Colab often stay blank, so the tunnel link is the supported path." + ) + elif cloudflare_requested: + embed_note = ( + "Could not open a Cloudflare tunnel, so Studio may be unreachable on Colab. " + "Check the logs above and re-run this cell. Pass " + '' + "cloudflare=True after fixing any tunnel errors." + ) + else: + embed_note = ( + "Colab proxy links cannot be opened in a new tab (they 404 outside this " + 'notebook). Re-run with start(cloudflare=True) for a working link.' + ) + return f""" +
+

+ + Unsloth Studio is Ready! +

+

+ {embed_note} +

+

+ {short_url} +

+
+ """ + + return f"""

""" - display(HTML(html)) + + +def show_link( + port: int = 8888, + *, + _url: "str | None" = None, + has_cloudflare_link: bool = False, + cloudflare_requested: bool = False, +): + """Display a styled ready card for the UI. + + Colab proxy URLs are informational only (no new-tab open; they 404 outside the cell); + non-proxy URLs keep a clickable open button. *_url* is an optional pre-fetched proxy + URL to avoid a second eval_js round-trip. + """ + from IPython.display import display, HTML + + url = _url if _url is not None else get_colab_url(port) + logger.info(f"๐ŸŒ Unsloth Studio URL: {url}") + display( + HTML( + _ready_card_html( + url, + port, + has_cloudflare_link = has_cloudflare_link, + cloudflare_requested = cloudflare_requested, + ) + ) + ) + + +def _warn_colab_cloudflare_missing(*, use_cloudflare: bool, cloudflare_url: "str | None") -> None: + """Log a prominent warning when Colab expected a tunnel but none was opened.""" + if not use_cloudflare or cloudflare_url or not _is_colab_runtime(): + return + logger.warning( + "Colab Cloudflare tunnel unavailable โ€” Studio is unlikely to be reachable in this " + "notebook. Check the logs above for tunnel or auth errors, then re-run start()." + ) def _bootstrap_password_pending() -> bool: """True while the default admin still owes a bootstrap-password change. - While pending, main.py injects that password into same-origin GETs, and a public - tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin - access. Fails safe to pending if the state cannot be read. + While pending, a public tunnel GET (no Origin) reads as same-origin and gets the + injected password, so sharing the link would leak admin access. Fails safe to pending. """ try: from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME @@ -121,9 +370,8 @@ def _bootstrap_password_pending() -> bool: def start_cloudflare_tunnel(port: int) -> "str | None": """Open a shareable Cloudflare quick tunnel to localhost:*port*, or None. - run_server suppresses the tunnel on Colab by design, so we start it directly. - Refused while the bootstrap password is pending; any failure collapses to None - and the Colab proxy still works. + run_server suppresses the tunnel on Colab, so we start it directly. Refused while the + bootstrap password is pending; any failure collapses to None (Colab proxy still works). """ if _bootstrap_password_pending(): logger.warning( @@ -152,9 +400,9 @@ def start_cloudflare_tunnel(port: int) -> "str | None": def _publish_cloudflare_url(cloudflare_url: "str | None") -> None: """Publish a directly-started tunnel URL onto app.state so /api/health advertises it. - run_server only sets this when it opens the tunnel itself, which it skips on Colab, - so we set it here. Otherwise the frontend's API examples fall back to an - unreachable server_url. Best-effort. + run_server sets this only when it opens the tunnel itself (skipped on Colab), so we + set it here; otherwise the frontend's API examples fall back to an unreachable + server_url. Best-effort. """ if not cloudflare_url: return @@ -183,8 +431,7 @@ def _stop_cloudflare_tunnel() -> None: def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: """True only if Unsloth Studio (not some other app) answers /api/health on *port*. - The service-marker check stops the reuse path reusing or tunneling a foreign - process that merely serves /api/health. + The service-marker check stops the reuse path reusing or tunneling a foreign process. """ import json, urllib.request try: @@ -222,31 +469,45 @@ def _shareable_link_html(cloudflare_url: str) -> str: """ -def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None): - """Render the Unsloth header + iframe for *port*, with a shareable-link card above - when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe.""" - url = get_colab_url(port) - logger.info(f"๐ŸŒ Unsloth Studio URL: {url}") - if cloudflare_url: - logger.info(f"๐Ÿ”— Shareable Cloudflare link: {cloudflare_url}") +# Height for serve_kernel_port_as_iframe (~82vh on a 1080p screen, clamped). +_COLAB_IFRAME_HEIGHT = 900 + +def _embed_kernel_port_iframe(port: int) -> bool: + """Embed Studio via Colab's native kernel-port iframe helper. + + Only trusted on a real Colab runtime: colabtools can import ``google.colab`` and + queue browser-side JS without appending an iframe, so callers outside Colab must use + the HTML iframe path instead. + """ + if not _is_colab_runtime(): + return False + try: + from google.colab import output as colab_output + except ImportError: + return False + try: + colab_output.serve_kernel_port_as_iframe( + port, + height = _COLAB_IFRAME_HEIGHT, + width = "100%", + ) + return True + except Exception as e: + logger.info(f"serve_kernel_port_as_iframe failed ({e}); trying HTML iframe.") + return False + + +def _embed_html_iframe(url: str, port: int) -> bool: + """Fallback embed: raw HTML iframe when the Colab helper is unavailable.""" try: from IPython.display import HTML, display + except ImportError: + return False - iframe_id = f"unsloth-studio-{port}" - - # Truncated header URL โ€” best-effort, falls back to full URL. - try: - port_prefix = f"{port}-" - idx = url.index(port_prefix) - next_dash = url.index("-", idx + len(port_prefix)) - short_url = url[: next_dash + 1] + "..." - except (ValueError, IndexError): - short_url = url - - if cloudflare_url: - display(HTML(_shareable_link_html(cloudflare_url))) - + short_url = _short_colab_url(url, port) + iframe_id = f"unsloth-studio-{port}" + try: display( HTML(f"""
""") ) - except Exception: - # Fallback: Colab's built-in helper. + return True + except Exception as e: + logger.info(f"HTML iframe embed failed ({e}).") + return False + + +def _show_and_embed( + port: int, + *, + cloudflare_url: "str | None" = None, + colab_login: "tuple[str, str] | None" = None, + cloudflare_requested: bool = False, +): + """Render the Unsloth ready card + iframe for *port*. + + Prefer Colab's ``serve_kernel_port_as_iframe`` on real Colab; raw HTML iframe is the + fallback. Cloudflare cards stay clickable. + """ + url = get_colab_url(port) + logger.info(f"๐ŸŒ Unsloth Studio URL: {url}") + if cloudflare_url: + logger.info(f"๐Ÿ”— Shareable Cloudflare link: {cloudflare_url}") + + _warn_colab_cloudflare_missing( + use_cloudflare = cloudflare_requested, + cloudflare_url = cloudflare_url, + ) + + if cloudflare_url: try: - from google.colab import output as colab_output - colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") - except ImportError: - pass + from IPython.display import HTML, display + display(HTML(_shareable_link_html(cloudflare_url))) + except Exception as e: + logger.info(f"Could not render Cloudflare link card ({e}).") + + if colab_login: + try: + _show_colab_login_credentials(*colab_login) + except Exception as e: + logger.info(f"Could not render Colab login card ({e}).") + + try: + show_link( + port, + _url = url, + has_cloudflare_link = bool(cloudflare_url), + cloudflare_requested = cloudflare_requested, + ) + except Exception as e: + logger.info(f"Could not render Unsloth link card ({e}).") + + # On Colab with a working tunnel, skip the in-cell proxy embed (often blank). + if _is_colab_runtime() and cloudflare_url: + return + + # Real Colab: kernel helper needs only the port (works when eval_js failed). + if _is_colab_runtime(): + if _embed_kernel_port_iframe(port): + return + _embed_html_iframe(url, port) -def start(port: int = 8888, *, cloudflare: bool = False): +def start(port: int = 8888, *, cloudflare: "bool | None" = None): """Start Unsloth Studio in Colab and display the URL. Args: port: Port to bind/serve on. - cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any - device (default OFF). It exposes Unsloth's login page beyond Colab, so it - stays an explicit opt-in; the default shows only the in-tab proxy iframe. + cloudflare: Shareable Cloudflare HTTPS link. ``None`` (default) auto-enables on + real Colab because the in-cell proxy embed is often blank; pass ``False`` to + skip the tunnel or ``True`` to force it on other runtimes. Usage: - start() # Colab-proxy iframe only (default) - start(cloudflare=True) # also open a shareable Cloudflare link + start() # Cloudflare link on Colab (auto); proxy iframe elsewhere + start(cloudflare=False) # Colab proxy iframe only (often blank on current Colab) + start(cloudflare=True) # force Cloudflare link on any runtime """ import time logger.info("๐Ÿฆฅ Starting Unsloth Studio...") + use_cloudflare = _colab_wants_cloudflare(cloudflare) - # Fast path: Unsloth already running (cell re-run). Re-launching would collide on - # the port, so just re-show the link and iframe. + # Fast path: already running (cell re-run); re-show link/iframe instead of rebinding the port. if _is_studio_healthy(port): logger.info(f" Unsloth is already running on port {port} โ€” reusing existing server.") # try/finally: tear the tunnel down even if interrupted mid-start/render. try: - cf_url = start_cloudflare_tunnel(port) if cloudflare else None + colab_login = _finalize_colab_admin_password() if use_cloudflare else None + cf_url = start_cloudflare_tunnel(port) if use_cloudflare else None _publish_cloudflare_url(cf_url) - _show_and_embed(port, cloudflare_url = cf_url) + _show_and_embed( + port, + cloudflare_url = cf_url, + colab_login = colab_login, + cloudflare_requested = use_cloudflare, + ) for _ in range(10000): time.sleep(300) print("=", end = "", flush = True) @@ -313,7 +634,6 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Loading backend...") from run import run_server - # Auto-detect frontend path repo_root = Path(__file__).parent.parent frontend_path = repo_root / "frontend" / "dist" @@ -323,8 +643,7 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Starting server...") try: - # cloudflare=False: this helper owns the tunnel (Colab's own - # start(cloudflare=...) drives it), so pin it off explicitly. + # cloudflare=False: this helper owns the tunnel (via start(cloudflare=...)), so pin it off. app = run_server( host = "0.0.0.0", port = port, @@ -339,14 +658,12 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.error(f"โŒ Unsloth Studio failed to start: {exc}") return - # run_server auto-increments the port if in use; read back the bound port so the - # proxy URL and iframe point at the right place. + # run_server may auto-increment the port; read back the bound port for the proxy URL/iframe. actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port logger.info(f" Server started on port {actual_port}!") - # Poll health endpoint before showing the link โ€” avoids the race where ready_event - # fires but the process hasn't finished binding. + # Poll health before showing the link: avoids the race where ready_event fires pre-bind. import urllib.request server_ready = False @@ -365,12 +682,17 @@ def start(port: int = 8888, *, cloudflare: bool = False): ) return - # Open the tunnel now the server is healthy, publish its URL for /api/health, and - # tear it down on interrupt (try/finally) rather than orphan the process. + # Server healthy: finalize Colab auth, open the tunnel, publish URL, tear down on interrupt. try: - cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None + colab_login = _finalize_colab_admin_password() if use_cloudflare else None + cf_url = start_cloudflare_tunnel(actual_port) if use_cloudflare else None _publish_cloudflare_url(cf_url) - _show_and_embed(actual_port, cloudflare_url = cf_url) + _show_and_embed( + actual_port, + cloudflare_url = cf_url, + colab_login = colab_login, + cloudflare_requested = use_cloudflare, + ) # Keep kernel alive so the daemon server thread runs. for _ in range(10000): diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 0e0044702e..135c9fccf6 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -27,7 +27,6 @@ from .constants import ( ) from .parse import apply_update, coerce_event, parse_log_message from .types import Job -from .worker import run_job_process from loggers import get_logger logger = get_logger(__name__) @@ -169,12 +168,18 @@ class JobManager: native_path_secret_removed_for_child_start, run_without_native_path_secret, ) + from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths - with native_path_secret_removed_for_child_start(): + cache_env = get_hf_cache_paths().child_env({}) + + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): mp_q = _CTX.Queue() proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_job_process,), + args = ("core.data_recipe.jobs.worker", "run_job_process", cache_env), kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload}, daemon = True, ) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index c8be50b08b..1b24c46e65 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -1048,6 +1048,21 @@ class ExportBackend: "Use the safetensors adapter instead.", None, ) + # llama.cpp's convert_lora_to_gguf.py has no concept of DoRA's + # lora_magnitude_vector tensors: it only reads the standard + # lora_A/lora_B delta, so exporting a DoRA adapter would silently + # drop the magnitude rescaling and produce a GGUF LoRA file that + # loads fine but no longer matches the trained model. + _peft_config = getattr(self.current_model, "peft_config", {}).get("default") + if getattr(_peft_config, "use_dora", False): + return ( + False, + "GGUF LoRA export is not supported for DoRA adapters: the GGUF LoRA " + "format has no way to represent DoRA's magnitude vectors, so the " + "exported file would silently lose the DoRA behavior. Use the " + "safetensors adapter instead, or merge to a full GGUF model.", + None, + ) outtype = str(gguf_outtype).lower() if outtype not in _GGUF_LORA_OUTTYPES: return ( diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 6d1a928f2e..aaf48615f0 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -230,16 +230,20 @@ class ExportOrchestrator: native_path_secret_removed_for_child_start, run_without_native_path_secret, ) + from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths - from .worker import run_export_process + cache_env = get_hf_cache_paths().child_env({}) - with native_path_secret_removed_for_child_start(): + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_export_process,), + args = ("core.export.worker", "run_export_process", cache_env), kwargs = { "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index 93c7da72cb..b59f2bcce0 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -76,8 +76,14 @@ class AudioCodecManager: if self._snac_model is not None: return from snac import SNAC + from utils.hf_cache_settings import active_hf_hub_cache - self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval() + # Route weights to the selected cache; this can run in the main process. + self._snac_model = ( + SNAC.from_pretrained("hubertsiuzdak/snac_24khz", cache_dir = active_hf_hub_cache()) + .to(device) + .eval() + ) logger.info("Loaded SNAC codec (24kHz)") def _load_bicodec( diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 8d262bbb0f..2f46470091 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -8,6 +8,7 @@ from unsloth.chat_templates import get_chat_template from transformers import TextIteratorStreamer, TextStreamer from peft import PeftModel, PeftModelForCausalLM +import contextlib import json import sys import torch @@ -1942,8 +1943,30 @@ class InferenceBackend: + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" ) + with torch.inference_mode(): - with torch.amp.autocast("cuda", dtype = model.dtype): + # Derive the autocast device from the loaded model, not from the + # global backend: a CPU-fallback DAC on an XPU/CUDA host must not + # open a GPU autocast context around CPU tensors. + device_type = ( + model.device.type + if hasattr(model.device, "type") + else str(model.device).split(":", 1)[0] + ) + # Clamp to autocast-supported backends so exotic devices + # (e.g. "meta" during accelerate offloaded loading) do not raise. + # MPS is autocast-supported since torch 2.3, keep it in the set. + if device_type not in ("cuda", "xpu", "mps", "cpu"): + device_type = "cpu" + # CPU and XPU autocast only accept bfloat16/float16. For a + # float32 model, skip autocast entirely to avoid raising or + # producing a warning on every generate call. + autocast_dtype_supported = model.dtype in (torch.bfloat16, torch.float16) + if device_type in ("cpu", "xpu") and not autocast_dtype_supported: + autocast_ctx = contextlib.nullcontext() + else: + autocast_ctx = torch.amp.autocast(device_type, dtype = model.dtype) + with autocast_ctx: inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2c7433f7a4..1fe134c3f9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -247,6 +247,59 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out +def _bundled_hip_present(binary_dir: str) -> bool: + """True when a prebuilt bundle ships its own HIP backend library.""" + if not binary_dir: + return False + try: + # Glob the version suffix (libggml-hip.so, .so.0, .so.0.11.1) the same + # way the installer's runtime health check matches libggml-hip.so*. + return any(Path(str(binary_dir)).glob("libggml-hip.so*")) + except OSError: + return False + + +def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]": + """System ROCm lib dir(s) to prepend before a prebuilt's bundled HIP, on native Linux. + + The bundled bare-metal HIP runtime can mismatch the host amdkfd driver and crash + in hsa_init(); prepending the whole system ROCm lib dir loads a driver-matched, + version-consistent stack (libhsa-runtime64 / libamdhip64 / librocblas) ahead of it. + The whole dir is deliberate: mixing the bundle's rocBLAS with a different-version + system HIP/ROCR risks missing symbols. UNSLOTH_LLAMA_NO_SYSTEM_ROCM=1 keeps the pure + bundle (for a host whose system ROCm lacks this arch); no-op on WSL / non-Linux. + """ + if os.environ.get("UNSLOTH_LLAMA_NO_SYSTEM_ROCM") == "1": + return [] + if sys.platform != "linux" or os.path.exists("/dev/dxg"): + return [] + if not os.path.exists("/dev/kfd"): + return [] + if not _bundled_hip_present(binary_dir): + return [] + # Env-configured ROCm root first; /opt/rocm only as a fallback so a stale + # /opt/rocm doesn't shadow the driver-matching install these vars point at. + candidates = [] + for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + val = os.environ.get(var) + if val: + candidates.append(val) + candidates.append("/opt/rocm") + out: "list[str]" = [] + seen: "set[str]" = set() + for base in candidates: + for lib_sub in ("lib", "lib64"): + d = os.path.join(base, lib_sub) + if d in seen: + continue + seen.add(d) + if os.path.exists(os.path.join(d, "libhsa-runtime64.so")) or os.path.exists( + os.path.join(d, "libhsa-runtime64.so.1") + ): + out.append(d) + return out + + # Plan-without-action re-prompt state now lives in tool_call_parser (imported above). # Default max_tokens to the effective context when known. The floor is high @@ -526,7 +579,14 @@ def _swa_entry_from_layer_types(lt) -> Optional[object]: def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: try: from huggingface_hub import hf_hub_download - cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model") + from utils.hf_cache_settings import active_hf_hub_cache + + cfg_path = hf_hub_download( + repo_id, + "config.json", + repo_type = "model", + cache_dir = active_hf_hub_cache(), + ) with open(cfg_path) as f: cfg = json.load(f) except Exception: @@ -928,6 +988,7 @@ def _cached_hf_snapshot_file( filename: str, *, expected_size: Optional[int] = None, + cache_dir: Optional[str] = None, ) -> Optional[str]: """Return a cached snapshot file even when HF's current-ref probe misses it.""" if not filename: @@ -936,8 +997,22 @@ def _cached_hf_snapshot_file( if not parts or any(part in (".", "..") for part in parts): return None try: - from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(repo_id): + if cache_dir is None: + from utils.models.model_config import _iter_hf_cache_snapshots + snapshots = _iter_hf_cache_snapshots(repo_id) + else: + from hub.utils.hf_cache_state import iter_active_repo_cache_dirs + snapshots = ( + snapshot + for repo_dir in iter_active_repo_cache_dirs( + "model", + repo_id, + root = Path(cache_dir), + ) + for snapshot in (repo_dir / "snapshots").glob("*") + if snapshot.is_dir() + ) + for snap in snapshots: candidate = snap.joinpath(*parts) if not candidate.is_file(): continue @@ -1179,6 +1254,16 @@ def _snapshot_dir_of(path: str) -> Optional[Path]: return None +def _hub_cache_dir_for_snapshot_path(path: Optional[str]) -> Optional[str]: + """Return the HF Hub cache root that owns a snapshot-contained path.""" + if not path: + return None + snapshot = _snapshot_dir_of(path) + if snapshot is None or snapshot.parent.name != "snapshots": + return None + return str(snapshot.parent.parent.parent) + + def _companion_snapshot_sibling( near_path: str, pick: Callable[[list[str]], Optional[str]] ) -> Optional[str]: @@ -1938,6 +2023,10 @@ class LlamaCppBackend: self._tensor_split: Optional[List[float]] = None # User-picked physical GPU indices (None = automatic selection). self._gpu_ids: Optional[List[int]] = None + # RAW requested GPU pin, before the fit narrowed it. self._gpu_ids records the + # EFFECTIVE (fit-narrowed) pin for /status; dedupe compares this raw value so a + # [0, 1] narrowed to [0] and re-sent as [0, 1] still matches (#7239). + self._requested_gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -2409,6 +2498,46 @@ class LlamaCppBackend: """User-picked physical GPU indices, or None for automatic selection.""" return self._gpu_ids + @property + def requested_gpu_ids(self) -> Optional[List[int]]: + """RAW requested GPU pin (before the fit narrowed it), or None for auto. + gpu_ids echoes the EFFECTIVE pin for /status.""" + return self._requested_gpu_ids + + def matches_gpu_ids(self, gpu_ids: Optional[List[int]]) -> bool: + """Whether a requested pin is already satisfied by the active runner. + + A regular GGUF load may narrow the requested placement pool to the + smallest fitting subset. Accept both the original request and the + effective status-echoed subset so either can round-trip without a + needless reload. Diffusion drives one device and keeps its existing + lowest-device normalization. + """ + if self._is_diffusion: + requested = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + return requested == (self._gpu_ids or None) + + requested = sorted(int(x) for x in gpu_ids) if gpu_ids else None + raw = self._requested_gpu_ids or None + effective = self._gpu_ids or None + return requested == raw or requested == effective + + def _record_matching_gpu_request(self, gpu_ids: Optional[List[int]]) -> None: + """Adopt the caller's explicit pool after a full already-loaded match. + + Matching an effective subset avoids a reload, but the incoming request + is still the user's latest placement intent. Record it so status and a + later reload do not restore GPUs the user just removed. + """ + if self._is_diffusion: + self._requested_gpu_ids = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + else: + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + if self._last_load_kwargs is not None: + self._last_load_kwargs["gpu_ids"] = ( + list(self._requested_gpu_ids) if self._requested_gpu_ids else None + ) + @property def n_layers(self) -> Optional[int]: """Model layer count (GGUF block_count), or None if unknown.""" @@ -2652,6 +2781,7 @@ class LlamaCppBackend: "found": False, "mtp_token": None, "supports_mtp": False, + "mtp_probe_inconclusive": True, "ngram_mod_flavor": None, "supports_ngram_mod": False, "spec_draft_n_max_flag": None, @@ -2684,6 +2814,9 @@ class LlamaCppBackend: supports_no_cache_prompt = False supports_metrics = False supports_slot_save = False + saw_spec_type = False + probe_ok = False + help_text = "" try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( @@ -2695,6 +2828,7 @@ class LlamaCppBackend: check = False, env = probe_env, ) + probe_ok = result.returncode == 0 help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented # continuation), so the "argument has been removed" description @@ -2739,17 +2873,19 @@ class LlamaCppBackend: return False return "argument has been removed" not in desc - # MTP token from the --spec-type line. - spec_line = "" - for line in help_text.splitlines(): - if "--spec-type" in line: - spec_line = line - break - # PR #22673 used draft-mtp; later renamed to mtp. - if "draft-mtp" in spec_line: - mtp_token = "draft-mtp" - elif re.search(r"[|,\[]mtp[|,\]]", spec_line): - mtp_token = "mtp" + # MTP token from the full --spec-type help block (decl + indented + # continuation). First-line-only probing missed builds putting the + # enum on the next line (#7302). Prefer draft-mtp (PR #22673) over mtp. + spec_help = blocks.get("--spec-type") or "" + if not spec_help: + # Fallback: join --spec-type lines, avoiding incidental "mtp" in --help. + spec_help = "\n".join( + line for line in help_text.splitlines() if "--spec-type" in line + ) + mtp_token = cls._mtp_token_from_spec_help(spec_help) + # Only a resolved --spec-type block confirms missing MTP; empty/crash + # leaves saw_spec_type False so supports_mtp fails open. + saw_spec_type = bool(spec_help.strip()) and "--spec-type" in spec_help # ngram-mod flag flavor. Post-rename builds advertise both new # args (real) and legacy ones (stubs); pre-rename builds only @@ -2785,11 +2921,29 @@ class LlamaCppBackend: supports_slot_save = _is_real("--slot-save-path") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") + saw_spec_type = False + probe_ok = False + help_text = "" + + help_nonempty = bool(help_text.strip()) + # Confirmed only when a successful --help lists a --spec-type block with + # mtp/draft-mtp; nonempty --help without it is a definitive pre-spec + # binary; failed/empty probes stay inconclusive (#7302). + if saw_spec_type and probe_ok: + supports_mtp = mtp_token is not None + mtp_probe_inconclusive = False + elif help_nonempty and probe_ok: + supports_mtp = False + mtp_probe_inconclusive = False + else: + supports_mtp = False + mtp_probe_inconclusive = True info = { "found": True, "mtp_token": mtp_token, - "supports_mtp": mtp_token is not None, + "supports_mtp": supports_mtp, + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": ngram_mod_flavor, "supports_ngram_mod": ngram_mod_flavor is not None, "spec_draft_n_max_flag": spec_draft_n_max_flag, @@ -2805,6 +2959,21 @@ class LlamaCppBackend: cls._capability_cache[cache_key] = info return info + @staticmethod + def _mtp_token_from_spec_help(spec_help: str) -> Optional[str]: + """Extract ``draft-mtp`` / ``mtp`` from a ``--spec-type`` help snippet. + + Prefers ``draft-mtp`` (llama.cpp PR #22673) over the later bare ``mtp`` + rename. Returns ``None`` when neither token appears as an enum value. + """ + text = spec_help or "" + if "draft-mtp" in text: + return "draft-mtp" + # Bare `mtp` enum token (`|mtp|`, `,mtp,`, ...), not a substring. + if re.search(r"(?physical mapping.""" try: import torch - is_rocm = getattr(torch.version, "hip", None) is not None + + # Same ROCm detection as _emit_child_gpu_visibility: AMD SDK wheels + # leave version.hip unset but encode "rocm" in __version__. The two + # must agree, else an inherited ROCR mask reads back as "no mask", + # ordinal 0 is labelled physical 0, and the child's new ROCR pin + # re-exposes the GPU the inherited mask was hiding. + is_rocm = ( + getattr(torch.version, "hip", None) is not None + or "rocm" in getattr(torch, "__version__", "").lower() + ) except Exception: is_rocm = False if is_rocm: hip_v = os.environ.get("HIP_VISIBLE_DEVICES") - rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") + # ROCR_VISIBLE_DEVICES is a Linux ROCr variable; Windows HIP has no + # ROCr layer, so a stray ROCR var there does not mask the runtime and + # must not be read as the ordinal->physical mapping (mirrors the + # Windows gate in _emit_child_gpu_visibility). + rocr_v = None if sys.platform == "win32" else os.environ.get("ROCR_VISIBLE_DEVICES") cvd = ( hip_v if hip_v is not None @@ -2882,20 +3064,52 @@ class LlamaCppBackend: return None @staticmethod - def _emit_child_gpu_visibility(env: dict, pinned: str) -> None: - """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on - ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child - seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP - mask at different layers, so the same indices apply twice -- ROCR reduces - and re-indexes from 0, then a non-zero HIP pin points out of range, HIP - enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone - narrows correctly; clear any inherited ROCR mask so it can't double up.""" + def _emit_child_gpu_visibility( + env: dict, + pinned: str, + *, + prefer_rocr: bool = False, + ) -> None: + """Write the child's GPU visibility mask: CUDA, plus a ROCm mirror on AMD + (masking only CUDA_VISIBLE_DEVICES leaves an AMD child seeing every GPU). + + Default: HIP_VISIBLE_DEVICES, clearing any inherited ROCR mask so the two + can't stack (ROCR re-indexes from 0, then a non-zero HIP pin points out of + range, HIP sees 0 devices, and llama.cpp falls back to CPU). + + prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask + filters only AFTER the HSA runtime enumerates every agent, and that + enumeration segfaults at startup on a GPU the build has no kernels for + (e.g. a gfx1103 iGPU under a gfx110X prebuilt), before llama-server logs a + line. ROCR drops the device at the driver layer, consuming physical ids. + The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps + the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a + Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin + would be dead there while the cleared HIP mask stops selecting.""" env["CUDA_VISIBLE_DEVICES"] = pinned try: import torch as _torch - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - env.pop("ROCR_VISIBLE_DEVICES", None) + + # torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may + # leave it unset but encode "rocm" in __version__ (mirrors detect_hardware). + if ( + getattr(_torch.version, "hip", None) is not None + or "rocm" in getattr(_torch, "__version__", "").lower() + ): + if prefer_rocr and pinned != "-1" and sys.platform != "win32": + env["ROCR_VISIBLE_DEVICES"] = pinned + env.pop("HIP_VISIBLE_DEVICES", None) + # ROCR re-indexes the visible agents from 0, and with HIP + # cleared HIP honours CUDA_VISIBLE_DEVICES -- so it must carry + # the post-ROCR ordinals (0..N-1), not the physical ids, else a + # non-zero pick points out of range and HIP sees 0 devices (the + # same stacking the default path avoids by clearing ROCR). + env["CUDA_VISIBLE_DEVICES"] = ",".join( + str(i) for i in range(len(pinned.split(","))) + ) + else: + env["HIP_VISIBLE_DEVICES"] = pinned + env.pop("ROCR_VISIBLE_DEVICES", None) except Exception as e: logger.debug("Failed to set ROCm visibility env vars for child: %s", e) @@ -2930,7 +3144,21 @@ class LlamaCppBackend: logger.debug("Could not read reported GPU order for split pin: %s", e) if order is None: order = sorted(inherited) - LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order)) + # Re-emit at the layer that produced the mapping. A parent masked only + # via ROCR_VISIBLE_DEVICES hides agents at the driver layer, and the + # default HIP re-emission clears that mask -- HSA then enumerates every + # agent again and can segfault at startup on an unsupported GPU the + # parent was hiding (the crash prefer_rocr exists to avoid). Linux-only, + # mirroring _resolve_visible_physical_ids: on Windows a stray ROCR var + # is dead and was not the mapping's source. + prefer_rocr = ( + sys.platform != "win32" + and env.get("HIP_VISIBLE_DEVICES") is None + and env.get("ROCR_VISIBLE_DEVICES") is not None + ) + LlamaCppBackend._emit_child_gpu_visibility( + env, ",".join(str(i) for i in order), prefer_rocr = prefer_rocr + ) @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: @@ -3592,6 +3820,9 @@ class LlamaCppBackend: lib_dirs.extend(_wsl_system_rocm_lib_dirs()) if lib_dirs: env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + # Native Linux AMD: system ROCm libs before the bundle's HIP runtime, + # which can be incompatible with the host amdkfd driver. + lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir)) lib_dirs.append(binary_dir) _arch = platform.machine() # x86_64, aarch64, etc. @@ -4434,6 +4665,14 @@ class LlamaCppBackend: LlamaCppBackend._gguf_skip_value(f, atype) return None + @classmethod + def _gguf_path_is_diffusion(cls, gguf_path: str, model_identifier: str) -> bool: + """Classify a downloaded GGUF without mutating the active backend.""" + probe = object.__new__(cls) + probe._model_identifier = model_identifier + probe._read_gguf_metadata(gguf_path) + return probe._is_diffusion + def _read_gguf_metadata(self, gguf_path: str) -> None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -4885,11 +5124,14 @@ class LlamaCppBackend: # the unload reset) so /status doesn't misreport TP and an identical # re-Apply doesn't reload against stale tensor-parallel state. self._tensor_parallel = False - # Record only the single device the runner actually uses (the lowest - # selected GPU, chosen above) -- not the whole pick. The diffusion runner - # is single-device, so echoing a multi-GPU list would misreport placement - # in /status and let a re-Apply dedup against GPUs the runner never used. + # The single-device runner records only the lowest selected GPU (chosen + # above), not the whole pick, and clears any explicit pin from a prior + # chat load; a multi-GPU list would misreport placement and mis-dedup. self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None + # The frontend prefers requested_gpu_ids when hydrating the picker. + # Diffusion uses only one device, so echo the collapsed effective pin, + # not unused members of the original request. + self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -4955,6 +5197,9 @@ class LlamaCppBackend: touching the shared one; defaults to the shared event. """ cancel_event = cancel_event if cancel_event is not None else self._cancel_event + from utils.hf_cache_settings import get_hf_cache_paths + + download_cache_dir = str(get_hf_cache_paths().hub_cache) try: import huggingface_hub # noqa: F401 -- presence check only except ImportError: @@ -5050,7 +5295,11 @@ class LlamaCppBackend: if not p.size: continue try: - cached_path = try_to_load_from_cache(hf_repo, p.path) + cached_path = try_to_load_from_cache( + hf_repo, + p.path, + cache_dir = download_cache_dir, + ) except Exception: cached_path = None if ( @@ -5061,6 +5310,7 @@ class LlamaCppBackend: hf_repo, p.path, expected_size = p.size, + cache_dir = download_cache_dir, ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: @@ -5074,12 +5324,8 @@ class LlamaCppBackend: total_download_bytes = max(0, total_bytes - already_cached_bytes) if total_download_bytes > 0: - cache_dir = os.environ.get( - "HF_HUB_CACHE", - str(Path.home() / ".cache" / "huggingface" / "hub"), - ) - Path(cache_dir).mkdir(parents = True, exist_ok = True) - free_bytes = shutil.disk_usage(cache_dir).free + Path(download_cache_dir).mkdir(parents = True, exist_ok = True) + free_bytes = shutil.disk_usage(download_cache_dir).free total_gb = total_download_bytes / (1024**3) free_gb = free_bytes / (1024**3) @@ -5097,7 +5343,7 @@ class LlamaCppBackend: # surface the disk shortfall for the requested variant. raise RuntimeError( f"Not enough disk space to download {gguf_filename}. " - f"Only {free_gb:.1f} GB free in {cache_dir}" + f"Only {free_gb:.1f} GB free in {download_cache_dir}" ) smaller = self._find_smallest_fitting_variant( hf_repo, @@ -5128,7 +5374,7 @@ class LlamaCppBackend: else: raise RuntimeError( f"Not enough disk space to download any variant. " - f"Only {free_gb:.1f} GB free in {cache_dir}" + f"Only {free_gb:.1f} GB free in {download_cache_dir}" ) except RuntimeError: raise @@ -5151,6 +5397,7 @@ class LlamaCppBackend: cancel_event = cancel_event, on_status = lambda m: logger.info(m), force_download = force, + cache_dir = download_cache_dir, ) for shard in gguf_extra_shards: if cancel_event.is_set(): @@ -5162,6 +5409,7 @@ class LlamaCppBackend: hf_token, cancel_event = cancel_event, force_download = force, + cache_dir = download_cache_dir, ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): @@ -5207,6 +5455,12 @@ class LlamaCppBackend: logger.info("Reusing cached %s: %s", label, cached) return cached + from utils.hf_cache_settings import get_hf_cache_paths + + companion_cache_dir = _hub_cache_dir_for_snapshot_path(near_path) or str( + get_hf_cache_paths().hub_cache + ) + if _hub_download_in_flight(hf_repo): logger.info("Skipping %s download while a hub download is active", label) return None @@ -5241,7 +5495,7 @@ class LlamaCppBackend: if target is None: try: from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(hf_repo): + for snap in _iter_hf_cache_snapshots(hf_repo, companion_cache_dir): rel_files = _gguf_snapshot_files(snap) target = pick(rel_files) if target is not None: @@ -5259,7 +5513,11 @@ class LlamaCppBackend: # hf_hub_download with hf_repo would miss the canonical file and silently # drop the companion. _cached_hf_snapshot_file scans every case variant. if _hf_env_offline(): - cached = _cached_hf_snapshot_file(hf_repo, target) + cached = _cached_hf_snapshot_file( + hf_repo, + target, + cache_dir = companion_cache_dir, + ) if cached: logger.info("Resolved %s from local HF cache: %s", label, cached) return cached @@ -5272,6 +5530,7 @@ class LlamaCppBackend: target, hf_token, cancel_event = cancel_event, + cache_dir = companion_cache_dir, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") @@ -5302,7 +5561,12 @@ class LlamaCppBackend: near_path = near_path, ) - def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: + def _cached_repo_mtp_drafter( + self, + hf_repo: str, + *, + cache_dir: Optional[str] = None, + ) -> Optional[str]: """A drafter already in this repo's local HF cache, reused offline when a fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all cached snapshots; else an existing ``MTP/`` copy (any precision -- the @@ -5312,7 +5576,12 @@ class LlamaCppBackend: roots: list[Path] = [] subdirs: list[Path] = [] - for snap in _iter_hf_cache_snapshots(hf_repo): # newest first + snapshots = ( + _iter_hf_cache_snapshots(hf_repo) + if cache_dir is None + else _iter_hf_cache_snapshots(hf_repo, cache_dir) + ) + for snap in snapshots: # newest first for f in sorted(_gguf_snapshot_files(snap)): if _is_companion_gguf_path(f) and "mmproj" not in f.lower(): (roots if "/" not in f else subdirs).append(snap / f) @@ -5365,7 +5634,10 @@ class LlamaCppBackend: # current cached file and refetch a changed one, so skip the probe here # rather than pair new weights with a stale draft. if _hf_env_offline(): - cached = self._cached_repo_mtp_drafter(hf_repo) + cached = self._cached_repo_mtp_drafter( + hf_repo, + cache_dir = _hub_cache_dir_for_snapshot_path(near_path), + ) if cached: logger.info(f"Reusing cached MTP drafter (offline): {cached}") return cached @@ -5783,6 +6055,24 @@ class LlamaCppBackend: and ("unknown" in text or "unsupported" in text or "not supported" in text) ) + @staticmethod + def _mmproj_retry_failure_message(*, projector_confirmed: bool, detail: str) -> str: + """User-facing error when the text-only --mmproj strip retry also fails. + + Confirmed projector-format mismatches keep the historical wording. + Bare signal crashes (common on some ROCm/driver paths) must not be + reported as "Vision projector incompatible" โ€” that misled #7302. + """ + if projector_confirmed: + return ( + "Vision projector incompatible with this llama.cpp " + "build, and the text-only retry also failed: " + detail + ) + return ( + "Vision model failed to start (llama-server crashed with " + "--mmproj), and the text-only retry also failed: " + detail + ) + @staticmethod def _output_has_nonprojector_diagnostic(output: str) -> bool: """True when the output already names a concrete non-projector cause (out @@ -5984,6 +6274,8 @@ class LlamaCppBackend: gpu_layers: int = -1, n_cpu_moe: int = 0, tensor_split: Optional[List[float]] = None, + # Explicit GPU placement pool (issue #7164). None/[] = auto-select; + # the fitter may pin the smallest subset of this pool that fits. gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused @@ -6081,15 +6373,63 @@ class LlamaCppBackend: self._cancel_event.clear() - # โ”€โ”€ Phase 1: kill old process (under lock, fast) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - with self._lock: - self._kill_process() - # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() is_vulkan_backend = self._is_vulkan_backend(binary) + # โ”€โ”€ Vulkan-ordinal preflight (BEFORE the Phase 1 kill) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # An explicit Vulkan pin the ggml probe never enumerated cannot be honored. + # Validate it ABOVE the kill so an invalid selection leaves the live model + # untouched: CUDA ids are range-checked at the route, but Vulkan ordinals are + # not, so a stale gpu_ids=[99] used to kill the server then 400, leaving + # nothing running (#7239). _get_gpu_memory needs only the binary (safe pre- + # download) and reuses the later fit's issubset logic. Guarded on a found + # Vulkan build + a pin so a deferred not-found stays deferred for diffusion. + if is_vulkan_backend and gpu_ids and binary: + _pf_wanted = {int(x) for x in gpu_ids} + _pf_probed = {g[0] for g in self._get_gpu_memory(binary)} + if not _pf_wanted.issubset(_pf_probed): + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {sorted(_pf_wanted)} not " + f"present. Available Vulkan devices: {sorted(_pf_probed)}." + ) + + # A remote uncached GGUF may only reveal that it needs the + # single-device diffusion runner after download. On Vulkan, an + # explicit gpu_ids request cannot be mapped from ggml ordinals to + # that runner's CUDA physical index. Download and classify the main + # file before killing the healthy server so this late rejection is + # non-destructive. The Phase 2 call below reuses this cached path. + _preflight_model_path = None + if is_vulkan_backend and gpu_ids and hf_repo: + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo + with _hf_offline_if_dns_dead(): + _preflight_model_path = self._download_gguf( + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + ) + if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier): + raise ValueError( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ) + + # โ”€โ”€ Phase 1: kill old process (under lock, fast) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + with self._lock: + self._kill_process() + # โ”€โ”€ Phase 2: download (NO lock held, so cancel can proceed) โ”€โ”€ # mtp_draft_path arrives set for local Gemma loads (detected # sibling); for -hf loads it's None here and resolved just below. @@ -6111,7 +6451,7 @@ class LlamaCppBackend: ) hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): - model_path = self._download_gguf( + model_path = _preflight_model_path or self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, hf_token = hf_token, @@ -6161,6 +6501,18 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: + # The diffusion runner pins its child by CUDA visibility mask, so a + # ggml Vulkan ordinal cannot be honored (wrong GPU / CPU fallback). + # Route and remote-download preflights reject before teardown; keep + # this as a final defense if classification ever disagrees. + if is_vulkan_backend and gpu_ids: + raise ValueError( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ) # Not a tensor/layer GGUF: clear any preserved-fallback flag from a # prior load (this path skips the command builder that clears it). self._layer_preserves_tensor_intent = False @@ -6381,6 +6733,12 @@ class LlamaCppBackend: # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). _layer_min_gpus = 1 + # An explicit Vulkan ordinal absent from the ggml probe cannot be + # honored; flag it in the fit and reject after the try (raising inside + # would be swallowed into the --fit-on fallback). Bound before the try. + _vulkan_explicit_unmatched = False + _vulkan_requested_ids: list[int] = [] + _vulkan_available_ordinals: list[int] = [] try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -6393,6 +6751,28 @@ class LlamaCppBackend: # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] + # Restrict the fit (and thus the layer plan + pin env) to the + # selected GPUs; fail-open if none match so a stale UI choice + # can't strand the load on CPU (issue #7164). + if gpu_ids: + # A Vulkan build indexes by ggml ordinal. An explicit ordinal + # absent from the probe can't be pinned, so reject after the try + # rather than fail-open onto a device the user didn't pick. + _wanted_ids = {int(x) for x in gpu_ids} + # Reject if ANY requested ordinal is absent, not only when none + # match: [0, 99] against {0, 1} silently drops 99. Comparing the + # full requested set (before filter narrows) still lets the fitter + # pick a valid subset later -- that is narrowing, not absence. + _probed_ordinals = {g[0] for g in gpus} + if is_vulkan_backend and not _wanted_ids.issubset(_probed_ordinals): + _vulkan_explicit_unmatched = True + _vulkan_requested_ids = sorted(_wanted_ids) + _vulkan_available_ordinals = sorted(_probed_ordinals) + # Restrict the probed pool to the selection; fail-open (keep the + # full pool) if none match so a stale UI choice can't strand the + # load on CPU (issue #7164). + _sel_gpus = [g for g in gpus if g[0] in _wanted_ids] + gpus = _sel_gpus if _sel_gpus else gpus total_by_idx = {idx: total for idx, _f, total in _gpu_mem} # GPU picker: restrict every mode to the chosen devices, so # auto selection only considers them and manual mask to @@ -7219,6 +7599,17 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # An unenumerated explicit Vulkan ordinal can't be pinned; fail loudly + # instead of fitting onto an unselected device. Clear the raw selection + # the early state-publish recorded so it never leaks into gpu_ids (#7239). + if _vulkan_explicit_unmatched: + self._gpu_ids = None + self._requested_gpu_ids = None + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {_vulkan_requested_ids} not " + f"present. Available Vulkan devices: {_vulkan_available_ordinals}." + ) + # GPU picker: when no narrower subset was chosen (manual, or # a failed/file-size selection), pin the whole picked set so the # model can't spill onto an unpicked GPU. @@ -7582,11 +7973,45 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) - # Vulkan pins via --device (a cmd arg, unlike the env-based - # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's - # last-wins parsing lets a user --device override Unsloth's pick. - if is_vulkan_backend and gpu_indices is not None: - cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # Vulkan pins via --device (a cmd arg), before user extras so a user + # --device wins. Fall back to raw ids when the fit did not narrow. + _vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None) + + # Record the pin actually applied (fit-narrowed gpu_indices, else the raw + # request) for the keep-warm loop, dedupe, and /status, so an explicit + # [0, 1] narrowed to [0] records [0] and /status never echoes an ordinal + # the child never saw. Auto selection (no gpu_ids) stays None (#7239). + if is_vulkan_backend: + # Only record an EXPLICIT Vulkan pin: an auto pick still narrows + + # pins below, but recording it would misreport an explicit pin and + # make dedupe miss the loaded server; mirrors the CUDA/ROCm branch. + self._gpu_ids = ( + sorted(int(x) for x in _vulkan_pin_ids) + if (gpu_ids and _vulkan_pin_ids) + else None + ) + elif gpu_ids: + # Physical pin: the fit-selected subset when the fit ran, else the raw + # user selection so an explicit choice is honoured even when the fit + # could not size the model. + _effective_pin_ids = ( + [int(x) for x in gpu_indices] + if gpu_indices is not None + else [int(x) for x in gpu_ids] + ) + self._gpu_ids = ( + sorted(int(x) for x in _effective_pin_ids) if _effective_pin_ids else None + ) + else: + self._gpu_ids = None + + # Also record the RAW requested pin (before the fit narrowed it). Load + # dedupe compares this so a [0, 1] narrowed to [0] and re-sent as [0, 1] + # still matches, while /status keeps echoing the effective pin (#7239). + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + + if is_vulkan_backend and _vulkan_pin_ids is not None: + cmd += LlamaCppBackend._vulkan_pin_args(_vulkan_pin_ids) # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Unsloth's auto-set flags. Already @@ -7655,10 +8080,10 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so - # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device - # (above), not here. + # Pin to selected GPU(s) (issue #7164; resolved above into gpu_indices). + # On ROCm, narrowing only CUDA_VISIBLE_DEVICES leaves the AMD child + # seeing the full set, so set HIP_VISIBLE_DEVICES too. Vulkan is pinned + # via --device (above), not here. # A deliberate zero-offload load with no GPU companions runs # entirely on CPU, yet a visible CUDA device still costs the child # ~0.5 GB (context + compute scratch) that the CPU-only @@ -7684,7 +8109,12 @@ class LlamaCppBackend: # default FASTEST_FIRST order (#5025). if gpu_ids: env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices)) + # Mask on AMD at the ROCr/HSA layer: HIP-only masking still + # enumerates every agent first, which segfaults on a deselected + # unsupported GPU (e.g. gfx1103 iGPU under a gfx110X prebuilt). + self._emit_child_gpu_visibility( + env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True + ) elif manual_tensor_split_emitted and not is_vulkan_backend: # A manual per-GPU ratio across ALL GPUs (no explicit pick, so # no CUDA_VISIBLE_DEVICES mask above): the UI built the @@ -8016,23 +8446,29 @@ class LlamaCppBackend: self._kill_process() # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). + _projector_msg = self._is_projector_incompatibility(out) + _signal_mmproj_guess = self._is_signal_crash( + _crash_rc + ) and not self._output_has_nonprojector_diagnostic(out) if ( launched_with_mmproj and not self._cancel_event.is_set() - and ( - self._is_projector_incompatibility(out) - or ( - self._is_signal_crash(_crash_rc) - and not self._output_has_nonprojector_diagnostic(out) - ) - ) + and (_projector_msg or _signal_mmproj_guess) ): - logger.warning( - "llama-server could not load this model's vision " - "projector (--mmproj). The installed llama.cpp build is " - "likely too old for it. Loading text-only for this " - "session; run 'unsloth studio update' to enable vision." - ) + if _projector_msg: + logger.warning( + "llama-server could not load this model's vision " + "projector (--mmproj). The installed llama.cpp build is " + "likely too old for it. Loading text-only for this " + "session; run 'unsloth studio update' to enable vision." + ) + else: + logger.warning( + "llama-server crashed while loading this model's vision " + "projector (--mmproj). Retrying text-only for this " + "session; if this persists, run 'unsloth studio update' " + "or check GPU/driver logs." + ) cmd = self._strip_mmproj_args(_last_spawn_cmd) # This retry bypasses _spawn_and_wait, so refresh the # launched-argv snapshot itself -- the zero-offload @@ -8046,14 +8482,30 @@ class LlamaCppBackend: # an OS-killed text-only retry still gets the OOM message. _retry_rc = self._process.poll() if self._process is not None else None self._kill_process() + # If the text-only retry ALSO hard-crashed (a signal, not + # OOM/timeout), the vision projector was never the cause: + # llama-server is faulting during GPU/driver init. Say so + # -- with the ROCm fix -- instead of blaming the mmproj. + if self._is_signal_crash(_retry_rc): + raise RuntimeError( + "llama-server crashed at startup on both the vision " + "and text-only attempts -- a GPU driver/runtime " + "initialization crash, not a model or vision-projector " + "problem. This often means an unsupported secondary " + "GPU; on AMD/ROCm, hide it with ROCR_VISIBLE_DEVICES " + "(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first " + "GPU) before launching Unsloth Studio." + ) + _retry_detail = self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + self._model_identifier, + _retry_rc, + ) raise RuntimeError( - "Vision projector incompatible with this llama.cpp " - "build, and the text-only retry also failed: " - + self._classify_llama_start_failure( - "\n".join(self._stdout_lines[-50:]), - gguf_path, - self._model_identifier, - _retry_rc, + self._mmproj_retry_failure_message( + projector_confirmed = _projector_msg, + detail = _retry_detail, ) ) else: @@ -8283,18 +8735,29 @@ class LlamaCppBackend: caps = self.probe_server_capabilities(binary) mtp_token = caps.get("mtp_token") if caps else None if not mtp_token: - logger.warning( - "Requested MTP speculative decoding but " - "llama-server lacks --spec-type mtp/draft-mtp; " - "run `unsloth studio update`. Loading without " - "speculative decoding." - ) + inconclusive = bool(caps.get("mtp_probe_inconclusive")) if caps else True + if inconclusive: + logger.info( + "Requested MTP speculative decoding but llama-server MTP " + "capability probe was inconclusive; loading without " + "speculative decoding." + ) + else: + logger.warning( + "Requested MTP speculative decoding but " + "llama-server lacks --spec-type mtp/draft-mtp; " + "run `unsloth studio update`. Loading without " + "speculative decoding." + ) # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins # over env) so the child matches the binary-capability gate and # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. flags.append("--spec-default") self._speculative_type = "default" - self._spec_fallback_reason = "binary_no_mtp" + if inconclusive: + self._spec_fallback_reason = None + else: + self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max" @@ -8560,16 +9023,10 @@ class LlamaCppBackend: ) ): return False - # A changed GPU pick must reload (compare order-insensitively; None/[] - # both mean automatic). The diffusion runner collapses a multi-GPU pick - # to its single lowest device, so self._gpu_ids holds just that device; - # normalize the request the same way, or a multi-GPU pick that resolves - # to the same device needlessly reloads. - if self._is_diffusion: - requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None - else: - requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None - if (self._gpu_ids or None) != requested_gpu_pick: + # A changed GPU pick must reload. Regular GGUF accepts either the raw + # requested placement pool or the effective status-echoed subset; + # diffusion compares its normalized single-device pick. + if not self.matches_gpu_ids(gpu_ids): return False # Compare on the canonical requested mode. With --spec-type in @@ -8627,6 +9084,7 @@ class LlamaCppBackend: current = list(self._extra_args) if self._extra_args is not None else [] if list(extra_args) != current: return False + self._record_matching_gpu_request(gpu_ids) return True def _classify_gpu_offload( @@ -8758,12 +9216,15 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + # GPU-pin state describes the active runner only; clear it so an explicit + # pin never leaks into the next (or diffusion) runner. + self._gpu_ids = None + self._requested_gpu_ids = None self._tensor_parallel = False self._gpu_memory_mode = "auto" self._gpu_layers = -1 self._n_cpu_moe = 0 self._tensor_split = None - self._gpu_ids = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py index 64ab38ec75..9e3eaeda3f 100644 --- a/studio/backend/core/inference/local_model_resolver.py +++ b/studio/backend/core/inference/local_model_resolver.py @@ -146,6 +146,7 @@ def _build_index() -> dict[str, _LocalGgufEntry]: _is_hidden_model, ) from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + from utils.hf_cache_settings import known_hf_hub_caches index: dict[str, _LocalGgufEntry] = {} seen_hf: set[str] = set() @@ -174,7 +175,12 @@ def _build_index() -> dict[str, _LocalGgufEntry]: except Exception as exc: logger.debug("auto-switch: ./models scan failed: %s", exc) try: - for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + for hf_dir in ( + *known_hf_hub_caches(), + _resolve_hf_cache_dir(), + legacy_hf_cache_dir(), + hf_default_cache_dir(), + ): found += _scan_hf_once(hf_dir) except Exception as exc: logger.debug("auto-switch: HF cache scan failed: %s", exc) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index e78c93b6f3..d19c67a01a 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -181,19 +181,27 @@ def _vlm_messages_have_tool_history(messages): ) -def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): +def _build_generation_stats( + prompt_n, + prompt_tps, + gen_n, + gen_tps, + cached_n = 0, +): """Map mlx stream stats onto the usage/timings shape llama-server emits.""" prompt_n = int(prompt_n or 0) gen_n = int(gen_n or 0) + cached_n = int(cached_n or 0) prompt_tps = float(prompt_tps or 0.0) gen_tps = float(gen_tps or 0.0) prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0 predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0 + total_prompt_n = prompt_n + cached_n return { "usage": { - "prompt_tokens": prompt_n, + "prompt_tokens": total_prompt_n, "completion_tokens": gen_n, - "total_tokens": prompt_n + gen_n, + "total_tokens": total_prompt_n + gen_n, }, "timings": { "prompt_n": prompt_n, @@ -204,11 +212,123 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): "predicted_ms": predicted_ms, "predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0, "predicted_per_second": gen_tps, - "cache_n": 0, + "cache_n": cached_n, }, } +PROMPT_CACHE_ENTRIES = 6 +PROMPT_CACHE_MEMORY_FRACTION = 0.15 +PROMPT_CACHE_FALLBACK_BYTES = 2 * 1024**3 + + +def _mlx_prompt_cache_api(): + try: + from mlx_lm.models.cache import ( + LRUPromptCache, + can_trim_prompt_cache, + make_prompt_cache, + trim_prompt_cache, + ) + except ImportError: + return None + return LRUPromptCache, make_prompt_cache, can_trim_prompt_cache, trim_prompt_cache + + +def _prompt_cache_max_bytes(recommended_gb = None): + override = os.environ.get("UNSLOTH_MLX_PROMPT_CACHE_BYTES") + if override: + try: + return max(int(override), 0) + except ValueError: + logger.warning("Ignoring non-integer UNSLOTH_MLX_PROMPT_CACHE_BYTES=%r", override) + if recommended_gb: + return int(recommended_gb * 1e9 * PROMPT_CACHE_MEMORY_FRACTION) + return PROMPT_CACHE_FALLBACK_BYTES + + +def _flatten_kv_entries(cache): + for entry in cache: + nested = getattr(entry, "caches", None) + if nested is None: + yield entry + else: + yield from _flatten_kv_entries(nested) + + +def _kv_prefix_coverage(cache): + covered = None + for entry in _flatten_kv_entries(cache): + offset = getattr(entry, "offset", None) + if offset is None: + return None + if getattr(entry, "start_position", 0): + return None + window = getattr(entry, "max_size", None) + if window is not None and offset > window: + return None + if covered is None: + covered = offset + elif covered != offset: + return None + return covered + + +class _MLXPromptCacheHistory: + def __init__(self, max_entries, max_bytes): + api = _mlx_prompt_cache_api() + if api is None: + raise RuntimeError("mlx-lm is too old for LRUPromptCache") + lru_cls, make, can_trim, trim = api + self._make_prompt_cache = make + self._can_trim = can_trim + self._trim = trim + self._max_bytes = max_bytes + self._lru = lru_cls(max_size = max_entries, max_bytes = max_bytes) + + def fetch(self, model, key, tokens): + cache, rest = self._lru.fetch_nearest_cache(key, list(tokens)) + if cache is not None: + if rest: + return cache, list(rest) + if self._can_trim(cache) and self._trim(cache, 1) == 1: + return cache, list(tokens[-1:]) + if len(tokens) > 1: + head = list(tokens[:-1]) + cache, rest = self._lru.fetch_nearest_cache(key, head) + if cache is not None: + covered = len(head) - len(rest) + return cache, list(tokens[covered:]) + return self._make_prompt_cache(model), list(tokens) + + def insert(self, key, tokens, cache): + # An over-budget entry evicts itself and every other conversation. + nbytes = sum(getattr(entry, "nbytes", 0) for entry in cache) + if nbytes > self._max_bytes: + logger.debug( + "MLX prompt cache: skipping %.2f GB entry over the %.2f GB budget", + nbytes / 1e9, + self._max_bytes / 1e9, + ) + return + covered = _kv_prefix_coverage(cache) + if covered is None: + logger.debug("MLX prompt cache: skipping cache with unverifiable prefix coverage") + return + tokens = list(tokens) + if covered > len(tokens): + logger.debug( + "MLX prompt cache: cache covers %d tokens but only %d were tracked", + covered, + len(tokens), + ) + return + tokens = tokens[:covered] + if not tokens: + return + self._lru.insert_cache(key, tokens, cache) + + def _mlx_distributed_rank_size(group = None): """Return ``(rank, world_size)`` for an optional MLX distributed group.""" if group is None: @@ -313,6 +433,55 @@ class MLXInferenceBackend: # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} + self._prompt_cache_history = None + self._prompt_cache_unavailable = False + + def _prompt_cache(self): + if self._prompt_cache_history is not None or self._prompt_cache_unavailable: + return self._prompt_cache_history + max_bytes = _prompt_cache_max_bytes(self._memory_limits_applied.get("recommended_gb")) + if max_bytes <= 0: + self._prompt_cache_unavailable = True + logger.info("MLX prompt cache disabled by budget") + return None + try: + self._prompt_cache_history = _MLXPromptCacheHistory( + PROMPT_CACHE_ENTRIES, + max_bytes, + ) + except Exception as exc: + self._prompt_cache_unavailable = True + logger.info("MLX prompt cache unavailable (%s); prefilling every request", exc) + return None + logger.info( + "MLX prompt cache: %d entries, %.2f GB budget", + PROMPT_CACHE_ENTRIES, + max_bytes / 1e9, + ) + return self._prompt_cache_history + + def _clear_prompt_cache(self): + self._prompt_cache_history = None + self._prompt_cache_unavailable = False + + def _prepare_prompt_cache(self, prompt, adapter_state): + history = self._prompt_cache() + if history is None: + return prompt, None, None, None, 0 + try: + tokenizer = self._tokenizer + bos = getattr(tokenizer, "bos_token", None) + add_special_tokens = bos is None or not prompt.startswith(bos) + tokens = list(tokenizer.encode(prompt, add_special_tokens = add_special_tokens)) + if not tokens: + return prompt, None, None, None, 0 + key = f"{self.active_model_name}|{adapter_state!r}" + cache, rest = history.fetch(self._model, key, tokens) + except Exception as exc: + logger.debug("MLX prompt cache lookup failed: %s", exc) + return prompt, None, None, None, 0 + return rest, cache, key, tokens, len(tokens) - len(rest) + def _configure_memory_limits(self): """Apply Metal memory caps before loading a model. @@ -535,6 +704,7 @@ class MLXInferenceBackend: self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None + self._clear_prompt_cache() gc.collect() mx.clear_cache() @@ -731,24 +901,34 @@ class MLXInferenceBackend: # prefix on every native-protocol snapshot just as the normal # decoding path does below. normalized_output = think_prefix - logger.info( - "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s", - len(prompt), - max_new_tokens, - type(self._model).__name__, - type(self._tokenizer).__name__, - ) with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): + ( + gen_prompt, + prompt_cache, + cache_key, + prompt_tokens, + cached_n, + ) = self._prepare_prompt_cache(prompt, _adapter_state) + logger.info( + "Generating: prompt_len=%d, cached=%d, max_tokens=%d, model=%s, tokenizer=%s", + len(prompt), + cached_n, + max_new_tokens, + type(self._model).__name__, + type(self._tokenizer).__name__, + ) final_response = None try: # Enter request-scoped model state before yielding any response. if think_prefix: yield think_prefix gen_kwargs = dict( - prompt = prompt, + prompt = gen_prompt, max_tokens = max_new_tokens, sampler = sampler, ) + if prompt_cache is not None: + gen_kwargs["prompt_cache"] = prompt_cache if logits_processors is not None: gen_kwargs["logits_processors"] = logits_processors for response in stream_generate( @@ -757,6 +937,7 @@ class MLXInferenceBackend: **gen_kwargs, ): final_response = response + token_ids.append(response.token) if preserve_native_channels: piece = getattr(response, "text", None) or "" delta = normalizer.feed(piece) @@ -764,7 +945,6 @@ class MLXInferenceBackend: normalized_output += delta yield normalized_output else: - token_ids.append(response.token) cumulative = self._tokenizer.decode( token_ids, skip_special_tokens = True, @@ -773,6 +953,13 @@ class MLXInferenceBackend: if cancel_event and cancel_event.is_set(): break + if prompt_cache is not None and prompt_tokens is not None: + history = self._prompt_cache_history + if history is not None: + try: + history.insert(cache_key, prompt_tokens + token_ids, prompt_cache) + except Exception as exc: + logger.debug("MLX prompt cache insert failed: %s", exc) except Exception as e: import traceback logger.error("stream_generate failed:\n%s", traceback.format_exc()) @@ -785,6 +972,7 @@ class MLXInferenceBackend: getattr(final_response, "prompt_tps", 0.0), getattr(final_response, "generation_tokens", 0), getattr(final_response, "generation_tps", 0.0), + cached_n, ) if normalizer is not None: cancelled = cancel_event is not None and cancel_event.is_set() diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 75ef9c2399..616384386d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -27,7 +27,7 @@ import uuid from io import BytesIO from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union -from utils.hardware import prepare_gpu_selection +from utils.hardware import get_device, prepare_gpu_selection # Re-exported from the shared helper so GGUF, training, and inference share one # type; kept importable here for backwards compatibility. @@ -217,10 +217,14 @@ class InferenceOrchestrator: native_path_secret_removed_for_child_start, run_without_native_path_secret, ) + from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths - from .worker import run_inference_process + cache_env = get_hf_cache_paths().child_env({}) - with native_path_secret_removed_for_child_start(): + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._cancel_event = _CTX.Event() @@ -228,7 +232,7 @@ class InferenceOrchestrator: self._proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_inference_process,), + args = ("core.inference.worker", "run_inference_process", cache_env), kwargs = { "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, @@ -1008,6 +1012,8 @@ class InferenceOrchestrator: ) sub_config["resolved_gpu_ids"] = resolved_gpu_ids sub_config["gpu_selection"] = gpu_selection + # Parent-detected backend for the worker's apply_gpu_ids(). + sub_config["device_backend"] = get_device().value # Recheck the sidecar reservation BEFORE tearing the old worker down, # for REPAIRS only: an install holds this same lifecycle gate, so it diff --git a/studio/backend/core/inference/stt_ggml_sidecar.py b/studio/backend/core/inference/stt_ggml_sidecar.py new file mode 100644 index 0000000000..02b376dea5 --- /dev/null +++ b/studio/backend/core/inference/stt_ggml_sidecar.py @@ -0,0 +1,876 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""whisper.cpp (GGML/GGUF) speech-to-text sidecar for Studio dictation. + +Runs the same curated Whisper checkpoints as the Transformers sidecar +(stt_sidecar.py) through whisper.cpp's `whisper-server`, ~2.5x faster at +identical quality on Apple Silicon and CPU because its Metal/CPU kernels run +the weights in f16 where PyTorch MPS requires fp32. + +Owns a single `whisper-server` subprocess bound to 127.0.0.1 on an ephemeral +port; the model loads on demand, stays warm between dictations, and unloads +after the same keep-alive as the Transformers sidecar. Curated GGML checkpoints +are single files from `unslothai/whisper-*-GGUF`, downloaded directly rather +than through the Model Hub (whose variant planner only handles `.gguf` chat +layouts). + +Binary discovery mirrors `_find_llama_server_binary`: env override, then managed +Studio home, then PATH. With no binary the engine is unavailable and dictation +falls back to the Transformers sidecar; `scripts/build_whisper_cpp.sh` installs +the binary. +""" + +from __future__ import annotations + +import io +import json +import os +import re +import shutil +import socket +import subprocess +import sys +import threading +import time +import urllib.request +import uuid +import wave +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator, Optional + +from loggers import get_logger + +from core.inference.stt_sidecar import ( + STT_KEEP_ALIVE_SECONDS, + SttAudioDecodeError, + SttLanguageError, + SttLoadCancelledError, + SttModelIdError, + SttModelNotDownloadedError, + SttUnavailableError, + _decode_audio_bounded, + _known_whisper_languages, + _TARGET_SAMPLE_RATE, + _training_active, + normalize_whisper_language, +) +from utils.prebuilt.child_env import isolate_home, scrub_env, wsl_system_rocm_lib_dirs +from utils.prebuilt.runtime_libs import dedupe_existing_dirs +from utils.prebuilt.whisper_layout import lookup_marker +from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid + +logger = get_logger(__name__) + +# Curated GGML checkpoints, one repo per model. Keys match the Transformers +# sidecar's ids so the frontend reuses one picker; values are the single file +# inside each repo. +GGML_STT_REPOS: dict[str, str] = { + "tiny": "unslothai/whisper-tiny-GGUF", + "base": "unslothai/whisper-base-GGUF", + "small": "unslothai/whisper-small-GGUF", + "large-v3-turbo": "unslothai/whisper-large-v3-turbo-GGUF", + "large-v3": "unslothai/whisper-large-v3-GGUF", +} +GGML_STT_MODELS: dict[str, str] = { + "tiny": "whisper-tiny.bin", + "base": "whisper-base.bin", + "small": "whisper-small.bin", + "large-v3-turbo": "whisper-large-v3-turbo.bin", + "large-v3": "whisper-large-v3.bin", +} +DEFAULT_GGML_STT_MODEL = "small" + +_SERVER_START_TIMEOUT_SECONDS = 120.0 +_TRANSCRIBE_TIMEOUT_SECONDS = 600.0 + + +class SttEngineUnavailableError(SttUnavailableError): + """whisper-server is not installed; the GGUF dictation engine is off.""" + + +def resolve_ggml_model_id(model: Optional[str]) -> str: + """Validate a curated GGML model id. Custom repos are not supported here.""" + if model is None or not str(model).strip(): + return DEFAULT_GGML_STT_MODEL + normalized = str(model).strip() + if normalized in GGML_STT_MODELS: + return normalized + raise SttModelIdError( + f"STT model '{model}' is not a curated GGUF dictation model. " + f"Choose one of: {', '.join(GGML_STT_MODELS)}." + ) + + +def _managed_whisper_cpp_dir() -> Path: + """`/whisper.cpp` in custom mode, else `~/.unsloth/whisper.cpp`. + + Mirrors `managed_node_dir` / `_find_llama_server_binary` so managed runtimes + share one parent directory. + """ + legacy = Path.home() / ".unsloth" / "whisper.cpp" + try: + from utils.paths.storage_roots import studio_root + + resolved = studio_root() + legacy_studio = Path.home() / ".unsloth" / "studio" + try: + is_legacy = resolved.resolve() == legacy_studio.resolve() + except (OSError, ValueError): + is_legacy = resolved == legacy_studio + return legacy if is_legacy else (resolved / "whisper.cpp") + except (ImportError, OSError, ValueError): + override = ( + os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") or "" + ).strip() + if override: + try: + return Path(override).expanduser().resolve() / "whisper.cpp" + except (OSError, ValueError): + return Path(override).expanduser() / "whisper.cpp" + return legacy + + +def find_whisper_server_binary() -> Optional[str]: + """Locate the whisper-server binary. + + Search order: + 1. WHISPER_SERVER_PATH environment variable (direct path to binary) + 2. UNSLOTH_WHISPER_CPP_PATH env var (custom whisper.cpp install dir) + 3. managed dir: /whisper.cpp/{,build/bin/}whisper-server + 4. whisper-server on PATH + """ + binary_name = "whisper-server.exe" if sys.platform == "win32" else "whisper-server" + + def _layout_candidates(d: Path) -> list[Path]: + cands = [d / binary_name, d / "build" / "bin" / binary_name] + if sys.platform == "win32": + cands.append(d / "build" / "bin" / "Release" / binary_name) + return cands + + env_path = os.environ.get("WHISPER_SERVER_PATH") + if env_path: + p = Path(env_path) + if _is_runnable(p): + return str(p) + + custom_dir = os.environ.get("UNSLOTH_WHISPER_CPP_PATH") + if custom_dir: + for p in _layout_candidates(Path(custom_dir)): + if _is_runnable(p): + return str(p) + + for p in _layout_candidates(_managed_whisper_cpp_dir()): + if _is_runnable(p): + return str(p) + + return shutil.which(binary_name) + + +def _is_runnable(p: Path) -> bool: + """A real whisper-server is an executable file. On Windows os.access(X_OK) is + effectively an existence check; on Unix it rejects a non-executable stub so a + half-written or wrong-mode file isn't mistaken for the server.""" + return p.is_file() and (sys.platform == "win32" or os.access(p, os.X_OK)) + + +def _whisper_install_marker(binary: str) -> Optional[dict]: + """The prebuilt install marker above ``binary``, or None (source/custom builds).""" + return lookup_marker(binary).marker + + +def slim_runtime_intact(binary: str) -> bool: + """True unless the marker says slim and the linked ggml runtime is missing + beside the server. New markers record the exact wired filenames + (linked_libraries), all of which must be present; legacy markers without the + field fall back to the per-OS core ggml name globs. A broken slim install + reads as engine-unavailable (reinstall via `unsloth studio update`), never a + crash at load.""" + lookup = lookup_marker(binary) + marker = lookup.marker + if lookup.invalid or marker is None: + return not lookup.slim_collision + if not marker or marker.get("install_kind") != "slim": + return True + if lookup.authoritative: + valid = marker.get("component") == "whisper.cpp" + valid = valid and isinstance(marker.get("schema_version"), int) + valid = valid and all( + isinstance(marker.get(key), str) and marker[key] + for key in ("release_tag", "backend", "paired_llama_tag") + ) + valid = valid and isinstance(marker.get("linked_libraries"), list) + valid = valid and bool(marker.get("linked_libraries")) + valid = valid and all( + isinstance(name, str) and name and Path(name).name == name + for name in marker["linked_libraries"] + ) + if not valid: + return False + bin_dir = Path(binary).parent + linked = marker.get("linked_libraries") + if isinstance(linked, list) and linked and all(isinstance(name, str) for name in linked): + intact = all((bin_dir / name).is_file() for name in linked) + else: + if sys.platform == "win32": + required = ("ggml.dll", "ggml-base.dll") + elif sys.platform == "darwin": + required = ("libggml*.dylib", "libggml-base*.dylib") + else: + required = ("libggml.so*", "libggml-base.so*") + intact = all(any(p.is_file() for p in bin_dir.glob(pattern)) for pattern in required) + runtime_dirs = marker.get("linked_runtime_directories") + if intact and isinstance(runtime_dirs, list) and runtime_dirs: + intact = all( + isinstance(name, str) + and name + and (bin_dir / name).is_dir() + and any(path.is_file() for path in (bin_dir / name).rglob("*")) + for name in runtime_dirs + ) + if intact and marker.get("backend") == "rocm": + expected_runtime_dirs = set() if sys.platform == "win32" else {"hipblaslt", "rocblas"} + intact = ( + marker.get("runtime_wiring_version") == 2 + and isinstance(runtime_dirs, list) + and set(runtime_dirs) == expected_runtime_dirs + ) + if not intact: + logger.warning( + "slim whisper install is missing its linked ggml runtime at " + f"{bin_dir}; run `unsloth studio update` to reinstall it" + ) + return intact + + +def is_available() -> bool: + binary = find_whisper_server_binary() + if binary is None: + return False + if not slim_runtime_intact(binary): + return False + try: + import av # noqa: F401 + except Exception: + # No PyAV means every transcription 501s on decode. + return False + return True + + +def ensure_engine_available() -> str: + binary = find_whisper_server_binary() + if binary is None: + raise SttEngineUnavailableError( + "The local transcription runtime is not installed. Run " + "`unsloth studio update` to install it." + ) + if not slim_runtime_intact(binary): + raise SttEngineUnavailableError( + "The local transcription runtime is missing its paired ggml " + "libraries. Run `unsloth studio update` to reinstall it." + ) + return binary + + +# --------------------------------------------------------------------------- +# whisper-server child-process environment +# --------------------------------------------------------------------------- +# Build the whisper-server env: prepend the binary dir (co-located libs win, and +# a backstop where the loader ignores the rpath) and scrub secret-bearing vars the +# binary never needs. On WSL2 ROCm the system HIP libs go first, since a bundle's +# bare-metal HIP cannot drive /dev/dxg. A CUDA bundle ships libggml-cuda.so but not +# libcudart/libcublas (paired with the user's PyTorch), so add the +# CUDA-from-PyTorch runtime dirs the selection gated on, else the backend cannot +# resolve a runtime that lives only in wheels. Mirrors llama's binary_env(); the +# scrub/WSL/dedupe helpers live in utils.prebuilt. + +# Module-level aliases keep the historical patch points for tests and callers. +_wsl_system_rocm_lib_dirs = wsl_system_rocm_lib_dirs +_dedupe_existing_dirs = dedupe_existing_dirs + + +def _whisper_server_child_env(binary: str) -> dict[str, str]: + """Env for the whisper-server subprocess: secrets scrubbed, home/profile vars + repointed at a managed scratch dir (a downloaded binary must not see the real + home's token caches), co-located libs on the loader path, WSL system HIP first + on WSL2 ROCm.""" + env = scrub_env(os.environ) + isolate_home(env, str(_managed_whisper_cpp_dir() / ".child_home")) + bin_dir = str(Path(binary).parent) + # A CUDA bundle needs the CUDA-from-PyTorch wheel dirs so libcudart/libcublas + # resolve at launch when they live only in site-packages/nvidia/*/lib. Placed + # after bin_dir so co-located libs still win; empty for other bundles. + cuda_runtime_dirs: list[str] = [] + bundle_dir = Path(bin_dir) + has_cuda_module = any( + path.is_file() + for pattern in ("libggml-cuda.so*", "ggml-cuda*.dll") + for path in bundle_dir.glob(pattern) + ) + if has_cuda_module: + try: + from utils.prebuilt.runtime_libs import python_runtime_dirs + cuda_runtime_dirs = python_runtime_dirs() + except Exception: + cuda_runtime_dirs = [] + if sys.platform == "win32": + var, lead = "PATH", [bin_dir, *cuda_runtime_dirs] + elif sys.platform == "darwin": + var, lead = "DYLD_LIBRARY_PATH", [bin_dir] + else: + var, lead = "LD_LIBRARY_PATH", [bin_dir, *cuda_runtime_dirs] + wsl_rocm = _wsl_system_rocm_lib_dirs() + if wsl_rocm: + lead = [*wsl_rocm, bin_dir, *cuda_runtime_dirs] + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + existing = [p for p in env.get(var, "").split(os.pathsep) if p] + env[var] = os.pathsep.join(_dedupe_existing_dirs([*lead, *existing])) + return env + + +# --------------------------------------------------------------------------- +# Model file download (single files; deliberately outside the Model Hub flow) +# --------------------------------------------------------------------------- + + +def _cached_model_path(model_id: str) -> Optional[str]: + """Path of a fully downloaded GGML file in the shared HF cache, else None.""" + from huggingface_hub import hf_hub_download + try: + return hf_hub_download( + repo_id = GGML_STT_REPOS[model_id], + filename = GGML_STT_MODELS[model_id], + local_files_only = True, + ) + except Exception: + return None + + +class _GgmlDownloadState: + """Tracks one background hf_hub_download of a curated GGML file.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._thread: Optional[threading.Thread] = None + self._model_id: Optional[str] = None + self._error: Optional[str] = None + self._total_bytes: Optional[int] = None + self._etag: Optional[str] = None + + def status(self) -> dict: + with self._lock: + downloading = self._thread is not None and self._thread.is_alive() + return { + "downloading": downloading, + "model": self._model_id if downloading else None, + "error": self._error, + "bytes_total": self._total_bytes if downloading else None, + "bytes_done": self._incomplete_bytes() if downloading else None, + } + + def _incomplete_bytes(self) -> Optional[int]: + """Best-effort progress: size of the in-flight blob in the HF cache. + + hf_hub_download writes ``blobs/.incomplete``; prefer this file's + etag, else the largest in-flight blob. + """ + try: + from huggingface_hub.constants import HF_HUB_CACHE + + # Caller may hold the non-reentrant self._lock; bare reads are safe. + model_id = self._model_id + if not model_id: + return None + repo_dir = ( + Path(HF_HUB_CACHE) + / f"models--{GGML_STT_REPOS[model_id].replace('/', '--')}" + / "blobs" + ) + if not repo_dir.is_dir(): + return None + etag = self._etag + if etag: + target = repo_dir / f"{etag}.incomplete" + if target.is_file(): + return target.stat().st_size + sizes = [p.stat().st_size for p in repo_dir.glob("*.incomplete") if p.is_file()] + return max(sizes) if sizes else None + except Exception: + return None + + def start( + self, + model_id: str, + hf_token: Optional[str] = None, + ) -> None: + model_id = resolve_ggml_model_id(model_id) + with self._lock: + if self._thread is not None and self._thread.is_alive(): + if self._model_id == model_id: + return + raise SttModelIdError( + f"Another GGUF dictation model ('{self._model_id}') is still " + "downloading; wait for it to finish." + ) + self._model_id = model_id + self._error = None + self._total_bytes = None + self._etag = None + thread = threading.Thread(target = self._run, args = (model_id, hf_token), daemon = True) + self._thread = thread + thread.start() + + def _run(self, model_id: str, hf_token: Optional[str]) -> None: + repo_id = GGML_STT_REPOS[model_id] + filename = GGML_STT_MODELS[model_id] + try: + from huggingface_hub import ( + get_hf_file_metadata, + hf_hub_download, + hf_hub_url, + ) + try: + # One HEAD request for the total and etag. + meta = get_hf_file_metadata(hf_hub_url(repo_id, filename), token = hf_token or None) + with self._lock: + self._total_bytes = meta.size + self._etag = meta.etag + except Exception: + pass + hf_hub_download( + repo_id = repo_id, + filename = filename, + token = hf_token or None, + ) + except Exception as exc: + logger.warning("GGUF STT download failed for %s: %s", model_id, exc) + with self._lock: + self._error = f"Download failed for '{model_id}'." + + +_download_state = _GgmlDownloadState() + + +def start_model_download(model: Optional[str], hf_token: Optional[str] = None) -> None: + _download_state.start(resolve_ggml_model_id(model), hf_token) + + +def download_status() -> dict: + return _download_state.status() + + +# --------------------------------------------------------------------------- +# WAV packaging +# --------------------------------------------------------------------------- + + +def _pcm_to_wav_bytes(decoded_audio) -> bytes: + """Wrap decoded float32 mono 16 kHz PCM into an in-memory 16-bit WAV.""" + import numpy as np + + clipped = np.clip(decoded_audio, -1.0, 1.0) + pcm16 = (clipped * 32767.0).astype(" None: + self._lock = threading.RLock() + self._process: Optional[subprocess.Popen] = None + self._port: Optional[int] = None + self._model_id: Optional[str] = None + self._idle_timer: Optional[threading.Timer] = None + self._idle_generation = 0 + self._keep_alive_seconds = keep_alive_seconds + # Set while whisper-server starts so training admission can account for + # the accelerator memory it is about to bind. Read without the lock. + self._loading = False + # A still-starting whisper-server is cancellable so training can preempt + # it before it binds accelerator memory. Assigned inside self._lock but + # acted on without it: cancel_pending_load() runs while load() holds the + # lock, so the event is the source of truth and terminating the process + # is a best-effort fast path. + self._load_cancel_event: Optional[threading.Event] = None + self._starting_process: Optional[subprocess.Popen] = None + # Set before the updater waits for _lock, then kept set while it owns + # the lock and atomically replaces the managed install tree. New loads + # fail fast instead of starting a process from files being swapped. + self._update_in_progress = False + + @property + def loaded_model(self) -> Optional[str]: + # Lock-free status read (like stt_sidecar.py): transcribe() holds + # self._lock for the whole inference call (up to + # _TRANSCRIBE_TIMEOUT_SECONDS), and status polls plus training admission + # must not block behind it. _process_alive() snapshots self._process + # before poll(), which subprocess guards with _waitpid_lock, so a + # concurrent unload is safe. + return self._model_id if self._process_alive() else None + + @property + def device(self) -> Optional[str]: + return "whisper.cpp" if self._process_alive() else None + + def is_loading(self) -> bool: + # True only while whisper-server is starting (seconds to bind its GPU + # backend); load() sets and clears the flag around that window. + return self._loading + + @property + def keep_alive_seconds(self) -> float: + return self._keep_alive_seconds + + def _process_alive(self) -> bool: + # Snapshot self._process once: a concurrent unload() nulls it under the + # lock, so lock-free readers would otherwise re-read None between the + # truthiness check and .poll(). + process = self._process + return process is not None and process.poll() is None + + # -- idle unload ------------------------------------------------------ + + def _cancel_idle_unload_locked(self) -> None: + self._idle_generation += 1 + if self._idle_timer is not None: + self._idle_timer.cancel() + self._idle_timer = None + + def _schedule_idle_unload_locked(self) -> None: + self._cancel_idle_unload_locked() + if not self._process_alive(): + return + generation = self._idle_generation + timer = threading.Timer(self._keep_alive_seconds, self._idle_unload, args = (generation,)) + timer.daemon = True + self._idle_timer = timer + timer.start() + + def _idle_unload(self, generation: int) -> None: + with self._lock: + if generation != self._idle_generation: + return + logger.info("Unloading idle GGUF STT model %s", self._model_id) + self._release_locked() + + # -- process lifecycle ------------------------------------------------- + + def _release_locked(self) -> None: + self._cancel_idle_unload_locked() + process = self._process + self._process = None + self._port = None + self._model_id = None + if process is not None and process.poll() is None: + process.terminate() + try: + process.wait(timeout = 10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout = 10) + if process is not None: + forget_pid(process.pid) + + def unload(self) -> None: + with self._lock: + self._release_locked() + + def _raise_if_update_in_progress(self) -> None: + if self._update_in_progress: + raise SttEngineUnavailableError( + "The local transcription runtime is being updated. Try dictation again shortly." + ) + + @contextmanager + def update_maintenance(self) -> Iterator[bool]: + """Block new loads while the managed whisper.cpp tree is replaced. + + The flag is published before waiting for an existing transcription to + release ``_lock``. Holding that lock across the yielded installer phase + prevents Windows from relocking the executable and prevents every host + from starting a process against a partially swapped tree. The yielded + value records whether a warm model had to be unloaded. + """ + self._update_in_progress = True + try: + with self._lock: + model_was_active = self._process_alive() + self._release_locked() + yield model_was_active + finally: + self._update_in_progress = False + + def cancel_pending_load(self) -> bool: + # Preempt a starting whisper-server so training does not launch while it + # binds accelerator memory. load() holds self._lock for the whole startup, + # so act without the lock: signal abort and terminate the starting + # process. _wait_for_server observes the event and raises, then load() + # reaps the process and releases the lock. + if not self._loading: + return False + event = self._load_cancel_event + if event is None: + return False + event.set() + process = self._starting_process + if process is not None and process.poll() is None: + try: + process.terminate() + except Exception: + pass + return True + + def wait_for_load_to_settle(self) -> None: + # load() holds self._lock across startup and cancel cleanup, so acquiring + # it blocks until a cancelled server is killed, reaped, and its + # accelerator memory released. + with self._lock: + pass + + @staticmethod + def _reserve_free_port() -> tuple[socket.socket, int]: + """Bind an ephemeral port and keep the socket held. + + The caller closes the reservation immediately before spawning + whisper-server, shrinking the window in which another local process + could bind the port. SO_REUSEADDR lets the child rebind right after. + """ + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("127.0.0.1", 0)) + return s, s.getsockname()[1] + + def _ensure_model_downloaded(self, model_id: str) -> str: + path = _cached_model_path(model_id) + if path is None: + raise SttModelNotDownloadedError( + f"STT model '{model_id}' (GGUF) is not downloaded. " + "Download it in Settings, then Voice, before loading it." + ) + return path + + def load(self, model: Optional[str] = None) -> None: + """Start (or switch) whisper-server for the requested curated model.""" + self._raise_if_update_in_progress() + model_id = resolve_ggml_model_id(model) + with self._lock: + self._raise_if_update_in_progress() + binary = ensure_engine_available() + if self._process_alive() and self._model_id == model_id: + self._schedule_idle_unload_locked() + return + model_path = self._ensure_model_downloaded(model_id) + self._release_locked() + reservation, port = self._reserve_free_port() + command = [binary, "-m", model_path, "--host", "127.0.0.1", "--port", str(port)] + marker = _whisper_install_marker(binary) + if _training_active(): + # Keep whisper.cpp off the accelerator during training (like the + # Transformers sidecar's CPU choice) so a mid-training dictation + # cannot reclaim the VRAM training just freed. + command.append("--no-gpu") + elif marker is not None and marker.get("backend") == "cpu": + # A deliberate CPU install must stay CPU: the slim wiring links + # every llama ggml backend (including CUDA/ROCm), so without + # this flag a cpu-selected install would still grab the GPU. + command.append("--no-gpu") + logger.info( + "Starting whisper-server for STT model %s on 127.0.0.1:%s", + model_id, + port, + ) + cancel_event = threading.Event() + self._load_cancel_event = cancel_event + self._loading = True + try: + # Release the reservation as late as possible: whisper-server + # binds the port moments after this close. + reservation.close() + process = subprocess.Popen( + command, + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, + stdin = subprocess.DEVNULL, + # Co-located GPU libs on the loader path (WSL system HIP first), + # secrets scrubbed from the downloaded binary's env. + env = _whisper_server_child_env(binary), + # Die with Studio (Linux PDEATHSIG, Windows job) so a crash + # never orphans a server holding the model. + **child_popen_kwargs(), + ) + self._starting_process = process + adopt_pid(process.pid) # terminate_all backstop for graceful exits + try: + self._wait_for_server(process, port, cancel_event) + except Exception: + if process.poll() is None: + process.kill() + process.wait(timeout = 10) + forget_pid(process.pid) + raise + self._process = process + self._port = port + self._model_id = model_id + self._schedule_idle_unload_locked() + finally: + reservation.close() # no-op when already released before spawn + self._loading = False + self._load_cancel_event = None + self._starting_process = None + + @staticmethod + def _wait_for_server( + process: subprocess.Popen, + port: int, + cancel_event: Optional[threading.Event] = None, + ) -> None: + deadline = time.monotonic() + _SERVER_START_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if cancel_event is not None and cancel_event.is_set(): + raise SttLoadCancelledError( + "GGUF STT model loading was cancelled so training could start." + ) + if process.poll() is not None: + raise SttEngineUnavailableError( + "The local transcription runtime exited before becoming " + "ready; the model file may be corrupt or unsupported." + ) + # Require a whisper-server-specific response twice, with the managed + # child alive around each probe. An arbitrary local process that won + # the bind race would otherwise be mistaken for the sidecar and + # receive the user's microphone audio. + if GgmlSttSidecar._probe_is_whisper_server(process, port) and ( + GgmlSttSidecar._probe_is_whisper_server(process, port) + ): + return + time.sleep(0.2) + raise SttEngineUnavailableError("The local transcription runtime did not start in time.") + + @staticmethod + def _probe_is_whisper_server(process: subprocess.Popen, port: int) -> bool: + """One readiness probe: our child is alive and the responder looks like + whisper.cpp's server (its index page and errors identify whisper).""" + if process.poll() is not None: + return False + try: + req = urllib.request.Request(f"http://127.0.0.1:{port}/", method = "GET") + with urllib.request.urlopen(req, timeout = 2) as response: + body = response.read(65536) + except Exception: + return False + if process.poll() is not None: + return False + return b"whisper" in body.lower() + + # -- transcription ------------------------------------------------------ + + def transcribe( + self, + audio: bytes, + model: Optional[str] = None, + language: Optional[str] = None, + fast: bool = False, + ) -> dict: + """Transcribe encoded audio bytes via whisper-server. + + Accepts any container PyAV can decode (same validation and caps as the + Transformers sidecar). Returns {text, language, duration, model}. + """ + self._raise_if_update_in_progress() + ensure_engine_available() + model_id = resolve_ggml_model_id(model) + lang = normalize_whisper_language(language) + known_languages = _known_whisper_languages() + if lang is not None and known_languages is not None and lang not in known_languages: + raise SttLanguageError( + f"Language '{language}' is not supported by STT model '{model_id}'." + ) + # Reject a missing model before decoding so a long clip does not burn CPU + # only to 409 (matches the Transformers sidecar's preflight). + self._ensure_model_downloaded(model_id) + decoded_audio = _decode_audio_bounded(audio) + wav_bytes = _pcm_to_wav_bytes(decoded_audio) + with self._lock: + try: + self.load(model_id) + text = self._post_inference(wav_bytes, lang, fast) + finally: + self._schedule_idle_unload_locked() + duration = (len(decoded_audio) / _TARGET_SAMPLE_RATE) if len(decoded_audio) else None + return { + "text": text, + "language": lang, + "duration": duration, + "model": model_id, + } + + def _post_inference(self, wav_bytes: bytes, lang: Optional[str], fast: bool) -> str: + boundary = uuid.uuid4().hex + fields = { + "temperature": "0.0", + "response_format": "json", + # Match the Transformers sidecar: 5-way beam search, greedy for fast. + "beam_size": "1" if fast else "5", + "language": lang or "auto", + } + parts: list[bytes] = [] + for name, value in fields.items(): + parts.append( + ( + f"--{boundary}\r\nContent-Disposition: form-data; " + f'name="{name}"\r\n\r\n{value}\r\n' + ).encode() + ) + parts.append( + ( + f"--{boundary}\r\nContent-Disposition: form-data; " + 'name="file"; filename="dictation.wav"\r\n' + "Content-Type: audio/wav\r\n\r\n" + ).encode() + + wav_bytes + + b"\r\n" + ) + parts.append(f"--{boundary}--\r\n".encode()) + body = b"".join(parts) + req = urllib.request.Request( + f"http://127.0.0.1:{self._port}/inference", + data = body, + headers = {"Content-Type": f"multipart/form-data; boundary={boundary}"}, + ) + try: + with urllib.request.urlopen(req, timeout = _TRANSCRIBE_TIMEOUT_SECONDS) as resp: + payload = json.load(resp) + except SttAudioDecodeError: + raise + except Exception as exc: + raise SttEngineUnavailableError( + "The local transcription runtime did not answer the request." + ) from exc + text = payload.get("text") + if not isinstance(text, str): + raise SttAudioDecodeError("Could not decode the audio.") + # whisper.cpp joins segments with newlines; dictation wants one line. + return " ".join(part.strip() for part in text.splitlines() if part.strip()).strip() + + +_sidecar: Optional[GgmlSttSidecar] = None + + +def get_ggml_stt_sidecar() -> GgmlSttSidecar: + global _sidecar + if _sidecar is None: + _sidecar = GgmlSttSidecar() + return _sidecar diff --git a/studio/backend/core/inference/stt_sidecar.py b/studio/backend/core/inference/stt_sidecar.py new file mode 100644 index 0000000000..edf57c16e3 --- /dev/null +++ b/studio/backend/core/inference/stt_sidecar.py @@ -0,0 +1,1142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Standalone speech-to-text (STT) sidecar for dictation. + +Loads a Whisper model (via Transformers) in the backend process, separate from +the chat model's inference subprocess, so dictation works with any chat model +without evicting it. Curated defaults plus any Transformers-compatible Whisper +repo; weights come through Studio's Model Hub and stay warm briefly between +dictations. CUDA runs float16; MPS and CPU run float32. +""" + +from __future__ import annotations + +import gc +import hashlib +import io +import json +import os +import re +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +# Multilingual Whisper defaults: stable API/UI id -> Hub repository. A request +# may instead pass a validated Hugging Face `owner/model` id. +STT_MODELS: dict[str, str] = { + "tiny": "unsloth/whisper-tiny", + "base": "unsloth/whisper-base", + "small": "unsloth/whisper-small", + "large-v3-turbo": "unsloth/whisper-large-v3-turbo", + "large-v3": "unsloth/whisper-large-v3", +} +DEFAULT_STT_MODEL = "small" +STT_KEEP_ALIVE_SECONDS = 5 * 60 +_HF_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$") +_HF_COMMIT_SHA = re.compile(r"^[0-9a-f]{40}$") + +# Bound decoded PCM length so a crafted upload cannot exhaust memory (callers +# also cap the encoded bytes). +_MAX_AUDIO_SECONDS = 30 * 60 +_TARGET_SAMPLE_RATE = 16000 + +# Non-weight files WhisperProcessor/WhisperForConditionalGeneration may load. +# Weight selection is built from pinned Hub metadata. A custom repo id is +# attacker-controllable, so only safetensors weights are accepted: a +# pytorch_model.bin is a pickle and executes code while Transformers +# deserializes it (see utils/security/file_security.py), and this path skips +# the malware gate the normal model loader applies. +_STT_SNAPSHOT_SUPPORT_FILES = ( + "config.json", + "generation_config.json", + "preprocessor_config.json", + "processor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "normalizer.json", + "special_tokens_map.json", + "added_tokens.json", +) +_STT_SAFETENSORS_INDEX = "model.safetensors.index.json" +_STT_SAFETENSORS_WEIGHTS = "model.safetensors" +_STT_REVISION_RECORD_VERSION = 1 + + +@dataclass(frozen = True) +class _SelectedHubFile: + path: str + size: int + blob_key: Optional[str] + + +@dataclass(frozen = True) +class _CachedSttSnapshot: + path: Optional[Path] + is_multilingual: Optional[bool] + + +class SttUnavailableError(RuntimeError): + """The STT backend (PyTorch/Transformers or PyAV) is not installed.""" + + +class SttLoadCancelledError(RuntimeError): + """An in-flight STT model load was cancelled for training.""" + + +class SttModelNotDownloadedError(RuntimeError): + """The selected model is not complete in the shared Hub cache.""" + + +class SttModelIdError(ValueError): + """The requested custom model is not a valid Hugging Face repository id.""" + + +class SttModelCompatibilityError(ValueError): + """The requested repository is not a Transformers Whisper checkpoint.""" + + +class SttAudioDecodeError(ValueError): + """The uploaded bytes could not be decoded as audio.""" + + +class SttAudioTooLongError(ValueError): + """The decoded audio exceeds the bounded transcription duration.""" + + +class SttLanguageError(ValueError): + """The requested language is not supported by the selected STT model.""" + + +_WHISPER_LANGUAGE_ALIASES = { + # Legacy/browser BCP-47 primaries whose Whisper code differs. + "cmn": "zh", + "fil": "tl", + "in": "id", + "iw": "he", + "ji": "yi", + "nb": "no", + "nn": "no", +} + + +def normalize_whisper_language(language: Optional[str]) -> Optional[str]: + """Convert a BCP-47 locale into the short code Whisper expects.""" + if not language: + return None + normalized = language.strip().replace("_", "-").lower() + if not normalized or normalized == "auto": + return None + primary = normalized.split("-", 1)[0] + return _WHISPER_LANGUAGE_ALIASES.get(primary, primary) + + +def _known_whisper_languages() -> Optional[frozenset[str]]: + """Return Whisper's language codes without constructing/loading a model.""" + try: + from transformers.models.whisper.tokenization_whisper import LANGUAGES + except Exception: + # Transformers unavailable or the constant moved: skip the check. + return None + return frozenset(LANGUAGES) + + +def ensure_stt_available() -> None: + """Raise when the complete local Whisper backend cannot be imported.""" + try: + import av # noqa: F401 + import torch # noqa: F401 + import transformers # noqa: F401 + except Exception as exc: + raise SttUnavailableError( + "Speech-to-text needs PyTorch, Transformers, and PyAV. " + "Run `unsloth studio update` to install them." + ) from exc + + +def is_available() -> bool: + """True when the complete local Whisper backend can be imported.""" + try: + ensure_stt_available() + except SttUnavailableError: + return False + return True + + +def resolve_model_id(model: Optional[str]) -> str: + """Resolve a curated id or validate a custom Hugging Face repository.""" + if not model: + return DEFAULT_STT_MODEL + normalized = model.strip() + if normalized in STT_MODELS: + return normalized + if _HF_REPO_ID.fullmatch(normalized): + return normalized + raise SttModelIdError( + "STT model must be one of Studio's defaults or a Hugging Face " + "repository in 'owner/model' form." + ) + + +def resolve_model_repo(model_id: str) -> str: + """Return the Hub repository for a curated or custom model id.""" + resolved = resolve_model_id(model_id) + return STT_MODELS.get(resolved, resolved) + + +def _is_whisper_config(config: object) -> bool: + """True when Hub/local config metadata identifies a Whisper ASR model.""" + if not isinstance(config, dict): + return False + model_type = config.get("model_type") + if isinstance(model_type, str) and model_type.strip().lower() == "whisper": + return True + architectures = config.get("architectures") + return isinstance(architectures, list) and any( + isinstance(name, str) and name == "WhisperForConditionalGeneration" + for name in architectures + ) + + +def _read_json_object(path: Path) -> dict: + try: + with open(path, "r", encoding = "utf-8") as file: + value = json.load(file) + return value if isinstance(value, dict) else {} + except Exception: + return {} + + +def _active_hf_hub_cache() -> Path: + """Return the active Hub cache while respecting runtime test overrides.""" + explicit = (os.environ.get("HF_HUB_CACHE") or "").strip() + if explicit: + return Path(explicit).expanduser() + hf_home = (os.environ.get("HF_HOME") or "").strip() + if hf_home: + return Path(hf_home).expanduser() / "hub" + from huggingface_hub.constants import HF_HUB_CACHE + + return Path(HF_HUB_CACHE) + + +def _repo_cache_dir(repo: str) -> Path: + return _active_hf_hub_cache() / f"models--{repo.replace('/', '--')}" + + +def _revision_record_path(repo: str) -> Path: + from utils.paths.storage_roots import cache_root + digest = hashlib.sha256(repo.encode("utf-8")).hexdigest() + return cache_root() / "stt-revisions" / f"{digest}.json" + + +def _write_revision_record(repo: str, revision: str) -> None: + """Persist immutable identity only, never an HF-cache absolute path.""" + if not _HF_COMMIT_SHA.fullmatch(revision): + return + path = _revision_record_path(repo) + tmp = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex[:8]}") + try: + path.parent.mkdir(parents = True, exist_ok = True) + with tmp.open("w", encoding = "utf-8") as handle: + json.dump( + { + "version": _STT_REVISION_RECORD_VERSION, + "repo": repo, + "revision": revision, + }, + handle, + ) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + except OSError as exc: + logger.debug("Could not persist STT revision for %s: %s", repo, exc) + try: + tmp.unlink(missing_ok = True) + except OSError: + pass + + +def _read_revision_record(repo: str) -> Optional[str]: + payload = _read_json_object(_revision_record_path(repo)) + if payload.get("version") != _STT_REVISION_RECORD_VERSION or payload.get("repo") != repo: + return None + revision = payload.get("revision") + return revision if isinstance(revision, str) and _HF_COMMIT_SHA.fullmatch(revision) else None + + +def _safe_snapshot_for_revision(repo: str, revision: str) -> Optional[Path]: + """Resolve a canonical SHA below this repository's active snapshots dir.""" + if not _HF_COMMIT_SHA.fullmatch(revision): + return None + snapshots = _repo_cache_dir(repo) / "snapshots" + candidate = snapshots / revision + try: + snapshots_resolved = snapshots.resolve() + candidate_resolved = candidate.resolve() + except (OSError, RuntimeError): + return None + if snapshots_resolved not in candidate_resolved.parents or not candidate_resolved.is_dir(): + return None + return candidate_resolved + + +def _snapshot_usable(model_id: str, snapshot: Path) -> bool: + if not _snapshot_is_complete(snapshot): + return False + if model_id not in STT_MODELS: + return _is_whisper_config(_read_json_object(snapshot / "config.json")) + return True + + +def _find_complete_cached_snapshot(model: Optional[str]) -> Optional[Path]: + """Find one complete local snapshot without contacting the Hub.""" + model_id = resolve_model_id(model) + repo = resolve_model_repo(model_id) + + recorded = _read_revision_record(repo) + if recorded: + snapshot = _safe_snapshot_for_revision(repo, recorded) + if snapshot is not None and _snapshot_usable(model_id, snapshot): + return snapshot + + ref = _repo_cache_dir(repo) / "refs" / "main" + try: + revision = ref.read_text(encoding = "utf-8").strip() + except OSError: + revision = "" + snapshot = _safe_snapshot_for_revision(repo, revision) + if snapshot is not None and _snapshot_usable(model_id, snapshot): + _write_revision_record(repo, revision) + return snapshot + + snapshots = _repo_cache_dir(repo) / "snapshots" + try: + revisions = sorted( + ( + (path.stat().st_mtime_ns, path.name) + for path in snapshots.iterdir() + if path.is_dir() and _HF_COMMIT_SHA.fullmatch(path.name) + ), + reverse = True, + ) + except OSError: + return None + for _mtime, revision in revisions: + snapshot = _safe_snapshot_for_revision(repo, revision) + if snapshot is not None and _snapshot_usable(model_id, snapshot): + _write_revision_record(repo, revision) + return snapshot + return None + + +def _selected_file_from_sibling(sibling) -> _SelectedHubFile: + lfs = getattr(sibling, "lfs", None) + blob_key = getattr(lfs, "sha256", None) or getattr(sibling, "blob_id", None) + return _SelectedHubFile( + path = sibling.rfilename, + size = max(0, int(getattr(sibling, "size", 0) or 0)), + blob_key = blob_key if isinstance(blob_key, str) and blob_key else None, + ) + + +def _select_snapshot_files(info, load_index) -> tuple[_SelectedHubFile, ...]: + """Select support files and one complete safetensors weight set. Pickle + (pytorch_model.bin) weights are never selected: they are an RCE sink on a + custom repo id (see _STT_SNAPSHOT_SUPPORT_FILES).""" + siblings = { + sibling.rfilename: sibling + for sibling in (getattr(info, "siblings", None) or []) + if isinstance(getattr(sibling, "rfilename", None), str) + } + selected = {name for name in _STT_SNAPSHOT_SUPPORT_FILES if name in siblings} + + index_name: Optional[str] = None + if _STT_SAFETENSORS_INDEX in siblings: + index_name = _STT_SAFETENSORS_INDEX + elif _STT_SAFETENSORS_WEIGHTS in siblings: + selected.add(_STT_SAFETENSORS_WEIGHTS) + else: + raise SttModelCompatibilityError( + "The STT repository has no safetensors model weights. Only safetensors " + "checkpoints are supported; convert the model with save_pretrained(safe_serialization=True)." + ) + + if index_name is not None: + weight_map = load_index(index_name).get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise SttModelCompatibilityError(f"Invalid checkpoint index '{index_name}'.") + shards = set(weight_map.values()) + if not all(isinstance(shard, str) and shard in siblings for shard in shards): + raise SttModelCompatibilityError(f"Checkpoint index '{index_name}' has missing shards.") + # The index JSON is attacker-controlled: a safetensors index can name + # pytorch_model-*.bin shards, which Transformers still loads through + # torch.load (pickle) since it dispatches per shard by file extension. + # Require every shard to be safetensors so no pickle file is selected. + if not all(shard.endswith(".safetensors") for shard in shards): + raise SttModelCompatibilityError( + f"Checkpoint index '{index_name}' references non-safetensors shards." + ) + selected.add(index_name) + selected.update(shards) + + return tuple(_selected_file_from_sibling(siblings[name]) for name in sorted(selected)) + + +def validate_remote_model(model: Optional[str], hf_token: Optional[str] = None) -> dict: + """Verify a custom Hub repository is Whisper-compatible without downloading weights.""" + model_id = resolve_model_id(model) + repo = resolve_model_repo(model_id) + if model_id in STT_MODELS: + return {"model": model_id, "repo": repo} + + try: + from huggingface_hub import HfApi + info = HfApi(token = hf_token or False).model_info( + repo, + expand = ["config", "sha"], + timeout = 10, + ) + except Exception as exc: + raise SttModelCompatibilityError( + f"Could not verify STT model '{model_id}'. " + "Check that the repository exists and your Hugging Face token can access it." + ) from exc + + if not _is_whisper_config(getattr(info, "config", None)): + raise SttModelCompatibilityError( + f"STT model '{model_id}' is not a compatible Transformers Whisper model." + ) + revision = getattr(info, "sha", None) + if not isinstance(revision, str) or not _HF_COMMIT_SHA.fullmatch(revision): + raise SttModelCompatibilityError( + f"Could not resolve an immutable revision for STT model '{model_id}'." + ) + # The commit that was validated; the download pins to it so the repo cannot + # be swapped between validation and snapshot_download (TOCTOU). + return {"model": model_id, "repo": repo, "revision": revision} + + +def _is_missing_local_model_error(exc: BaseException) -> bool: + """Recognize a local-cache-only miss by name/message, without importing HF + internals (tolerates huggingface_hub/Transformers moving the exception).""" + current: Optional[BaseException] = exc + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + if type(current).__name__ in ("LocalEntryNotFoundError", "EntryNotFoundError"): + return True + message = str(current).lower() + if "local_files_only" in message or "does not appear to have a file" in message: + return True + current = current.__cause__ or current.__context__ + return False + + +def _snapshot_is_complete(snapshot: Path) -> bool: + """True when a cached snapshot holds every file loading needs. + + An aborted download can leave only metadata behind, and an offline lookup + cannot know the repo's full file list, so verify config, preprocessor, + tokenizer, and weights directly. is_file() follows cache symlinks, so a + link from an interrupted blob download does not count. + """ + # Safetensors only: a cached pytorch_model.bin is a pickle load path and is + # never treated as a usable snapshot (a repo shipping only pickle weights + # re-resolves and fails closed in _select_snapshot_files). + index = snapshot / _STT_SAFETENSORS_INDEX + if index.is_file(): + # Sharded safetensors checkpoint: every shard must exist and be + # safetensors (a safe index naming .bin shards would still pickle-load + # them, matching the _select_snapshot_files guard). + weight_map = _read_json_object(index).get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + return False + shards = set(weight_map.values()) + if not all(isinstance(shard, str) and shard.endswith(".safetensors") for shard in shards): + return False + has_weights = all((snapshot / shard).is_file() for shard in shards) + else: + has_weights = (snapshot / _STT_SAFETENSORS_WEIGHTS).is_file() + # WhisperProcessor needs the tokenizer: either the fast tokenizer.json or + # the slow vocab.json + merges.txt pair. + has_tokenizer = (snapshot / "tokenizer.json").is_file() or ( + (snapshot / "vocab.json").is_file() and (snapshot / "merges.txt").is_file() + ) + return ( + has_weights + and has_tokenizer + and (snapshot / "config.json").is_file() + and (snapshot / "preprocessor_config.json").is_file() + ) + + +def is_model_downloaded(model: Optional[str]) -> bool: + """True when a usable Whisper snapshot exists in the local HF cache.""" + try: + return _find_complete_cached_snapshot(model) is not None + except Exception: + return False + + +class _SnapshotDownloadState: + """Tracks one background snapshot_download of a dictation repository. + + Like stt_ggml_sidecar's tracker, but a Transformers checkpoint is a whole + repo, so progress is the byte count of its cache blobs. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._thread: Optional[threading.Thread] = None + self._model_id: Optional[str] = None + self._repo: Optional[str] = None + self._error: Optional[str] = None + self._total_bytes: Optional[int] = None + self._selected_files: tuple[_SelectedHubFile, ...] = () + self._complete = False + + def status(self) -> dict: + with self._lock: + downloading = self._thread is not None and self._thread.is_alive() + show_progress = downloading or self._complete + return { + "downloading": downloading, + "model": self._model_id if downloading else None, + "error": self._error, + "bytes_total": self._total_bytes if show_progress else None, + "bytes_done": self._blob_bytes() if show_progress else None, + } + + def _blob_bytes(self) -> Optional[int]: + """Best-effort progress: bytes in the repo's HF cache blobs. + + Counts only the selected support files and one selected weight format, + including in-progress ``.incomplete`` blobs. + """ + try: + # Caller may hold the non-reentrant self._lock; a bare read is safe. + repo = self._repo + selected_files = self._selected_files + if not repo or not selected_files: + return None + blobs = _repo_cache_dir(repo) / "blobs" + if not blobs.is_dir(): + return 0 + done = 0 + for selected in selected_files: + if not selected.blob_key: + continue + complete = blobs / selected.blob_key + incomplete = blobs / f"{selected.blob_key}.incomplete" + candidate = complete if complete.is_file() else incomplete + if candidate.is_file(): + done += min(candidate.stat().st_size, selected.size) + total = self._total_bytes + return min(done, total) if total is not None else done + except Exception: + return None + + def start( + self, + model_id: str, + hf_token: Optional[str] = None, + revision: Optional[str] = None, + ) -> None: + model_id = resolve_model_id(model_id) + with self._lock: + if self._thread is not None and self._thread.is_alive(): + if self._model_id == model_id: + return + raise SttModelIdError( + f"Another dictation model ('{self._model_id}') is still " + "downloading; wait for it to finish." + ) + self._model_id = model_id + self._repo = resolve_model_repo(model_id) + self._error = None + self._total_bytes = None + self._selected_files = () + self._complete = False + thread = threading.Thread( + target = self._run, args = (self._repo, hf_token, revision), daemon = True + ) + self._thread = thread + thread.start() + + def _run( + self, + repo: str, + hf_token: Optional[str], + revision: Optional[str] = None, + ) -> None: + try: + from huggingface_hub import HfApi, hf_hub_download, snapshot_download + + info = HfApi(token = hf_token or None).model_info( + repo, + revision = revision, + files_metadata = True, + timeout = 30, + ) + if not revision: + revision = getattr(info, "sha", None) + if not isinstance(revision, str) or not _HF_COMMIT_SHA.fullmatch(revision): + raise SttModelCompatibilityError( + f"Could not resolve an immutable revision for STT model '{repo}'." + ) + + def load_index(filename: str) -> dict: + path = hf_hub_download( + repo_id = repo, + filename = filename, + revision = revision, + token = hf_token or None, + ) + return _read_json_object(Path(path)) + + selected_files = _select_snapshot_files(info, load_index) + total = sum(selected.size for selected in selected_files) + with self._lock: + self._selected_files = selected_files + self._total_bytes = total or None + snapshot = Path( + snapshot_download( + repo_id = repo, + revision = revision, + allow_patterns = [selected.path for selected in selected_files], + token = hf_token or None, + ) + ) + if not _snapshot_is_complete(snapshot): + raise SttModelCompatibilityError( + f"Downloaded STT snapshot for '{repo}' is incomplete." + ) + _write_revision_record(repo, revision) + with self._lock: + self._complete = True + except Exception as exc: + logger.warning("STT snapshot download failed for %s: %s", repo, exc) + with self._lock: + self._error = f"Download failed for '{repo}'." + + +_download_state = _SnapshotDownloadState() + + +def start_model_download( + model: Optional[str], + hf_token: Optional[str] = None, + revision: Optional[str] = None, +) -> None: + _download_state.start(resolve_model_id(model), hf_token, revision = revision) + + +def download_status() -> dict: + return _download_state.status() + + +def _training_active() -> bool: + try: + from core.training import get_training_backend + return bool(get_training_backend().is_training_active()) + except Exception: + return False + + +def _clear_device_cache(device: Optional[str]) -> None: + gc.collect() + try: + import torch + if device == "cuda": + torch.cuda.empty_cache() + elif device == "mps": + torch.mps.empty_cache() + except Exception: + pass + + +def _pick_device(): + """Return (device, torch_dtype) for the Whisper model. + + CUDA uses float16. MPS and CPU use float32: Whisper's decoder is unstable in + float16 on MPS and degenerates into repeated tokens. + """ + try: + import torch + + # New loads use CPU during training; a resident GPU model may stay put + # when the training admission check confirms enough headroom. + training_active = _training_active() + if not training_active and torch.cuda.is_available(): + return "cuda", torch.float16 + if ( + not training_active + and getattr(torch.backends, "mps", None) is not None + and torch.backends.mps.is_available() + ): + return "mps", torch.float32 + return "cpu", torch.float32 + except Exception as exc: + logger.debug("STT device detection failed, using CPU: %s", exc) + import torch + return "cpu", torch.float32 + + +def _decode_audio_bounded(audio: bytes): + """Decode to 16 kHz mono PCM without buffering unbounded audio. + + A small, highly-compressed upload can expand far past the encoded request + limit once decoded, so decode frame-by-frame and enforce the sample cap as + frames arrive, then hand the array straight to Whisper. + """ + try: + import av + import numpy as np + from av.error import FFmpegError, InvalidDataError + except ImportError as exc: + raise SttUnavailableError( + "Speech-to-text needs the PyAV package to decode audio. " + "Run `unsloth studio update` to install it." + ) from exc + + max_samples = _MAX_AUDIO_SECONDS * _TARGET_SAMPLE_RATE + sample_count = 0 + raw_buffer = io.BytesIO() + resampler = av.audio.resampler.AudioResampler( + format = "s16", + layout = "mono", + rate = _TARGET_SAMPLE_RATE, + ) + # Group frames before resampling so short clips need one resampler call + # rather than one per codec frame. + fifo = av.audio.fifo.AudioFifo() + + def write_frame(frame) -> None: + nonlocal sample_count + array = frame.to_ndarray() + sample_count += array.size + if sample_count > max_samples: + max_minutes = _MAX_AUDIO_SECONDS // 60 + unit = "minute" if max_minutes == 1 else "minutes" + raise SttAudioTooLongError(f"Audio must be {max_minutes} {unit} or shorter.") + raw_buffer.write(array) + + try: + with av.open(io.BytesIO(audio), mode = "r", metadata_errors = "ignore") as container: + if not container.streams.audio: + raise SttAudioDecodeError("Could not decode the audio.") + frames = iter(container.decode(audio = 0)) + while True: + try: + frame = next(frames) + except StopIteration: + break + except InvalidDataError: + # Skip a corrupt frame rather than fail the whole transcription. + continue + frame.pts = None + fifo.write(frame) + if fifo.samples >= 500000: + for resampled in resampler.resample(fifo.read()): + write_frame(resampled) + if fifo.samples > 0: + for resampled in resampler.resample(fifo.read()): + write_frame(resampled) + for resampled in resampler.resample(None): + write_frame(resampled) + except (SttAudioDecodeError, SttAudioTooLongError): + raise + except (FFmpegError, ValueError, RuntimeError) as exc: + raise SttAudioDecodeError("Could not decode the audio.") from exc + finally: + del fifo, resampler + + if sample_count == 0: + raise SttAudioDecodeError("Could not decode the audio.") + decoded = np.frombuffer(raw_buffer.getbuffer(), dtype = np.int16).astype(np.float32) + decoded /= 32768.0 + return decoded + + +class WhisperSttSidecar: + """Lazily loaded Whisper model with idle eviction. Thread-safe.""" + + def __init__(self, keep_alive_seconds: float = STT_KEEP_ALIVE_SECONDS) -> None: + self._engine = None + self._model_id: Optional[str] = None + self._device: Optional[str] = None + self._lock = threading.RLock() + self._load_state_lock = threading.Lock() + self._loading = False + self._load_cancel_event: Optional[threading.Event] = None + self._keep_alive_seconds = max(0.0, keep_alive_seconds) + self._idle_timer: Optional[threading.Timer] = None + self._idle_generation = 0 + + @property + def loaded_model(self) -> Optional[str]: + return self._model_id + + @property + def device(self) -> Optional[str]: + return self._device + + def is_loading(self) -> bool: + with self._load_state_lock: + return self._loading + + def cancel_pending_load(self) -> bool: + """Cancel a model load without waiting for the model lock.""" + with self._load_state_lock: + event = self._load_cancel_event + if not self._loading or event is None: + return False + event.set() + return True + + def wait_for_load_to_settle(self) -> None: + """Block until any in-flight load() has exited and freed its memory. + + load() holds self._lock throughout, including the from_pretrained()/ + .to(device) allocation and cancel cleanup, so acquiring the lock here + waits for that memory to be freed. + """ + with self._lock: + pass + + def _begin_load(self) -> threading.Event: + event = threading.Event() + with self._load_state_lock: + self._load_cancel_event = event + self._loading = True + return event + + def _end_load(self, event: threading.Event) -> None: + with self._load_state_lock: + if self._load_cancel_event is event: + self._load_cancel_event = None + self._loading = False + + @staticmethod + def _raise_if_load_cancelled(event: threading.Event) -> None: + if event.is_set(): + raise SttLoadCancelledError("STT model loading was cancelled so training could start.") + + @property + def keep_alive_seconds(self) -> float: + return self._keep_alive_seconds + + def _cancel_idle_unload_locked(self) -> None: + self._idle_generation += 1 + timer = self._idle_timer + self._idle_timer = None + if timer is not None: + timer.cancel() + + def _schedule_idle_unload_locked(self) -> None: + self._cancel_idle_unload_locked() + if self._engine is None or self._keep_alive_seconds <= 0: + return + generation = self._idle_generation + timer = threading.Timer( + self._keep_alive_seconds, + self._idle_unload, + args = (generation,), + ) + timer.daemon = True + self._idle_timer = timer + timer.start() + + def _idle_unload(self, generation: int) -> None: + with self._lock: + if generation != self._idle_generation or self._engine is None: + return + logger.info("Unloading idle STT model %s", self._model_id) + self._release_engine_locked() + + def _release_engine_locked(self) -> None: + self._cancel_idle_unload_locked() + engine = self._engine + device = self._device + self._engine = None + self._model_id = None + self._device = None + del engine + _clear_device_cache(device) + + def _build_model(self, snapshot_path: str, device: str, dtype, cancel_event: threading.Event): + """Load a Whisper model + processor from the local Hub cache. + + local_files_only keeps the Model Hub the only download path; a cache + miss raises so the caller can surface SttModelNotDownloadedError. + """ + import torch + from transformers import WhisperForConditionalGeneration, WhisperProcessor + + processor = None + model = None + try: + processor = WhisperProcessor.from_pretrained(snapshot_path, local_files_only = True) + self._raise_if_load_cancelled(cancel_event) + # use_safetensors forces the pickle-free load path even if a + # pytorch_model.bin somehow reached the cache; the selector and the + # completeness check already exclude pickle weights upstream. + model = WhisperForConditionalGeneration.from_pretrained( + snapshot_path, torch_dtype = dtype, local_files_only = True, use_safetensors = True + ) + self._raise_if_load_cancelled(cancel_event) + model.to(torch.device(device)) + self._raise_if_load_cancelled(cancel_event) + model.eval() + return model, processor + except SttLoadCancelledError: + model = None + processor = None + _clear_device_cache(device) + raise + + def _ensure_model_downloaded(self, model_id: str) -> _CachedSttSnapshot: + """Validate the local snapshot before decode or model replacement. + + Returns the checkpoint's multilingual flag when local metadata provides + it. Curated defaults are known multilingual. + """ + model_id = resolve_model_id(model_id) + with self._lock: + if self._engine is not None and self._model_id == model_id: + resident_model = ( + self._engine[0] if isinstance(self._engine, (tuple, list)) else self._engine + ) + generation_config = getattr(resident_model, "generation_config", None) + is_multilingual = getattr(generation_config, "is_multilingual", None) + return _CachedSttSnapshot( + path = None, + is_multilingual = is_multilingual if isinstance(is_multilingual, bool) else None, + ) + snapshot_path = _find_complete_cached_snapshot(model_id) + if snapshot_path is None: + raise SttModelNotDownloadedError( + f"STT model '{model_id}' is not downloaded. " + "Download it in Settings, then Voice, before loading it." + ) + + if model_id in STT_MODELS: + return _CachedSttSnapshot(path = snapshot_path, is_multilingual = True) + + if not _is_whisper_config(_read_json_object(snapshot_path / "config.json")): + raise SttModelCompatibilityError( + f"STT model '{model_id}' is not a compatible Transformers Whisper model." + ) + generation_config = _read_json_object(snapshot_path / "generation_config.json") + is_multilingual = generation_config.get("is_multilingual") + if isinstance(is_multilingual, bool): + return _CachedSttSnapshot(path = snapshot_path, is_multilingual = is_multilingual) + if resolve_model_repo(model_id).lower().endswith(".en"): + return _CachedSttSnapshot(path = snapshot_path, is_multilingual = False) + return _CachedSttSnapshot(path = snapshot_path, is_multilingual = None) + + def load(self, model: Optional[str] = None): + """Load (or switch to) a model, reusing it if already resident. + + Returns a ``(model, processor)`` pair. + """ + model_id = resolve_model_id(model) + with self._lock: + ensure_stt_available() + if self._engine is not None and self._model_id == model_id: + self._schedule_idle_unload_locked() + return self._engine + import torch + + cancel_event = self._begin_load() + candidate = None + device: Optional[str] = None + try: + cached = self._ensure_model_downloaded(model_id) + snapshot_path = cached.path + if snapshot_path is None: + raise SttModelNotDownloadedError( + f"STT model '{model_id}' is not downloaded. " + "Download it in Settings, then Voice, before loading it." + ) + self._raise_if_load_cancelled(cancel_event) + device, dtype = _pick_device() + self._release_engine_locked() + logger.info("Loading STT model %s (%s) on %s", model_id, snapshot_path, device) + + def not_downloaded(cause: BaseException) -> SttModelNotDownloadedError: + return SttModelNotDownloadedError( + f"STT model '{model_id}' is not downloaded. " + "Download it in Settings, then Voice, before loading it." + ) + + retry_on_cpu = False + try: + candidate = self._build_model(str(snapshot_path), device, dtype, cancel_event) + self._raise_if_load_cancelled(cancel_event) + except SttLoadCancelledError: + raise + except Exception as exc: + if _is_missing_local_model_error(exc): + raise not_downloaded(exc) from exc + if device == "cpu": + raise + logger.warning("STT load on %s failed (%s); retrying on CPU", device, exc) + retry_on_cpu = True + if retry_on_cpu: + # Retry outside the handler: live exception state pins frames + # referencing the partly loaded model, so leave it before + # clearing the cache to release that memory. + _clear_device_cache(device) + try: + candidate = self._build_model( + str(snapshot_path), + "cpu", + torch.float32, + cancel_event, + ) + self._raise_if_load_cancelled(cancel_event) + except SttLoadCancelledError: + raise + except Exception as cpu_exc: + if _is_missing_local_model_error(cpu_exc): + raise not_downloaded(cpu_exc) from cpu_exc + raise + device = "cpu" + with self._load_state_lock: + self._raise_if_load_cancelled(cancel_event) + self._engine = candidate + self._model_id = model_id + self._device = device + self._load_cancel_event = None + self._loading = False + self._schedule_idle_unload_locked() + logger.info("STT model %s ready on %s", model_id, device) + return self._engine + except SttLoadCancelledError: + candidate = None + self._release_engine_locked() + _clear_device_cache(device) + raise + finally: + self._end_load(cancel_event) + + def _transcribe_decoded(self, model_id: str, decoded_audio, generate_kwargs: dict) -> str: + """Run Whisper on already-decoded 16 kHz mono PCM and return text. + + Feeds a pre-decoded array so nothing here touches the Transformers audio + path (torchcodec/ffmpeg). Splits into 30s windows (Whisper's receptive + field); short clips take one pass. + """ + import torch + + model, processor = self.load(model_id) + effective_generate_kwargs = dict(generate_kwargs) + generation_config = getattr(model, "generation_config", None) + if getattr(generation_config, "is_multilingual", None) is False: + # English-only checkpoints fix language and task in their generation + # config, and Transformers rejects passing them here. + effective_generate_kwargs.pop("task", None) + effective_generate_kwargs.pop("language", None) + window = 30 * _TARGET_SAMPLE_RATE + target_dtype = getattr(model, "dtype", None) + parts: list[str] = [] + with torch.no_grad(): + for start in range(0, max(len(decoded_audio), 1), window): + segment = decoded_audio[start : start + window] + if segment.size == 0: + continue + inputs = processor( + segment, + sampling_rate = _TARGET_SAMPLE_RATE, + return_tensors = "pt", + ) + features = inputs.input_features.to(model.device) + if target_dtype is not None: + features = features.to(target_dtype) + generated = model.generate(features, **effective_generate_kwargs) + text = processor.batch_decode(generated, skip_special_tokens = True) + parts.append(text[0] if text else "") + return " ".join(part.strip() for part in parts if part.strip()).strip() + + def transcribe( + self, + audio: bytes, + model: Optional[str] = None, + language: Optional[str] = None, + fast: bool = False, + ) -> dict: + """Transcribe encoded audio bytes to text. + + Accepts any container PyAV can decode: wav, mp3, opus/webm, ogg, + m4a/aac. Returns {text, language, duration, model}. + """ + # Reject a missing runtime up front, before the cache and bounded decode. + ensure_stt_available() + # A set language beats auto-detect. API takes BCP-47; Whisper wants short + # codes like en or fr. + lang = normalize_whisper_language(language) + # Pin the requested id: another request may switch the resident model + # mid-transcription, so sidecar state is not this request's identity. + model_id = resolve_model_id(model) + known_languages = _known_whisper_languages() + if lang is not None and known_languages is not None and lang not in known_languages: + raise SttLanguageError( + f"Language '{language}' is not supported by STT model '{model_id}'." + ) + cached = self._ensure_model_downloaded(model_id) + if cached.is_multilingual is False and lang not in (None, "en"): + raise SttLanguageError( + f"Language '{language}' is not supported by English-only STT model '{model_id}'." + ) + decoded_audio = _decode_audio_bounded(audio) + # condition_on_prev_tokens=False stops a fresh clip inheriting prior + # context, which causes runaway repeats. + generate_kwargs = { + "task": "transcribe", + "condition_on_prev_tokens": False, + "num_beams": 5, + } + if lang is not None: + generate_kwargs["language"] = lang + if fast: + # Short voiced clips: greedy decoding drops beam search for latency. + generate_kwargs["num_beams"] = 1 + # Serialize inference with model switches and unloads. + with self._lock: + try: + text = self._transcribe_decoded(model_id, decoded_audio, generate_kwargs) + finally: + self._schedule_idle_unload_locked() + duration = (len(decoded_audio) / _TARGET_SAMPLE_RATE) if len(decoded_audio) else None + return { + "text": text, + "language": lang, + "duration": duration, + "model": model_id, + } + + def unload(self) -> None: + with self._lock: + self._release_engine_locked() + + +_sidecar: Optional[WhisperSttSidecar] = None + + +def get_stt_sidecar() -> WhisperSttSidecar: + global _sidecar + if _sidecar is None: + _sidecar = WhisperSttSidecar() + return _sidecar diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index bc9ffe85c2..0ef6dd46cf 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -18,6 +18,7 @@ import queue import random import re import shlex +import shutil import ssl import subprocess import sys @@ -328,6 +329,7 @@ def _find_blocked_commands(command: str) -> set[str]: # Directory holding the sandbox ``sitecustomize.py`` shim (code-interpreter # path remap); placed on the sandboxed child's PYTHONPATH in _build_safe_env. _SANDBOX_SITE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sandbox_site") + # โ”€โ”€ "Approve for me" (permission_mode="auto") safety detection โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ # Auto mode pauses only calls classified here as potentially unsafe. The sandbox # and hard blocks (blocklist, rlimits) still apply at run time; this gate only @@ -2491,15 +2493,124 @@ def is_potentially_unsafe_tool_call(name: str, arguments: dict) -> bool: return True +def _canon_win_path(p: str) -> str: + """Canonical form for trust comparison: realpath (expands 8.3 aliases and + resolves junctions/symlinks) + normcase/normpath.""" + return os.path.normcase(os.path.normpath(os.path.realpath(p))) + + +def _augment_native_program_roots(roots: list[str]) -> list[str]: + """Add the native Program Files sibling for any x86 root by stripping the + `` (x86)`` suffix, so a 32-bit process (whose known-folder ids map only to + the x86 root) still trusts a 64-bit Git install.""" + out = list(roots) + for root in roots: + base = root.rstrip("\\/") + if base.lower().endswith(" (x86)"): + native = base[: -len(" (x86)")] + if native and native not in out: + out.append(native) + return out + + +def _windows_program_roots() -> list[str]: + """Program Files install roots, resolved ONLY from the Windows known-folder + API (SHGetKnownFolderPath). Fails closed (returns ``[]``) if the API is + unavailable: env vars (%ProgramFiles%, even %SystemDrive%) are caller- + overrideable and could relocate the trust boundary, so we never derive a + trusted root from them. On any real Windows host shell32 is present, so + this only returns empty in a broken/non-Windows environment where the + sandbox git-PATH feature is not needed anyway (#7317). + """ + roots: list[str] = [] + try: + import ctypes + from ctypes import wintypes + + # FOLDERID_ProgramFiles, _ProgramFilesX86, _ProgramFilesX64. The X64 + # id (Win10 1703+) yields the native root even from a 32-bit process, + # where the first two both map to Program Files (x86). + folder_ids = ( + "{905e63b6-c1bf-494e-b29c-65b732d3d21a}", + "{7C5A40EF-A0FB-4BFC-874A-C0F2E0B9FA8E}", + "{6D809377-6AF0-444b-8957-A3773F02200E}", + ) + _SHGet = ctypes.windll.shell32.SHGetKnownFolderPath + _CoTaskMemFree = ctypes.windll.ole32.CoTaskMemFree + for fid in folder_ids: + guid = ctypes.create_string_buffer(16) + ctypes.windll.ole32.CLSIDFromString(wintypes.LPCWSTR(fid), ctypes.byref(guid)) + ptr = ctypes.c_wchar_p() + if _SHGet(ctypes.byref(guid), 0, None, ctypes.byref(ptr)) == 0: + if ptr.value: + roots.append(ptr.value) + _CoTaskMemFree(ptr) + except Exception: + return [] + return _augment_native_program_roots(roots) + + +def _resolve_trusted_windows_git() -> tuple[str, str]: + """Find a git launcher in a TRUSTED Program Files dir. Returns + ``(canonical_dir, ext)`` or ``("", "")``. + + ``shutil.which`` returns only the first PATH match, which may be an + untrusted user shim; scan the remaining PATH entries for a later trusted + Git so bare ``git`` still resolves (#7317). + """ + exts = [e for e in (os.environ.get("PATHEXT") or ".EXE;.CMD;.BAT;.COM").split(os.pathsep)] + candidates: list[str] = [] + primary = shutil.which("git") + if primary: + candidates.append(primary) + for entry in (os.environ.get("PATH") or "").split(os.pathsep): + entry = entry.strip().strip('"') + if not entry or not os.path.isabs(entry): + continue + for ext in exts: + cand = os.path.join(entry, "git" + ext) + if os.path.isfile(cand): + candidates.append(cand) + for git_exe in candidates: + git_dir = os.path.dirname(git_exe) + if os.path.isabs(git_dir) and _is_trusted_windows_program_dir(git_dir): + return os.path.realpath(git_dir), os.path.splitext(git_exe)[1].upper() + return "", "" + + +def _is_trusted_windows_program_dir(path: str) -> bool: + """True when ``path`` sits under a system-managed Program Files root. + + Only the Program Files roots are trusted (admin-writable only), resolved + via the known-folder API so an overridden env var cannot relocate them, + never ``%SystemRoot%`` (Git does not install there and it holds + world-writable subdirs like ``Windows\\Temp``). Per-user managers + (Scoop/Choco shims under the profile) are refused. Paths are canonicalized + so 8.3 aliases and junctions still resolve to their real root (#7317). + """ + norm = _canon_win_path(path) + for root in _windows_program_roots(): + root_norm = _canon_win_path(root) + if norm == root_norm or norm.startswith(root_norm + os.sep): + return True + return False + + def _build_safe_env(workdir: str) -> dict[str, str]: """Build a minimal, credential-free environment for sandboxed subprocesses. Whitelist-built from scratch (parent env NOT inherited): only PATH/HOME/ TMPDIR/LANG/TERM/PYTHONIOENCODING/PYTHONPATH (+VIRTUAL_ENV or Windows - SystemRoot) reach the child; all credential vars (HF_TOKEN, AWS_*, etc.) - are absent. HOME points at the sandbox workdir so SDKs can't read the + SystemRoot and a minimal PATHEXT) reach the child; all credential vars + (HF_TOKEN, AWS_*, etc.) are absent. HOME points at the sandbox workdir so SDKs can't read the operator's cached creds. PYTHONPATH carries only the sandbox sitecustomize shim directory. + + PATH starts with the Studio interpreter / venv and OS system dirs so + ``python``/``pip`` stay pinned. On Windows only, Git-for-Windows install + dirs from the host PATH are appended so bare ``git`` resolves (#7317). + User-writable host PATH entries (venv, ``node_modules/.bin``, etc.) are + never inherited โ€” they could shadow auto-safe terminal commands. """ # Start from the running interpreter's dir so 'python'/'pip' resolve to the # same environment the Unsloth server runs in. @@ -2519,6 +2630,20 @@ def _build_safe_env(workdir: str) -> dict[str, str]: else: path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"]) + # Windows Git installs live outside System32; inherit the dir of the git + # the HOST shell resolves, but ONLY when it sits under a system install + # root (Program Files, windir). A user-writable dir (Scoop/Choco shims) + # is refused: it would let an attacker drop rg.exe/jq.exe beside git and + # have an auto-approved bare command execute it (#7317). + git_ext = "" + if sys.platform == "win32": + # Append the CANONICAL (realpath) trusted git dir, scanning past any + # untrusted user shim that sorts first on PATH; the canonical path + # cannot be retargeted via a junction after the trust check. + _trusted_git_dir, git_ext = _resolve_trusted_windows_git() + if _trusted_git_dir: + path_entries.append(_trusted_git_dir) + # Deduplicate, preserving order. deduped = list(dict.fromkeys(p for p in path_entries if p)) @@ -2538,6 +2663,15 @@ def _build_safe_env(workdir: str) -> dict[str, str]: # Windows needs SystemRoot for Python/subprocess to work. if sys.platform == "win32": env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows") + # Restrict PATHEXT so cwd .BAT/.CMD cannot hijack bare names (#7317). + pathext = ".EXE;.COM" + if git_ext and git_ext not in (".EXE", ".COM"): + # Keep the host git launcher (e.g. a .CMD shim) resolvable. + pathext += ";" + git_ext + env["PATHEXT"] = pathext + # cmd/CreateProcess search cwd before PATH for bare names; disable so + # a workdir rg.exe/git.exe cannot shadow auto-approved commands. + env["NoDefaultCurrentDirectoryInExePath"] = "1" return env diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 9f301ba37e..367de196f7 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -794,7 +794,7 @@ def run_inference_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) - apply_gpu_ids(config.get("resolved_gpu_ids")) + apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) model_name = config["model_name"] diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index b141e59422..facd989b27 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -188,7 +188,14 @@ class LlamaServerBackend: match = [f for f in files if variant in f.lower()] or files filename = sorted(match, key = len)[0] logger.info("resolving GGUF embedder %s/%s", repo, filename) - self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token) + from utils.hf_cache_settings import active_hf_hub_cache + + self._model_path = hf_hub_download( + repo_id = repo, + filename = filename, + token = token, + cache_dir = active_hf_hub_cache(), + ) self._model_repo = desired self._dim = None return self._model_path diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 15be7f1249..3354585d2a 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -22,6 +22,7 @@ from typing import Callable from utils.hardware.hardware import DeviceType, get_device from utils.transformers_dtype import dtype_kwargs +from utils.utils import hf_env_offline from . import config @@ -103,9 +104,15 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: else: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError + from utils.hf_cache_settings import active_hf_hub_cache try: - local = hf_hub_download(name, "modules.json", token = token or None) + local = hf_hub_download( + name, + "modules.json", + token = token or None, + cache_dir = active_hf_hub_cache(), + ) except EntryNotFoundError: return () data = json.loads(open(local).read()) @@ -119,30 +126,55 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: return () -def _guard_model_security(name: str) -> None: +def _guard_model_security(name: str, local_only: bool = False) -> None: """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside SentenceTransformer regardless of trust_remote_code. Defense in depth behind the /settings gate (a name can also arrive via env/default); local paths and unreachable scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error. + + ``local_only`` (offline) inspects the local cache; subdir probes are skipped (they'd hit the + network and hang, and the offline gate walks the whole snapshot anyway). """ try: from utils.security import evaluate_file_security, security_load_subdirs token = _ambient_hf_token() - # Union the audio-model load roots with the ST module dirs so a flagged pickle - # directly under a Transformer module dir (0_Transformer/) blocks instead of - # passing as an unreferenced nested shard. - load_subdirs = tuple( - dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) - ) - blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + if local_only: + load_subdirs = () + else: + # Union audio-model load roots with ST module dirs so a flagged pickle under a + # Transformer module dir blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + (*security_load_subdirs(name, token), *_st_module_subdirs(name, token)) + ) + ) + blocked = evaluate_file_security( + name, hf_token = token, load_subdirs = load_subdirs, local_only_load = local_only + ).blocked except Exception: return if blocked: - raise UnsafeEmbeddingModelError( - f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security " - "scan; refusing to load. Set a different RAG embedding model." + reason = ( + "has cached pickle weights that cannot be security-scanned offline and no " + "safetensors alternative" + if local_only + else "is flagged as unsafe by Hugging Face's security scan" ) + raise UnsafeEmbeddingModelError( + f"Embedding model {name!r} {reason}; refusing to load. " + "Set a different RAG embedding model." + ) + + +def _st_accepts_local_files_only(st_cls) -> bool: + """Whether this SentenceTransformer version accepts local_files_only; passing it to an + older constructor raises, so gate on the signature.""" + try: + import inspect + return "local_files_only" in inspect.signature(st_cls.__init__).parameters + except Exception: + return False def _get(model_name: str | None = None): @@ -150,15 +182,35 @@ def _get(model_name: str | None = None): for a ~1.5x speedup at negligible accuracy loss.""" global _model, _name name = model_name or config.effective_embedding_model() + # Capture offline state once so the gate and the load agree (no window where the gate is + # skipped as offline but the constructor then reaches the network). + local_only = hf_env_offline() with _lock: if _model is None or _name != name: _install_torchao_stub_once() from sentence_transformers import SentenceTransformer + from utils.hf_cache_settings import active_hf_hub_cache device = _device() logger.info("loading embedding model %s on %s", name, device) - _guard_model_security(name) - _model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16")) + _guard_model_security(name, local_only) + st_kwargs = dict( + device = device, + cache_folder = active_hf_hub_cache(), + model_kwargs = dtype_kwargs("float16"), + ) + load_target = name + if local_only: + from utils.utils import hf_cache_snapshot_dir + snapshot = hf_cache_snapshot_dir(name) + if snapshot is not None: + # Load from the local snapshot dir: a local path never touches the Hub, so + # this is offline-safe on ANY sentence-transformers version (even ones + # predating local_files_only). + load_target = str(snapshot) + elif _st_accepts_local_files_only(SentenceTransformer): + st_kwargs["local_files_only"] = True + _model = SentenceTransformer(load_target, **st_kwargs) _name = name return _model diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index bbd9a895ab..17183484c5 100644 --- a/studio/backend/core/training/resume.py +++ b/studio/backend/core/training/resume.py @@ -4,6 +4,8 @@ """Helpers for validating resumable training outputs.""" import json +import pickletools +import zipfile from pathlib import Path from typing import Optional @@ -33,21 +35,158 @@ def _checkpoint_step(path: Path) -> int: return -1 -def get_resume_checkpoint_path(path_value: str) -> Optional[str]: +_MODEL_FILES = ( + "adapter_model.safetensors", + "adapter_model.bin", + "model.safetensors", + "pytorch_model.bin", +) +_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json") + + +def _valid_state_file(path: Path, require_tensor: bool = True) -> bool: + try: + if not path.is_file() or path.stat().st_size == 0: + return False + if path.suffix == ".safetensors": + try: + from safetensors import SafetensorError, safe_open + except ImportError: + return False + try: + with safe_open(str(path), framework = "np") as state: + return bool(state.keys()) + except SafetensorError: + return False + if path.suffix in {".bin", ".pt"}: + with zipfile.ZipFile(path) as state: + infos = state.infolist() + names = [info.filename for info in infos] + data_name = next( + (name for name in names if name == "data.pkl" or name.endswith("/data.pkl")), + None, + ) + if data_name is None: + return False + data_prefix = data_name.removesuffix("data.pkl") + "data/" + operations = list(pickletools.genops(state.read(data_name))) + if not operations or operations[-1][0].name != "STOP": + return False + if not require_tensor: + return True + # Require a non-empty tensor record; a zero-byte one fails torch.load. + return any( + info.filename.startswith(data_prefix) + and not info.is_dir() + and info.file_size > 0 + for info in infos + ) + # Unrecognized state-file formats are not usable resume state. + return False + except (OSError, ValueError, zipfile.BadZipFile): + return False + + +def _checkpoint_state(path: Path) -> Optional[int]: + try: + state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8")) + step = state.get("global_step") if isinstance(state, dict) else None + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + return None + if isinstance(step, bool) or not isinstance(step, int) or step < 0: + return None + directory_step = _checkpoint_step(path) + return step if directory_step < 0 or step == directory_step else None + + +_INDEX_SHARD_SUFFIX = { + "model.safetensors.index.json": ".safetensors", + "pytorch_model.bin.index.json": ".bin", +} + + +def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool: + # Shard must be a relative, in-format path contained in the checkpoint dir. + if not isinstance(shard, str) or not shard: + return False + if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix: + return False + try: + root = checkpoint.resolve(strict = True) + candidate = (checkpoint / shard).resolve(strict = True) + candidate.relative_to(root) + except (OSError, ValueError): + return False + return _valid_state_file(candidate) + + +def _has_model_state(path: Path) -> bool: + if any(_valid_state_file(path / name) for name in _MODEL_FILES): + return True + for name in _MODEL_INDEXES: + try: + index = json.loads((path / name).read_text(encoding = "utf-8")) + shards = set(index["weight_map"].values()) + except ( + AttributeError, + OSError, + KeyError, + TypeError, + UnicodeDecodeError, + json.JSONDecodeError, + ): + continue + expected_suffix = _INDEX_SHARD_SUFFIX[name] + if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards): + return True + return False + + +def is_resume_checkpoint_valid( + path: Path, + expected_step: Optional[int] = None, + backend: Optional[str] = None, +) -> bool: + step = _checkpoint_state(path) if path.is_dir() else None + step_valid = step is not None and (expected_step is None or step == expected_step) + if backend == "mlx": + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + else: + valid_bundle = ( + _has_model_state(path) + # optimizer/scheduler state can be validly tensor-free (e.g. SGD without + # momentum); _has_model_state still requires real model tensors. + and _valid_state_file(path / "optimizer.pt", require_tensor = False) + and _valid_state_file(path / "scheduler.pt", require_tensor = False) + ) + if backend is None and not valid_bundle: + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + return step_valid and valid_bundle + + +def get_resume_checkpoint_path( + path_value: str, expected_step: Optional[int] = None +) -> Optional[str]: path = resolve_output_dir(path_value) if not _is_under_outputs(path) or not path.is_dir(): return None - if (path / "trainer_state.json").is_file(): + if is_resume_checkpoint_valid(path, expected_step): return str(path) - checkpoints = [ - child - for child in path.glob("checkpoint-*") - if child.is_dir() and (child / "trainer_state.json").is_file() - ] - if not checkpoints: - return None - return str(max(checkpoints, key = _checkpoint_step)) + checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True) + return next( + ( + str(checkpoint) + for checkpoint in checkpoints + if _checkpoint_step(checkpoint) >= 0 + and is_resume_checkpoint_valid(checkpoint, expected_step) + ), + None, + ) def normalize_resume_output_dir(path_value: str) -> str: @@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool: def can_resume_run(run: dict) -> bool: if run.get("resumed_later"): return False + # Set when a stop-and-save failed to write a current-step checkpoint. + if run.get("resume_blocked"): + return False if _uses_s3_dataset(run): return False + status = run.get("status") + if status == "error": + # A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability. + return has_resume_state(run.get("output_dir")) + final_step = run.get("final_step") total_steps = run.get("total_steps") has_remaining_steps = ( @@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool: or total_steps <= 0 or final_step < total_steps ) - return ( - run.get("status") == "stopped" - and has_remaining_steps - and has_resume_state(run.get("output_dir")) - ) + return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir")) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 8e419849cb..b858fe6f17 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -891,6 +891,7 @@ class UnslothTrainer: use_gradient_checkpointing: str = "unsloth", use_rslora: bool = False, use_loftq: bool = False, + use_dora: bool = False, modules_to_save: list = None, ) -> bool: """ @@ -993,6 +994,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, ) # Audio VLM models support VLM-style layer selection @@ -1023,6 +1025,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, task_type = None, ) @@ -1042,6 +1045,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, ) @@ -1067,6 +1071,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, modules_to_save = modules_to_save, ) @@ -1087,6 +1092,7 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, + use_dora = use_dora, loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, modules_to_save = modules_to_save, ) @@ -1481,6 +1487,9 @@ class UnslothTrainer: SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" SNAC_SAMPLE_RATE = 24000 + + # SNAC codec unvalidated on Intel XPU; keep the pre-PR CPU + # fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" max_length = self.max_seq_length or 2048 tokenizer = self.tokenizer @@ -1642,7 +1651,8 @@ class UnslothTrainer: del snac_model gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: @@ -1669,6 +1679,8 @@ class UnslothTrainer: import numpy as np import torchaudio.transforms as T + # Spark-TTS BiCodec unvalidated on Intel XPU; keep the pre-PR CPU + # fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" # sparktts lives in the SparkAudio/Spark-TTS GitHub repo, not the HF model @@ -1857,7 +1869,8 @@ class UnslothTrainer: del audio_tokenizer gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: @@ -1894,6 +1907,8 @@ class UnslothTrainer: from datasets import Dataset as HFDataset from utils.paths import ensure_dir, tmp_root + # OuteTTS DAC/Whisper preprocess unvalidated on Intel XPU; keep the + # pre-PR CPU fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" # Clone OuteTTS repo (same as audio_codecs._load_dac) @@ -2065,7 +2080,8 @@ class UnslothTrainer: del prompt_processor gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index b407ba39a5..8592dabbfe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -30,7 +30,7 @@ from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt -from utils.hardware import prepare_gpu_selection +from utils.hardware import get_device, prepare_gpu_selection from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -196,6 +196,7 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), "use_rslora": values.get("use_rslora", False), "use_loftq": values.get("use_loftq", False), + "use_dora": values.get("use_dora", False), "train_on_completions": values.get("train_on_completions", False), "finetune_vision_layers": values.get("finetune_vision_layers", True), "finetune_language_layers": values.get("finetune_language_layers", True), @@ -219,6 +220,9 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: config[key] = values.get(key) if config["training_type"] == "Full Finetuning": config["load_in_4bit"] = False + # The parent's detected backend: the worker's apply_gpu_ids() targets the + # right visibility env var from this, without probing torch pre-mask. + config["device_backend"] = get_device().value return config @@ -452,6 +456,7 @@ class _MLXTrainerAdapter: use_gradient_checkpointing: Union[str, bool] = "unsloth", use_rslora: bool = False, use_loftq: bool = False, + use_dora: bool = False, ) -> bool: self._peft_config = { "use_lora": bool(use_lora), @@ -462,6 +467,7 @@ class _MLXTrainerAdapter: "gradient_checkpointing": use_gradient_checkpointing, "use_rslora": bool(use_rslora), "use_loftq": bool(use_loftq), + "use_dora": bool(use_dora), "finetune_vision_layers": bool(finetune_vision_layers), "finetune_language_layers": bool(finetune_language_layers), "finetune_attention_modules": bool(finetune_attention_modules), @@ -569,6 +575,7 @@ class _MLXTrainerAdapter: "gradient_checkpointing": "unsloth", "use_rslora": False, "use_loftq": False, + "use_dora": False, "finetune_vision_layers": True, "finetune_language_layers": True, "finetune_attention_modules": True, @@ -754,6 +761,9 @@ class TrainingBackend: def __init__(self): # Subprocess state self._proc: Optional[mp.Process] = None + # True from the sidecar-swap handshake until the worker is recorded, so + # installs and STT loads treat the startup window as active. + self._spawn_in_progress: bool = False self._event_queue: Any = None self._stop_queue: Any = None self._pump_thread: Optional[threading.Thread] = None @@ -761,6 +771,7 @@ class TrainingBackend: # Left True after an abnormal death so _ensure_pump_alive spots a crash. self._pump_running: bool = False self._lock = threading.Lock() + self._run_intent_lock = threading.RLock() # Stop watchdog: after a stop is requested, escalates to force_terminate() # if the worker does not exit on its own within a bounded time. The watched @@ -773,6 +784,7 @@ class TrainingBackend: self._progress = TrainingProgress() self._should_stop = False self._cancel_requested = False # True only for stop(save=False) + self._cancel_cleanup_output_dir: Optional[str] = None # Throttled training-status logging to the server log (not one line/step). self._last_progress_log_ts: float = 0.0 @@ -792,6 +804,8 @@ class TrainingBackend: # Job metadata self.current_job_id: Optional[str] = None self._output_dir: Optional[str] = None + self._resume_source_run_id: Optional[str] = None + self._terminal_finalize_payload: Optional[dict] = None # DB persistence self._metric_buffer: list[dict] = [] @@ -819,6 +833,7 @@ class TrainingBackend: job_id: str, *, before_spawn = None, + resume_source_run_id: Optional[str] = None, **kwargs, ) -> bool: """Spawn a subprocess to run the full training pipeline. @@ -924,16 +939,21 @@ class TrainingBackend: config["resolved_gpu_ids"] = resolved_gpu_ids config["gpu_selection"] = gpu_selection - from .worker import run_training_process + from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths + + cache_env = get_hf_cache_paths().child_env({}) try: - with native_path_secret_removed_for_child_start(): + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): event_queue = _CTX.Queue() stop_queue = _CTX.Queue() proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_training_process,), + args = ("core.training.worker", "run_training_process", cache_env), kwargs = { "event_queue": event_queue, "stop_queue": stop_queue, @@ -956,6 +976,7 @@ class TrainingBackend: self.current_job_id = job_id self._should_stop = False self._cancel_requested = False + self._cancel_cleanup_output_dir = None self._complete_seen.clear() self._progress = TrainingProgress( is_training = True, status_message = "Initializing training..." @@ -972,7 +993,10 @@ class TrainingBackend: self.eval_loss_history.clear() self.eval_step_history.clear() self.eval_enabled = False - self._output_dir = None + self._output_dir = config.get("output_dir") if resume_source_run_id else None + self._progress.output_dir = self._output_dir + self._resume_source_run_id = resume_source_run_id + self._terminal_finalize_payload = None self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False @@ -982,6 +1006,7 @@ class TrainingBackend: self._db_started_at = datetime.now(timezone.utc).isoformat() # Start each job Xet-first; keep config so a stall can respawn over HTTP. self._last_full_config = config + self._last_hf_cache_env = cache_env self._in_model_load = False self._xet_fallback_used = False self._needs_xet_respawn = False @@ -990,6 +1015,17 @@ class TrainingBackend: # in history during model loading and a fast terminal worker can't race the # pump into a duplicate create/finalize. From here the pump only finalizes. self._ensure_db_run_created() + if resume_source_run_id and not self._db_run_created: + if proc.is_alive(): + proc.terminate() + proc.join(timeout = 5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout = 2.0) + self._progress.is_training = False + self._progress.error = "Resume checkpoint is no longer available." + self._spawn_in_progress = False + return False # Assign handles and start the pump together under the lock so a concurrent # poll can't see a live _proc with no pump and spawn a duplicate. @@ -1011,28 +1047,75 @@ class TrainingBackend: def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" - self._should_stop = True - if not save: - self._cancel_requested = True - with self._lock: - if self._stop_queue is not None: - try: - self._stop_queue.put({"type": "stop", "save": save}) - except (OSError, ValueError): - pass - # Update progress immediately for responsive UI. - self._progress.status_message = ( - "Stopping training and saving checkpoint..." if save else "Cancelling training..." - ) - # Guarantee the run finalizes even if the worker wedges after saving. - self._start_stop_watchdog(cancel = not save) + with self._run_intent_lock: + with self._lock: + run_id = self.current_job_id + if not save and run_id: + persist_error: Optional[Exception] = None + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import mark_run_cancel_requested + + self._ensure_db_run_created() + with self._lock: + terminal_payload = self._terminal_finalize_payload + if ( + terminal_payload + and terminal_payload.get("expected_job_id") == run_id + ): + return False + if not mark_run_cancel_requested(run_id): + if self._db_run_created: + return False + raise RuntimeError( + "Training run disappeared before cancellation persisted" + ) + if self.current_job_id != run_id: + return False + self._should_stop = self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + persist_error = None + break + except Exception as exc: + persist_error = exc + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + if persist_error is not None: + raise RuntimeError("Failed to persist Stop-without-Save") from persist_error + with self._lock: + if self.current_job_id != run_id: + return False + if save or not run_id: + self._should_stop = True + if not save and not run_id: + self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + if self._stop_queue is not None: + try: + self._stop_queue.put({"type": "stop", "save": save}) + except (OSError, ValueError): + pass + self._progress.status_message = ( + "Stopping training and saving checkpoint..." + if save + else "Cancelling training..." + ) + self._start_stop_watchdog(cancel = not save, expected_job_id = run_id) return True - def _start_stop_watchdog(self, cancel: bool) -> None: + def _start_stop_watchdog( + self, + cancel: bool, + expected_job_id: Optional[str] = None, + ) -> None: """Start a daemon that force-terminates the worker if a requested stop does not exit on its own. No-op if no worker is alive or a live watchdog already watches this proc (a stale watchdog on an old proc never blocks a new run's watcher).""" with self._lock: + if expected_job_id is not None and self.current_job_id != expected_job_id: + return proc = self._proc if proc is None or not proc.is_alive(): return @@ -1113,8 +1196,9 @@ class TrainingBackend: watched_job_id: Optional[str] = None, ) -> None: """Finalize parent state after a force-terminate so the UI leaves "Stopping..." - even if the worker is wedged in driver teardown; preserves output_dir so a saved - checkpoint is kept. No-ops if a new run already replaced the watched worker, so a + even if the worker is wedged in driver teardown; preserves output_dir on a save so + the checkpoint is kept, and clears it on a cancel (Stop without saving must not + offer resume/export). No-ops if a new run already replaced the watched worker, so a stale watchdog never marks a fresh run stopped or drops its handle. Supersession is checked on both the watched proc and job id: start_training sets @@ -1134,7 +1218,18 @@ class TrainingBackend: return # a new run is already starting up; leave its state alone run_id = self.current_job_id # == watched_job_id self._progress.is_training = False - self._progress.status_message = "Training stopped." + terminal_payload = self._terminal_finalize_kwargs() + status = terminal_payload["status"] + error_message = terminal_payload.get("error_message") + output_dir = terminal_payload["output_dir"] + clear_output_dir = terminal_payload["clear_output_dir"] + resume_blocked = bool(terminal_payload.get("resume_blocked")) + with self._lock: + if self.current_job_id != run_id: + return + self._progress.status_message = error_message or "Training stopped." + if error_message: + self._progress.error = error_message # Create the row if a start-time create failed (no-op otherwise; skips when the pump # is mid-create, in which case its create-then-finalize records the run instead). self._ensure_db_run_created() @@ -1148,7 +1243,8 @@ class TrainingBackend: batch: list = [] final_step = final_loss = duration = None loss_history: list = [] - output_dir = self._output_dir + if clear_output_dir: + self._output_dir = self._progress.output_dir = None if claim: self._run_finalized = True # claim this run's finalize batch = list(self._metric_buffer) @@ -1161,7 +1257,17 @@ class TrainingBackend: loss_history = list(self.loss_history) if claim: self._finish_stopped_run( - run_id, output_dir, batch, final_step, final_loss, duration, loss_history + run_id, + output_dir, + batch, + final_step, + final_loss, + duration, + loss_history, + status = status, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) with self._lock: if target_proc is None or self._proc is target_proc: @@ -1176,6 +1282,10 @@ class TrainingBackend: final_loss: Optional[float], duration: Optional[float], loss_history: list, + status: str = "stopped", + error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: """Record a force-stopped run finished by its captured id, from state snapshotted under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE, @@ -1194,14 +1304,16 @@ class TrainingBackend: sparkline = downsample(loss_history, 50) finish_run( id = run_id, - status = "stopped", + status = status, ended_at = datetime.now(timezone.utc).isoformat(), final_step = final_step, final_loss = final_loss, duration_seconds = duration, loss_sparkline = _json.dumps(sparkline), output_dir = output_dir, - error_message = None, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) return except Exception: @@ -1231,7 +1343,7 @@ class TrainingBackend: logger.info("Force-terminating training subprocess (pid=%s)", proc.pid) proc.terminate() cancelled = self._cancel_requested - output_dir = self._output_dir + output_dir = self._cancel_cleanup_output_dir or self._output_dir if proc is not None: proc.join(timeout = 5.0) @@ -1304,7 +1416,11 @@ class TrainingBackend: self._last_full_config = config logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall") - from .worker import run_training_process + cache_env = getattr(self, "_last_hf_cache_env", None) + if not cache_env: + from utils.hf_cache_settings import get_hf_cache_paths + cache_env = get_hf_cache_paths().child_env({}) + from utils.hf_cache_settings import child_environment_for_spawn # This run is active, so an install request 409s rather than proceeds: a reservation seen here # is transient (an aborting install or short lazy repair). Wait it out instead of stranding the @@ -1336,12 +1452,15 @@ class TrainingBackend: # crashed respawn cannot wedge is_training_active until restart. try: try: - with native_path_secret_removed_for_child_start(): + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): event_queue = _CTX.Queue() stop_queue = _CTX.Queue() new_proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_training_process,), + args = ("core.training.worker", "run_training_process", cache_env), kwargs = { "event_queue": event_queue, "stop_queue": stop_queue, @@ -1595,17 +1714,60 @@ class TrainingBackend: ) self._ensure_db_run_created() - self._finalize_run_in_db( - status = "stopped" if self._should_stop else "error", - error_message = None - if self._should_stop - else "Training process terminated unexpectedly", - ) + terminal_payload = self._terminal_finalize_kwargs() + with self._lock: + if terminal_payload["clear_output_dir"]: + self._output_dir = self._progress.output_dir = None + if terminal_payload.get("error_message"): + self._progress.error = terminal_payload["error_message"] + self._progress.status_message = terminal_payload["error_message"] + self._finalize_run_in_db(**terminal_payload) except Exception: logger.exception("Training event pump: finalization after worker exit failed") self._pump_running = False return + def _has_current_resume_checkpoint(self, output_dir, step) -> bool: + # A valid checkpoint at the current step means the stop-and-save landed on + # disk even if the worker died before confirming it. + if not output_dir or not isinstance(step, int) or step <= 0: + return False + from core.training.resume import get_resume_checkpoint_path + return get_resume_checkpoint_path(output_dir, expected_step = step) is not None + + def _terminal_finalize_kwargs(self) -> dict: + with self._lock: + job_id = self.current_job_id + payload = self._terminal_finalize_payload + if payload and payload.get("expected_job_id") == job_id: + return dict(payload) + cancel, stopped = self._cancel_requested, self._should_stop + output_dir = None if cancel else self._output_dir + step = self._progress.step + existing_error = self._progress.error + status, error, blocked = ( + ("stopped", None, cancel) + if stopped + else ( + "error", + existing_error or "Training process terminated unexpectedly", + False, + ) + ) + # Block only when no valid current-step checkpoint actually landed. + if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step): + status = "error" + error = "Stop and Save ended before a valid current-step checkpoint was written." + blocked = True + return { + "status": status, + "error_message": error, + "output_dir": output_dir, + "clear_output_dir": cancel, + "resume_blocked": blocked, + "expected_job_id": job_id, + } + def _handle_event(self, event: dict) -> None: """Apply a subprocess event to local state. @@ -1764,6 +1926,15 @@ class TrainingBackend: elif etype == "eval_configured": self.eval_enabled = True + elif etype == "output_dir": + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = self._progress.output_dir = None + else: + self._output_dir = event_output_dir + db_action = "persist_output_dir" + elif etype == "status": self._progress.status_message = event.get("message", "") self._progress.is_training = True @@ -1778,7 +1949,12 @@ class TrainingBackend: self._complete_seen.set() self._progress.is_training = False self._progress.is_completed = not stopped - self._output_dir = event.get("output_dir") + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = None + else: + self._output_dir = event_output_dir self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: @@ -1788,11 +1964,16 @@ class TrainingBackend: db_action_kwargs = { "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) elif etype == "error": self._progress.is_training = False self._progress.error = event.get("error", "Unknown error") + if self._cancel_requested: + self._output_dir = self._progress.output_dir = None logger.error("Training error: %s", event.get("error")) stack = event.get("stack", "") if stack: @@ -1801,29 +1982,36 @@ class TrainingBackend: db_action = "create_and_finalize" else: db_action = "finalize" + stop_save_failed = ( + self._should_stop + and not self._cancel_requested + and not self._has_current_resume_checkpoint( + self._output_dir, self._progress.step + ) + ) db_action_kwargs = { - "status": "stopped" if self._should_stop else "error", + "status": "stopped" + if self._should_stop + and not stop_save_failed + and not event.get("keep_error_status") + else "error", "error_message": event.get("error", "Unknown error"), + "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "resume_blocked": stop_save_failed or bool(event.get("resume_blocked")), + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) # --- DB I/O outside the lock --- if db_action == "create_run": - try: - from storage.studio_db import create_run - - create_run( - id = db_action_kwargs["job_id"], - model_name = db_action_kwargs["model_name"], - dataset_name = db_action_kwargs["dataset_name"], - config_json = db_action_kwargs["config_json"], - started_at = db_action_kwargs["started_at"], - total_steps = db_action_kwargs["total_steps"], - ) - self._db_run_created = True + self._ensure_db_run_created() + if self._db_run_created: if db_action_kwargs["total_steps"]: self._db_total_steps_set = True - except Exception: - logger.warning("Failed to create DB run record", exc_info = True) + self._persist_output_dir() + elif db_action == "persist_output_dir": + self._persist_output_dir() elif db_action == "create_and_finalize": self._ensure_db_run_created() self._finalize_run_in_db(**db_action_kwargs) @@ -1842,6 +2030,22 @@ class TrainingBackend: if etype == "progress": self._log_training_progress() + def _persist_output_dir(self) -> None: + with self._lock: + if ( + not self._output_dir + or not self.current_job_id + or not self._db_run_created + or self._cancel_requested + ): + return + run_id, output_dir = self.current_job_id, self._output_dir + try: + from storage.studio_db import update_run_output_dir + update_run_output_dir(run_id, output_dir) + except Exception: + logger.warning("Failed to persist output_dir", exc_info = True) + def _log_training_progress(self) -> None: """One throttled training-status line to the server log (the per-step stream still goes to the UI via SSE): first step, then at most every 30s, plus the @@ -1875,6 +2079,7 @@ class TrainingBackend: caller create at a time, and ``_db_run_created`` is published only after ``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running).""" + self._run_intent_lock.acquire() with self._lock: if ( self._db_run_created @@ -1882,6 +2087,7 @@ class TrainingBackend: or not self.current_job_id or not self._db_config ): + self._run_intent_lock.release() return self._db_create_in_progress = True # only one caller creates job_id = self.current_job_id @@ -1898,6 +2104,12 @@ class TrainingBackend: or _s3_dataset_name(db_config.get("s3_dataset")) or "unknown" ) + with self._lock: + if self.current_job_id != job_id: + return + output_dir = self._output_dir + cancel_requested = self._cancel_requested + resumed_from_run_id = self._resume_source_run_id create_run( id = job_id, model_name = db_config["model_name"], @@ -1905,6 +2117,9 @@ class TrainingBackend: config_json = _json.dumps(db_config), started_at = started_at, total_steps = total_steps, + output_dir = output_dir, + cancel_requested = cancel_requested, + resumed_from_run_id = resumed_from_run_id, ) created = True except Exception: @@ -1919,12 +2134,15 @@ class TrainingBackend: if created: self._db_run_created = True # publish only after the insert commits self._db_create_in_progress = False + self._run_intent_lock.release() def _finalize_run_in_db( self, status: str, error_message: Optional[str] = None, output_dir: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, expected_job_id: Optional[str] = None, ) -> None: """Flush remaining metrics and mark a run finished in the DB. Claims the finalize @@ -1947,26 +2165,33 @@ class TrainingBackend: duration = self._progress.elapsed_seconds loss_history = list(self.loss_history) self._flush_metrics_to_db(run_id = run_id) - try: - from storage.studio_db import finish_run - from utils.downsample import downsample + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import finish_run + from utils.downsample import downsample - sparkline = downsample(loss_history, 50) - finish_run( - id = run_id, - status = status, - ended_at = datetime.now(timezone.utc).isoformat(), - final_step = final_step, - final_loss = final_loss, - duration_seconds = duration, - loss_sparkline = _json.dumps(sparkline), - output_dir = output_dir, - error_message = error_message, - ) - except Exception: - with self._lock: - self._run_finalized = False # unclaim so a later flush can retry - logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) + finish_run( + id = run_id, + status = status, + ended_at = datetime.now(timezone.utc).isoformat(), + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, + loss_sparkline = _json.dumps(downsample(loss_history, 50)), + output_dir = output_dir, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, + ) + return + except Exception: + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + continue + with self._lock: + if self.current_job_id == run_id: + self._run_finalized = False + logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None: """Flush buffered metrics to the DB and update live progress. The target run id, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 111f4fdd0f..03327d3320 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -90,6 +90,79 @@ _FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS" # run_training_process() and isn't GC'd mid-run. _WINDOWS_ROCM_GROUPED_MM_LIB = None + +def _install_grouped_mm_cpu_fallback(torch_mod, logger, label): + """Register a Python mm/bmm fallback for torch._grouped_mm and return the Library. + + RDNA4 (gfx1200/gfx1201) ships a null HIP _grouped_mm kernel on ROCm <= 7.12 + (fixed in 7.13; ROCm/TheRock #5284). JitDecomp dispatches _grouped_mm to the + null kernel and crashes; overriding the CUDA dispatch key bypasses it. Shared + by the Windows and Linux ROCm guards. Keep the returned Library referenced so + the registration outlives the caller. + """ + import warnings as _warnings + + _gm_lib = torch_mod.library.Library("aten", "IMPL") + + def _grouped_mm_safe_impl( + self, + mat2, + offs = None, + bias = None, + out_dtype = None, + ): + """Python mm/bmm fallback for _grouped_mm on gfx120X (null HIP kernel, ROCm <= 7.12).""" + _t = torch_mod + if offs is None: + # No offsets: 2-D -> mm, 3-D batched -> bmm (unconditional mm broke 3-D MoE). + if self.dim() == 3 and mat2.dim() == 3: + result = _t.bmm(self.contiguous(), mat2.contiguous()) + elif self.dim() == 3 and mat2.dim() == 2: + result = _t.matmul(self.contiguous(), mat2.contiguous()) + elif self.dim() == 2 and mat2.dim() == 3: + result = _t.matmul(self.contiguous(), mat2.contiguous()) + else: + result = _t.mm(self.contiguous(), mat2.contiguous()) + else: + # Grouped: offs[i] is the exclusive end-row of group i. + offs_list = offs.tolist() + pieces = [] + prev = 0 + for idx, end in enumerate(offs_list): + end = int(end) + a_part = self[prev:end].contiguous() + b_part = mat2[idx].contiguous() if mat2.dim() == 3 else mat2.contiguous() + pieces.append(_t.mm(a_part, b_part)) + prev = end + # Include trailing rows not covered by offs. + if prev < self.shape[0]: + a_tail = self[prev:].contiguous() + b_tail = mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous() + pieces.append(_t.mm(a_tail, b_tail)) + result = ( + _t.cat(pieces, dim = 0) + if pieces + else _t.zeros(0, mat2.shape[-1], device = self.device, dtype = self.dtype) + ) + if bias is not None: + result = result + bias + if out_dtype is not None: + result = result.to(out_dtype) + elif result.dtype != self.dtype: + result = result.to(self.dtype) + return result + + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA") + logger.info( + "%s: patched _grouped_mm CUDA dispatch (null HIP kernel on gfx120X, " + "ROCm <= 7.12 -- bypassed with Python mm fallback)", + label, + ) + return _gm_lib + + # Subprocesses don't inherit os.add_dll_directory registrations. Replicate # main.py's Windows ROCm DLL setup so the first `import torch` finds # amdhip64.dll. Handles retained at module scope so they aren't GC'd. @@ -702,8 +775,9 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: 3. Device-name substring match (last resort when all arch attrs absent; AMD SDK / Radeon wheels may not populate them): - gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M`` - - gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395), - ``Radeon 8050S`` (cut-down SKU) + - gfx1151 Strix Halo / Gorgon Halo: ``Radeon 8065S`` (Ryzen AI + Max+ 495), ``Radeon 8060S`` (Ryzen AI MAX+ + 395), ``Radeon 8050S`` (cut-down SKU) """ gcn_arch = "" for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"): @@ -728,7 +802,11 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: # Arch attrs absent โ€” fall back to device-name matching. dev_lower = (getattr(props, "name", "") or "").lower() is_unified = ( - "890m" in dev_lower or "880m" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower + "890m" in dev_lower + or "880m" in dev_lower + or "8065s" in dev_lower + or "8060s" in dev_lower + or "8050s" in dev_lower ) return gcn_arch, is_unified @@ -1469,6 +1547,10 @@ def _run_mlx_training(event_queue, stop_queue, config): message = "LoftQ is not supported for MLX training yet." _send("error", error = message) raise NotImplementedError(message) + if config.get("use_dora"): + message = "DoRA is not supported for MLX training yet." + _send("error", error = message) + raise NotImplementedError(message) if config.get("is_embedding"): message = "Embedding model training is not supported for MLX training yet." _send("error", error = message) @@ -1840,8 +1922,15 @@ def _run_mlx_training(event_queue, stop_queue, config): # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it from utils.paths import ensure_dir - output_dir = _resolve_mlx_output_dir(config, model_name) + # Resume must land in the original run dir even when config lacks output_dir. + resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint( + resume_from_checkpoint + ) + output_dir = _resolve_mlx_output_dir( + {**config, "output_dir": resume_dir} if resume_dir else config, model_name + ) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) # โ”€โ”€ 6. Create trainer โ”€โ”€ eval_steps_val = config.get("eval_steps", 0) or 0 @@ -2067,6 +2156,17 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.add_eval_callback(_on_eval) + _opt_ref = [None] + _orig_build_optimizer = getattr(trainer, "_build_optimizer", None) + + if callable(_orig_build_optimizer): + + def _capture_optimizer(total_steps): + _opt_ref[0] = _orig_build_optimizer(total_steps) + return _opt_ref[0] + + trainer._build_optimizer = _capture_optimizer + # โ”€โ”€ 11. Run training โ”€โ”€ gc.collect() mx.synchronize() @@ -2082,31 +2182,58 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.save_model = _save_model # โ”€โ”€ 12. Save and finalize โ”€โ”€ - if trainer.stop_requested: - if not _stop_save[0]: - # Cancel (save=False): skip saving. - _send("complete", output_dir = None, status_message = "Training cancelled") + def _finish_tracking() -> None: + # Runs on every save/finalize exit so TB/W&B never leak on early return. + if tb_writer is not None: + try: + tb_writer.close() + except Exception: + pass + if wandb_run is not None: + try: + wandb_run.finish() + except Exception: + pass + + def _stop_checkpoint_ok() -> bool: + if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir): + return True + _send( + "error", + error = ( + "Failed to save a resumable checkpoint after stop. " + "Model files were saved, but this run cannot be resumed." + ), + # A user stop finalizes as 'stopped'; keep this failure's error status so history explains it. + keep_error_status = True, + # Older checkpoints are stale; resuming would roll back past this stop. + resume_blocked = True, + ) + return False + + try: + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + # Stop-and-save promises a resumable checkpoint, not just model files. + if not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: - _send("status", status_message = "Saving stopped model...") + _send("status", status_message = "Saving model...") mx.synchronize() trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training stopped") - else: - _send("status", status_message = "Saving model...") - mx.synchronize() - trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training completed") - - if tb_writer is not None: - try: - tb_writer.close() - except Exception: - pass - if wandb_run is not None: - try: - wandb_run.finish() - except Exception: - pass + # A save-stop can race the natural final save; it made the same promise. + if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training completed") + finally: + _finish_tracking() def _is_current_process_apple_silicon() -> bool: @@ -2250,7 +2377,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> env = os.getenv("ENVIRONMENT_TYPE", "production"), ) - apply_gpu_ids(config.get("resolved_gpu_ids")) + apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) model_name = config["model_name"] @@ -2644,80 +2771,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # so 7.13+ uses the real GPU kernel. if not _hip_ver_at_least(7, 13): try: - import warnings as _warnings - - _gm_lib = _torch_for_rocm.library.Library("aten", "IMPL") - - def _grouped_mm_safe_impl( - self, - mat2, - offs = None, - bias = None, - out_dtype = None, - ): - """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm โ‰ค 7.12).""" - _t = _torch_for_rocm - if offs is None: - # No offsets: 2-D -> mm, 3-D batched -> bmm - # (unconditional mm broke 3-D MoE). - if self.dim() == 3 and mat2.dim() == 3: - result = _t.bmm(self.contiguous(), mat2.contiguous()) - elif self.dim() == 3 and mat2.dim() == 2: - # Broadcast 2-D mat2 across the batch dim. - result = _t.matmul(self.contiguous(), mat2.contiguous()) - elif self.dim() == 2 and mat2.dim() == 3: - # Broadcast 2-D self across batch via matmul. - result = _t.matmul(self.contiguous(), mat2.contiguous()) - else: - result = _t.mm(self.contiguous(), mat2.contiguous()) - else: - # Grouped: offs[i] is the exclusive end-row of group i. - offs_list = offs.tolist() - pieces = [] - prev = 0 - for idx, end in enumerate(offs_list): - end = int(end) - a_part = self[prev:end].contiguous() - if mat2.dim() == 3: - b_part = mat2[idx].contiguous() - else: - b_part = mat2.contiguous() - pieces.append(_t.mm(a_part, b_part)) - prev = end - # Include trailing rows not covered by offs. - if prev < self.shape[0]: - a_tail = self[prev:].contiguous() - b_tail = ( - mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous() - ) - pieces.append(_t.mm(a_tail, b_tail)) - result = ( - _t.cat(pieces, dim = 0) - if pieces - else _t.zeros( - 0, - mat2.shape[-1], - device = self.device, - dtype = self.dtype, - ) - ) - if bias is not None: - result = result + bias - if out_dtype is not None: - result = result.to(out_dtype) - elif result.dtype != self.dtype: - result = result.to(self.dtype) - return result - - with _warnings.catch_warnings(): - _warnings.simplefilter("ignore") - _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA") - - _WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC - logger.info( - "Windows ROCm: patched _grouped_mm CUDA dispatch " - "(null HIP kernel on gfx1200, ROCm โ‰ค 7.12 โ€” " - "bypassed with Python mm fallback)" + _WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback( + _torch_for_rocm, logger, "Windows ROCm" ) except Exception as _patch_exc: logger.warning( @@ -2731,6 +2786,44 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> "skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)" ) + # โ”€โ”€ 1f-linux. Linux ROCm RDNA4 _grouped_mm null kernel โ”€โ”€ + # The win32 guard above misses Linux: RDNA4 (gfx1200/gfx1201) hits the same null + # HIP _grouped_mm kernel at ROCm <= 7.12 (fixed 7.13, ROCm/TheRock #5284). Gate on + # arch + HIP < 7.13 so NVIDIA/CUDA and non-RDNA4 AMD are untouched; no-op if fixed. + if sys.platform.startswith("linux") and _hw.IS_ROCM: + try: + _torch_lin = sys.modules.get("torch") + if _torch_lin is not None and _torch_lin.cuda.is_available(): + # Prefer torch.version.hip, else rocmX.Y from torch.__version__ (AMD + # SDK / Radeon wheels leave version.hip unset). Unknown version on a + # gfx120X build -> assume affected unless it is a post-fix rocmsdk wheel. + _hip_str = str(getattr(getattr(_torch_lin, "version", None), "hip", "") or "") + _ver = getattr(_torch_lin, "__version__", "").lower() + _m = re.match(r"(\d+)\.(\d+)", _hip_str) or re.search(r"rocm(\d+)\.(\d+)", _ver) + if _m: + _hip_lt_713 = (int(_m.group(1)), int(_m.group(2))) < (7, 13) + else: + _hip_lt_713 = "rocmsdk" not in _ver + # Scan every visible GPU (device_map="balanced" can place layers on a + # later RDNA4 card, so device 0 is not enough). Match gfx120X by arch, + # or by RX 9000 / R9700 name when the wheel omits gcnArchName. + _rdna4 = False + for _i in range(_torch_lin.cuda.device_count()): + _props = _torch_lin.cuda.get_device_properties(_i) + _lin_arch, _ = _rocm_classify_unified_memory(_props) + _lin_name = (getattr(_props, "name", "") or "").lower() + if _lin_arch.lower() in ("gfx1200", "gfx1201") or ( + not _lin_arch and re.search(r"rx\s*90[0-9]0|r9700", _lin_name) + ): + _rdna4 = True + break + if _rdna4 and _hip_lt_713: + _WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback( + _torch_lin, logger, "Linux ROCm gfx120X" + ) + except Exception as _gm_lin_exc: + logger.warning("Linux ROCm gfx120X: could not patch _grouped_mm: %s", _gm_lin_exc) + # โ”€โ”€ 1g. ROCm OOM guard โ”€โ”€ # On ROCm, exhausting VRAM can hang the HIP driver instead of raising. # set_per_process_memory_fraction caps the allocator so PyTorch raises @@ -3097,6 +3190,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), use_rslora = config.get("use_rslora", False), use_loftq = config.get("use_loftq", False), + use_dora = config.get("use_dora", False), ) elif use_lora: _send_status(event_queue, "Configuring LoRA adapters...") @@ -3113,6 +3207,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), use_rslora = config.get("use_rslora", False), use_loftq = config.get("use_loftq", False), + use_dora = config.get("use_dora", False), ) else: _send_status(event_queue, "Preparing model for full finetuning...") @@ -3177,6 +3272,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) tensorboard_dir = config.get("tensorboard_dir") if config.get("enable_tensorboard", False): @@ -3296,6 +3392,61 @@ def _send_status(event_queue: Any, message: str) -> None: ) +def _emit_output_dir(event_queue: Any, output_dir: str) -> None: + try: + event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()}) + except Exception: + pass + + +def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool: + if step <= 0: + return False + from core.training.resume import is_resume_checkpoint_valid + return is_resume_checkpoint_valid( + Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx" + ) + + +def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool: + """Write a full resume checkpoint for a stopped MLX run. + + Returns True when a checkpoint for the current training step exists. + """ + step = int(getattr(trainer, "_global_step", 0) or 0) + # A periodic save or a resumed run may already cover the current step. + if _mlx_has_checkpoint_at_step(output_dir, step): + return True + if step <= 0 or optimizer is None: + return False + ckpt_dir = Path(output_dir) / f"checkpoint-{step}" + if ckpt_dir.is_symlink(): + # Refuse a symlinked dir: it could redirect writes outside output_dir. + logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir) + return False + try: + ckpt_dir.mkdir(parents = True, exist_ok = True) + from unsloth_zoo.mlx.utils import ( + save_optimizer_state, + save_trainable_adapters, + save_trainer_state, + ) + + save_trainable_adapters(trainer.model, str(ckpt_dir)) + save_optimizer_state(optimizer, str(ckpt_dir)) + save_trainer_state( + { + "global_step": step, + "train_loss_history": list(getattr(trainer, "_train_loss_history", [])), + }, + str(ckpt_dir), + ) + logger.info("Saved stop checkpoint to %s", ckpt_dir) + except Exception: + logger.exception("Failed to write stop checkpoint under %s", output_dir) + return _mlx_has_checkpoint_at_step(output_dir, step) + + def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None: """Self-contained embedding model training pipeline. @@ -3485,6 +3636,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> use_gradient_checkpointing = gradient_checkpointing, random_state = config.get("random_seed", 3407), use_rslora = config.get("use_rslora", False), + use_dora = config.get("use_dora", False), loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if config.get("use_loftq") else None, @@ -3660,6 +3812,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> config.get("project_name"), ) output_dir = str(resolve_output_dir(output_dir)) + _emit_output_dir(event_queue, output_dir) num_epochs = config.get("num_epochs", 2) batch_size = config.get("batch_size", 256) diff --git a/studio/backend/hub/routes/datasets.py b/studio/backend/hub/routes/datasets.py index edf4f36ac0..7c7cc274d3 100644 --- a/studio/backend/hub/routes/datasets.py +++ b/studio/backend/hub/routes/datasets.py @@ -61,9 +61,11 @@ async def list_cached_datasets(current_subject: str = Depends(get_current_subjec @router.delete("/cached", response_model = DeleteCachedDatasetResponse) async def delete_cached_dataset( - repo_id: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject) + repo_id: str = Body(..., embed = True), + cache_path: Optional[str] = Body(None, embed = True), + current_subject: str = Depends(get_current_subject), ): - return await cache_inventory.delete_cached_dataset_response(repo_id) + return await cache_inventory.delete_cached_dataset_response(repo_id, cache_path) @router.get("/download-progress", response_model = DownloadProgressResponse) diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py index 4b6c179a2b..dc3e3641bc 100644 --- a/studio/backend/hub/routes/inventory.py +++ b/studio/backend/hub/routes/inventory.py @@ -28,6 +28,7 @@ from hub.schemas.inventory import ( CachedModelsResponse, DeleteCachedModelResponse, GgufVariantsResponse, + HiddenModelsResponse, LocalModelListResponse, ModelsFolderResponse, RecommendedFoldersResponse, @@ -214,6 +215,16 @@ async def list_cached_models( return await cache_inventory.list_cached_models_response(hf_token) +@router.get("/hidden-models", response_model = HiddenModelsResponse) +async def list_hidden_models(current_subject: str = Depends(get_current_subject)): + import asyncio + + from routes.models import hidden_model_matchers + + needles, exact_ids, exact_paths = await asyncio.to_thread(hidden_model_matchers) + return HiddenModelsResponse(needles = needles, exact_ids = exact_ids, exact_paths = exact_paths) + + @router.delete( "/delete-cached", response_model = DeleteCachedModelResponse, @@ -222,7 +233,8 @@ async def list_cached_models( async def delete_cached_model( repo_id: str = Body(...), variant: Optional[str] = Body(None), + cache_path: Optional[str] = Body(None), hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - return await deletion.delete_cached_model_response(repo_id, variant, hf_token) + return await deletion.delete_cached_model_response(repo_id, variant, hf_token, cache_path) diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index ef95efe2f2..ca0f4658a3 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -99,6 +99,10 @@ class LocalModelInfo(BaseModel): None, description = "HF repo id for cached models, e.g. org/model", ) + active_cache: Optional[bool] = Field( + None, + description = "Whether this HF entry belongs to the current download cache.", + ) base_model: Optional[str] = Field( None, description = "Base model from adapter_config.json when this is an adapter", @@ -160,6 +164,7 @@ class CachedRepoBase(BaseModel): repo_id: str size_bytes: int = 0 cache_path: Optional[str] = None + last_modified: Optional[float] = None partial: bool = False partial_transport: Optional[str] = None inventory_id: Optional[str] = None @@ -189,6 +194,12 @@ class CachedModelsResponse(BaseModel): cached: List[CachedModelRepo] = Field(default_factory = list) +class HiddenModelsResponse(BaseModel): + needles: List[str] = Field(default_factory = list) + exact_ids: List[str] = Field(default_factory = list) + exact_paths: List[str] = Field(default_factory = list) + + class AddScanFolderRequest(BaseModel): """Request body for adding a custom scan folder.""" diff --git a/studio/backend/hub/services/datasets/cache_inventory.py b/studio/backend/hub/services/datasets/cache_inventory.py index a180c9df58..c106b62ab7 100644 --- a/studio/backend/hub/services/datasets/cache_inventory.py +++ b/studio/backend/hub/services/datasets/cache_inventory.py @@ -20,12 +20,11 @@ from hub.utils import inventory_scan as hf_cache_scan from hub.utils.hf_cache_state import ( purge_partial_repo, purge_repo_cache_dirs, + resolve_delete_target_root, resolve_destructive_case_matches, ) from hub.utils.paths import ( - hf_default_cache_dir, is_valid_repo_id as _is_valid_repo_id, - legacy_hf_cache_dir, resolve_cached_repo_id_case, ) @@ -43,38 +42,8 @@ def _collect_hf_cache_scans() -> tuple[list, set[str]]: def _hf_hub_cache_roots() -> list[Path]: - roots: list[Path] = [] - seen: set[str] = set() - - def _add(path: Optional[Path]) -> None: - if path is None or not path.is_dir(): - return - try: - resolved = str(path.resolve()) - except OSError: - return - if resolved in seen: - return - seen.add(resolved) - roots.append(path) - - try: - from huggingface_hub.constants import HF_HUB_CACHE - _add(Path(HF_HUB_CACHE)) - except Exception: - pass - - hf_hub_cache = os.environ.get("HF_HUB_CACHE") - if hf_hub_cache: - _add(Path(hf_hub_cache).expanduser()) - - hf_home = os.environ.get("HF_HOME") - if hf_home: - _add(Path(hf_home).expanduser() / "hub") - - _add(legacy_hf_cache_dir()) - _add(hf_default_cache_dir()) - return roots + from hub.utils.hf_cache_state import hf_cache_roots + return hf_cache_roots() def _repo_id_from_hub_dataset_dir(name: str) -> str | None: @@ -207,6 +176,21 @@ def _repo_id_from_datasets_cache_dir(name: str) -> str | None: return repo_id if _is_valid_repo_id(repo_id) else None +def _is_processed_dataset_cache_path(repo_id: str, cache_path: str) -> bool: + """True when *cache_path* is this repo's processed Arrow cache dir + (``___`` directly under an HF_DATASETS_CACHE root). Such rows + have no Hub ``datasets--`` layout, so they are deleted via the processed + path and must not be rejected as an invalid cache_path.""" + try: + resolved = Path(cache_path).expanduser().resolve(strict = False) + except (OSError, RuntimeError, ValueError): + return False + if resolved.name.lower() != repo_id.replace("/", "___").lower(): + return False + roots = {r.resolve(strict = False) for r in _hf_datasets_cache_roots()} + return resolved.parent.resolve(strict = False) in roots + + def _processed_dataset_cache_size(path: Path) -> int: total = 0 try: @@ -361,7 +345,7 @@ async def list_cached_datasets_response() -> dict: ) from exc -async def delete_cached_dataset_response(repo_id: str) -> dict: +async def delete_cached_dataset_response(repo_id: str, cache_path: Optional[str] = None) -> dict: """Remove a cached dataset repo from the HF cache.""" if not _is_valid_repo_id(repo_id): raise HTTPException(status_code = 400, detail = "Invalid repo_id format") @@ -373,22 +357,40 @@ async def delete_cached_dataset_response(repo_id: str) -> dict: detail = "Cancel the active download before deleting.", ) try: - return await asyncio.to_thread(_delete_cached_dataset_blocking, repo_key) + return await asyncio.to_thread(_delete_cached_dataset_blocking, repo_key, cache_path) finally: downloads.registry.end_delete(repo_key) hf_cache_scan.invalidate_hf_cache_scans() -def _delete_cached_dataset_blocking(repo_id: str) -> dict: +def _delete_cached_dataset_blocking(repo_id: str, cache_path: Optional[str] = None) -> dict: scans, _seen_roots = _collect_hf_cache_scans() - candidate_entries = [] + # Group this dataset's copies by owning cache root, then target exactly one + # cache so a delete never removes copies in other, previously selected caches. + owners: dict = {} for hf_cache in scans: for repo_info in hf_cache.repos: if str(repo_info.repo_type) != "dataset": continue - if repo_info.repo_id.lower() == repo_id.lower(): - candidate_entries.append((hf_cache, repo_info)) + if repo_info.repo_id.lower() != repo_id.lower(): + continue + try: + owner = Path(repo_info.repo_path).parent.resolve(strict = False) + except (OSError, RuntimeError, ValueError): + continue + owners.setdefault(owner, []).append((hf_cache, repo_info)) + + target_root = resolve_delete_target_root("dataset", repo_id, cache_path, owners.keys()) + # A processed-only dataset row sends its Arrow cache path (___ + # under HF_DATASETS_CACHE), which is not a Hub datasets-- dir, so + # resolve_delete_target_root returns None. Accept it and fall through to the + # processed-cache delete rather than rejecting a legitimate row. + if target_root is None and not ( + cache_path and _is_processed_dataset_cache_path(repo_id, cache_path) + ): + raise HTTPException(status_code = 400, detail = "Invalid cache_path") + candidate_entries = owners.get(target_root, []) if target_root is not None else [] matched_repo_ids = resolve_destructive_repo_ids( repo_id, [str(repo_info.repo_id) for _hf_cache, repo_info in candidate_entries], @@ -414,7 +416,26 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict: exc_info = True, ) - processed_deleted, processed_failures = _delete_processed_dataset_cache(repo_id) + # Restrict the processed Arrow-cache delete to the selected cache's datasets + # root so it never removes copies under other cache homes. A processed + # cache_path scopes to its own root; a Hub target scopes to the datasets root + # sharing its cache home; an unspecified cache_path stays global (legacy). + processed_roots: Optional[set[Path]] + if not cache_path: + processed_roots = None + elif _is_processed_dataset_cache_path(repo_id, cache_path): + processed_roots = {Path(cache_path).expanduser().resolve(strict = False).parent} + else: + home = target_root.parent if target_root is not None else None + processed_roots = { + root.resolve(strict = False) + for root in _hf_datasets_cache_roots() + if home is not None and root.resolve(strict = False).parent == home + } + + processed_deleted, processed_failures = _delete_processed_dataset_cache( + repo_id, only_roots = processed_roots + ) failures.extend(processed_failures) if failures: raise HTTPException( @@ -427,15 +448,23 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict: # ``scan_cache_dir()`` skips blob-only/corrupt repos the revision delete # can't touch, yet the fallback scanner shows them; purge the whole dir. - cache_purged = purge_repo_cache_dirs("dataset", repo_id) - partial_purged = purge_partial_repo("dataset", repo_id) - state_purged = download_manifest.purge_all_state_for_repo("dataset", repo_id) > 0 + # Only for a Hub cache target; a processed-only path has no Hub dir/state. + cache_purged = partial_purged = state_purged = False + if target_root is not None: + cache_purged = purge_repo_cache_dirs("dataset", repo_id, root = target_root) + partial_purged = purge_partial_repo("dataset", repo_id, root = target_root) + state_purged = ( + download_manifest.purge_all_state_for_repo("dataset", repo_id, hub_cache = target_root) + > 0 + ) if not (deleted or processed_deleted or cache_purged or partial_purged or state_purged): raise HTTPException(status_code = 404, detail = "Dataset not found in cache") return {"status": "deleted", "repo_id": repo_id} -def _delete_processed_dataset_cache(repo_id: str) -> tuple[bool, list[str]]: +def _delete_processed_dataset_cache( + repo_id: str, only_roots: Optional[set[Path]] = None +) -> tuple[bool, list[str]]: import shutil target = repo_id.replace("/", "___") @@ -443,6 +472,10 @@ def _delete_processed_dataset_cache(repo_id: str) -> tuple[bool, list[str]]: deleted = False failures: list[str] = [] for root in _hf_datasets_cache_roots(): + # Scope to the selected cache's datasets root(s): a delete must not remove + # processed copies living under other, previously selected cache homes. + if only_roots is not None and root.resolve(strict = False) not in only_roots: + continue try: entries = [ entry diff --git a/studio/backend/hub/services/datasets/downloads.py b/studio/backend/hub/services/datasets/downloads.py index 5efac562fa..b412a339e9 100644 --- a/studio/backend/hub/services/datasets/downloads.py +++ b/studio/backend/hub/services/datasets/downloads.py @@ -159,12 +159,18 @@ async def download_dataset_response( use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet) transport = download_lifecycle.resolve_transport(use_xet) + from utils.hf_cache_settings import get_hf_cache_paths + + cache_paths = get_hf_cache_paths() + cache_env = cache_paths.child_env({}) claimed, claim_state = _registry.claim( key, transport, repo_type = "dataset", repo_id = repo_id, + hub_cache = str(cache_paths.hub_cache), + xet_cache = str(cache_paths.xet_cache), ) generation = _registry.current_generation(key) if not claimed: @@ -176,7 +182,12 @@ async def download_dataset_response( "accepted": _registry.adoptable(key), "generation": generation, } - download_manifest.clear_cancel_marker("dataset", repo_id, None) + download_manifest.clear_cancel_marker( + "dataset", + repo_id, + None, + hub_cache = cache_paths.hub_cache, + ) state = download_lifecycle.launch_worker( _registry, @@ -185,6 +196,7 @@ async def download_dataset_response( ["--repo-id", repo_id, "--dataset"], hf_token, use_xet = use_xet, + cache_env = cache_env, ), hf_token = hf_token, label = repo_id, diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index 23f8c7c911..8e14427a56 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -11,7 +11,7 @@ import sys import time import threading from pathlib import Path -from typing import Callable, Optional +from typing import Callable, Mapping, Optional from fastapi import HTTPException @@ -57,6 +57,7 @@ def spawn_worker( *, use_xet: bool, protected_blob_hashes: Optional[frozenset[str]] = None, + cache_env: Optional[Mapping[str, str]] = None, ) -> subprocess.Popen: """Spawn the download worker. @@ -68,7 +69,11 @@ def spawn_worker( """ cwd = backend_dir() mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP - env = os.environ.copy() + from utils.hf_cache_settings import get_hf_cache_paths + + env = get_hf_cache_paths().child_env() + if cache_env is not None: + env.update(cache_env) if protected_blob_hashes: env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes)) else: @@ -230,6 +235,7 @@ def finalize_worker_exit( (stderr_data or b"").decode("utf-8", "replace").strip(), hf_token = hf_token, ) + metadata = registry.get_job_metadata(key) state = classify_exit(rc, cancel_requested = cancel_requested) if state == "complete": registry.set_job(key, "complete") @@ -252,13 +258,13 @@ def finalize_worker_exit( repo_type, repo_id, download_registry.variant_from_key(key), + hub_cache = metadata.hub_cache if metadata is not None else None, ) except Exception as exc: logger.debug(f"clear_cancel_marker failed for {repo_id} (rc=0): {exc}") elif state == "cancelled": # Read metadata before the terminal set_job so a concurrent eviction # can't drop it; the job key is the fallback variant label. - metadata = registry.get_job_metadata(key) registry.set_job(key, "cancelled") logger.info(f"{log_prefix} cancelled: {label} (rc={rc})") download_registry.persist_cancel_marker( @@ -268,6 +274,7 @@ def finalize_worker_exit( if metadata is not None and metadata.variant else download_registry.variant_from_key(key), cancel_marker_transport or transport, + hub_cache = metadata.hub_cache if metadata is not None else None, logger = logger, ) else: @@ -303,6 +310,7 @@ def _set_retry_failure_state( metadata.transport if metadata is not None and metadata.transport else fallback_transport, + hub_cache = metadata.hub_cache if metadata is not None else None, logger = logger, ) return state @@ -371,6 +379,7 @@ def _try_http_retry( repo_type, repo_id, progress_blob_hashes, + root = Path(original_metadata.hub_cache) if original_metadata.hub_cache else None, ) if progress_blob_hashes else 0 @@ -403,6 +412,8 @@ def _try_http_retry( generation = generation, replace_active = True, cancel_marker_transport = original_metadata.transport, + hub_cache = original_metadata.hub_cache, + xet_cache = original_metadata.xet_cache, ) if claimed: break @@ -446,11 +457,24 @@ def _try_http_retry( label, ) try: + cache_env = ( + { + "HF_HUB_CACHE": original_metadata.hub_cache, + "HF_XET_CACHE": original_metadata.xet_cache, + } + if original_metadata.hub_cache and original_metadata.xet_cache + else None + ) + spawn_kwargs = { + "use_xet": False, + "protected_blob_hashes": peer_hashes or None, + } + if cache_env is not None: + spawn_kwargs["cache_env"] = cache_env proc = spawn_worker( args, hf_token, - use_xet = False, - protected_blob_hashes = peer_hashes or None, + **spawn_kwargs, ) except Exception as exc: scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token) diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 54a25482f2..807ec70991 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -31,9 +31,9 @@ from hub.services.models.common import ( _is_checkpoint_weight_name, _is_gguf_filename, _is_main_gguf_filename, + _is_mmproj_filename, _is_transformers_safetensors_weight_name, _local_inventory_id, - _prefer_complete_larger, _runtime_for_format, ) @@ -132,6 +132,39 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _blob_mtime(file_obj) -> float: + ts = getattr(file_obj, "blob_last_modified", None) + if isinstance(ts, (int, float)) and ts > 0: + return float(ts) + blob_path = getattr(file_obj, "blob_path", None) + if blob_path: + try: + return float(Path(blob_path).stat().st_mtime) + except OSError: + pass + return 0.0 + + +def _repo_gguf_last_modified(repo_info) -> float: + latest = 0.0 + for revision in repo_info.revisions: + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + latest = max(latest, _blob_mtime(f)) + return latest + + +def _repo_has_mmproj(repo_info) -> bool: + # An mmproj file only makes a repo vision-capable when it is an actual GGUF + # projector; a non-GGUF sidecar (e.g. mmproj_config.json) does not, and the + # runtime's projector detection is GGUF-only. + return any( + _is_gguf_filename(f.file_name) and _is_mmproj_filename(f.file_name) + for revision in repo_info.revisions + for f in revision.files + ) + + def _cached_repo_file_name(file_obj) -> str: file_path = getattr(file_obj, "file_path", None) if file_path: @@ -216,24 +249,46 @@ def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[ def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool: if existing is None: return True - return _prefer_complete_larger( - bool(candidate.get("partial")), - int(candidate.get("size_bytes") or 0), - bool(existing.get("partial")), - int(existing.get("size_bytes") or 0), - ) + candidate_partial = bool(candidate.get("partial")) + existing_partial = bool(existing.get("partial")) + if candidate_partial != existing_partial: + return not candidate_partial + candidate_active = bool(candidate.get("active_cache")) + existing_active = bool(existing.get("active_cache")) + if candidate_active != existing_active: + return candidate_active + return int(candidate.get("size_bytes") or 0) > int(existing.get("size_bytes") or 0) def _cache_inventory_fields( repo_id: str, model_format: ModelFormat, *, + repo_path: Optional[Path] = None, + snapshot_path: Optional[Path] = None, + active_hub_cache: Optional[Path] = None, partial: bool = False, requires_variant: bool = False, ) -> dict: + load_id = repo_id + active_cache = True + if repo_path is not None: + try: + if active_hub_cache is None: + from utils.hf_cache_settings import get_hf_cache_paths + active_hub_cache = get_hf_cache_paths().hub_cache + active_root = active_hub_cache.resolve(strict = False) + cached_root = repo_path.parent.resolve(strict = False) + if cached_root != active_root: + active_cache = False + load_id = str(snapshot_path or repo_path.resolve(strict = False)) + except (OSError, RuntimeError, ValueError): + active_cache = False + load_id = str(snapshot_path or repo_path) return { "inventory_id": _local_inventory_id("cache", model_format, repo_id), - "load_id": repo_id, + "load_id": load_id, + "active_cache": active_cache, "model_format": model_format, "runtime": _runtime_for_format(model_format), "format_variant": None, @@ -260,6 +315,9 @@ def _is_hidden_infra_repo(*values: str | None) -> bool: def _scan_cached_gguf() -> list[dict]: """Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread.""" cache_scans = all_hf_cache_scans() + from utils.hf_cache_settings import get_hf_cache_paths + + active_hub_cache = get_hf_cache_paths().hub_cache seen_lower: dict[str, dict] = {} for hf_cache in cache_scans: @@ -271,7 +329,10 @@ def _scan_cached_gguf() -> list[dict]: repo_path = Path(repo_info.repo_path) snapshot_path = _cached_model_snapshot_path(repo_path) total_size = _repo_gguf_size_bytes(repo_info) - has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id) + has_variant_state, variant_state_size = _gguf_variant_state_summary( + repo_id, + hub_cache = repo_path.parent, + ) is_hidden_infra = _is_hidden_infra_repo( repo_id, str(repo_path), @@ -291,6 +352,7 @@ def _scan_cached_gguf() -> list[dict]: continue key = repo_id.lower() existing = seen_lower.get(key) + last_modified = _repo_gguf_last_modified(repo_info) row = { "repo_id": repo_id, "size_bytes": max(total_size, variant_state_size), @@ -300,19 +362,34 @@ def _scan_cached_gguf() -> list[dict]: # per-variant detail lives on GgufVariantDetail. "partial_transport": None, } + last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if last_modified > 0: + row["last_modified"] = last_modified row.update( _cache_inventory_fields( repo_id, "gguf", + repo_path = repo_path, + snapshot_path = snapshot_path, + active_hub_cache = active_hub_cache, partial = bool(row["partial"]), requires_variant = True, ) ) + if _repo_has_mmproj(repo_info): + row["capabilities"]["supports_vision"] = True # Visible infra variants remain management-only. if is_hidden_infra: row["capabilities"]["can_chat"] = False if _prefer_cache_row(row, existing): + if existing and existing["capabilities"].get("supports_vision"): + row["capabilities"]["supports_vision"] = True seen_lower[key] = row + else: + if last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified + if row["capabilities"].get("supports_vision"): + existing["capabilities"]["supports_vision"] = True except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") @@ -340,13 +417,14 @@ class _CachedNonGgufPayload(NamedTuple): size_bytes: int has_runnable_weights: bool model_format: ModelFormat + last_modified: float def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: - all_weight_blobs: dict[str, int] = {} - adapter_blobs: dict[str, int] = {} - safetensors_blobs: dict[str, int] = {} - checkpoint_blobs: dict[str, int] = {} + all_weight_blobs: dict[str, tuple[int, float]] = {} + adapter_blobs: dict[str, tuple[int, float]] = {} + safetensors_blobs: dict[str, tuple[int, float]] = {} + checkpoint_blobs: dict[str, tuple[int, float]] = {} has_config = False has_adapter_config = False has_adapter_weights = False @@ -354,12 +432,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: has_transformers_safetensors = False has_checkpoint = False - def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None: + def _record_blob( + target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str + ) -> None: blob_path = getattr(file_obj, "blob_path", None) size = int(file_obj.size_on_disk or 0) key = str(blob_path) if blob_path else f"{rev_id}:{file_name}" - target[key] = size - all_weight_blobs[key] = size + value = (size, _blob_mtime(file_obj)) + target[key] = value + all_weight_blobs[key] = value for revision in repo_info.revisions: rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) @@ -403,18 +484,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: or "unknown" ) if model_format == "adapter": - size_bytes = sum(adapter_blobs.values()) + selected_blobs = adapter_blobs elif model_format == "safetensors": - size_bytes = sum(safetensors_blobs.values()) + selected_blobs = safetensors_blobs elif model_format == "checkpoint": - size_bytes = sum(checkpoint_blobs.values()) + selected_blobs = checkpoint_blobs else: - size_bytes = sum(all_weight_blobs.values()) + selected_blobs = all_weight_blobs return _CachedNonGgufPayload( - size_bytes = size_bytes, + size_bytes = sum(size for size, _mtime in selected_blobs.values()), has_runnable_weights = model_format != "unknown", model_format = model_format, + last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0), ) @@ -435,6 +517,19 @@ def _read_json_object(path: Path) -> dict: return {} +def _is_whisper_model_config(config: object) -> bool: + if not isinstance(config, dict): + return False + model_type = config.get("model_type") + if isinstance(model_type, str) and model_type.strip().lower() == "whisper": + return True + architectures = config.get("architectures") + return isinstance(architectures, list) and any( + isinstance(name, str) and name == "WhisperForConditionalGeneration" + for name in architectures + ) + + def _read_model_card_frontmatter(path: Path) -> dict: try: text = path.read_text(encoding = "utf-8") @@ -465,6 +560,8 @@ def _cached_model_local_metadata(repo_path: Path) -> dict: result: dict = {} config = _read_json_object(snapshot / "config.json") + if _is_whisper_model_config(config): + result["_hidden_stt"] = True quant_method = ( config.get("quantization_config", {}).get("quant_method") if isinstance(config.get("quantization_config"), dict) @@ -491,11 +588,15 @@ def _cached_model_local_metadata(repo_path: Path) -> dict: def _scan_cached_models() -> list[dict]: """Synchronous HF-cache disk walk for non-GGUF model repos; runs in a worker thread.""" cache_scans = all_hf_cache_scans() + from utils.hf_cache_settings import get_hf_cache_paths + + active_hub_cache = get_hf_cache_paths().hub_cache seen_lower: dict[str, dict] = {} inspected = 0 skipped_gguf = 0 skipped_no_weights = 0 + skipped_stt = 0 for hf_cache in cache_scans: for repo_info in hf_cache.repos: inspected += 1 @@ -523,6 +624,10 @@ def _scan_cached_models() -> list[dict]: continue key = repo_id.lower() existing = seen_lower.get(key) + local_metadata = _cached_model_local_metadata(repo_path) + if local_metadata.pop("_hidden_stt", False): + skipped_stt += 1 + continue snapshot_partial = hf_cache_scan.is_snapshot_partial( "model", repo_id, @@ -542,27 +647,40 @@ def _scan_cached_models() -> list[dict]: if snapshot_partial else None ), - **_cached_model_local_metadata(repo_path), + **local_metadata, } + last_modified = max( + payload.last_modified, + (existing or {}).get("last_modified", 0.0), + ) + if last_modified > 0: + row["last_modified"] = last_modified row.update( _cache_inventory_fields( repo_id, payload.model_format, + repo_path = repo_path, + snapshot_path = snapshot_path, + active_hub_cache = active_hub_cache, partial = bool(row["partial"]), ) ) if _prefer_cache_row(row, existing): seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") continue cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) logger.info( - "Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d returned=%d", + "Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d " + "skipped_stt=%d returned=%d", inspected, skipped_gguf, skipped_no_weights, + skipped_stt, len(cached), ) return cached diff --git a/studio/backend/hub/services/models/common.py b/studio/backend/hub/services/models/common.py index f381bffe9c..4c0e296fdc 100644 --- a/studio/backend/hub/services/models/common.py +++ b/studio/backend/hub/services/models/common.py @@ -150,7 +150,9 @@ def _prefer_complete_larger( return candidate_size_bytes > existing_size_bytes -def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]: +def _gguf_variant_state_summary( + repo_id: str, *, hub_cache: Optional[str | Path] = None +) -> tuple[bool, int]: """Whether GGUF variant-scoped state exists and its expected size; a cancelled/in-progress variant may have only manifests/markers/`.incomplete` blobs, which inventory needs to avoid a generic fallback row.""" from hub.utils import download_manifest @@ -159,10 +161,16 @@ def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]: for variant, _path in download_manifest.iter_variant_manifests( "model", repo_id, + hub_cache = hub_cache, ): key = variant.lower() variant_keys.add(key) - manifest = download_manifest.read_manifest("model", repo_id, variant) + manifest = download_manifest.read_manifest( + "model", + repo_id, + variant, + hub_cache = hub_cache, + ) if manifest is None: continue size_by_variant[key] = max( @@ -172,6 +180,7 @@ def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]: for variant, _path in download_manifest.iter_variant_markers( "model", repo_id, + hub_cache = hub_cache, ): variant_keys.add(variant.lower()) return bool(variant_keys), sum(size_by_variant.values()) @@ -432,8 +441,13 @@ def _local_model_info( base_model_source: Optional[str] = None, adapter_type: Optional[str] = None, training_method: Optional[str] = None, + active_cache: Optional[bool] = None, ) -> LocalModelInfo: - load_id = model_id if source == "hf_cache" and model_id else str(load_path) + load_id = ( + model_id + if source == "hf_cache" and model_id and active_cache is not False + else str(load_path) + ) semantic_id = model_id or str(load_path) return LocalModelInfo( id = load_id, @@ -445,6 +459,7 @@ def _local_model_info( ), load_id = load_id, model_id = model_id, + active_cache = active_cache if source == "hf_cache" else None, display_name = display_name or (scan_path.stem if scan_path.is_file() else scan_path.name), path = str(load_path), size_bytes = max(0, int(size_bytes or 0)), @@ -476,6 +491,7 @@ def _classify_local_path( model_id: Optional[str] = None, updated_at: Optional[float] = None, partial: bool = False, + active_cache: Optional[bool] = None, ) -> list[LocalModelInfo]: load_path = load_path or scan_path files = ( @@ -512,6 +528,7 @@ def _classify_local_path( requires_variant = scan_path.is_dir(), format_variant = variant, size_bytes = gguf_size_bytes, + active_cache = active_cache, ) ) @@ -574,6 +591,7 @@ def _classify_local_path( ), adapter_type = adapter_type if model_format == "adapter" else None, training_method = training_method if model_format == "adapter" else None, + active_cache = active_cache, ) ) elif not rows: @@ -592,6 +610,7 @@ def _classify_local_path( updated_at = updated_at, partial = partial or trusted_hf_cache_repo, size_bytes = size_bytes, + active_cache = active_cache, ) ) diff --git a/studio/backend/hub/services/models/deletion.py b/studio/backend/hub/services/models/deletion.py index 636a223d4e..c736908058 100644 --- a/studio/backend/hub/services/models/deletion.py +++ b/studio/backend/hub/services/models/deletion.py @@ -19,8 +19,10 @@ from hub.utils import inventory_scan as hf_cache_scan from hub.utils.gguf import extract_quant_label, extract_quant_token from hub.utils.hf_cache_state import ( INCOMPLETE_SUFFIX, + iter_repo_cache_dirs, purge_partial_repo, purge_repo_cache_dirs, + resolve_delete_target_root, ) from hub.utils.paths import ( is_valid_gguf_variant as _is_valid_gguf_variant, @@ -184,6 +186,7 @@ def _delete_gguf_variant_from_repos( hf_token: Optional[str], *, sibling_active: bool = False, + root: Optional[Path] = None, ) -> dict: failures: list[str] = [] removed_snapshots = 0 @@ -265,6 +268,7 @@ def _delete_gguf_variant_from_repos( hf_token, extra_hashes = frozenset(completed_hashes), companions = not sibling_active, + root = root, ) if incomplete_result.unresolved: raise HTTPException( @@ -276,7 +280,7 @@ def _delete_gguf_variant_from_repos( ), ) - state_purged = download_manifest.purge_state("model", repo_id, variant) + state_purged = download_manifest.purge_state("model", repo_id, variant, hub_cache = root) # Reclaim the empty quant folder so it stops 404ing on delete. removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant) removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos) @@ -316,6 +320,8 @@ def reclaim_replaced_gguf_variant( variant: str, keep_main_hashes: frozenset[str], hf_token: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> dict: """Prune stale main-GGUF files for a variant after a replacement verified. @@ -366,12 +372,22 @@ def reclaim_replaced_gguf_variant( "reason": "scan_failed", } + if hub_cache is None: + from utils.hf_cache_settings import get_hf_cache_paths + hub_cache = get_hf_cache_paths().hub_cache + try: + target_hub_cache = Path(hub_cache).expanduser().resolve(strict = False) + except (OSError, RuntimeError, ValueError): + target_hub_cache = Path(hub_cache).expanduser() + candidate_repos = [ repo_info for hf_cache in cache_scans for repo_info in hf_cache.repos if str(getattr(repo_info, "repo_type", "")) == "model" and str(getattr(repo_info, "repo_id", "")).lower() == repo_id.lower() + and getattr(repo_info, "repo_path", None) + and Path(repo_info.repo_path).parent.resolve(strict = False) == target_hub_cache ] try: matched_repo_ids = resolve_destructive_repo_ids( @@ -493,10 +509,24 @@ def reclaim_replaced_gguf_variant( def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool: - """True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``.""" + """Match a loaded repo ID or an on-disk path inside any copy of the repo.""" rid = repo_id.lower() lid = loaded_id.lower() - return lid == rid or lid.startswith(f"{rid}/") + if lid == rid or lid.startswith(f"{rid}/"): + return True + + try: + loaded_path = Path(loaded_id).expanduser().resolve(strict = False) + except (OSError, RuntimeError, ValueError): + return False + for repo_dir in iter_repo_cache_dirs("model", repo_id): + try: + resolved_repo = repo_dir.resolve(strict = False) + if loaded_path == resolved_repo or loaded_path.is_relative_to(resolved_repo): + return True + except (OSError, RuntimeError, ValueError): + continue + return False def _loaded_repo_variant_blocks_delete( @@ -560,6 +590,7 @@ async def delete_cached_model_response( repo_id: str, variant: Optional[str] = None, hf_token: Optional[str] = None, + cache_path: Optional[str] = None, ): """Delete a cached model repo (or a specific GGUF variant) from the HF cache. @@ -603,14 +634,19 @@ async def delete_cached_model_response( ) raise HTTPException(status_code = 400, detail = detail) try: - return await asyncio.to_thread(_delete_cached_model_blocking, repo_id, variant, hf_token) + return await asyncio.to_thread( + _delete_cached_model_blocking, repo_id, variant, hf_token, cache_path + ) finally: downloads.registry.end_delete(repo_key, variant) cache_inventory.invalidate_hf_cache_scans() def _delete_cached_model_blocking( - repo_id: str, variant: Optional[str], hf_token: Optional[str] + repo_id: str, + variant: Optional[str], + hf_token: Optional[str], + cache_path: Optional[str] = None, ) -> dict: try: # If a sibling quant is downloading concurrently, restrict this delete to @@ -621,13 +657,26 @@ def _delete_cached_model_blocking( cache_scans = cache_inventory.all_hf_cache_scans() - candidate_entries = [] + # A repo can live in several remembered caches. Group its copies by the + # cache root that owns each, then target exactly one cache so a delete + # never removes copies in other, previously selected caches. + owners: dict = {} for hf_cache in cache_scans: for repo_info in hf_cache.repos: if str(repo_info.repo_type) != "model": continue - if repo_info.repo_id.lower() == repo_id.lower(): - candidate_entries.append((hf_cache, repo_info)) + if repo_info.repo_id.lower() != repo_id.lower(): + continue + try: + owner = Path(repo_info.repo_path).parent.resolve(strict = False) + except (OSError, RuntimeError, ValueError): + continue + owners.setdefault(owner, []).append((hf_cache, repo_info)) + + target_root = resolve_delete_target_root("model", repo_id, cache_path, owners.keys()) + if target_root is None: + raise HTTPException(status_code = 400, detail = "Invalid cache_path") + candidate_entries = owners.get(target_root, []) matched_repo_ids = resolve_destructive_repo_ids( repo_id, @@ -642,10 +691,15 @@ def _delete_cached_model_blocking( if not target_entries: if variant is None: - cache_purged = purge_repo_cache_dirs("model", repo_id) or purge_partial_repo( - "model", repo_id + cache_purged = purge_repo_cache_dirs( + "model", repo_id, root = target_root + ) or purge_partial_repo("model", repo_id, root = target_root) + state_purged = ( + download_manifest.purge_all_state_for_repo( + "model", repo_id, hub_cache = target_root + ) + > 0 ) - state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0 if cache_purged or state_purged: return {"status": "deleted", "repo_id": repo_id} if variant: @@ -654,6 +708,7 @@ def _delete_cached_model_blocking( variant, hf_token, companions = not sibling_active, + root = target_root, ) if incomplete_result.unresolved: raise HTTPException( @@ -668,6 +723,7 @@ def _delete_cached_model_blocking( "model", repo_id, variant, + hub_cache = target_root, ) if incomplete_result.deleted > 0 or state_purged: return { @@ -684,6 +740,7 @@ def _delete_cached_model_blocking( [repo for _cache, repo in target_entries], hf_token, sibling_active = sibling_active, + root = target_root, ) deleted_revisions = False @@ -702,9 +759,11 @@ def _delete_cached_model_blocking( delete_strategy.execute() deleted_revisions = True - cache_purged = purge_repo_cache_dirs("model", repo_id) - partial_purged = purge_partial_repo("model", repo_id) - state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0 + cache_purged = purge_repo_cache_dirs("model", repo_id, root = target_root) + partial_purged = purge_partial_repo("model", repo_id, root = target_root) + state_purged = ( + download_manifest.purge_all_state_for_repo("model", repo_id, hub_cache = target_root) > 0 + ) if not (deleted_revisions or cache_purged or partial_purged or state_purged): raise HTTPException(status_code = 404, detail = "No revisions found for model") diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index 862af0141a..c93b21c082 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -90,6 +90,7 @@ def _spawn_download_worker( hf_token: Optional[str], use_xet: bool = True, protected_blob_hashes: Optional[frozenset[str]] = None, + cache_env: Optional[dict[str, str]] = None, ) -> subprocess.Popen: args = ["--repo-id", repo_id] if variant: @@ -99,6 +100,7 @@ def _spawn_download_worker( hf_token, use_xet = use_xet, protected_blob_hashes = protected_blob_hashes, + cache_env = cache_env, ) @@ -125,6 +127,10 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional key = _download_job_key(repo_id, variant) use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet) transport = download_lifecycle.resolve_transport(use_xet) + from utils.hf_cache_settings import get_hf_cache_paths + + cache_paths = get_hf_cache_paths() + cache_env = cache_paths.child_env({}) variant_blob_hashes = frozenset() variant_progress_blob_hashes = frozenset() completed_baseline_bytes = 0 @@ -175,6 +181,8 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional progress_blob_hashes = variant_progress_blob_hashes, completed_baseline_bytes = completed_baseline_bytes, admission_check = lambda: not _load_in_flight(repo_id), + hub_cache = str(cache_paths.hub_cache), + xet_cache = str(cache_paths.xet_cache), ) generation = _registry.current_generation(key) if not claimed: @@ -189,7 +197,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional "accepted": _registry.adoptable(key), "generation": generation, } - download_manifest.clear_cancel_marker("model", repo_id, variant) + download_manifest.clear_cancel_marker( + "model", + repo_id, + variant, + hub_cache = cache_paths.hub_cache, + ) # Blobs a concurrent same-repo variant is already writing (e.g. a shared # mmproj). The worker must not purge these during cache preparation. protected_blob_hashes = _registry.peer_blob_hashes(key) if variant else frozenset() @@ -204,6 +217,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional hf_token, use_xet = use_xet, protected_blob_hashes = protected_blob_hashes, + cache_env = cache_env, ), hf_token = hf_token, label = label, diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index 7d9c3ac665..effb6a32ae 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -30,6 +30,7 @@ from hub.utils.paths import ( ) from utils.paths.external_media import ( linux_run_media_mount_roots, + macos_volume_roots, windows_drive_roots, ) from hub.services.models.common import _safe_is_dir @@ -187,7 +188,7 @@ def _build_browse_allowlist( _add(Path.home()) if media_roots is None: - media_roots = linux_run_media_mount_roots() + media_roots = [*linux_run_media_mount_roots(), *macos_volume_roots()] if drive_roots is None: drive_roots = windows_drive_roots() for p in media_roots: @@ -195,6 +196,12 @@ def _build_browse_allowlist( for p in drive_roots: _add(p) _add(_resolve_hf_cache_dir()) + try: + from utils.hf_cache_settings import known_hf_cache_homes + for cache_home in known_hf_cache_homes(): + _add(cache_home) + except Exception: # noqa: BLE001 -- best-effort + pass try: _add(hf_default_cache_dir()) except Exception: # noqa: BLE001 -- best-effort @@ -431,7 +438,7 @@ def browse_folders_response( # Probe removable-media and Windows drive roots once; the allowlist and # chips reuse the result so a disconnected mapped drive isn't scanned twice. - media_roots = linux_run_media_mount_roots() + media_roots = [*linux_run_media_mount_roots(), *macos_volume_roots()] drive_roots = windows_drive_roots() # Build the allowlist once -- the sandbox check and suggestion chips share # it so chips are always navigable. diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py index 33f0297ff5..533fd2ca5a 100644 --- a/studio/backend/hub/services/models/gguf_variants.py +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -9,6 +9,7 @@ import asyncio import threading import time from collections import OrderedDict +from pathlib import Path from typing import NamedTuple, Optional from fastapi import HTTPException @@ -22,6 +23,7 @@ from hub.utils.hf_errors import hf_error_status from hub.utils.hf_cache_state import ( INCOMPLETE_SUFFIX, iter_destructive_repo_cache_dirs, + repo_cache_dir_name, ) from hub.utils.gguf import ( extract_quant_label, @@ -233,8 +235,14 @@ def _manifest_variant_blob_hashes( variant: str, *, include_companions: bool = True, + repo_cache_dir: Optional[Path] = None, ) -> frozenset[str]: - manifest = download_manifest.read_manifest("model", repo_id, variant) + manifest = download_manifest.read_manifest( + "model", + repo_id, + variant, + hub_cache = repo_cache_dir.parent if repo_cache_dir is not None else None, + ) if manifest is None: return frozenset() variant_key = variant.lower() @@ -257,6 +265,7 @@ def gguf_variant_blob_hashes( *, include_companions: bool = True, allow_remote: bool = True, + repo_cache_dir: Optional[Path] = None, ) -> frozenset[str]: key = _variant_blob_hash_cache_key( repo_id, @@ -271,9 +280,9 @@ def gguf_variant_blob_hashes( repo_id, variant, include_companions = include_companions, + repo_cache_dir = repo_cache_dir, ) if hashes: - _variant_hash_cache_set(key, hashes) return hashes requirement_key = _variant_hash_cache_key(repo_id, variant, hf_token) requirement = _variant_requirement_cache_get(requirement_key) @@ -287,11 +296,22 @@ def gguf_variant_blob_hashes( return frozenset() -def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]: - return hf_cache_scan.partial_transport_for("model", repo_id, variant) +def _partial_transport_for_variant( + repo_id: str, + variant: str, + repo_cache_dir: Optional[Path] = None, +) -> Optional[str]: + return hf_cache_scan.partial_transport_for( + "model", + repo_id, + variant, + repo_cache_dir, + ) -def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str]]]: +def _local_main_gguf_blobs_by_quant( + repo_id: str, repo_cache_dir: Optional[Path] = None +) -> dict[str, dict[str, set[str]]]: """Map quant -> repo-relative expected GGUF filename -> cached blob hashes. Shared companions are copied into each main-quant bucket so update checks can @@ -313,6 +333,14 @@ def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str continue if str(getattr(repo_info, "repo_id", "")).lower() != target_lower: continue + if repo_cache_dir is not None: + try: + if Path(repo_info.repo_path).resolve(strict = False) != repo_cache_dir.resolve( + strict = False + ): + continue + except (AttributeError, OSError, RuntimeError, ValueError): + continue for path, hashes in cache_inventory._repo_gguf_blob_map( repo_info, include_companions = True, @@ -388,6 +416,7 @@ def delete_variant_incomplete_blobs_result( *, extra_hashes: frozenset[str] = frozenset(), companions: bool = True, + root: Optional[Path] = None, ) -> VariantIncompleteDeleteResult: # With a sibling still downloading, ``companions=False`` keeps a shared mmproj # from being unlinked out from under it; the repo's last delete reclaims it. @@ -409,8 +438,9 @@ def delete_variant_incomplete_blobs_result( ) deleted = 0 # Destructive iterator: only the exact-case match (or abort if ambiguous), - # so a case-variant sibling repo's partials are never unlinked. - for entry in iter_destructive_repo_cache_dirs("model", repo_id): + # so a case-variant sibling repo's partials are never unlinked. ``root`` scopes + # the purge to one cache so a delete never touches another cache's partials. + for entry in iter_destructive_repo_cache_dirs("model", repo_id, root = root): blobs_dir = entry / "blobs" if not blobs_dir.is_dir(): continue @@ -425,15 +455,37 @@ def delete_variant_incomplete_blobs_result( return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False) +def _repo_cache_dir_for_request(repo_id: str, local_path: Optional[str]) -> Path: + """Resolve the one Hub repo cache represented by this variant request.""" + expected_name = repo_cache_dir_name("model", repo_id).lower() + if local_path: + try: + local = Path(local_path).expanduser().resolve(strict = False) + for candidate in (local, *local.parents): + if candidate.name.lower() == expected_name: + return candidate + except (OSError, RuntimeError, ValueError): + pass + from utils.hf_cache_settings import get_hf_cache_paths + + return get_hf_cache_paths().hub_cache / repo_cache_dir_name("model", repo_id) + + def _mark_empty_dir_cleanables( - repo_id: str, response: GgufVariantsResponse + repo_id: str, + response: GgufVariantsResponse, + repo_cache_dir: Optional[Path] = None, ) -> GgufVariantsResponse: """Surface empty leftover ``/`` folders (interrupted downloads) as partial so the UI can delete them -- on local/offline paths too, not just a remote listing. A listed quant is flipped to partial; an unlisted one is appended as a zero-byte cleanable entry.""" try: - empty_labels = list_empty_gguf_variant_dirs(repo_id) + empty_labels = ( + list_empty_gguf_variant_dirs(repo_id, root = repo_cache_dir.parent) + if repo_cache_dir is not None + else list_empty_gguf_variant_dirs(repo_id) + ) except Exception as e: logger.warning(f"Failed to scan empty GGUF variant folders for {repo_id}: {e}") return response @@ -468,6 +520,11 @@ async def get_gguf_variants_response( """ def _compute() -> GgufVariantsResponse: + repo_cache_dir = ( + None if is_local_path(repo_id) else _repo_cache_dir_for_request(repo_id, local_path) + ) + hub_cache = repo_cache_dir.parent if repo_cache_dir is not None else None + def _local_response( response_repo_id: str, variants, has_vision: bool ) -> GgufVariantsResponse: @@ -511,6 +568,7 @@ async def get_gguf_variants_response( partial_transport = _partial_transport_for_variant( response_repo_id, v.quant, + repo_cache_dir, ), ) for v in variants @@ -532,7 +590,7 @@ async def get_gguf_variants_response( local_only = prefer_local_cache or offline if local_only: - cached = list_gguf_variants_from_hf_cache(repo_id) + cached = list_gguf_variants_from_hf_cache(repo_id, root = hub_cache) if cached is not None: variants, has_vision = cached return _local_response(repo_id, variants, has_vision) @@ -540,7 +598,7 @@ async def get_gguf_variants_response( variants, has_vision = list_local_gguf_variants(local_path) if variants or has_vision: return _local_response(repo_id, variants, has_vision) - partial = list_partial_gguf_variants_from_state(repo_id) + partial = list_partial_gguf_variants_from_state(repo_id, hub_cache = hub_cache) if partial is not None: variants, has_vision = partial return _partial_local_response(repo_id, variants, has_vision) @@ -560,11 +618,11 @@ async def get_gguf_variants_response( try: variants, has_vision, siblings = list_gguf_variants(repo_id, hf_token = hf_token) except Exception: - cached = list_gguf_variants_from_hf_cache(repo_id) + cached = list_gguf_variants_from_hf_cache(repo_id, root = hub_cache) if cached is not None: variants, has_vision = cached return _local_response(repo_id, variants, has_vision) - partial = list_partial_gguf_variants_from_state(repo_id) + partial = list_partial_gguf_variants_from_state(repo_id, hub_cache = hub_cache) if partial is not None: variants, has_vision = partial return _partial_local_response(repo_id, variants, has_vision) @@ -581,7 +639,7 @@ async def get_gguf_variants_response( cached_filenames_by_snapshot: list[dict[str, int]] = [] cached_quant_bytes_by_snapshot: list[dict[str, int]] = [] if _is_valid_repo_id(repo_id): - for snap in iter_hf_cache_snapshots(repo_id): + for snap in iter_hf_cache_snapshots(repo_id, root = hub_cache): try: gguf_paths = list(_iter_gguf_paths(snap)) except (OSError, RuntimeError, ValueError) as e: @@ -694,11 +752,20 @@ async def get_gguf_variants_response( partial_quants: set[str] = set() partial_quant_transports: dict[str, Optional[str]] = {} try: - incomplete_hashes = download_registry.incomplete_blob_hashes("model", repo_id) + incomplete_hashes = download_registry.incomplete_blob_hashes( + "model", + repo_id, + active_only = True, + root = hub_cache, + ) except Exception as e: logger.warning(f"Failed to compute partial GGUF variants for {repo_id}: {e}") incomplete_hashes = set() - scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan("model", repo_id) + scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan( + "model", + repo_id, + repo_cache_dir, + ) # Manifest + marker + main incomplete-blob check: catches variants whose # download was cancelled or whose expected shards are missing/undersized. for variant in variants: @@ -711,6 +778,7 @@ async def get_gguf_variants_response( variant.quant, hf_token, include_companions = False, + repo_cache_dir = repo_cache_dir, ) if hf_cache_scan.is_variant_partial( repo_id, @@ -718,11 +786,13 @@ async def get_gguf_variants_response( scan_snapshot_dir, incomplete_blob_hashes = incomplete_hashes, variant_blob_hashes = variant_hashes, + repo_cache_dir = repo_cache_dir, ): partial_quants.add(variant.quant) partial_quant_transports[variant.quant] = _partial_transport_for_variant( repo_id, variant.quant, + repo_cache_dir, ) except Exception as e: logger.warning( @@ -744,10 +814,14 @@ async def get_gguf_variants_response( partial_quants.add(variant.quant) partial_quant_transports.setdefault( variant.quant, - _partial_transport_for_variant(repo_id, variant.quant), + _partial_transport_for_variant( + repo_id, + variant.quant, + repo_cache_dir, + ), ) - local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id) + local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id, repo_cache_dir) def _variant_detail(v) -> GgufVariantDetail: is_partial = v.quant in partial_quants @@ -790,14 +864,20 @@ async def get_gguf_variants_response( if skip: raise enriched = _mark_empty_dir_cleanables( - repo_id, GgufVariantsResponse(repo_id = repo_id, variants = []) + repo_id, + GgufVariantsResponse(repo_id = repo_id, variants = []), + _repo_cache_dir_for_request(repo_id, local_path), ) if enriched.variants: return enriched raise if skip: return response - return _mark_empty_dir_cleanables(repo_id, response) + return _mark_empty_dir_cleanables( + repo_id, + response, + _repo_cache_dir_for_request(repo_id, local_path), + ) try: return await asyncio.to_thread(_compute_with_cleanables) diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py index b34532fa35..9cf260b157 100644 --- a/studio/backend/hub/services/models/local_inventory.py +++ b/studio/backend/hub/services/models/local_inventory.py @@ -106,11 +106,8 @@ def _is_model_directory_for_scan(path: Path, *, entry_limit: int | None) -> bool def _resolve_hf_cache_dir() -> Path: - try: - from huggingface_hub.constants import HF_HUB_CACHE - return Path(HF_HUB_CACHE) - except Exception: - return Path.home() / ".cache" / "huggingface" / "hub" + from utils.hf_cache_settings import get_hf_cache_paths + return get_hf_cache_paths().hub_cache def _scan_models_dir( @@ -202,7 +199,12 @@ def _hf_repo_dir_has_content(repo_dir: Path) -> bool: return False -def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]: +def _scan_hf_cache( + cache_dir: Path, + *, + entry_limit: int | None = None, + active_cache: bool = True, +) -> List[LocalModelInfo]: if not _safe_is_dir(cache_dir): return [] @@ -240,7 +242,10 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L repo_dir, ) gguf_partial = hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir) - has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(model_id) + has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary( + model_id, + hub_cache = cache_dir, + ) snapshot_partial_transport = ( hf_cache_scan.partial_transport_for( "model", @@ -252,23 +257,25 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L ) resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_dir) scan_path = Path(resolved) if resolved else repo_dir + load_path = repo_dir if active_cache else scan_path # partial=False here; _apply_format_aware_partial below rewrites per-row # so a hybrid repo's gguf row doesn't taint its safetensors row. rows = _classify_local_path( scan_path, "hf_cache", - load_path = repo_dir, + load_path = load_path, display_name = model_id.split("/")[-1], model_id = model_id, updated_at = updated_at, partial = False, + active_cache = active_cache, ) if not rows: if has_gguf_variant_state and gguf_partial: rows = [ _local_model_info( scan_path = repo_dir, - load_path = repo_dir, + load_path = load_path, source = "hf_cache", model_format = "gguf", display_name = model_id.split("/")[-1], @@ -277,6 +284,7 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L partial = True, requires_variant = True, size_bytes = gguf_variant_state_size, + active_cache = active_cache, ) ] else: @@ -285,13 +293,14 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L rows = [ _local_model_info( scan_path = repo_dir, - load_path = repo_dir, + load_path = load_path, source = "hf_cache", model_format = "unknown", display_name = model_id.split("/")[-1], model_id = model_id, updated_at = updated_at, partial = snapshot_partial or gguf_partial, + active_cache = active_cache, ) ] elif ( @@ -302,7 +311,7 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L rows.append( _local_model_info( scan_path = repo_dir, - load_path = repo_dir, + load_path = load_path, source = "hf_cache", model_format = "gguf", display_name = model_id.split("/")[-1], @@ -311,6 +320,7 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L partial = True, requires_variant = True, size_bytes = gguf_variant_state_size, + active_cache = active_cache, ) ) rows = _apply_format_aware_partial( @@ -515,14 +525,39 @@ async def _collect_models_from_default_sources( local_models += await _scan_source("HF cache", _scan_hf_cache, hf_cache_dir) if _safe_is_dir(legacy_hf) and legacy_hf.resolve() != hf_cache_dir.resolve(): - local_models += await _scan_source("legacy HF cache", _scan_hf_cache, legacy_hf) + local_models += await _scan_source( + "legacy HF cache", + lambda path: _scan_hf_cache(path, active_cache = False), + legacy_hf, + ) if ( _safe_is_dir(hf_default) and hf_default.resolve() != hf_cache_dir.resolve() and hf_default.resolve() != legacy_hf.resolve() ): - local_models += await _scan_source("default HF cache", _scan_hf_cache, hf_default) + local_models += await _scan_source( + "default HF cache", + lambda path: _scan_hf_cache(path, active_cache = False), + hf_default, + ) + + from utils.hf_cache_settings import known_hf_hub_caches + + seen_hf = { + os.path.normcase(str(path.resolve(strict = False))) + for path in (hf_cache_dir, legacy_hf, hf_default) + } + for previous_cache in known_hf_hub_caches(): + key = os.path.normcase(str(previous_cache.resolve(strict = False))) + if key in seen_hf: + continue + seen_hf.add(key) + local_models += await _scan_source( + "previous HF cache", + lambda path: _scan_hf_cache(path, active_cache = False), + previous_cache, + ) for lm_dir in lm_dirs: local_models += await _scan_source("LM Studio", _scan_lmstudio_dir, lm_dir) @@ -543,7 +578,11 @@ def _scan_custom_folder(folder_path: Path) -> List[LocalModelInfo]: limit = _MAX_MODELS_PER_CUSTOM_FOLDER, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES, ) - + _scan_hf_cache(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES) + + _scan_hf_cache( + folder_path, + entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES, + active_cache = False, + ) + _scan_lmstudio_dir(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES) ) if m.model_format in supported_formats @@ -610,12 +649,20 @@ def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelI row_key = model.inventory_id or model.id key = f"{row_key}\x00custom" if model.source == "custom" else row_key existing = deduped.get(key) - if existing is None or _prefer_complete_larger( - model.partial, - model.size_bytes, - existing.partial, - existing.size_bytes, - ): + prefer_candidate = existing is None + if existing is not None: + if model.partial != existing.partial: + prefer_candidate = not model.partial + elif (model.active_cache is True) != (existing.active_cache is True): + prefer_candidate = model.active_cache is True + else: + prefer_candidate = _prefer_complete_larger( + model.partial, + model.size_bytes, + existing.partial, + existing.size_bytes, + ) + if prefer_candidate: deduped[key] = model return sorted( deduped.values(), diff --git a/studio/backend/hub/services/snapshot_progress.py b/studio/backend/hub/services/snapshot_progress.py index 1fdf05e2e5..c3db6fed7a 100644 --- a/studio/backend/hub/services/snapshot_progress.py +++ b/studio/backend/hub/services/snapshot_progress.py @@ -86,9 +86,20 @@ def _snapshot_complete_on_disk( return False if variant is None and hf_cache_scan.repo_cache_dir_has_incomplete_blobs(entry): return False - if download_manifest.has_cancel_marker(repo_type, repo_id, variant): + hub_cache = entry.parent + if download_manifest.has_cancel_marker( + repo_type, + repo_id, + variant, + hub_cache = hub_cache, + ): return False - manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + manifest = download_manifest.read_manifest( + repo_type, + repo_id, + variant, + hub_cache = hub_cache, + ) if manifest is None: return False return download_manifest.verify_against_disk(manifest, snapshot_dir).ok @@ -118,6 +129,8 @@ def compute_snapshot_progress( 0, int(getattr(metadata, "completed_baseline_bytes", 0) or 0), ) + metadata_hub_cache = getattr(metadata, "hub_cache", None) + active_root = Path(metadata_hub_cache) if metadata_hub_cache else None expected_total = max(expected_bytes, 0) # Always resolve the revision's blob hashes so stale blobs from a superseded @@ -134,11 +147,17 @@ def compute_snapshot_progress( count_finalized_unscoped = variant is None readings: list[tuple[int, int, Optional[str], bool]] = [] - for entry in preferred_repo_cache_dirs( - repo_type, - repo_id, - force_active = force_active, - ): + cache_dirs = ( + preferred_repo_cache_dirs( + repo_type, + repo_id, + force_active = force_active, + active_root = active_root, + ) + if active_root is not None + else preferred_repo_cache_dirs(repo_type, repo_id, force_active = force_active) + ) + for entry in cache_dirs: completed_bytes = 0 in_progress_bytes = 0 cache_path = hf_cache_scan.resolve_hf_cache_realpath(entry) diff --git a/studio/backend/hub/tests/test_dataset_services.py b/studio/backend/hub/tests/test_dataset_services.py index 4890714cd0..6aab07cc46 100644 --- a/studio/backend/hub/tests/test_dataset_services.py +++ b/studio/backend/hub/tests/test_dataset_services.py @@ -72,57 +72,115 @@ def test_dataset_cache_scan_merges_raw_and_processed_rows(monkeypatch): assert rows[0]["partial"] is False -def test_delete_cached_dataset_attempts_all_roots_before_raising(monkeypatch): +def test_delete_cached_dataset_scopes_delete_to_selected_root(monkeypatch, tmp_path): + """A dataset present in the active cache and a previously selected cache is + deleted only from the selected root, so the other cache's copy survives.""" calls = [] - purged_state = [] + target_hub = tmp_path / "active" / "hub" + other_hub = tmp_path / "previous" / "hub" + for hub in (target_hub, other_hub): + (hub / "datasets--Org--Data").mkdir(parents = True) class _DeleteStrategy: - def __init__(self, label: str, fail: bool): + def __init__(self, label: str): self.label = label - self.fail = fail def execute(self): calls.append(self.label) - if self.fail: - raise RuntimeError(f"{self.label} failed") - class _Cache: - def __init__(self, label: str, fail: bool): - self.cache_dir = label - self.repos = [ + def _cache(label: str, hub): + return SimpleNamespace( + cache_dir = label, + repos = [ SimpleNamespace( repo_type = "dataset", repo_id = "Org/Data", + repo_path = str(hub / "datasets--Org--Data"), revisions = [SimpleNamespace(commit_hash = f"{label}-rev")], ) - ] - self.fail = fail - - def delete_revisions(self, *_revisions): - return _DeleteStrategy(self.cache_dir, self.fail) + ], + delete_revisions = lambda *_revs, _label = label: _DeleteStrategy(_label), + ) monkeypatch.setattr( cache_inventory, "_collect_hf_cache_scans", - lambda: ([_Cache("first", True), _Cache("second", False)], set()), + lambda: ([_cache("active", target_hub), _cache("previous", other_hub)], set()), ) monkeypatch.setattr( cache_inventory, "_delete_processed_dataset_cache", - lambda _repo_id: (True, []), + lambda _repo_id, **_kwargs: (False, []), ) monkeypatch.setattr( cache_inventory.download_manifest, "purge_all_state_for_repo", - lambda *_args: purged_state.append(True) or 1, + lambda *_args, **_kwargs: 0, + ) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = target_hub), + ) + monkeypatch.setattr( + "hub.utils.hf_cache_state.hf_cache_roots", + lambda: [target_hub, other_hub], ) - with pytest.raises(HTTPException) as exc_info: - cache_inventory._delete_cached_dataset_blocking("Org/Data") + result = cache_inventory._delete_cached_dataset_blocking("Org/Data") - assert exc_info.value.status_code == 500 - assert calls == ["first", "second"] - assert purged_state == [] + assert result == {"status": "deleted", "repo_id": "Org/Data"} + # Only the selected (active) cache's revision is deleted; the previous + # cache's copy is never touched. + assert calls == ["active"] + assert not (target_hub / "datasets--Org--Data").exists() + assert (other_hub / "datasets--Org--Data").exists() + + +def test_delete_processed_only_dataset_accepts_processed_cache_path(monkeypatch, tmp_path): + """A processed-only dataset row sends its Arrow cache path (___ + under HF_DATASETS_CACHE), which is not a Hub datasets-- dir. The delete must + accept it and run the processed-cache delete instead of raising 400.""" + datasets_root = tmp_path / "datasets" + processed_dir = datasets_root / "Org___Data" + processed_dir.mkdir(parents = True) + + # No Hub-cache copy exists; only the processed Arrow cache holds this repo. + monkeypatch.setattr(cache_inventory, "_collect_hf_cache_scans", lambda: ([], set())) + monkeypatch.setattr(cache_inventory, "_hf_datasets_cache_roots", lambda: [datasets_root]) + processed_calls: list[str] = [] + monkeypatch.setattr( + cache_inventory, + "_delete_processed_dataset_cache", + lambda repo_id, **_kwargs: (processed_calls.append(repo_id) or True, []), + ) + + result = cache_inventory._delete_cached_dataset_blocking("Org/Data", str(processed_dir)) + + assert result == {"status": "deleted", "repo_id": "Org/Data"} + assert processed_calls == ["Org/Data"] + + +def test_delete_processed_dataset_scopes_to_selected_root(monkeypatch, tmp_path): + """A dataset processed under two HF_DATASETS_CACHE roots is deleted only from + the selected root; the copy under the other cache home survives (real delete, + not stubbed).""" + selected_root = tmp_path / "selected" / "datasets" + other_root = tmp_path / "other" / "datasets" + for root in (selected_root, other_root): + (root / "Org___Data").mkdir(parents = True) + + monkeypatch.setattr(cache_inventory, "_collect_hf_cache_scans", lambda: ([], set())) + monkeypatch.setattr( + cache_inventory, "_hf_datasets_cache_roots", lambda: [selected_root, other_root] + ) + + result = cache_inventory._delete_cached_dataset_blocking( + "Org/Data", str(selected_root / "Org___Data") + ) + + assert result == {"status": "deleted", "repo_id": "Org/Data"} + assert not (selected_root / "Org___Data").exists() # the selected copy is deleted + assert (other_root / "Org___Data").exists() # the other cache home is untouched def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch): @@ -139,22 +197,22 @@ def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch): monkeypatch.setattr( cache_inventory, "_delete_processed_dataset_cache", - lambda _repo_id: (False, []), + lambda _repo_id, **_kwargs: (False, []), ) monkeypatch.setattr( cache_inventory, "purge_repo_cache_dirs", - lambda _repo_type, repo_id: purged_dirs.append(repo_id) or True, + lambda _repo_type, repo_id, **_kwargs: purged_dirs.append(repo_id) or True, ) monkeypatch.setattr( cache_inventory, "purge_partial_repo", - lambda *_args: False, + lambda *_args, **_kwargs: False, ) monkeypatch.setattr( cache_inventory.download_manifest, "purge_all_state_for_repo", - lambda *_args: 0, + lambda *_args, **_kwargs: 0, ) result = cache_inventory._delete_cached_dataset_blocking("Org/Data") @@ -172,22 +230,22 @@ def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch): monkeypatch.setattr( cache_inventory, "_delete_processed_dataset_cache", - lambda _repo_id: (False, []), + lambda _repo_id, **_kwargs: (False, []), ) monkeypatch.setattr( cache_inventory, "purge_repo_cache_dirs", - lambda *_args: False, + lambda *_args, **_kwargs: False, ) monkeypatch.setattr( cache_inventory, "purge_partial_repo", - lambda *_args: False, + lambda *_args, **_kwargs: False, ) monkeypatch.setattr( cache_inventory.download_manifest, "purge_all_state_for_repo", - lambda *_args: 0, + lambda *_args, **_kwargs: 0, ) with pytest.raises(HTTPException) as exc_info: diff --git a/studio/backend/hub/tests/test_download_manifest_scoping.py b/studio/backend/hub/tests/test_download_manifest_scoping.py new file mode 100644 index 0000000000..966eeaf0c2 --- /dev/null +++ b/studio/backend/hub/tests/test_download_manifest_scoping.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +from types import SimpleNamespace + +from hub.utils import download_manifest, state_dir + + +def _write_manifest(path, payload): + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(json.dumps(payload), encoding = "utf-8") + + +def test_purge_state_preserves_active_legacy_when_deleting_inactive_cache(monkeypatch, tmp_path): + """A scoped delete of an inactive cache must not erase the unscoped legacy + state, which _legacy_state_applies attributes to the active cache.""" + active = tmp_path / "active" / "hub" + previous = tmp_path / "previous" / "hub" + for path in (active, previous): + path.mkdir(parents = True) + + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = str(active)), + ) + + # Unowned legacy manifest -> belongs to the active cache. + legacy = state_dir.manifest_path("model", "Org/Model") + _write_manifest(legacy, {"version": 1}) + # The inactive cache's own scoped copy is the one being deleted. + scoped = state_dir.manifest_path("model", "Org/Model", hub_cache = str(previous)) + _write_manifest(scoped, {"version": 1, "hub_cache": str(previous)}) + + removed = download_manifest.purge_state("model", "Org/Model", hub_cache = str(previous)) + + assert removed is True + assert not scoped.is_file() # the inactive cache's copy is gone + assert legacy.is_file() # the active cache's legacy state survives + + +def test_purge_state_removes_legacy_owned_by_the_deleted_cache(monkeypatch, tmp_path): + """A legacy file that recorded the deleted cache as its owner is purged.""" + active = tmp_path / "active" / "hub" + previous = tmp_path / "previous" / "hub" + for path in (active, previous): + path.mkdir(parents = True) + + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = str(active)), + ) + + legacy = state_dir.manifest_path("model", "Org/Model") + _write_manifest(legacy, {"version": 1, "hub_cache": str(previous)}) + + removed = download_manifest.purge_state("model", "Org/Model", hub_cache = str(previous)) + + assert removed is True + assert not legacy.is_file() # owned by the deleted cache -> purged diff --git a/studio/backend/hub/tests/test_empty_variant_folder.py b/studio/backend/hub/tests/test_empty_variant_folder.py index 33bf6c6819..3ed8e69e0d 100644 --- a/studio/backend/hub/tests/test_empty_variant_folder.py +++ b/studio/backend/hub/tests/test_empty_variant_folder.py @@ -120,10 +120,16 @@ def _force_compute_to_raise(monkeypatch): monkeypatch.setattr(gguf_variants, "list_gguf_variants", _boom, raising = False) monkeypatch.setattr( - gguf_variants, "list_gguf_variants_from_hf_cache", lambda repo_id: None, raising = False + gguf_variants, + "list_gguf_variants_from_hf_cache", + lambda repo_id, root = None: None, + raising = False, ) monkeypatch.setattr( - gguf_variants, "list_partial_gguf_variants_from_state", lambda repo_id: None, raising = False + gguf_variants, + "list_partial_gguf_variants_from_state", + lambda repo_id, hub_cache = None: None, + raising = False, ) @@ -133,7 +139,11 @@ def test_get_variants_surfaces_cleanable_when_metadata_fails(monkeypatch): import asyncio _force_compute_to_raise(monkeypatch) - monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + monkeypatch.setattr( + gguf_variants, + "list_empty_gguf_variant_dirs", + lambda repo_id, root = None: {"UD-IQ1_S"}, + ) resp = asyncio.run( gguf_variants.get_gguf_variants_response( @@ -152,7 +162,11 @@ def test_get_variants_reraises_when_no_cleanable(monkeypatch): from fastapi import HTTPException _force_compute_to_raise(monkeypatch) - monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set()) + monkeypatch.setattr( + gguf_variants, + "list_empty_gguf_variant_dirs", + lambda repo_id, root = None: set(), + ) try: asyncio.run( diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 693d945ee1..fa5862a13c 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -2,6 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import asyncio +import json import sys from pathlib import Path from types import SimpleNamespace @@ -102,6 +103,236 @@ def test_big_endian_detection_ignores_model_name_be_token(): ) +def _cached_model_row(tmp_path: Path, *, partial: bool, active_cache: bool | None, size_bytes: int): + path = tmp_path / f"cache-{active_cache}-{partial}-{size_bytes}" + return model_common._local_model_info( + scan_path = path, + load_path = path, + source = "hf_cache", + model_format = "safetensors", + model_id = "Org/Model", + partial = partial, + active_cache = active_cache, + size_bytes = size_bytes, + ) + + +@pytest.mark.parametrize("reverse", [False, True]) +def test_local_inventory_prefers_complete_previous_cache_copy(tmp_path, reverse): + active_partial = _cached_model_row( + tmp_path, + partial = True, + active_cache = True, + size_bytes = 20, + ) + previous_complete = _cached_model_row( + tmp_path, + partial = False, + active_cache = False, + size_bytes = 10, + ) + rows = [active_partial, previous_complete] + if reverse: + rows.reverse() + + result = local_inventory._dedupe_local_models(rows) + + assert result == [previous_complete] + + +def test_local_inventory_compares_all_non_active_cache_copies(tmp_path): + inactive_partial = _cached_model_row( + tmp_path, + partial = True, + active_cache = False, + size_bytes = 20, + ) + custom_complete = _cached_model_row( + tmp_path, + partial = False, + active_cache = None, + size_bytes = 10, + ) + + assert local_inventory._dedupe_local_models([inactive_partial, custom_complete]) == [ + custom_complete + ] + + +def test_local_inventory_prefers_active_cache_when_copies_are_equally_complete(tmp_path): + previous = _cached_model_row( + tmp_path, + partial = False, + active_cache = False, + size_bytes = 20, + ) + active = _cached_model_row( + tmp_path, + partial = False, + active_cache = True, + size_bytes = 10, + ) + + assert local_inventory._dedupe_local_models([previous, active]) == [active] + + +def test_loaded_repo_match_accepts_previous_cache_snapshot_path(monkeypatch, tmp_path): + repo_dir = tmp_path / "old-hub" / "models--Org--Model" + snapshot = repo_dir / "snapshots" / "revision" + snapshot.mkdir(parents = True) + monkeypatch.setattr(deletion, "iter_repo_cache_dirs", lambda *_args: iter([repo_dir])) + + assert deletion._loaded_id_matches_repo(str(snapshot), "Org/Model") is True + assert deletion._loaded_id_matches_repo(str(snapshot / "model.gguf"), "Org/Model") is True + assert deletion._loaded_id_matches_repo(str(tmp_path / "other"), "Org/Model") is False + + +def test_cached_inventory_loads_previous_cache_copy_by_snapshot(monkeypatch, tmp_path): + active_hub = tmp_path / "active-hub" + previous_repo = tmp_path / "previous-hub" / "models--Org--Model" + snapshot = previous_repo / "snapshots" / "revision" + snapshot.mkdir(parents = True) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = active_hub), + ) + + fields = cache_inventory._cache_inventory_fields( + "Org/Model", + "safetensors", + repo_path = previous_repo, + snapshot_path = snapshot, + ) + + assert fields["load_id"] == str(snapshot) + + +def test_cached_inventory_keeps_repo_id_for_active_cache(monkeypatch, tmp_path): + active_hub = tmp_path / "active-hub" + active_repo = active_hub / "models--Org--Model" + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = active_hub), + ) + + fields = cache_inventory._cache_inventory_fields( + "Org/Model", + "safetensors", + repo_path = active_repo, + ) + + assert fields["load_id"] == "Org/Model" + + +def test_cached_inventory_prefers_active_copy_when_completeness_matches(): + previous = {"partial": False, "active_cache": False, "size_bytes": 200} + active = {"partial": False, "active_cache": True, "size_bytes": 100} + + assert cache_inventory._prefer_cache_row(active, previous) is True + assert cache_inventory._prefer_cache_row(previous, active) is False + + +def test_cached_inventory_prefers_complete_copy_before_active_cache(): + previous = {"partial": False, "active_cache": False, "size_bytes": 100} + active_partial = {"partial": True, "active_cache": True, "size_bytes": 200} + + assert cache_inventory._prefer_cache_row(previous, active_partial) is True + assert cache_inventory._prefer_cache_row(active_partial, previous) is False + + +def test_inventory_scans_every_dynamic_cache_root(monkeypatch, tmp_path): + first = tmp_path / "first-hub" + second = tmp_path / "second-hub" + unreadable = tmp_path / "unreadable-hub" + first.mkdir() + second.mkdir() + unreadable.mkdir() + scanned = [] + + monkeypatch.setattr( + inventory_scan, + "hf_cache_roots", + lambda: [first, unreadable, second], + ) + + def scan_cache(cache_dir): + path = Path(cache_dir) + scanned.append(path) + if path == unreadable: + raise PermissionError("unreadable") + return SimpleNamespace(cache_dir = cache_dir) + + monkeypatch.setattr("huggingface_hub.scan_cache_dir", scan_cache) + + result = inventory_scan._compute_all_hf_cache_scans() + + assert scanned == [first, unreadable, second] + assert [Path(scan.cache_dir) for scan in result] == [first, second] + + +def test_inventory_applies_download_state_to_its_owning_cache(monkeypatch, tmp_path): + state_root = tmp_path / "state" + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + repo_id = "Org/Model" + repo_name = "models--Org--Model" + repo_a = cache_a / repo_name + repo_b = cache_b / repo_name + snapshot_a = repo_a / "snapshots" / "revision" + snapshot_b = repo_b / "snapshots" / "revision" + snapshot_a.mkdir(parents = True) + snapshot_b.mkdir(parents = True) + (snapshot_a / "config.json").write_bytes(b"x") + (snapshot_b / "config.json").write_bytes(b"xx") + + monkeypatch.setattr(state_dir, "cache_root", lambda: state_root) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = cache_b), + ) + assert download_manifest.write_manifest( + "model", + repo_id, + None, + [download_manifest.ExpectedFile(path = "config.json", size = 2)], + "http", + hub_cache = cache_a, + ) + + assert inventory_scan.is_snapshot_partial("model", repo_id, repo_a) is True + assert inventory_scan.is_snapshot_partial("model", repo_id, repo_b) is False + assert inventory_scan.partial_transport_for("model", repo_id, None, repo_a) == "http" + assert inventory_scan.partial_transport_for("model", repo_id, None, repo_b) is None + + +def test_inventory_scopes_cancel_markers_to_their_owning_cache(monkeypatch, tmp_path): + state_root = tmp_path / "state" + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + repo_id = "Org/Model" + repo_name = "models--Org--Model" + repo_a = cache_a / repo_name + repo_b = cache_b / repo_name + repo_a.mkdir(parents = True) + repo_b.mkdir(parents = True) + + monkeypatch.setattr(state_dir, "cache_root", lambda: state_root) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = cache_b), + ) + assert download_manifest.write_cancel_marker( + "model", + repo_id, + "Q4_K_M", + "xet", + hub_cache = cache_a, + ) + + assert inventory_scan.is_variant_partial(repo_id, "Q4_K_M", repo_cache_dir = repo_a) is True + assert inventory_scan.is_variant_partial(repo_id, "Q4_K_M", repo_cache_dir = repo_b) is False + + def test_list_local_gguf_variants_skips_big_endian_sibling(tmp_path): (tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 100) (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"y" * 10) @@ -163,8 +394,19 @@ def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path "http", ) - marker_path = state_dir.marker_path("model", repo_id, variant) - manifest_path = state_dir.manifest_path("model", repo_id, variant) + hub_cache = download_manifest._canonical_hub_cache() + marker_path = state_dir.marker_path( + "model", + repo_id, + variant, + hub_cache = hub_cache, + ) + manifest_path = state_dir.manifest_path( + "model", + repo_id, + variant, + hub_cache = hub_cache, + ) assert marker_path is not None assert manifest_path is not None @@ -181,6 +423,97 @@ def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path ] +def test_download_state_isolated_across_hub_cache_switches(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + selected = SimpleNamespace(hub_cache = cache_a) + + from utils import hf_cache_settings + + monkeypatch.setattr(hf_cache_settings, "get_hf_cache_paths", lambda: selected) + expected_a = [download_manifest.ExpectedFile(path = "a.gguf", size = 1)] + expected_b = [download_manifest.ExpectedFile(path = "b.gguf", size = 2)] + + assert download_manifest.write_manifest("model", "Owner/Repo", "Q4_K_M", expected_a) + assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http") + + selected.hub_cache = cache_b + assert download_manifest.write_manifest("model", "Owner/Repo", "Q4_K_M", expected_b) + + manifest_b = download_manifest.read_manifest("model", "Owner/Repo", "Q4_K_M") + manifest_a = download_manifest.read_manifest( + "model", + "Owner/Repo", + "Q4_K_M", + hub_cache = cache_a, + ) + + assert manifest_b is not None and manifest_b.expected_files == tuple(expected_b) + assert manifest_a is not None and manifest_a.expected_files == tuple(expected_a) + assert not download_manifest.has_cancel_marker("model", "Owner/Repo", "Q4_K_M") + assert download_manifest.has_cancel_marker( + "model", + "Owner/Repo", + "Q4_K_M", + hub_cache = cache_a, + ) + assert len(list((tmp_path / "hub-state" / "manifests").rglob("*.json"))) == 2 + + +def test_legacy_unscoped_download_state_falls_back_only_for_selected_cache(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = cache_a), + ) + manifest = state_dir.manifest_path("model", "Owner/Repo", "Q4_K_M") + marker = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M") + assert manifest is not None and marker is not None + manifest.write_text( + json.dumps( + { + "version": 1, + "repo_id": "Owner/Repo", + "variant": "Q4_K_M", + "expected_files": [{"path": "model.gguf", "size": 10}], + "transport": "http", + } + ), + encoding = "utf-8", + ) + marker.write_text( + json.dumps({"version": 1, "repo_id": "Owner/Repo", "variant": "Q4_K_M"}), + encoding = "utf-8", + ) + + assert download_manifest.read_manifest("model", "Owner/Repo", "Q4_K_M") is not None + assert download_manifest.has_cancel_marker("model", "Owner/Repo", "Q4_K_M") + assert list(download_manifest.iter_variant_manifests("model", "Owner/Repo")) == [ + ("Q4_K_M", manifest) + ] + assert list(download_manifest.iter_variant_markers("model", "Owner/Repo")) == [ + ("Q4_K_M", marker) + ] + assert ( + download_manifest.read_manifest( + "model", + "Owner/Repo", + "Q4_K_M", + hub_cache = cache_b, + ) + is None + ) + assert not download_manifest.has_cancel_marker( + "model", + "Owner/Repo", + "Q4_K_M", + hub_cache = cache_b, + ) + + class _RecordingLogger: def __init__(self): self.warnings = [] @@ -416,8 +749,15 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa "Q4_K_M", [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)], "http", + hub_cache = repo_path.parent, + ) + assert download_manifest.write_cancel_marker( + "model", + "Org/PartialGguf", + "Q4_K_M", + "http", + hub_cache = repo_path.parent, ) - assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http") monkeypatch.setattr( cache_inventory, "all_hf_cache_scans", @@ -484,6 +824,7 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa "Q8_0", [download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)], "http", + hub_cache = Path(embedder.repo_path).parent, ) monkeypatch.setattr( cache_inventory, @@ -1206,6 +1547,7 @@ def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, ), ], "http", + hub_cache = entry.parent, ) requirement = gguf_variants._GgufVariantRequirement( @@ -1292,6 +1634,7 @@ def test_gguf_progress_subtracts_new_job_completed_baseline(monkeypatch, tmp_pat ), ], "http", + hub_cache = entry.parent, ) requirement = gguf_variants._GgufVariantRequirement( @@ -1462,6 +1805,7 @@ def test_gguf_progress_complete_on_disk_ignores_full_baseline(monkeypatch, tmp_p ), ], "http", + hub_cache = entry.parent, ) requirement = gguf_variants._GgufVariantRequirement( @@ -1861,8 +2205,15 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_ "Q4_K_M", [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)], "http", + hub_cache = cache_dir, + ) + assert download_manifest.write_cancel_marker( + "model", + "Org/PartialGguf", + "Q4_K_M", + "http", + hub_cache = cache_dir, ) - assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http") monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []) monkeypatch.setattr( local_inventory.hf_cache_scan, @@ -2117,7 +2468,7 @@ def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch monkeypatch.setattr( gguf_variants, "iter_hf_cache_snapshots", - lambda _repo_id: [snapshot], + lambda _repo_id, root = None: [snapshot], ) monkeypatch.setattr( gguf_variants, @@ -2136,6 +2487,70 @@ def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch assert result.variants[0].partial is True +def test_gguf_variants_scopes_partial_state_to_requested_cache(monkeypatch, tmp_path): + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + repo_id = "Org/SharedRepo" + repo_name = "models--Org--SharedRepo" + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + repo_a = cache_a / repo_name + snapshot_a = repo_a / "snapshots" / "revision" + snapshot_a.mkdir(parents = True) + (snapshot_a / "model-Q8_0.gguf").write_bytes(b"complete") + blobs_b = cache_b / repo_name / "blobs" + blobs_b.mkdir(parents = True) + (blobs_b / "q8-hash.incomplete").write_bytes(b"partial") + + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = cache_b), + ) + assert download_manifest.write_cancel_marker( + "model", + repo_id, + "Q8_0", + "http", + hub_cache = cache_b, + ) + monkeypatch.setattr( + gguf_variants, + "list_gguf_variants", + lambda *_args, **_kwargs: ( + [ + SimpleNamespace( + filename = "model-Q8_0.gguf", + quant = "Q8_0", + display_label = None, + size_bytes = 8, + ) + ], + False, + [ + SimpleNamespace( + rfilename = "model-Q8_0.gguf", + size = 8, + lfs = SimpleNamespace(sha256 = "q8-hash"), + ) + ], + ), + ) + monkeypatch.setattr(cache_inventory, "all_hf_cache_scans", lambda: []) + + result = asyncio.run( + gguf_variants.get_gguf_variants_response( + repo_id, + local_path = str(repo_a), + ) + ) + + assert result.variants[0].downloaded is True + assert result.variants[0].partial is False + + def test_download_registry_repo_keys_are_case_insensitive(): registry = download_registry.DownloadRegistry() @@ -2444,6 +2859,34 @@ def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, t assert (blobs / "shared-mmproj.incomplete").exists() +def test_prepare_cache_for_transport_uses_captured_root(monkeypatch, tmp_path): + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + repo_name = "models--Org--Repo" + partial_a = cache_a / repo_name / "blobs" / "blob.incomplete" + partial_b = cache_b / repo_name / "blobs" / "blob.incomplete" + partial_a.parent.mkdir(parents = True) + partial_b.parent.mkdir(parents = True) + partial_a.write_bytes(b"a") + partial_b.write_bytes(b"b") + monkeypatch.setattr( + download_registry, + "hf_cache_root", + lambda create = False, root = None: root or cache_b, + ) + + purged = download_registry.prepare_cache_for_transport( + "model", + "Org/Repo", + download_registry.TRANSPORT_HTTP, + root = cache_a, + ) + + assert purged == 1 + assert not partial_a.exists() + assert partial_b.exists() + + def _vision_cache_root(monkeypatch, tmp_path): root = tmp_path / "hub" blobs = root / "models--Org--Vision" / "blobs" @@ -2802,6 +3245,47 @@ def test_shutdown_skips_marker_for_worker_that_exits_cleanly(monkeypatch): assert markers == ["Org/Cut"] +def test_orphan_reaper_uses_worker_cache_root_after_setting_changes(monkeypatch, tmp_path): + workers = tmp_path / "workers" + workers.mkdir() + cache_a = tmp_path / "cache-a" / "hub" + cache_b = tmp_path / "cache-b" / "hub" + partial = cache_a / "models--Org--Model" / "blobs" / "abc.incomplete" + partial.parent.mkdir(parents = True) + partial.write_bytes(b"partial") + cache_b.mkdir(parents = True) + monkeypatch.setattr(state_dir, "workers_dir", lambda: workers) + monkeypatch.setattr(download_registry, "_process_alive", lambda _pid: False) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = cache_b), + ) + markers = [] + monkeypatch.setattr( + download_registry, + "persist_cancel_marker", + lambda *args, **kwargs: markers.append(args), + ) + metadata = download_registry.DownloadMetadata( + repo_type = "model", + repo_id = "Org/Model", + variant = None, + transport = download_registry.TRANSPORT_HTTP, + hub_cache = str(cache_a), + xet_cache = str(tmp_path / "cache-a" / "xet"), + ) + download_registry.write_worker_breadcrumb("org/model", 1234, metadata) + [breadcrumb] = list(workers.iterdir()) + payload = json.loads(breadcrumb.read_text(encoding = "utf-8")) + assert payload["hub_cache"] == str(cache_a) + assert payload["xet_cache"] == str(tmp_path / "cache-a" / "xet") + + download_registry.reap_orphan_workers() + + assert markers == [("model", "Org/Model", None, "http")] + assert list(workers.iterdir()) == [] + + def test_model_claim_register_cancel_uses_registry_marker_owner(monkeypatch): killed = [] @@ -3125,12 +3609,19 @@ def _build_variant_cache_repo(repo_dir, blob_specs, snapshot_links): return repo -def _patch_variant_delete_side_effects(monkeypatch): +def _patch_variant_delete_side_effects(monkeypatch, hub_cache = None): monkeypatch.setattr( deletion.download_manifest, "purge_state", lambda *_args, **_kwargs: False, ) + # The repo under test lives in this cache; make it the active one so the + # delete scopes to it (default target root is the active hub cache). + if hub_cache is not None: + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = hub_cache), + ) def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path): @@ -3308,7 +3799,7 @@ def test_delete_variant_keeps_blob_shared_with_other_snapshot(monkeypatch, tmp_p "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])], ) - _patch_variant_delete_side_effects(monkeypatch) + _patch_variant_delete_side_effects(monkeypatch, tmp_path) result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None) @@ -3335,7 +3826,7 @@ def test_delete_variant_unlinks_unshared_blob(monkeypatch, tmp_path): "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])], ) - _patch_variant_delete_side_effects(monkeypatch) + _patch_variant_delete_side_effects(monkeypatch, tmp_path) result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None) @@ -3361,7 +3852,7 @@ def test_delete_variant_surfaces_locked_file_as_conflict(monkeypatch, tmp_path): "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])], ) - _patch_variant_delete_side_effects(monkeypatch) + _patch_variant_delete_side_effects(monkeypatch, tmp_path) real_unlink = Path.unlink diff --git a/studio/backend/hub/utils/download_manifest.py b/studio/backend/hub/utils/download_manifest.py index 5366689296..ac0ccb5490 100644 --- a/studio/backend/hub/utils/download_manifest.py +++ b/studio/backend/hub/utils/download_manifest.py @@ -77,6 +77,7 @@ class Manifest: started_at: str expected_files: tuple[ExpectedFile, ...] transport: Optional[str] = None + hub_cache: Optional[str] = None @dataclass(frozen = True) @@ -86,6 +87,78 @@ class VerifyResult: size_mismatched: tuple[str, ...] +def _canonical_hub_cache(hub_cache: Optional[str | Path] = None) -> Optional[str]: + if hub_cache is None: + try: + from utils.hf_cache_settings import get_hf_cache_paths + hub_cache = get_hf_cache_paths().hub_cache + except Exception: + return None + try: + return str(Path(hub_cache).expanduser().resolve(strict = False)) + except (OSError, RuntimeError, ValueError): + return str(hub_cache) + + +def _read_state_payload(path: Path) -> Optional[dict]: + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + logger.debug("Could not read Hub state %s: %s", path, exc) + return None + return data if isinstance(data, dict) else None + + +def _legacy_state_applies( + path: Path, + requested_hub_cache: Optional[str], + *, + fail_closed: bool = False, +) -> bool: + """Whether an old unscoped state file belongs to the requested cache. + + Transitional files that recorded their cache keep that ownership. Older + files with no ownership can only be attributed to the currently selected + cache, which matches the single-cache behavior under which they were + written without leaking them into remembered inactive caches. + """ + data = _read_state_payload(path) + if data is not None: + recorded = data.get("hub_cache") + if isinstance(recorded, str) and recorded: + return _canonical_hub_cache(recorded) == requested_hub_cache + elif not fail_closed: + return False + return requested_hub_cache == _canonical_hub_cache() + + +def _state_read_path( + path_factory, + repo_type: RepoType, + repo_id: str, + variant: Optional[str], + hub_cache: Optional[str | Path], + *, + fail_closed: bool = False, +) -> Optional[Path]: + requested = _canonical_hub_cache(hub_cache) + scoped = path_factory(repo_type, repo_id, variant, hub_cache = requested) + try: + if scoped is not None and scoped.is_file(): + return scoped + except OSError: + pass + legacy = path_factory(repo_type, repo_id, variant) + if legacy is None or legacy == scoped: + return None + try: + if not legacy.is_file(): + return None + except OSError: + return None + return legacy if _legacy_state_applies(legacy, requested, fail_closed = fail_closed) else None + + def _atomic_write_json(path: Path, payload: dict) -> bool: # Per-write uuid suffix so a concurrent caller or a stale tmp from a # previous crash cannot collide with the in-flight write. @@ -124,6 +197,8 @@ def write_manifest( variant: Optional[str], expected_files: Sequence[ExpectedFile], transport: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> bool: """Write/overwrite the manifest for this triple. Best-effort. @@ -131,7 +206,13 @@ def write_manifest( worst-case fallback is the pre-fix scanner behavior (one missed partial detection), which is no regression. """ - path = manifest_path(repo_type, repo_id, variant) + recorded_hub_cache = _canonical_hub_cache(hub_cache) + path = manifest_path( + repo_type, + repo_id, + variant, + hub_cache = recorded_hub_cache, + ) if path is None: return False payload = { @@ -149,6 +230,7 @@ def write_manifest( for f in expected_files ], "transport": transport, + "hub_cache": recorded_hub_cache, } return _atomic_write_json(path, payload) @@ -157,6 +239,8 @@ def read_manifest( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> Optional[Manifest]: """Return the manifest if present and parseable; ``None`` otherwise. @@ -171,15 +255,17 @@ def read_manifest( ``_MANIFEST_VERSION`` and widen this check) or live under a different filename, so an incompatible payload can never mis-classify rows. """ - path = manifest_path(repo_type, repo_id, variant) + path = _state_read_path( + manifest_path, + repo_type, + repo_id, + variant, + hub_cache, + ) if path is None or not path.is_file(): return None - try: - data = json.loads(path.read_text(encoding = "utf-8")) - except (OSError, ValueError) as exc: - logger.debug("Could not read manifest %s: %s", path, exc) - return None - if not isinstance(data, dict): + data = _read_state_payload(path) + if data is None: return None if data.get("version") != _MANIFEST_VERSION: logger.debug( @@ -216,6 +302,7 @@ def read_manifest( started_at = str(data.get("started_at", "")), expected_files = tuple(expected), transport = transport if transport in ("http", "xet") else None, + hub_cache = data.get("hub_cache") if isinstance(data.get("hub_cache"), str) else None, ) @@ -289,6 +376,8 @@ def write_cancel_marker( repo_id: str, variant: Optional[str] = None, transport: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> bool: """Record that this triple was cancelled. Idempotent across repeated cancels. @@ -296,7 +385,13 @@ def write_cancel_marker( inventory rows so the UI labels HTTP retries as continuable and XET retries as full redownloads. None is accepted for forward-compat. """ - path = marker_path(repo_type, repo_id, variant) + recorded_hub_cache = _canonical_hub_cache(hub_cache) + path = marker_path( + repo_type, + repo_id, + variant, + hub_cache = recorded_hub_cache, + ) if path is None: return False payload = { @@ -306,6 +401,7 @@ def write_cancel_marker( "variant": variant, "transport": transport, "cancelled_at": datetime.now(timezone.utc).isoformat(), + "hub_cache": recorded_hub_cache, } return _atomic_write_json(path, payload) @@ -314,6 +410,8 @@ def read_cancel_marker_transport( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> Optional[str]: """Return the transport recorded in the cancel marker, or ``None`` if no marker exists or it is unreadable. @@ -330,15 +428,17 @@ def read_cancel_marker_transport( ``None`` keeps the neutral "Retry" label. * Unknown future versions โ†’ ``None`` (unknown layout, unknown transport). """ - path = marker_path(repo_type, repo_id, variant) + path = _state_read_path( + marker_path, + repo_type, + repo_id, + variant, + hub_cache, + ) if path is None or not path.is_file(): return None - try: - data = json.loads(path.read_text(encoding = "utf-8")) - except (OSError, ValueError) as exc: - logger.debug("Could not read cancel marker %s: %s", path, exc) - return None - if not isinstance(data, dict): + data = _read_state_payload(path) + if data is None: return None version = data.get("version") if version == _LEGACY_MARKER_VERSION: @@ -351,10 +451,30 @@ def read_cancel_marker_transport( return None +def _all_matching_state_paths( + parent: Optional[Path], repo_type: RepoType, repo_id: str, variant: Optional[str] +) -> tuple[Path, ...]: + if parent is None: + return () + legacy_path = ( + manifest_path(repo_type, repo_id, variant) + if parent.name == "manifests" + else marker_path(repo_type, repo_id, variant) + ) + if legacy_path is None: + return () + try: + return tuple(path for path in parent.rglob(legacy_path.name) if path.is_file()) + except OSError: + return () + + def clear_cancel_marker( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> None: """Remove the cancel marker for this triple if present. @@ -362,31 +482,48 @@ def clear_cancel_marker( download-start (a fresh attempt supersedes prior cancel state) and again at successful completion (cleans up if the start clear failed). """ - path = marker_path(repo_type, repo_id, variant) - if path is None: - return - try: - path.unlink(missing_ok = True) - except OSError as exc: - logger.debug("Could not clear cancel marker %s: %s", path, exc) + requested = _canonical_hub_cache(hub_cache) + path = marker_path( + repo_type, + repo_id, + variant, + hub_cache = requested, + ) + legacy = marker_path(repo_type, repo_id, variant) + paths = [path] + if ( + legacy is not None + and legacy != path + and _legacy_state_applies(legacy, requested, fail_closed = True) + ): + paths.append(legacy) + for target in paths: + if target is None: + continue + try: + target.unlink(missing_ok = True) + except OSError as exc: + logger.debug("Could not clear cancel marker %s: %s", target, exc) def has_cancel_marker( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> bool: - """File-existence check only. Body is never read. - - Fail-closed: a corrupt marker still returns ``True`` because the - file's existence is the signal (the user once cancelled this - triple, even if the body is unreadable). - """ - path = marker_path(repo_type, repo_id, variant) - if path is None: - return False + """Return whether a cancel marker applies to the selected cache.""" + path = _state_read_path( + marker_path, + repo_type, + repo_id, + variant, + hub_cache, + fail_closed = True, + ) try: - return path.is_file() + return path is not None and path.is_file() except OSError: return False @@ -395,48 +532,124 @@ def delete_manifest( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> bool: - path = manifest_path(repo_type, repo_id, variant) - if path is None: - return False - try: - if not path.is_file(): - return False - path.unlink() - return True - except OSError as exc: - logger.debug("Could not delete manifest %s: %s", path, exc) - return False + requested = _canonical_hub_cache(hub_cache) + path = manifest_path( + repo_type, + repo_id, + variant, + hub_cache = requested, + ) + legacy = manifest_path(repo_type, repo_id, variant) + paths = [path] + if legacy is not None and legacy != path and _legacy_state_applies(legacy, requested): + paths.append(legacy) + removed = False + for target in paths: + if target is None: + continue + try: + if target.is_file(): + target.unlink() + removed = True + except OSError as exc: + logger.debug("Could not delete manifest %s: %s", target, exc) + return removed def purge_state( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> bool: """Remove manifest + cancel marker for this triple. Returns ``True`` - when anything was present on disk before the call. Idempotent.""" - marker_existed = has_cancel_marker(repo_type, repo_id, variant) - manifest_removed = delete_manifest(repo_type, repo_id, variant) - clear_cancel_marker(repo_type, repo_id, variant) - return marker_existed or manifest_removed + when anything was present on disk before the call. Idempotent. + + With ``hub_cache`` set, only that cache's scoped state (plus any legacy + unscoped file that belongs to it) is removed, so purging one cache's copy + never clears another cache's resumable/cancel state.""" + if hub_cache is None: + paths = ( + *_all_matching_state_paths(manifests_dir(), repo_type, repo_id, variant), + *_all_matching_state_paths(cancelled_dir(), repo_type, repo_id, variant), + ) + else: + requested = _canonical_hub_cache(hub_cache) + candidates = [ + manifest_path(repo_type, repo_id, variant, hub_cache = hub_cache), + marker_path(repo_type, repo_id, variant, hub_cache = hub_cache), + ] + # Legacy unscoped state is shared: an unowned file belongs to the active + # cache (per _legacy_state_applies), so only purge it when it belongs to + # the cache being deleted -- else deleting an inactive cache would erase + # the active cache's resume/cancel state. + for path_factory in (manifest_path, marker_path): + legacy = path_factory(repo_type, repo_id, variant) + if legacy is not None and _legacy_state_applies(legacy, requested): + candidates.append(legacy) + paths = tuple(p for p in candidates if p is not None) + removed = False + for path in paths: + try: + if path.is_file(): + path.unlink() + removed = True + except OSError as exc: + logger.debug("Could not purge Hub state %s: %s", path, exc) + return removed -def purge_all_state_for_repo(repo_type: RepoType, repo_id: str) -> int: +def purge_all_state_for_repo( + repo_type: RepoType, + repo_id: str, + *, + hub_cache: Optional[str | Path] = None, +) -> int: """Remove the snapshot-level manifest + marker AND every variant-keyed manifest + marker for this repo. Used by the route delete handlers so scanner state never outlives the cache it described. Returns the count - of (repo, variant) triples that had any state on disk.""" + of (repo, variant) triples that had any state on disk. + + With ``hub_cache`` set, only that cache's scoped state (plus any legacy + unscoped file) is enumerated and removed, so deleting one cache's copy does + not clear another cache's resumable/cancel state.""" removed = 0 - if purge_state(repo_type, repo_id, None): + if purge_state(repo_type, repo_id, None, hub_cache = hub_cache): removed += 1 variants: set[str] = set() - for variant, _ in iter_variant_manifests(repo_type, repo_id): - variants.add(variant) - for variant, _ in iter_variant_markers(repo_type, repo_id): - variants.add(variant) + prefix = variant_filename_prefix(repo_type, repo_id) + if hub_cache is None: + search = [(p, True) for p in (manifests_dir(), cancelled_dir()) if p is not None] + else: + # This cache's scoped dir (parent of its scoped path) plus the legacy + # unscoped base; glob (not rglob) so other caches' dirs are not swept. + search = [] + for scoped, base in ( + (manifest_path(repo_type, repo_id, None, hub_cache = hub_cache), manifests_dir()), + (marker_path(repo_type, repo_id, None, hub_cache = hub_cache), cancelled_dir()), + ): + if scoped is not None: + search.append((scoped.parent, False)) + if base is not None: + search.append((base, False)) + for parent, recursive in search: + try: + entries = tuple( + parent.rglob(f"{prefix}*.json") if recursive else parent.glob(f"{prefix}*.json") + ) + except OSError: + continue + for entry in entries: + if not entry.is_file(): + continue + fallback = entry.stem[len(prefix) :] + variants.add(_variant_from_state_file(entry, fallback)) for variant in variants: - if purge_state(repo_type, repo_id, variant): + if purge_state(repo_type, repo_id, variant, hub_cache = hub_cache): removed += 1 return removed @@ -453,35 +666,83 @@ def _variant_from_state_file(path: Path, fallback: str) -> str: def _iter_variant_state_files( - parent: Optional[Path], repo_type: RepoType, repo_id: str + parent: Optional[Path], + repo_type: RepoType, + repo_id: str, + hub_cache: Optional[str | Path], + *, + cancel_markers: bool, ) -> Iterator[tuple[str, Path]]: if parent is None: return - prefix = variant_filename_prefix(repo_type, repo_id) - try: - entries = list(parent.iterdir()) - except OSError: + path_factory = marker_path if cancel_markers else manifest_path + requested = _canonical_hub_cache(hub_cache) + scoped_probe = path_factory( + repo_type, + repo_id, + None, + hub_cache = requested, + ) + if scoped_probe is None: return - for entry in entries: - if not entry.is_file() or not entry.name.endswith(".json"): + prefix = variant_filename_prefix(repo_type, repo_id) + seen: set[str] = set() + for directory, legacy in ((scoped_probe.parent, False), (parent, True)): + if legacy and directory == scoped_probe.parent: continue - stem = entry.name[: -len(".json")] - if not stem.lower().startswith(prefix): + try: + entries = list(directory.iterdir()) + except OSError: continue - variant = stem[len(prefix) :] - if variant: - yield _variant_from_state_file(entry, variant), entry + for entry in entries: + if not entry.is_file() or not entry.name.endswith(".json"): + continue + stem = entry.name[: -len(".json")] + if not stem.lower().startswith(prefix) or entry.name in seen: + continue + if legacy and not _legacy_state_applies( + entry, + requested, + fail_closed = cancel_markers, + ): + continue + fallback = stem[len(prefix) :] + if fallback: + seen.add(entry.name) + yield _variant_from_state_file(entry, fallback), entry -def iter_variant_manifests(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]: +def iter_variant_manifests( + repo_type: RepoType, + repo_id: str, + *, + hub_cache: Optional[str | Path] = None, +) -> Iterator[tuple[str, Path]]: """Yield (variant, manifest_path) for every variant-keyed manifest written for this repo. Used by is_gguf_repo_partial to enumerate all variants present on disk so the all-variants-broken gate can run.""" - yield from _iter_variant_state_files(manifests_dir(), repo_type, repo_id) + yield from _iter_variant_state_files( + manifests_dir(), + repo_type, + repo_id, + hub_cache, + cancel_markers = False, + ) -def iter_variant_markers(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]: +def iter_variant_markers( + repo_type: RepoType, + repo_id: str, + *, + hub_cache: Optional[str | Path] = None, +) -> Iterator[tuple[str, Path]]: """Yield (variant, marker_path) for every variant-keyed cancel marker. Companion to iter_variant_manifests: catches variants cancelled before download-start ever wrote a manifest (very early failures).""" - yield from _iter_variant_state_files(cancelled_dir(), repo_type, repo_id) + yield from _iter_variant_state_files( + cancelled_dir(), + repo_type, + repo_id, + hub_cache, + cancel_markers = True, + ) diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index b6bdee3bce..9e2b7d1a6d 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -129,6 +129,8 @@ def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMeta "cancel_marker_transport": metadata.cancel_marker_transport if metadata is not None else None, + "hub_cache": metadata.hub_cache if metadata is not None else None, + "xet_cache": metadata.xet_cache if metadata is not None else None, } tmp = path.with_name(f".{path.name}.tmp-{pid}") try: @@ -236,6 +238,7 @@ def _settle_orphaned_download( repo_id: Optional[str], variant: Optional[str], transport: Optional[str], + hub_cache: Optional[str] = None, ) -> None: """Persist a cancel marker for a reaped orphan still mid-download so the next launch settles it to a resumable "cancelled" state instead of a phantom-running @@ -251,18 +254,42 @@ def _settle_orphaned_download( return from hub.utils import download_manifest - manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + cache_root = Path(hub_cache) if isinstance(hub_cache, str) and hub_cache else None + + manifest = download_manifest.read_manifest( + repo_type, + repo_id, + variant, + hub_cache = cache_root, + ) if repo_type == "model" and variant and manifest is None: return if manifest is None: - if not has_active_incomplete_blobs(repo_type, repo_id): + if not has_active_incomplete_blobs(repo_type, repo_id, root = cache_root): return else: - if _manifest_verifies_against_active_cache(repo_type, repo_id, manifest): + if _manifest_verifies_against_active_cache( + repo_type, + repo_id, + manifest, + root = cache_root, + ): return - if not _manifest_has_active_incomplete_blobs(repo_type, repo_id, manifest): + if not _manifest_has_active_incomplete_blobs( + repo_type, + repo_id, + manifest, + root = cache_root, + ): return - persist_cancel_marker(repo_type, repo_id, variant, transport, logger = logger) + persist_cancel_marker( + repo_type, + repo_id, + variant, + transport, + hub_cache = hub_cache, + logger = logger, + ) def reap_orphan_workers() -> None: @@ -309,6 +336,7 @@ def reap_orphan_workers() -> None: repo_id, data.get("variant"), data.get("cancel_marker_transport") or data.get("transport"), + data.get("hub_cache"), ) except Exception as exc: logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc) @@ -355,8 +383,13 @@ def _purge_incomplete_blobs( return removed -def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): +def _iter_active_snapshot_dirs( + repo_type: str, + repo_id: str, + *, + root: Optional[Path] = None, +) -> Iterator[Path]: + for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root): snapshots_dir = entry / "snapshots" if not snapshots_dir.is_dir(): continue @@ -369,24 +402,41 @@ def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: yield snapshot -def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manifest) -> bool: +def _manifest_verifies_against_active_cache( + repo_type: str, + repo_id: str, + manifest, + *, + root: Optional[Path] = None, +) -> bool: from hub.utils import download_manifest - for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id): + for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id, root = root): if download_manifest.verify_against_disk(manifest, snapshot_dir).ok: return True return False -def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest) -> bool: +def _manifest_has_active_incomplete_blobs( + repo_type: str, + repo_id: str, + manifest, + *, + root: Optional[Path] = None, +) -> bool: if not getattr(manifest, "variant", None): - return has_active_incomplete_blobs(repo_type, repo_id) + return has_active_incomplete_blobs(repo_type, repo_id, root = root) expected_hashes = frozenset( expected.sha256 for expected in manifest.expected_files if expected.sha256 ) if not expected_hashes: - return has_active_incomplete_blobs(repo_type, repo_id) + return has_active_incomplete_blobs(repo_type, repo_id, root = root) return bool( - incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection(expected_hashes) + incomplete_blob_hashes( + repo_type, + repo_id, + active_only = True, + root = root, + ).intersection(expected_hashes) ) @@ -459,6 +509,7 @@ def prepare_cache_for_transport( only_blob_hashes: Optional[frozenset[str]] = None, companion_blob_hashes: Optional[frozenset[str]] = None, protected_blob_hashes: Optional[frozenset[str]] = None, + root: Optional[Path] = None, ) -> int: """Guarantee any pre-existing ``.incomplete`` blobs are SAFE to resume under *mode*. Returns the number of partial blobs purged for untrusted provenance. @@ -485,14 +536,13 @@ def prepare_cache_for_transport( they are excluded from every purge so a shared companion is never deleted mid-write. - Scope: only the active ``HF_HUB_CACHE`` root is inspected. That suffices for - resume safety because ``snapshot_download`` runs without a ``cache_dir`` - override and so can only read or resume a ``.incomplete`` under this same - active root. Markers are written for the new mode before returning. + Scope: ``root`` selects the cache captured by the caller. It defaults to the + active ``HF_HUB_CACHE`` root for workers that inherit their cache through + the environment. Markers are written for the new mode before returning. """ if mode not in VALID_TRANSPORTS: raise ValueError(f"Invalid transport mode: {mode!r}") - root = hf_cache_root(create = True) + root = hf_cache_root(create = True) if root is None else hf_cache_root(create = True, root = root) if root is None: return 0 target = target_dir_name(repo_type, repo_id) @@ -618,10 +668,11 @@ def incomplete_blob_hashes( repo_id: str, *, active_only: bool = False, + root: Optional[Path] = None, ) -> set[str]: out: set[str] = set() entries = ( - iter_active_repo_cache_dirs(repo_type, repo_id) + iter_active_repo_cache_dirs(repo_type, repo_id, root = root) if active_only else iter_repo_cache_dirs(repo_type, repo_id) ) @@ -638,16 +689,24 @@ def incomplete_blob_hashes( return out -def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int: - """Sum finalized blob bytes for *blob_hashes* in the active HF cache root. +def completed_blob_bytes( + repo_type: str, + repo_id: str, + blob_hashes: frozenset[str], + *, + root: Optional[Path] = None, +) -> int: + """Sum finalized blob bytes for *blob_hashes* in a single HF cache root. - A worker only writes to the active ``HF_HUB_CACHE`` root, so a baseline must - ignore legacy/default roots that ``snapshot_download`` won't reuse this run. + A worker only writes to its captured ``HF_HUB_CACHE`` root, so a baseline + must be scoped to that root (``root``), not re-resolved to whatever cache is + active now; otherwise a runtime cache switch makes the retry baseline count + bytes from the wrong disk. """ if not blob_hashes: return 0 total = 0 - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root): blobs_dir = entry / "blobs" if not blobs_dir.is_dir(): continue @@ -712,6 +771,8 @@ class DownloadMetadata: # Bytes already complete before this job started; not counted as this run's # progress. completed_baseline_bytes: int = 0 + hub_cache: Optional[str] = None + xet_cache: Optional[str] = None @dataclass(frozen = True) @@ -752,6 +813,7 @@ def persist_cancel_marker( variant: Optional[str], transport: Optional[str], *, + hub_cache: Optional[str] = None, logger = logger, ) -> None: if not repo_type or not repo_id: @@ -763,6 +825,7 @@ def persist_cancel_marker( repo_id, variant, transport = transport, + hub_cache = hub_cache, ): logger.debug("write_cancel_marker returned False for %s", repo_id) except Exception as exc: @@ -971,6 +1034,7 @@ class DownloadRegistry: metadata_to_persist.repo_id, metadata_to_persist.variant, metadata_to_persist.transport, + hub_cache = metadata_to_persist.hub_cache, ) return False @@ -1033,6 +1097,8 @@ class DownloadRegistry: replace_active: bool = False, metadata_transport: Optional[str] = None, cancel_marker_transport: Optional[str] = None, + hub_cache: Optional[str] = None, + xet_cache: Optional[str] = None, ) -> tuple[bool, str]: key = normalize_job_key(key) repo = _repo_of_key(key) @@ -1106,6 +1172,8 @@ class DownloadRegistry: 0, int(completed_baseline_bytes or 0), ), + hub_cache = hub_cache, + xet_cache = xet_cache, ) if cancel_marker_transport is not None: self._cancel_marker_transports[key] = cancel_marker_transport @@ -1386,6 +1454,7 @@ class DownloadRegistry: metadata.repo_id, metadata.variant, metadata.cancel_marker_transport or metadata.transport, + hub_cache = metadata.hub_cache, ) reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = [] for key, proc, metadata in live: @@ -1401,6 +1470,7 @@ class DownloadRegistry: metadata.repo_id, metadata.variant, metadata.cancel_marker_transport or metadata.transport, + hub_cache = metadata.hub_cache, ) continue reaped.append((key, proc, metadata)) @@ -1421,6 +1491,7 @@ class DownloadRegistry: metadata.repo_id, metadata.variant, metadata.cancel_marker_transport or metadata.transport, + hub_cache = metadata.hub_cache, ) diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index 2e3de125f1..eb768db5d6 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -253,11 +253,16 @@ def _env_offline() -> bool: ) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes") -def iter_hf_cache_snapshots(repo_id: str): - from hub.utils.hf_cache_state import iter_repo_cache_dirs +def iter_hf_cache_snapshots(repo_id: str, root: Optional[Path] = None): + from hub.utils.hf_cache_state import iter_active_repo_cache_dirs, iter_repo_cache_dirs snapshots: list[Path] = [] - for repo_dir in iter_repo_cache_dirs("model", repo_id): + repo_dirs = ( + iter_active_repo_cache_dirs("model", repo_id, root = root) + if root is not None + else iter_repo_cache_dirs("model", repo_id) + ) + for repo_dir in repo_dirs: snapshots_dir = repo_dir / "snapshots" if not snapshots_dir.is_dir(): continue @@ -276,12 +281,17 @@ def iter_hf_cache_snapshots(repo_id: str): yield from snapshots -def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]: +def list_empty_gguf_variant_dirs(repo_id: str, root: Optional[Path] = None) -> set[str]: """Quant labels present only as an EMPTY snapshot ``/`` folder (an interrupted split download); a quant with shards in any snapshot is excluded.""" empty: dict[str, str] = {} nonempty: set[str] = set() - for snapshot in iter_hf_cache_snapshots(repo_id): + snapshots = ( + iter_hf_cache_snapshots(repo_id, root = root) + if root is not None + else iter_hf_cache_snapshots(repo_id) + ) + for snapshot in snapshots: try: entries = list(snapshot.iterdir()) except OSError: @@ -303,8 +313,15 @@ def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]: return {label for key, label in empty.items() if key not in nonempty} -def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - for snapshot in iter_hf_cache_snapshots(repo_id): +def list_gguf_variants_from_hf_cache( + repo_id: str, root: Optional[Path] = None +) -> Optional[tuple[list[GgufVariantInfo], bool]]: + snapshots = ( + iter_hf_cache_snapshots(repo_id, root = root) + if root is not None + else iter_hf_cache_snapshots(repo_id) + ) + for snapshot in snapshots: variants, has_vision = list_local_gguf_variants(str(snapshot)) if variants or has_vision: return variants, has_vision @@ -312,7 +329,7 @@ def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVa def list_partial_gguf_variants_from_state( - repo_id: str, + repo_id: str, hub_cache: Optional[Path] = None ) -> Optional[tuple[list[GgufVariantInfo], bool]]: """Reconstruct GGUF variants from download manifests/markers alone. @@ -328,10 +345,26 @@ def list_partial_gguf_variants_from_state( # original-casing label over a lowercased cancel marker for the same variant. seen: set[str] = set() ordered: list[str] = [] - for source in ( - download_manifest.iter_variant_manifests("model", repo_id), - download_manifest.iter_variant_markers("model", repo_id), - ): + sources = ( + ( + download_manifest.iter_variant_manifests("model", repo_id), + download_manifest.iter_variant_markers("model", repo_id), + ) + if hub_cache is None + else ( + download_manifest.iter_variant_manifests( + "model", + repo_id, + hub_cache = hub_cache, + ), + download_manifest.iter_variant_markers( + "model", + repo_id, + hub_cache = hub_cache, + ), + ) + ) + for source in sources: for variant, _path in source: key = variant.lower() if key not in seen: @@ -343,7 +376,16 @@ def list_partial_gguf_variants_from_state( variants: list[GgufVariantInfo] = [] has_vision = False for variant in ordered: - manifest = download_manifest.read_manifest("model", repo_id, variant) + manifest = ( + download_manifest.read_manifest("model", repo_id, variant) + if hub_cache is None + else download_manifest.read_manifest( + "model", + repo_id, + variant, + hub_cache = hub_cache, + ) + ) main_filename: Optional[str] = None size_bytes = 0 companion_bytes = 0 diff --git a/studio/backend/hub/utils/hf_cache_state.py b/studio/backend/hub/utils/hf_cache_state.py index 22c948b683..49a28c813c 100644 --- a/studio/backend/hub/utils/hf_cache_state.py +++ b/studio/backend/hub/utils/hf_cache_state.py @@ -29,12 +29,10 @@ def _safe_is_dir(path: Path) -> bool: return False -def hf_cache_root(*, create: bool = False) -> Optional[Path]: - try: - from huggingface_hub import constants as hf_constants - except ImportError: - return None - root = Path(hf_constants.HF_HUB_CACHE) +def hf_cache_root(*, create: bool = False, root: Optional[Path] = None) -> Optional[Path]: + from utils.hf_cache_settings import get_hf_cache_paths + + root = root or get_hf_cache_paths().hub_cache if create: try: root.mkdir(parents = True, exist_ok = True) @@ -46,6 +44,7 @@ def hf_cache_root(*, create: bool = False) -> Optional[Path]: def hf_cache_roots() -> list[Path]: from hub.utils.paths import hf_default_cache_dir, legacy_hf_cache_dir + from utils.hf_cache_settings import known_hf_hub_caches roots: list[Path] = [] seen: set[str] = set() @@ -62,7 +61,8 @@ def hf_cache_roots() -> list[Path]: seen.add(key) roots.append(path) - _add(hf_cache_root()) + for configured in known_hf_hub_caches(): + _add(configured) _add(legacy_hf_cache_dir()) _add(hf_default_cache_dir()) return roots @@ -181,12 +181,22 @@ def iter_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: continue -def iter_destructive_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: +def iter_destructive_repo_cache_dirs( + repo_type: str, + repo_id: str, + *, + root: Optional[Path] = None, +) -> Iterator[Path]: target = repo_cache_dir_name(repo_type, repo_id) folded_target = target.lower() - for root in hf_cache_roots(): + if root is not None: + scoped = hf_cache_root(root = root) + bases = [scoped] if scoped is not None else [] + else: + bases = hf_cache_roots() + for base in bases: try: - entries = [entry for entry in root.iterdir() if entry.name.lower() == folded_target] + entries = [entry for entry in base.iterdir() if entry.name.lower() == folded_target] except OSError: continue matched_names = resolve_destructive_case_matches( @@ -200,8 +210,13 @@ def iter_destructive_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[P yield entry -def iter_active_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: - root = hf_cache_root() +def iter_active_repo_cache_dirs( + repo_type: str, + repo_id: str, + *, + root: Optional[Path] = None, +) -> Iterator[Path]: + root = hf_cache_root(root = root) if root is None: return target = target_dir_name(repo_type, repo_id) @@ -218,12 +233,13 @@ def preferred_repo_cache_dirs( repo_id: str, *, force_active: bool = False, + active_root: Optional[Path] = None, ) -> list[Path]: - active_entries = list(iter_active_repo_cache_dirs(repo_type, repo_id)) + active_entries = list(iter_active_repo_cache_dirs(repo_type, repo_id, root = active_root)) if active_entries: return active_entries if force_active: - root = hf_cache_root() + root = hf_cache_root(root = active_root) if root is not None: canonical = repo_cache_dir_name(repo_type, repo_id) return [root / canonical] @@ -237,8 +253,13 @@ def has_incomplete_blobs(repo_type: str, repo_id: str) -> bool: return False -def has_active_incomplete_blobs(repo_type: str, repo_id: str) -> bool: - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): +def has_active_incomplete_blobs( + repo_type: str, + repo_id: str, + *, + root: Optional[Path] = None, +) -> bool: + for entry in iter_active_repo_cache_dirs(repo_type, repo_id, root = root): if repo_cache_dir_has_incomplete_blobs(entry): return True return False @@ -273,9 +294,14 @@ def _prune_empty_dirs(root: Path) -> bool: return removed -def purge_partial_repo(repo_type: str, repo_id: str) -> bool: +def purge_partial_repo( + repo_type: str, + repo_id: str, + *, + root: Optional[Path] = None, +) -> bool: removed = False - for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id): + for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id, root = root): blobs_dir = entry / "blobs" if blobs_dir.is_dir(): for blob in blobs_dir.iterdir(): @@ -290,9 +316,14 @@ def purge_partial_repo(repo_type: str, repo_id: str) -> bool: return removed -def purge_repo_cache_dirs(repo_type: str, repo_id: str) -> bool: +def purge_repo_cache_dirs( + repo_type: str, + repo_id: str, + *, + root: Optional[Path] = None, +) -> bool: removed = False - for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id): + for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id, root = root): try: if entry.is_symlink() or not entry.is_dir(): continue @@ -301,3 +332,59 @@ def purge_repo_cache_dirs(repo_type: str, repo_id: str) -> bool: except FileNotFoundError: continue return removed + + +def scoped_delete_root(repo_type: str, repo_id: str, cache_path: Optional[str]) -> Optional[Path]: + """Resolve the single cache root a delete of this repo may touch. + + Returns the active hub cache when *cache_path* is falsy, the owning cache + root when *cache_path* points inside a known cache, or ``None`` when + *cache_path* is set but not inside any known cache (caller should reject). + This keeps a delete of one inventory row from removing copies in other, + previously selected caches. + """ + from utils.hf_cache_settings import get_hf_cache_paths + + if not cache_path: + return Path(get_hf_cache_paths().hub_cache).resolve(strict = False) + try: + resolved = Path(cache_path).expanduser().resolve(strict = False) + except (OSError, RuntimeError, ValueError): + return None + expected = repo_cache_dir_name(repo_type, repo_id).lower() + repo_dir = next( + ( + candidate + for candidate in (resolved, *resolved.parents) + if candidate.name.lower() == expected + ), + None, + ) + if repo_dir is None: + return None + allowed = {r.resolve(strict = False) for r in hf_cache_roots()} + root = repo_dir.parent.resolve(strict = False) + return root if root in allowed else None + + +def resolve_delete_target_root( + repo_type: str, repo_id: str, cache_path: Optional[str], owner_roots +) -> Optional[Path]: + """Pick the single cache root a delete of this repo should target. + + An explicit *cache_path* wins (``None`` when it is not a known cache, so the + caller can reject it). Otherwise prefer the active cache when it holds a + copy, else the sole cache that does -- so a model that lives only in a + previously selected cache stays deletable while other caches are untouched. + """ + if cache_path: + return scoped_delete_root(repo_type, repo_id, cache_path) + from utils.hf_cache_settings import get_hf_cache_paths + + active = Path(get_hf_cache_paths().hub_cache).resolve(strict = False) + roots = list(owner_roots) + if active in roots: + return active + if len(roots) == 1: + return roots[0] + return active diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index 57ad7f6655..058fdf9b65 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -36,7 +36,7 @@ from hub.utils.state_dir import RepoType from hub.utils.hf_cache_state import ( INCOMPLETE_SUFFIX, has_incomplete_blobs, - hf_cache_root, + hf_cache_roots, iter_repo_cache_dirs, latest_snapshot_dir, repo_cache_dir_has_incomplete_blobs, @@ -127,33 +127,13 @@ def all_hf_cache_scans() -> list: def _compute_all_hf_cache_scans() -> list: from huggingface_hub import scan_cache_dir - from hub.utils.paths import legacy_hf_cache_dir, hf_default_cache_dir scans: list = [] - seen: set[str] = set() - try: - from huggingface_hub.constants import HF_HUB_CACHE - - active = Path(HF_HUB_CACHE).resolve() - seen.add(str(active)) - if active.is_dir(): - scans.append(scan_cache_dir()) - except Exception as exc: - logger.warning("Could not scan active HF cache: %s", exc) - - for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): + for cache_root in hf_cache_roots(): try: - extra = extra_fn() - # is_dir()/resolve() can raise on an inaccessible path; skip it. - if not extra.is_dir(): - continue - resolved = str(extra.resolve()) - if resolved in seen: - continue - seen.add(resolved) - scans.append(scan_cache_dir(cache_dir = str(extra))) + scans.append(scan_cache_dir(cache_dir = str(cache_root))) except Exception as exc: - logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc) + logger.warning("Could not scan HF cache %s: %s", cache_root, exc) return scans @@ -224,16 +204,8 @@ def _compose_partial(*signals: Callable[[], bool]) -> bool: return any(signal() for signal in signals) -def _state_applies_to_repo_cache_dir(repo_cache_dir: Optional[Path]) -> bool: - if repo_cache_dir is None: - return True - root = hf_cache_root() - if root is None: - return False - try: - return repo_cache_dir.resolve().parent == root.resolve() - except OSError: - return False +def _hub_cache_for_repo_dir(repo_cache_dir: Optional[Path]) -> Optional[Path]: + return repo_cache_dir.parent if repo_cache_dir is not None else None def _legacy_partial( @@ -285,12 +257,24 @@ def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir: Path) return False -def _gguf_variant_manifest_blob_hashes(repo_id: str) -> frozenset[str]: +def _gguf_variant_manifest_blob_hashes( + repo_id: str, repo_cache_dir: Optional[Path] = None +) -> frozenset[str]: from hub.utils import download_manifest hashes: set[str] = set() - for variant, _path in download_manifest.iter_variant_manifests("model", repo_id): - manifest = download_manifest.read_manifest("model", repo_id, variant) + hub_cache = _hub_cache_for_repo_dir(repo_cache_dir) + for variant, _path in download_manifest.iter_variant_manifests( + "model", + repo_id, + hub_cache = hub_cache, + ): + manifest = download_manifest.read_manifest( + "model", + repo_id, + variant, + hub_cache = hub_cache, + ) if manifest is None: continue for expected in manifest.expected_files: @@ -315,7 +299,7 @@ def _snapshot_legacy_partial( ) -> bool: if repo_type != "model": return _legacy_partial(repo_type, repo_id, repo_cache_dir) - ignored_hashes = _gguf_variant_manifest_blob_hashes(repo_id) + ignored_hashes = _gguf_variant_manifest_blob_hashes(repo_id, repo_cache_dir) if repo_cache_dir is not None: return _repo_cache_dir_has_snapshot_legacy_partial( repo_cache_dir, @@ -375,9 +359,12 @@ def _manifest_partial( ) -> bool: from hub.utils import download_manifest - if not _state_applies_to_repo_cache_dir(repo_cache_dir): - return False - manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + manifest = download_manifest.read_manifest( + repo_type, + repo_id, + variant, + hub_cache = _hub_cache_for_repo_dir(repo_cache_dir), + ) if manifest is None: return False resolved = ( @@ -452,10 +439,13 @@ def is_snapshot_partial( A manifest without a resolvable snapshot is partial: the worker got far enough to record expectations but did not leave a usable snapshot.""" from hub.utils import download_manifest - - state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) return _compose_partial( - lambda: state_applies and download_manifest.has_cancel_marker(repo_type, repo_id, None), + lambda: download_manifest.has_cancel_marker( + repo_type, + repo_id, + None, + hub_cache = _hub_cache_for_repo_dir(repo_cache_dir), + ), lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir), lambda: _manifest_partial( repo_type, @@ -484,10 +474,13 @@ def is_variant_partial( caller is checking many variants of the same repo (see is_gguf_repo_partial for that usage).""" from hub.utils import download_manifest - - state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) return _compose_partial( - lambda: state_applies and download_manifest.has_cancel_marker("model", repo_id, variant), + lambda: download_manifest.has_cancel_marker( + "model", + repo_id, + variant, + hub_cache = _hub_cache_for_repo_dir(repo_cache_dir), + ), lambda: bool( incomplete_blob_hashes and variant_blob_hashes @@ -526,22 +519,38 @@ def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> from hub.utils import download_manifest has_legacy_partial = _legacy_partial("model", repo_id, repo_cache_dir) - state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) snapshot_dir = resolve_snapshot_dir_for_scan( "model", repo_id, repo_cache_dir, ) variants: set[str] = set(_completed_gguf_variants(snapshot_dir)) - if state_applies: - for variant, _path in download_manifest.iter_variant_manifests( - "model", - repo_id, + hub_cache = _hub_cache_for_repo_dir(repo_cache_dir) + for variant, _path in download_manifest.iter_variant_manifests( + "model", + repo_id, + hub_cache = hub_cache, + ): + if ( + download_manifest.read_manifest( + "model", + repo_id, + variant, + hub_cache = hub_cache, + ) + is not None ): variants.add(variant) - for variant, _path in download_manifest.iter_variant_markers( + for variant, _path in download_manifest.iter_variant_markers( + "model", + repo_id, + hub_cache = hub_cache, + ): + if download_manifest.has_cancel_marker( "model", repo_id, + variant, + hub_cache = hub_cache, ): variants.add(variant) if not variants: @@ -576,14 +585,19 @@ def partial_transport_for( available.""" from hub.utils import download_manifest - if not _state_applies_to_repo_cache_dir(repo_cache_dir): - return None + hub_cache = _hub_cache_for_repo_dir(repo_cache_dir) marker_transport = download_manifest.read_cancel_marker_transport( repo_type, repo_id, variant, + hub_cache = hub_cache, ) if marker_transport is not None: return marker_transport - manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + manifest = download_manifest.read_manifest( + repo_type, + repo_id, + variant, + hub_cache = hub_cache, + ) return manifest.transport if manifest is not None else None diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index 5435202565..81621edcf9 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -277,12 +277,8 @@ def _memo_drop(memo_key: tuple[str, str]) -> None: def _hf_hub_cache_dir() -> Path: - try: - from huggingface_hub.constants import HF_HUB_CACHE - return Path(HF_HUB_CACHE) - except Exception as exc: - logger.debug("Could not read huggingface_hub HF_HUB_CACHE, using default: %s", exc) - return Path.home() / ".cache" / "huggingface" / "hub" + from utils.hf_cache_settings import get_hf_cache_paths + return get_hf_cache_paths().hub_cache def _hf_hub_cache_dirs() -> list[Path]: @@ -300,7 +296,10 @@ def _hf_hub_cache_dirs() -> list[Path]: seen.add(key) roots.append(resolved) - _add(_hf_hub_cache_dir()) + from utils.hf_cache_settings import known_hf_hub_caches + + for configured in known_hf_hub_caches(): + _add(configured) try: _add(legacy_hf_cache_dir()) _add(hf_default_cache_dir()) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index 898c03c87d..4650b97381 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -8,9 +8,10 @@ so it survives ``huggingface-cli delete-cache`` and any other HF-side cache lifecycle. Two subdirectories: /hub-state/ - manifests/ .json per-download expected-files manifest - cancelled/ .json per-download cancel marker + manifests/cache-/.json expected-files manifest + cancelled/cache-/.json cancel marker +The cache digest isolates state for the same repo across selectable Hub caches. The ```` mirrors HF's cache dir naming while the resulting manifest, cancel-marker, and atomic-write temp filenames fit common filesystem basename limits. Very long repo IDs use a stable hash in the state key: @@ -29,6 +30,7 @@ configuration failure. from __future__ import annotations import hashlib +import os import re from pathlib import Path from typing import Literal, Optional, get_args @@ -55,6 +57,7 @@ _STATE_EXTENSION = ".json" # _atomic_write_json writes "..tmp-<8hex>" beside the final file. _ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8 _MAX_VARIANT_FRAGMENT_LENGTH = 64 +_CACHE_SCOPE_DIGEST_LENGTH = 32 def state_root() -> Optional[Path]: @@ -130,13 +133,32 @@ def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str return f"{variant_filename_prefix(repo_type, repo_id)}{variant_fragment}" +def _cache_scope(parent: Path, hub_cache: Optional[str | Path]) -> Optional[Path]: + if hub_cache is None: + return parent + normalized = os.path.normcase(str(Path(hub_cache).expanduser())) + digest = hashlib.sha256(normalized.encode("utf-8")).hexdigest()[:_CACHE_SCOPE_DIGEST_LENGTH] + scoped = parent / f"cache-{digest}" + try: + scoped.mkdir(parents = True, exist_ok = True) + except OSError as exc: + logger.debug("Could not create cache-scoped Hub state dir %s: %s", scoped, exc) + return None + return scoped + + def manifest_path( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> Optional[Path]: """Path to the manifest file for this triple. May or may not exist.""" parent = _subdir(_MANIFESTS_SUBDIR) + if parent is None: + return None + parent = _cache_scope(parent, hub_cache) if parent is None: return None return parent / f"{_entry_key(repo_type, repo_id, variant)}.json" @@ -146,9 +168,14 @@ def marker_path( repo_type: RepoType, repo_id: str, variant: Optional[str] = None, + *, + hub_cache: Optional[str | Path] = None, ) -> Optional[Path]: """Path to the cancel-marker file for this triple. May or may not exist.""" parent = _subdir(_CANCELLED_SUBDIR) + if parent is None: + return None + parent = _cache_scope(parent, hub_cache) if parent is None: return None return parent / f"{_entry_key(repo_type, repo_id, variant)}.json" diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index e45357d311..9ff394b009 100644 --- a/studio/backend/hub/workers/hf_download.py +++ b/studio/backend/hub/workers/hf_download.py @@ -661,6 +661,7 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod variant, plan.main_hashes, hf_token, + hub_cache = Path(snapshot_path).parents[2], ) except Exception as e: print( diff --git a/studio/backend/main.py b/studio/backend/main.py index f686e29bf5..5af25efa74 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -289,6 +289,7 @@ from fastapi import Depends, FastAPI, HTTPException, Query, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, HTMLResponse, Response +from starlette.middleware.gzip import GZipMiddleware from pathlib import Path from datetime import datetime @@ -308,12 +309,14 @@ from routes import ( training_router, ) from routes.llama import router as llama_router +from routes.whisper import router as whisper_router from routes.preview import router as preview_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, token_router as hub_token_router, ) +from picker.routes import templates_router as picker_templates_router from hub.schemas.downloads import TransportCapabilities from hub.utils.download_registry import ( get_download_transport_capabilities, @@ -435,7 +438,11 @@ def _run_llama_cpp_startup_probes(app: FastAPI) -> None: import structlog as _structlog _log = _structlog.get_logger(__name__) - if _caps.get("found") and not _caps.get("supports_mtp"): + if ( + _caps.get("found") + and not _caps.get("supports_mtp") + and not _caps.get("mtp_probe_inconclusive") + ): _msg = ( "llama.cpp prebuilt lacks MTP support " "(--spec-type mtp/draft-mtp). Run `unsloth studio update`. " @@ -753,6 +760,8 @@ app.add_middleware(SecurityHeadersMiddleware) # headroom; non-upload routes keep the default body cap. import json as _json_for_413 # noqa: E402 from utils.upload_limits import ( # noqa: E402 + STT_AUDIO_JSON_MAX_BYTES, + STT_AUDIO_RAW_MAX_BYTES, UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES, default_request_body_limit_bytes, upload_request_limit_bytes, @@ -763,6 +772,7 @@ _BODY_PROTECTED_PREFIXES = ( "/v1/completions", "/p/", "/api/inference", + "/api/picker", "/api/data-recipe", "/api/datasets", "/api/hub", @@ -790,6 +800,14 @@ def _get_upload_passthrough_request_max_bytes(path: str) -> int: return default_request_body_limit_bytes() +def _get_request_body_max_bytes(path: str) -> int: + if path.startswith("/api/inference/audio/transcribe/raw"): + return STT_AUDIO_RAW_MAX_BYTES + if path.startswith("/api/inference/audio/transcribe"): + return STT_AUDIO_JSON_MAX_BYTES + return default_request_body_limit_bytes() + + async def _send_411(send) -> None: payload = _json_for_413.dumps( {"detail": "Content-Length required for upload requests."}, @@ -832,12 +850,14 @@ class MaxBodyMiddleware: app, max_bytes_getter, protected_prefixes: tuple, + request_max_bytes_getter = None, upload_passthrough_prefixes: tuple = (), upload_passthrough_max_bytes_getter = None, ): self.app = app self.max_bytes_getter = max_bytes_getter self.protected_prefixes = protected_prefixes + self.request_max_bytes_getter = request_max_bytes_getter self.upload_passthrough_prefixes = upload_passthrough_prefixes self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter @@ -854,6 +874,14 @@ class MaxBodyMiddleware: except Exception: return int(self.max_bytes_getter()) + def _request_max_bytes(self, path: str) -> int: + if self.request_max_bytes_getter is None: + return int(self.max_bytes_getter()) + try: + return int(self.request_max_bytes_getter(path)) + except Exception: + return int(self.max_bytes_getter()) + async def __call__(self, scope, receive, send): if scope["type"] != "http": await self.app(scope, receive, send) @@ -866,7 +894,7 @@ class MaxBodyMiddleware: await self.app(scope, receive, send) return - max_bytes = int(self.max_bytes_getter()) + max_bytes = self._request_max_bytes(path) declared = None for name, value in scope.get("headers", []): if name == b"content-length": @@ -931,6 +959,7 @@ app.add_middleware( MaxBodyMiddleware, max_bytes_getter = default_request_body_limit_bytes, protected_prefixes = _BODY_PROTECTED_PREFIXES, + request_max_bytes_getter = _get_request_body_max_bytes, upload_passthrough_prefixes = _BODY_UPLOAD_PASSTHROUGH_PREFIXES, upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes, ) @@ -989,11 +1018,13 @@ app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"]) app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"]) app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"]) +app.include_router(whisper_router, prefix = "/api/whisper", tags = ["whisper"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"]) app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"]) app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"]) app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"]) +app.include_router(picker_templates_router, prefix = "/api/picker", tags = ["picker"]) app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"]) # Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic @@ -1509,6 +1540,34 @@ def _should_inject_bootstrap(request: Request) -> bool: return _is_local_bootstrap_request(request) +_IMMUTABLE_ASSET_CACHE_CONTROL = "public, max-age=31536000, immutable" + + +class ImmutableStaticFiles(StaticFiles): + """Serve Vite's content-hashed assets without browser revalidation.""" + + def file_response( + self, + full_path, + stat_result, + scope, + status_code = 200, + ): + response = super().file_response(full_path, stat_result, scope, status_code) + response.headers["Cache-Control"] = _IMMUTABLE_ASSET_CACHE_CONTROL + return response + + +class _AssetGZipMiddleware(GZipMiddleware): + """Serve range requests uncompressed; gzip + 206 mislabels Content-Range.""" + + async def __call__(self, scope, receive, send): + if scope["type"] == "http" and any(key == b"range" for key, _ in scope["headers"]): + await self.app(scope, receive, send) + return + await super().__call__(scope, receive, send) + + def setup_frontend(app: FastAPI, build_path: Path): """Mount frontend static files (optional)""" if not build_path.exists(): @@ -1516,7 +1575,12 @@ def setup_frontend(app: FastAPI, build_path: Path): assets_dir = build_path / "assets" if assets_dir.exists(): - app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") + assets_app = _AssetGZipMiddleware( + ImmutableStaticFiles(directory = assets_dir), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") def _build_index_response(request: Request) -> Response: content = (build_path / "index.html").read_bytes() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d51d35189b..1758efe515 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -18,6 +18,8 @@ from pydantic import ( model_validator, ) +from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + class LoadRequest(BaseModel): """Request to load a model for inference""" @@ -54,17 +56,37 @@ class LoadRequest(BaseModel): @field_validator("chat_template_override") @classmethod def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]: - if value is not None and value.strip() == "": + if value is None: return None + # Char count is a lower bound on UTF-8 byte length: reject an oversized + # template before spending work encoding it. + if len(value) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + if value.strip() == "": + return None + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") return value cache_type_kv: Optional[str] = Field( None, - description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", + description = ( + "KV cache data type for both K and V " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32')" + ), ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", + description = ( + "GPU placement pool, for example [0, 1]. Omit or pass [] to use " + "automatic selection. CUDA/ROCm and Intel XPU values are physical " + "GPU indices; Vulkan values are ggml device ordinals. Explicit " + "physical IDs are unsupported when the parent visibility mask uses " + "non-numeric or subdevice entries, including CUDA_VISIBLE_DEVICES " + "with UUID/MIG entries and ZE_AFFINITY_MASK with subdevice tokens " + "(for example '0.0,0.1') or FLAT-hierarchy tile handles. For GGUF " + "models the fitter may pin the smallest subset of this pool that fits." + ), ) speculative_type: Optional[str] = Field( None, @@ -177,6 +199,32 @@ class UnloadRequest(BaseModel): model_path: str = Field(..., description = "Model identifier to unload") +class TranscribeRequest(BaseModel): + """Speech-to-text request for the dictation STT sidecar.""" + + audio: str = Field(..., description = "Base64-encoded audio (any common format)") + model: Optional[str] = Field(None, description = "STT model id; defaults server-side") + language: Optional[str] = Field(None, description = "BCP-47 language, or 'auto'/None to detect") + fast: bool = Field( + False, + description = "Use low-latency single-candidate decoding for dictation", + ) + engine: Optional[str] = Field( + None, + description = "STT engine: 'transformers' (default) or 'gguf' (whisper.cpp)", + ) + + +class SttLoadRequest(BaseModel): + """Warm the STT sidecar with a model without transcribing.""" + + model: Optional[str] = Field(None, description = "STT model id; defaults server-side") + engine: Optional[str] = Field( + None, + description = "STT engine: 'transformers' (default) or 'gguf' (whisper.cpp)", + ) + + class ValidateModelRequest(BaseModel): """Check whether an identifier resolves to a ModelConfig; does NOT load weights.""" @@ -206,6 +254,13 @@ class ValidateModelRequest(BaseModel): description = "Also read the native context length from the local GGUF header. " "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.", ) + include_chat_template: bool = Field( + False, + description = "Also read the embedded chat template from the local GGUF header, so a " + "native (picked / drag-drop) file's default template can be shown before it is loaded. " + "Opt-in and, like include_context_length, a metadata-only probe that skips the training " + "guard. Only the leased file's own embedded template is read, never sibling sidecars.", + ) class TransformersUpgradeInfo(BaseModel): @@ -266,6 +321,11 @@ class ValidateModelResponse(BaseModel): description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF " "header alongside context_length; 0 for dense models, None when not read.", ) + chat_template: Optional[str] = Field( + None, + description = "Embedded GGUF chat template, read from the header when include_chat_template " + "is set (native lease-backed picks); None for non-GGUF, over-cap, or not-read templates.", + ) # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. requires_transformers_upgrade: bool = Field( False, @@ -385,7 +445,10 @@ class LoadResponse(BaseModel): ) cache_type_kv: Optional[str] = Field( None, - description = "KV cache data type for K and V (e.g. 'f16', 'bf16', 'q8_0')", + description = ( + "KV cache data type for K and V " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32')" + ), ) chat_template: Optional[str] = Field( None, @@ -437,7 +500,14 @@ class LoadResponse(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.", + ) + requested_gpu_ids: Optional[List[int]] = Field( + None, + description = ( + "GPU placement pool requested by the user before fit-time narrowing, " + "or None for automatic selection." + ), ) @@ -538,7 +608,11 @@ class InferenceStatusResponse(BaseModel): ) cache_type_kv: Optional[str] = Field( None, - description = "KV cache quantization dtype (e.g. 'q8_0'), or None for default", + description = ( + "KV cache quantization dtype " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32'), " + "or None for default" + ), ) chat_template: Optional[str] = Field( None, description = "Model's default chat template (Jinja2 source), if any" @@ -601,7 +675,14 @@ class InferenceStatusResponse(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.", + ) + requested_gpu_ids: Optional[List[int]] = Field( + None, + description = ( + "GPU placement pool requested by the user before fit-time narrowing, " + "or None for automatic selection." + ), ) llama_cpp_supports_mtp: bool = Field( True, diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 54e88fed58..df6725c9c9 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -178,6 +178,14 @@ class LocalModelInfo(BaseModel): None, description = "HF repo id for cached models, e.g. org/model", ) + active_cache: Optional[bool] = Field( + None, + description = "Whether an HF model belongs to the current download cache.", + ) + partial: bool = Field( + False, + description = "Whether the cached model has an incomplete download.", + ) model_format: Optional[str] = Field( None, description = "Detected weights format ('gguf' when known). Lets the UI " diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 5a75246c07..4238403e00 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -47,6 +47,14 @@ class ProviderCreate(BaseModel): None, description = "Custom base URL (overrides registry default). Omit to use the default.", ) + models: list[str] = Field( + default_factory = list, + description = "Enabled model IDs for this connection", + ) + available_models: list[str] = Field( + default_factory = list, + description = "Discovered catalog model IDs last fetched for this connection", + ) class ProviderUpdate(BaseModel): @@ -55,6 +63,11 @@ class ProviderUpdate(BaseModel): display_name: Optional[str] = Field(None, description = "New display name") base_url: Optional[str] = Field(None, description = "New base URL") is_enabled: Optional[bool] = Field(None, description = "Enable or disable this provider") + models: Optional[list[str]] = Field(None, description = "Enabled model IDs for this connection") + available_models: Optional[list[str]] = Field( + None, + description = "Discovered catalog model IDs last fetched for this connection", + ) class ProviderResponse(BaseModel): @@ -65,6 +78,14 @@ class ProviderResponse(BaseModel): display_name: str = Field(..., description = "User-chosen label") base_url: str = Field(..., description = "API base URL") is_enabled: bool = Field(True, description = "Whether this provider is enabled") + models: list[str] = Field( + default_factory = list, + description = "Enabled model IDs for this connection", + ) + available_models: list[str] = Field( + default_factory = list, + description = "Discovered catalog model IDs last fetched for this connection", + ) created_at: str = Field(..., description = "ISO 8601 creation timestamp") updated_at: str = Field(..., description = "ISO 8601 last-update timestamp") diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0b50f63b95..0aca5da72c 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -470,6 +470,7 @@ class TrainingStartRequest(BaseModel): gradient_checkpointing: str = Field("", description = "Gradient checkpointing setting") use_rslora: bool = Field(False, description = "Use RSLoRA") use_loftq: bool = Field(False, description = "Use LoftQ") + use_dora: bool = Field(False, description = "Use DoRA") train_on_completions: bool = Field(False, description = "Train on completions only") # Vision-specific LoRA parameters @@ -496,7 +497,15 @@ class TrainingStartRequest(BaseModel): # GPU selection gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", + description = ( + "Physical GPU indices to use, for example [0, 1]. Omit or pass " + "[] to use automatic selection. Explicit gpu_ids are unsupported " + "when the parent visibility mask uses non-numeric or subdevice " + "entries -- this includes CUDA_VISIBLE_DEVICES with UUID/MIG " + "entries on NVIDIA, and ZE_AFFINITY_MASK with subdevice tokens " + "(e.g. '0.0,0.1') or FLAT-hierarchy (default) tile handles on " + "Intel XPU." + ), ) # S3 dataset source configuration @@ -505,6 +514,13 @@ class TrainingStartRequest(BaseModel): description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", ) + @field_validator("target_modules", mode = "before") + @classmethod + def _normalize_target_modules(cls, value: Any) -> Any: + # Sanitized non-LoRA history stores the unused value as null; treat it as a + # fresh request's omitted/default empty list on resume. + return [] if value is None else value + @model_validator(mode = "after") def _validate_streaming_splits(self) -> "TrainingStartRequest": # Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]" @@ -530,6 +546,37 @@ class TrainingStartRequest(BaseModel): raise ValueError("Either num_epochs or max_steps must be > 0; both cannot be 0.") return self + @model_validator(mode = "after") + def _validate_lora_variant_flags(self) -> "TrainingStartRequest": + # The frontend only ever sends one of these and never under Full + # Finetuning, but a direct API/YAML/CLI caller can bypass that. Nothing + # downstream breaks (full finetune ignores them, MLX rejects use_dora/ + # use_loftq outright), but reject early here for a clear error instead + # of a silently-ignored flag. + active = [ + name + for name, enabled in ( + ("use_rslora", self.use_rslora), + ("use_loftq", self.use_loftq), + ("use_dora", self.use_dora), + ) + if enabled + ] + if len(active) > 1: + raise ValueError( + f"Only one LoRA variant may be enabled at a time; got {active}. " + "use_rslora, use_loftq, and use_dora are mutually exclusive." + ) + # getattr, not self.training_type: model_construct() (used by tests that + # validate a single field in isolation) leaves required fields unset, and + # this is a mode="after" validator so it still runs on that partial instance. + if getattr(self, "training_type", None) == "Full Finetuning" and active: + raise ValueError( + f"{active[0]} requires an adapter method (LoRA/QLoRA or " + "Continued Pretraining); it has no effect under Full Finetuning." + ) + return self + class TrainingJobResponse(BaseModel): """Immediate response when training is initiated""" diff --git a/studio/backend/picker/__init__.py b/studio/backend/picker/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/picker/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/picker/routes/__init__.py b/studio/backend/picker/routes/__init__.py new file mode 100644 index 0000000000..c0e988c8bb --- /dev/null +++ b/studio/backend/picker/routes/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from .templates import router as templates_router + +__all__ = ["templates_router"] diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py new file mode 100644 index 0000000000..02b8bf7184 --- /dev/null +++ b/studio/backend/picker/routes/templates.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import asyncio +from typing import Optional + +from fastapi import APIRouter, Body, Depends, Query + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token + +from ..schemas import ( + MAX_CHAT_TEMPLATE_BYTES, + ModelTemplateResponse, + ValidateChatTemplateRequest, + ValidateChatTemplateResponse, +) +from ..service import read_default_chat_template, validate_chat_template + +router = APIRouter() + + +@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse) +async def validate_chat_template_route( + body: ValidateChatTemplateRequest = Body(...), + current_subject: str = Depends(get_current_subject), +) -> ValidateChatTemplateResponse: + return await asyncio.to_thread(validate_chat_template, body.template) + + +@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse) +async def get_default_chat_template_route( + model_name: str, + gguf_variant: Optional[str] = Query(None), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +) -> ModelTemplateResponse: + template = await asyncio.to_thread( + read_default_chat_template, model_name, hf_token, gguf_variant + ) + if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + template = None + return ModelTemplateResponse(model_name = model_name, chat_template = template) diff --git a/studio/backend/picker/schemas.py b/studio/backend/picker/schemas.py new file mode 100644 index 0000000000..b4f956188f --- /dev/null +++ b/studio/backend/picker/schemas.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from typing import Optional + +from pydantic import BaseModel, Field, field_validator + +# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at +# the API boundary so a direct caller cannot make Jinja parse an oversized +# template. MaxBodyMiddleware only caps the whole request body, not this field. +MAX_CHAT_TEMPLATE_BYTES = 65_536 + + +class ValidateChatTemplateRequest(BaseModel): + template: str = Field(default = "") + + @field_validator("template") + @classmethod + def _enforce_template_size(cls, value: str) -> str: + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + return value + + +class ValidateChatTemplateResponse(BaseModel): + valid: bool + error: Optional[str] = None + + +class ModelTemplateResponse(BaseModel): + model_name: str + chat_template: Optional[str] = None diff --git a/studio/backend/picker/service.py b/studio/backend/picker/service.py new file mode 100644 index 0000000000..ccf9c3e152 --- /dev/null +++ b/studio/backend/picker/service.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import logging +import os +import re +from pathlib import Path +from typing import Optional + +from hub.services.models.folder_browser import ( + _build_browse_allowlist, + _is_path_inside_allowlist, +) +from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots +from utils.models.gguf_metadata import read_gguf_chat_template +from utils.models.model_config import ( + _extract_quant_label, + _is_big_endian_gguf_path, + _is_mmproj, + _is_mtp_drafter, +) +from utils.hf_cache_settings import active_hf_hub_cache +from utils.paths.path_utils import ( + is_local_path, + normalize_path, + resolve_cached_repo_id_case, +) + +from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse + +logger = logging.getLogger(__name__) + +_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") + + +def _is_valid_repo_id(repo_id: str) -> bool: + return bool(_VALID_REPO_ID.fullmatch(repo_id)) + + +_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json") +_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja") +_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json") + +# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory +# before its template is size-checked. The JSON envelope may exceed a bare template +# (it carries other tokenizer metadata); the extracted template is still bounded by +# MAX_CHAT_TEMPLATE_BYTES downstream. +MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024 + + +def _read_bounded_text(path: Path, limit: int) -> Optional[str]: + """Read at most `limit` bytes of UTF-8 text; None if larger or unreadable.""" + try: + with path.open("rb") as f: + data = f.read(limit + 1) + except OSError: + return None + if len(data) > limit: + return None + try: + return data.decode("utf-8") + except UnicodeError: + return None + + +def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool: + # Block symlinked children from escaping the validated directory (realpath-checked). + # None = trusted caller (HF cache / remote download). + return allow_roots is None or _is_path_inside_allowlist(path, allow_roots) + + +def validate_chat_template(template: str) -> ValidateChatTemplateResponse: + text = (template or "").strip() + if not text: + return ValidateChatTemplateResponse(valid = True, error = None) + # Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a + # missing dependency must not crash API startup. + try: + from jinja2 import TemplateError + from jinja2.ext import Extension + from jinja2.sandbox import ImmutableSandboxedEnvironment + except ImportError: + return ValidateChatTemplateResponse(valid = True, error = None) + + class _GenerationTag(Extension): + # Accept Transformers' {% generation %} assistant-mask tag so a pasted HF + # chat template validates (we only parse it). + tags = {"generation"} + + def parse(self, parser): + next(parser.stream) + return parser.parse_statements(["name:endgeneration"], drop_needle = True) + + try: + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols", _GenerationTag], + ) + env.parse(text) + return ValidateChatTemplateResponse(valid = True, error = None) + except TemplateError as exc: + message = getattr(exc, "message", None) or str(exc) + lineno = getattr(exc, "lineno", None) + if lineno: + message = f"Line {lineno}: {message}" + return ValidateChatTemplateResponse(valid = False, error = message) + except Exception as exc: + return ValidateChatTemplateResponse(valid = False, error = str(exc)) + + +def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]: + if not isinstance(config, dict): + return None + raw = config.get("chat_template") + if isinstance(raw, str) and raw.strip(): + return raw + if isinstance(raw, list): + fallback: Optional[str] = None + for entry in raw: + if not isinstance(entry, dict): + continue + template = entry.get("template") + if not isinstance(template, str): + continue + if entry.get("name") == "default": + return template + if fallback is None: + fallback = template + return fallback + return None + + +def _chat_template_from_jinja_file( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + for rel in _JINJA_TEMPLATE_PATHS: + template_file = dir_path / rel + if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots): + continue + try: + if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES: + continue + template = template_file.read_text(encoding = "utf-8") + except Exception: + continue + if template.strip(): + return template + return None + + +def _chat_template_from_processor_payload(payload: object) -> Optional[str]: + # processor chat_template.json may be the template string itself or a + # {name: template} map, not only a tokenizer_config-shaped object. + if isinstance(payload, str): + return payload if payload.strip() else None + template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type] + if template: + return template + if isinstance(payload, dict): + # Named-template map: prefer "default", else the first non-empty entry + # (mirrors the tokenizer-config list fallback). + default = payload.get("default") + if isinstance(default, str) and default.strip(): + return default + for value in payload.values(): + if isinstance(value, str) and value.strip(): + return value + return None + + +def _chat_template_from_processor_json( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + for rel in _PROCESSOR_TEMPLATE_PATHS: + config_file = dir_path / rel + if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots): + continue + raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES) + if raw is None: + continue + try: + payload = json.loads(raw) + except Exception: + continue + template = _chat_template_from_processor_payload(payload) + if template: + return template + return None + + +def _chat_template_from_tokenizer_dir( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + jinja = _chat_template_from_jinja_file(dir_path, allow_roots) + if jinja: + return jinja + for rel in _TOKENIZER_CONFIG_PATHS: + config_file = dir_path / rel + if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots): + continue + raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES) + if raw is None: + continue + try: + config = json.loads(raw) + except Exception: + continue + template = _chat_template_from_tokenizer_config(config) + if template: + return template + return _chat_template_from_processor_json(dir_path, allow_roots) + + +_GGUF_SCAN_MAX_DEPTH = 2 + + +def _iter_ggufs(dir_path: Path) -> list[Path]: + if dir_path == dir_path.parent: + return [] + root = str(dir_path) + found: list[Path] = [] + for current, dirs, files in os.walk(root, followlinks = False): + rel = os.path.relpath(current, root) + depth = 0 if rel == os.curdir else rel.count(os.sep) + 1 + if depth >= _GGUF_SCAN_MAX_DEPTH: + dirs[:] = [] + for name in files: + if not name.lower().endswith(".gguf") or _is_mmproj(name): + continue + path = Path(current) / name + try: + rel = path.relative_to(dir_path).as_posix() + except ValueError: + rel = name + quant = _extract_quant_label(rel) + if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant): + continue + found.append(path) + return found + + +def _variant_matches(relative_path: str, needle: str) -> bool: + quant = _extract_quant_label(relative_path).lower() + if quant == needle: + return True + if extract_quant_label(relative_path).lower() == needle: + return True + prefix = f"{needle}-" + if not quant.startswith(prefix): + return False + suffix = quant[len(prefix) :] + if not suffix.endswith("bpw"): + return False + value = suffix[:-3] + return bool(value) and value.replace(".", "", 1).isdigit() + + +_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE) + + +def _is_nonfirst_gguf_split(path: Path) -> bool: + match = _GGUF_SPLIT_INDEX_RE.search(path.stem) + return match is not None and int(match.group(1)) != 1 + + +def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]: + try: + ggufs = sorted(_iter_ggufs(dir_path)) + except OSError: + return None + if not ggufs: + return None + needle = (gguf_variant or "").strip().lower() + if needle: + for path in ggufs: + try: + relative = path.relative_to(dir_path).as_posix() + except ValueError: + relative = path.name + if _variant_matches(relative, needle): + return path + return None + candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs + try: + return max(candidates, key = lambda path: path.stat().st_size) + except OSError: + return candidates[0] + + +def _chat_template_from_dir( + dir_path: Path, + gguf_variant: Optional[str] = None, + allow_roots: Optional[list[Path]] = None, +) -> Optional[str]: + def from_gguf() -> Optional[str]: + gguf = _find_gguf_in_dir(dir_path, gguf_variant) + if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots): + return None + return read_gguf_chat_template(str(gguf)) + + # Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the + # author's maintained template and supersede the GGUF's possibly-stale embedded + # copy. The variant only picks the GGUF fallback, so tokenizer-first precedence + # holds whether or not a variant is given. + return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf() + + +def read_default_chat_template( + model_name: str, + hf_token: Optional[str] = None, + gguf_variant: Optional[str] = None, +) -> Optional[str]: + if not isinstance(model_name, str) or not model_name.strip(): + return None + name = model_name.strip() + + if is_local_path(name): + try: + target = Path(normalize_path(name)).expanduser() + allow_roots = _build_browse_allowlist() + if not _is_path_inside_allowlist(target, allow_roots): + logger.debug("Refused chat template read outside allowed folders: %s", name) + return None + if name.lower().endswith(".gguf"): + # Prefer a maintained sidecar next to the file over the GGUF's + # embedded copy (tokenizer-first precedence, as elsewhere). + sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots) + if sidecar: + return sidecar + return read_gguf_chat_template(str(target)) + return _chat_template_from_dir(target, gguf_variant, allow_roots) + except Exception as exc: + logger.debug("Could not read local chat template for %s: %s", name, exc) + return None + + if not _is_valid_repo_id(name): + return None + + resolved = resolve_cached_repo_id_case(name) + + try: + # Resolve within each cached revision, newest first. A revision's sidecar + # supersedes its own embedded GGUF copy, but must not override a newer + # revision, so precedence stays per-snapshot rather than global. + for snapshot in iter_hf_cache_snapshots(resolved): + template = _chat_template_from_dir(snapshot, gguf_variant) + if template: + return template + except Exception as exc: + logger.debug("Could not read cached chat template for %s: %s", resolved, exc) + + try: + from huggingface_hub import HfApi, hf_hub_download + + _api = HfApi() + + def _remote_exceeds_cap(rel: str) -> bool: + # Best-effort: skip the download when the remote's advertised size + # exceeds the cap, so a maliciously large sidecar is never fetched. + try: + infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token) + except Exception: + return False + for info in infos: + size = getattr(info, "size", None) + if ( + getattr(info, "path", None) == rel + and isinstance(size, int) + and size > MAX_TEMPLATE_METADATA_BYTES + ): + return True + return False + + def _download_text(rel: str) -> Optional[str]: + if _remote_exceeds_cap(rel): + return None + try: + path = hf_hub_download( + resolved, + rel, + token = hf_token, + cache_dir = active_hf_hub_cache(), + ) + return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES) + except Exception: + return None + + for rel in _JINJA_TEMPLATE_PATHS: + template = _download_text(rel) + if not template or not template.strip(): + continue + # A raw Jinja sidecar is the whole template, so it must fit the route's + # response cap (the local path skips oversized .jinja too). Download stays + # bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small + # template still extracts below, but an over-cap Jinja is dropped so the + # search falls through to the tokenizer/processor template. + if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + continue + return template + + for rel in _TOKENIZER_CONFIG_PATHS: + raw = _download_text(rel) + if not raw: + continue + try: + config = json.loads(raw) + except Exception: + continue + template = _chat_template_from_tokenizer_config(config) + if template: + return template + + for rel in _PROCESSOR_TEMPLATE_PATHS: + raw = _download_text(rel) + if not raw: + continue + try: + payload = json.loads(raw) + except Exception: + continue + template = _chat_template_from_processor_payload(payload) + if template: + return template + + return None + except Exception as exc: + logger.debug("Could not fetch chat template for %s: %s", resolved, exc) + return None diff --git a/studio/backend/requirements/extras.txt b/studio/backend/requirements/extras.txt index 1baf2b6f2d..1601ccbfae 100644 --- a/studio/backend/requirements/extras.txt +++ b/studio/backend/requirements/extras.txt @@ -10,6 +10,7 @@ omegaconf einx pyloudnorm openai-whisper +av # PyAV: decode dictation audio (webm/opus/mp3/โ€ฆ) for the Whisper STT sidecar uroman # 4.0 MB - used for Outetts. MeCab # 19.9 MB - used for Outetts. inflect # number-to-words, required by OuteTTS diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 378fb33a60..847e89823b 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -7,7 +7,7 @@ # (current PyPI metadata still declares torch as a hard dep). # unsloth direct deps (from pyproject.toml [project].dependencies) -typer +typer>=0.12.0 # typer's full runtime dep tree. Required explicitly because this # file is installed with --no-deps. On Linux/Mac CI runners these # are often cached transitively; on a fresh windows-latest venv they diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index d779c8784e..1acc48e3a3 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -494,6 +494,11 @@ async def change_password( status_code = status.HTTP_401_UNAUTHORIZED, detail = "Current password is incorrect", ) + if any(ch.isspace() for ch in payload.new_password): + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, + detail = "New password cannot contain spaces", + ) if payload.current_password == payload.new_password: raise HTTPException( status_code = status.HTTP_400_BAD_REQUEST, diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 24b6dfb36d..6a0d49b47d 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -160,11 +160,26 @@ class ChatInferenceSettings(BaseModel): fastMode: Optional[bool] = None +class ChatPresetLoadConfig(BaseModel): + model_config = ConfigDict(extra = "forbid") + + customContextLength: Optional[int] = Field(default = None, gt = 0) + maxSeqLength: Optional[float] = None + kvCacheDtype: Optional[str] = None + speculativeType: Optional[str] = None + specDraftNMax: Optional[int] = Field(default = None, ge = 1, le = 16) + tensorParallel: Optional[bool] = None + gpuMemoryMode: Optional[Literal["manual"]] = None + gpuLayers: Optional[int] = None + nCpuMoe: Optional[int] = Field(default = None, ge = 0) + + class ChatPreset(BaseModel): model_config = ConfigDict(extra = "forbid") name: str params: ChatInferenceSettings + loadConfig: Optional[ChatPresetLoadConfig] = None class ChatSettingsPayload(BaseModel): diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 5456080f34..331eb5e1e0 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -23,40 +23,6 @@ def _is_valid_repo_id(repo_id: str) -> bool: return bool(_VALID_REPO_ID.fullmatch(repo_id)) -_dataset_size_cache: dict[str, int] = {} - - -def _get_dataset_size_cached(repo_id: str) -> int: - if repo_id in _dataset_size_cache: - return _dataset_size_cache[repo_id] - try: - from huggingface_hub import dataset_info as hf_dataset_info - - info = hf_dataset_info(repo_id, token = None, files_metadata = True) - total = sum(s.size for s in info.siblings if getattr(s, "size", None)) - _dataset_size_cache[repo_id] = total - return total - except Exception: - return 0 - - -def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: - """Resolved realpath for a HF cache repo dir: most-recent snapshot, else cache root. - - Mirrors routes/models.py; duplicated here to keep this module self-contained. - """ - try: - snapshots_dir = repo_dir / "snapshots" - if snapshots_dir.is_dir(): - snaps = [s for s in snapshots_dir.iterdir() if s.is_dir()] - if snaps: - latest = max(snaps, key = lambda s: s.stat().st_mtime) - return str(latest.resolve()) - return str(repo_dir.resolve()) - except Exception: - return None - - backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) @@ -64,6 +30,7 @@ if str(backend_path) not in sys.path: from utils.datasets import check_dataset_format from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token router = APIRouter() logger = get_logger(__name__) @@ -292,11 +259,13 @@ def _download_hf_metadata(*, repo_id: str, repo_files: list[str], token: str | N try: from huggingface_hub import hf_hub_download + from utils.hf_cache_settings import active_hf_hub_cache local_path = hf_hub_download( repo_id = repo_id, filename = metadata_file, repo_type = "dataset", token = token, + cache_dir = active_hf_hub_cache(), ) except Exception as exc: logger.warning(f"Could not read HF dataset metadata for {repo_id}: {exc}") @@ -525,77 +494,15 @@ def list_local_datasets( @router.get("/download-progress") async def get_dataset_download_progress( repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"), + hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - """Return download progress for a HuggingFace dataset repo. - - Mirrors ``GET /api/models/download-progress`` but scans the - ``datasets--owner--name`` cache dir under HF_HUB_CACHE, where in-progress - download bytes are visible. Returns ``cache_path`` so the UI can show it. - """ - _empty = { - "downloaded_bytes": 0, - "expected_bytes": 0, - "progress": 0, - "cache_path": None, - } - try: - if not _is_valid_repo_id(repo_id): - return _empty - - from huggingface_hub import constants as hf_constants - - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"datasets--{repo_id.replace('/', '--')}".lower() - completed_bytes = 0 - in_progress_bytes = 0 - cache_path: Optional[str] = None - - if cache_dir.is_dir(): - for entry in cache_dir.iterdir(): - if entry.name.lower() != target: - continue - cache_path = _resolve_hf_cache_realpath(entry) - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - break - for f in blobs_dir.iterdir(): - if not f.is_file(): - continue - if f.name.endswith(".incomplete"): - in_progress_bytes += f.stat().st_size - else: - completed_bytes += f.stat().st_size - break - - downloaded_bytes = completed_bytes + in_progress_bytes - if downloaded_bytes == 0: - return {**_empty, "cache_path": cache_path} - - expected_bytes = _get_dataset_size_cached(repo_id) - if expected_bytes <= 0: - return { - "downloaded_bytes": downloaded_bytes, - "expected_bytes": 0, - "progress": 0, - "cache_path": cache_path, - } - - # 95% threshold (as in the model endpoint): HF blob dedup makes - # completed_bytes drift under expected_bytes; inter-file gaps look "done". - if completed_bytes >= expected_bytes * 0.95: - progress = 1.0 - else: - progress = min(downloaded_bytes / expected_bytes, 0.99) - return { - "downloaded_bytes": downloaded_bytes, - "expected_bytes": expected_bytes, - "progress": round(progress, 3), - "cache_path": cache_path, - } - except Exception as e: - logger.warning(f"Error checking dataset download progress for {repo_id}: {e}") - return _empty + """Compatibility route backed by the shared multi-cache progress service.""" + from hub.services.datasets import downloads + return await downloads.get_dataset_download_progress_response( + repo_id, + hf_token = hf_token, + ) @router.post("/check-format", response_model = CheckFormatResponse) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index afd942e9a5..445a26f04d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -29,6 +29,8 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from utils.upload_limits import STT_AUDIO_B64_MAX_CHARS, STT_AUDIO_RAW_MAX_BYTES +from hub.dependencies import get_hf_token from core.inference.orchestrator import GenStreamError, GenStreamErrorRaised from core.inference.llama_admission import ( LlamaAdmissionCancelled, @@ -1692,6 +1694,8 @@ async def _aiter_llama_stream_items( from models.inference import ( LoadRequest, UnloadRequest, + TranscribeRequest, + SttLoadRequest, GenerateRequest, LoadResponse, LoadProgressResponse, @@ -3237,15 +3241,10 @@ def _request_matches_loaded_settings( ) ): return False - # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU - # request to its single lowest device (it drives one device only), so the - # backend records just that device; compare the request the same way, or a - # multi-GPU pick that resolves to the same device needlessly reloads. - if llama_backend.is_diffusion: - _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None - else: - _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None - if _req_gpu_ids != llama_backend.gpu_ids: + # A regular GGUF may narrow the requested placement pool. Accept either the + # original request or the effective status-echoed subset; diffusion keeps + # its single-device normalization. + if not llama_backend.matches_gpu_ids(request.gpu_ids): return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement @@ -3406,9 +3405,8 @@ async def _acquire_swap_gate() -> None: await asyncio.sleep(0.02) -# Counts in-flight auto-switch requests per (target, variant). The busy guard -# subtracts same-target waiters so concurrent requests for one model load once -# instead of each 409-ing the other. +# Counts auto-switch requests queued to load each (target, variant). They are not +# generating, so the drain wait below excludes them from the active inference count. _auto_switch_waiters: dict[tuple[str, str], int] = {} _auto_switch_waiters_guard = threading.Lock() @@ -3426,35 +3424,31 @@ def _note_switch_waiter(key: tuple[str, str], delta: int) -> None: _auto_switch_waiters.pop(key, None) -def _same_target_waiters(key: tuple[str, str]) -> int: +def _switch_waiter_count() -> int: with _auto_switch_waiters_guard: - return _auto_switch_waiters.get(key, 0) + return sum(max(0, count) for count in _auto_switch_waiters.values()) -# A second waiter map keyed by the raw requested model, registered before the -# (slow) resolve. The middleware counts a concurrent same-model request as -# in-flight before it resolves and joins _auto_switch_waiters, so without this -# the first request would see it as an unrelated request and 409. -_auto_switch_request_waiters: dict[str, int] = {} -_auto_switch_request_waiters_guard = threading.Lock() +async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None: + """Wait until a model replacement cannot interrupt active inference. - -def _request_waiter_key(requested_model: str) -> str: - return requested_model.strip().lower() - - -def _note_request_waiter(key: str, delta: int) -> None: - with _auto_switch_request_waiters_guard: - n = _auto_switch_request_waiters.get(key, 0) + delta - if n > 0: - _auto_switch_request_waiters[key] = n - else: - _auto_switch_request_waiters.pop(key, None) - - -def _same_request_waiters(key: str) -> int: - with _auto_switch_request_waiters_guard: - return _auto_switch_request_waiters.get(key, 0) + The caller holds ``inference_lifecycle_gate``, which prevents new inference + from starting while existing requests drain. Auto-switch requests that have + resolved their targets are scheduler waiters, not active generations, so + exclude them to avoid a queue deadlock. + """ + from core.inference.llama_keepwarm import other_inference_request_count + while True: + queued_switches = _switch_waiter_count() + if current_request_counted and queued_switches > 0: + queued_switches -= 1 + active_others = other_inference_request_count( + current_request_counted = current_request_counted, + include_pending = False, + ) + if active_others <= queued_switches: + return + await asyncio.sleep(0.02) def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: @@ -3582,7 +3576,6 @@ async def _maybe_auto_switch_model( from core.inference.local_model_resolver import resolve_local_gguf from core.inference.llama_keepwarm import ( get_last_unloaded_model, - other_inference_request_count, inference_lifecycle_gate, ) @@ -3603,12 +3596,7 @@ async def _maybe_auto_switch_model( if not auto_switch_on and get_auto_unload_idle_seconds() <= 0: return - # Register by the raw requested model before resolving (which can be slow): - # the middleware already counts a concurrent same-model request as in-flight, - # so the busy guard must know it shares this target even while it resolves. - request_key = _request_waiter_key(requested_model) - _note_request_waiter(request_key, 1) - try: + async def _resolve_and_switch() -> None: # Off the loop: a cold-cache rebuild walks several model dirs + HF caches. # With auto-switch off (or an omitted-model reload-only request), skip the # resolve so only the reload-stash path runs and no name is ever matched. @@ -3706,6 +3694,7 @@ async def _maybe_auto_switch_model( ) key = _switch_key(override_id, variant) _note_switch_waiter(key, 1) + waiter_noted = True try: async with _auto_switch_lock(): # The asyncio lock is per loop; add a process-wide gate so a swap on @@ -3718,31 +3707,6 @@ async def _maybe_auto_switch_model( if _already_serving(): _record_serving_alias() return - # Single slot: refuse a cross-model swap while another inference - # request is active rather than killing its response. Requests - # heading to this same target (by resolved id or raw name) are - # excluded, so concurrent requests for one model load once. A - # pending request is still in the middleware, not generating, so - # it is not counted here. - same_others = max( - _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 - ) - others = other_inference_request_count( - current_request_counted = True, include_pending = False - ) - # Not gated on the GGUF being loaded: _load_model_impl also - # tears down an active Unsloth backend before loading a GGUF, - # so refuse whenever any other inference request is in flight. - if others > same_others: - raise HTTPException( - status_code = 409, - detail = openai_error_body( - "Cannot switch models while another inference request is in progress.", - status = 409, - code = "model_switch_busy", - param = "model", - ), - ) # Apply this model's saved launch flags so the swap honors the config. override = get_model_override(override_id) load_kwargs = {"model_path": target_id, "gguf_variant": variant} @@ -3757,16 +3721,22 @@ async def _maybe_auto_switch_model( LoadRequest(**load_kwargs), fastapi_request, current_subject, + current_request_counted = True, ) # Advertise the repo id (not the concrete load path) as the loaded # model's public id and override key for /v1/models and idle stash. get_llama_cpp_backend()._openai_advertised_id = override_id finally: + # Deregister before releasing the gate: otherwise a swap on another + # loop counts this finished request as queued and unloads its model. + _note_switch_waiter(key, -1) + waiter_noted = False _auto_switch_process_lock.release() finally: - _note_switch_waiter(key, -1) - finally: - _note_request_waiter(request_key, -1) + if waiter_noted: + _note_switch_waiter(key, -1) + + await _resolve_and_switch() async def _auto_switch_from_request_body(request: Request, current_subject: str): @@ -3922,15 +3892,19 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: """Classify a GGUF as diffusion, normal, or unknown before it is loaded. ``None`` is important here: a remote GGUF whose header is not cached can - still be routed to the single-GPU diffusion runner after download. Treating - that case as normal would let Manual mode skip the training guard even - though the runner ignores Manual's llama-server placement controls. + still be routed to the single-GPU diffusion runner after download. Default + placement keeps that unknown case guarded until the header is available. """ identity = " ".join( str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") ).lower() - if "diffusion" in identity: - return True + # Name-only hint, used ONLY as a pre-download fallback, scoped to the + # DiffusionGemma runner family: a bare "diffusion" substring is common in + # ordinary text-model names/paths (e.g. "stable-diffusion-prompt"), and treating + # those as diffusion falsely rejects a valid Vulkan+gpu_ids GGUF (#7239). Normalize + # non-alphanumerics so "DiffusionGemma"/"diffusion-gemma" collapse to one token. + # The local header below stays authoritative. + name_says_diffusion = "diffusiongemma" in _re.sub(r"[^a-z0-9]+", "", identity) try: main = getattr(config, "gguf_file", None) @@ -3940,23 +3914,86 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: if repo and variant: from hub.utils.gguf import resolve_local_gguf_path main = resolve_local_gguf_path(repo, variant) - if not main or not Path(main).is_file(): - return None - - probe = LlamaCppBackend() - probe._read_gguf_metadata(str(main)) - if probe.is_diffusion: - return True - # A successfully decoded architecture proves that this is a normal - # llama-server GGUF. No architecture means the lightweight probe could - # not establish the routing decision, so preserve the unknown state. - if getattr(probe, "_architecture", None): - return False - return None + if main and Path(main).is_file(): + # The local GGUF header is authoritative (same probe the loader uses), so + # it can't be fooled by a "diffusion"-flavored name/path. + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A decoded architecture proves a normal llama-server GGUF; no architecture + # means the probe was inconclusive, so fall through to the name hint below. + if getattr(probe, "_architecture", None): + return False except Exception as e: logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + + # Header unavailable (remote uncached) or inconclusive: True only for the + # DiffusionGemma name family; otherwise None keeps an unknown remote GGUF guarded + # as potentially diffusion until its header proves otherwise. + return True if name_says_diffusion else None + + +async def _resolve_gguf_gpu_ids_for_request( + config: ModelConfig, gpu_ids: Optional[List[int]] +) -> Optional[List[int]]: + """Resolve and fully validate an explicit GGUF GPU placement pool. + + CUDA and ROCm use physical IDs. Vulkan uses ggml ordinals, so its device + existence check comes from the same ggml probe used by the loader. Both + /load and /validate call this before their training guard or any teardown. + """ + if not gpu_ids: return None + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + is_vulkan = LlamaCppBackend._is_vulkan_backend() + if get_device() == DeviceType.XPU and not is_vulkan: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + + if is_vulkan and _classify_diffusion_gguf(config) is True: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ), + ) + + try: + resolved = resolve_requested_gpu_ids(gpu_ids, is_vulkan = is_vulkan) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc + + if is_vulkan and resolved: + binary = LlamaCppBackend._find_llama_server_binary() + if binary: + probed = { + gpu[0] for gpu in await asyncio.to_thread(LlamaCppBackend._get_gpu_memory, binary) + } + wanted = {int(gpu_id) for gpu_id in resolved} + if not wanted.issubset(probed): + raise HTTPException( + status_code = 400, + detail = ( + f"Requested Vulkan GPU ordinal(s) {sorted(wanted)} not " + f"present. Available Vulkan devices: {sorted(probed)}." + ), + ) + + return resolved + def _guard_chat_load_against_training( config: ModelConfig, @@ -3996,8 +4033,18 @@ def _guard_chat_load_against_training( if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: return + # Vulkan GGUF pins are ggml ordinals, not CUDA physical IDs. Detect this + # before deriving a possible diffusion fallback device so an unknown remote + # GGUF never sends its ordinal through the CUDA single-device path. + is_vulkan = False + if is_gguf: + try: + is_vulkan = LlamaCppBackend._is_vulkan_backend() + except Exception as e: + logger.warning("Could not detect Vulkan backend for chat-load guard: %s", e) + diffusion_gpu = None - if is_gguf and diffusion_kind is not False: + if is_gguf and diffusion_kind is not False and not (is_vulkan and requested_gpu_ids): # Use the same token selection as the runner: an explicit pick wins, # followed by DG_GPU, the first parent-visible token, then GPU 0. diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( @@ -4024,6 +4071,7 @@ def _guard_chat_load_against_training( max_seq_length = max_seq_length, requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, + is_vulkan = is_vulkan, required_override_gb = required_override_gb, single_device_gpu = diffusion_gpu, ) @@ -4186,6 +4234,15 @@ def _maybe_unsupported_message(msg: str) -> str: return msg +def _raise_if_sidecar_swap_in_progress() -> None: + from utils.transformers_version import sidecar_swap_in_progress + if sidecar_swap_in_progress(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + + @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, @@ -4206,24 +4263,23 @@ async def load_model( # install can reserve while this request queues on the gate, so the pre-gate # check alone is only a fast path. from core.inference.llama_keepwarm import inference_lifecycle_gate - from utils.transformers_version import sidecar_swap_in_progress - _swap_409 = HTTPException( - status_code = 409, - detail = "A transformers installation is in progress. Retry when it completes.", - ) - if sidecar_swap_in_progress(): - raise _swap_409 + _raise_if_sidecar_swap_in_progress() # Hold the lifecycle gate across the load so idle auto-unload can't unload the # model mid-load. Auto-switch calls _load_model_impl directly since it already # holds this gate. async with inference_lifecycle_gate(): - if sidecar_swap_in_progress(): - raise _swap_409 + _raise_if_sidecar_swap_in_progress() return await _load_model_impl(request, fastapi_request, current_subject) -async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): +async def _load_model_impl( + request: LoadRequest, + fastapi_request: Request, + current_subject: str, + *, + current_request_counted: bool = False, +): from core.inference.llama_cpp import LlamaServerNotFoundError # A new load starts here; arm the progress throttle so this load's first @@ -4322,6 +4378,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) ): + llama_backend._record_matching_gpu_request(request.gpu_ids) logger.info( "Model already loaded (GGUF): " f"{model_log_label} variant={request.gguf_variant or llama_backend.hf_variant}, skipping reload" @@ -4368,6 +4425,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, ) else: if ( @@ -4434,41 +4492,12 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # GGUF supports gpu_ids: validate the pick up front (before the training - # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects - # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts - # are rejected outright: the picker's indices are torch-xpu ordinals neither - # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin - # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. - if config.is_gguf and effective_gpu_ids is not None: - from utils.hardware import DeviceType, get_device - from utils.hardware.hardware import resolve_requested_gpu_ids - - if get_device() == DeviceType.XPU: - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported on Intel XPU. " - "Omit gpu_ids to use all devices." - ), - ) - # Same reasoning for a Vulkan-only build: --device pins ggml's own - # Vulkan ordinals, so a physical pick can land on the wrong card on - # masked or non-contiguous hosts. - if LlamaCppBackend._is_vulkan_backend(): - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported with a Vulkan " - "llama.cpp build: physical GPU ids have no defined " - "mapping to Vulkan device ordinals. Omit gpu_ids to use " - "all devices." - ), - ) - try: - resolve_requested_gpu_ids(effective_gpu_ids) - except ValueError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc + # Validate the full GGUF placement pool before the training guard so an + # invalid physical ID or Vulkan ordinal is a clean 400, not a masked VRAM + # 409. The same helper is used by /validate. + gguf_gpu_ids: Optional[List[int]] = None + if config.is_gguf: + gguf_gpu_ids = await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids) if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, @@ -4557,6 +4586,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre ), ) + # Keep the resident model alive until every active generation finishes; + # the caller's lifecycle gate blocks new starts. + await _wait_for_model_switch_idle(current_request_counted = current_request_counted) + # A sidecar install can reserve the gate while inference drains, after the + # route-level checks above, so recheck before replacing either backend. + _raise_if_sidecar_swap_in_progress() + # Unload any active Unsloth model only after every hub conflict check. if unsloth_backend.active_model_name: logger.info( @@ -4585,8 +4621,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre gpu_layers = request.gpu_layers, n_cpu_moe = request.n_cpu_moe, tensor_split = request.tensor_split, - gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, + # Issue #7164: explicit GPU pin resolved to physical ids above. + gpu_ids = gguf_gpu_ids, ) if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server @@ -4760,6 +4797,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, ) # โ”€โ”€ Standard path: load via Unsloth/transformers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ @@ -4767,6 +4805,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Unload any active GGUF model first llama_backend = get_llama_cpp_backend() + await _wait_for_model_switch_idle(current_request_counted = current_request_counted) + _raise_if_sidecar_swap_in_progress() if llama_backend.is_loaded: logger.info("Unloading GGUF model before loading Unsloth model") llama_backend.unload_model() @@ -5051,36 +5091,8 @@ async def validate_model( # Apply the same training coexistence policy as /load before the frontend # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is - # a clean 400) before the guard sizes the model against training VRAM. - # XPU-host picks are rejected like /load (no defined mapping from the - # picker's torch-xpu ordinals to the launcher's device spaces). - if config.is_gguf and effective_gpu_ids is not None: - from utils.hardware import DeviceType, get_device - from utils.hardware.hardware import resolve_requested_gpu_ids - - if get_device() == DeviceType.XPU: - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported on Intel XPU. " - "Omit gpu_ids to use all devices." - ), - ) - if LlamaCppBackend._is_vulkan_backend(): - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported with a Vulkan " - "llama.cpp build: physical GPU ids have no defined " - "mapping to Vulkan device ordinals. Omit gpu_ids to use " - "all devices." - ), - ) - try: - resolve_requested_gpu_ids(effective_gpu_ids) - except ValueError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc + if config.is_gguf: + await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids) effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -5144,10 +5156,10 @@ async def validate_model( latest_tier_active_for, config.identifier, request.hf_token ): effective_load_in_4bit = False - # A metadata-only probe just reads the GGUF header and allocates no VRAM, - # so it must not be refused by the training guard. Real loads validate - # without include_context_length and /load applies the guard again. - if not request.include_context_length: + # A metadata-only probe reads the GGUF header and allocates no VRAM, so the + # training guard must not refuse it. Real loads omit include_context_length / + # include_chat_template, and /load applies the guard again. + if not (request.include_context_length or request.include_chat_template): # Match /load's inherited llama.cpp extras and parallel slot count so # validation cannot pass a smaller estimate than the subsequent load. effective_extra_args = _resolve_inherited_extra_args( @@ -5189,9 +5201,15 @@ async def validate_model( context_length: Optional[int] = None layer_count: Optional[int] = None moe_layer_count: Optional[int] = None - if request.include_context_length and is_gguf: + chat_template: Optional[str] = None + # Both header probes read the same local GGUF, so resolve it once. + if (request.include_context_length or request.include_chat_template) and is_gguf: from hub.utils.gguf import resolve_local_gguf_path - from utils.models.gguf_metadata import read_gguf_staged_dims + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + from utils.models.gguf_metadata import ( + read_gguf_chat_template, + read_gguf_staged_dims, + ) # Best-effort: a header-read failure must never fail validation of an # otherwise-valid model (the outer except turns it into a 400). @@ -5207,13 +5225,24 @@ async def validate_model( model_identifier, request.gguf_variant ) if local_gguf: - # Header walk reads tokenizer arrays for dense models (tens of - # ms); keep it off the event loop. - dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf) - if dims: - context_length = dims["context_length"] - layer_count = dims["layer_count"] - moe_layer_count = dims["moe_layer_count"] + if request.include_context_length: + # Header walk reads tokenizer arrays (tens of ms); keep it + # off the event loop. + dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf) + if dims: + context_length = dims["context_length"] + layer_count = dims["layer_count"] + moe_layer_count = dims["moe_layer_count"] + if request.include_chat_template: + # Read only the leased GGUF's own embedded template (the copy + # llama.cpp loads), never a sibling sidecar: the native grant + # authorizes just this path, so neighbours would be scope escalation. + raw_template = await asyncio.to_thread(read_gguf_chat_template, local_gguf) + if ( + raw_template is not None + and len(raw_template.encode("utf-8")) <= MAX_CHAT_TEMPLATE_BYTES + ): + chat_template = raw_template except Exception as e: logger.debug("Header probe failed for %s: %s", model_log_label, e) @@ -5232,6 +5261,7 @@ async def validate_model( context_length = context_length, layer_count = layer_count, moe_layer_count = moe_layer_count, + chat_template = chat_template, requires_transformers_upgrade = transformers_upgrade is not None, transformers_upgrade = transformers_upgrade, ) @@ -5805,10 +5835,15 @@ async def get_status(current_subject: str = Depends(get_current_subject)): try: _bin = type(llama_backend)._find_llama_server_binary() _caps = type(llama_backend).probe_server_capabilities(_bin) - _supports_mtp = bool(_caps.get("supports_mtp", False)) + # Fail open on inconclusive probes: False means a definitive + # "binary lacks MTP" to API consumers. + _supports_mtp = bool( + _caps.get("supports_mtp", False) + or (_caps.get("found", False) and _caps.get("mtp_probe_inconclusive", False)) + ) except Exception: _bin = None - _supports_mtp = True # fail open + _supports_mtp = False # no usable binary: MTP genuinely unavailable try: from utils.llama_cpp_freshness import check_prebuilt_freshness _freshness = check_prebuilt_freshness(_bin) @@ -5887,6 +5922,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, @@ -6051,6 +6087,7 @@ async def generate_audio( # Advertised repo id after an auto-switch load, else a clean public id, # never the absolute .gguf path. model_name = _llama_public_model_id(llama_backend) + _audio_model_id = getattr(llama_backend, "model_identifier", None) or model_name gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -6069,6 +6106,7 @@ async def generate_audio( if not model_info.get("is_audio"): raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") model_name = public_model_id(backend.active_model_name) + _audio_model_id = getattr(backend, "active_model_name", None) or model_name gen = lambda: backend.generate_audio_response( text = text, temperature = payload.temperature, @@ -6080,6 +6118,13 @@ async def generate_audio( use_adapter = payload.use_adapter, ) + # Apply per-model recommended sampling + any operator UNSLOTH_SAMPLING_* pin before + # generating, so `unsloth run --temperature` (and the other pins) and per-model + # recommendations reach audio (TTS) generation too, not just chat. The gen lambdas read + # payload.* lazily at call time, so filling here takes effect; this covers both the direct + # /audio/generate route and the chat-completions audio branches that delegate here. + _fill_recommended_sampling_openai(payload, _audio_model_id) + try: wav_bytes, sample_rate = await asyncio.to_thread(gen) except Exception as e: @@ -6107,6 +6152,342 @@ async def generate_audio( ) +# ===================================================================== +# Speech-to-text (STT) sidecar (/audio/transcribe, /audio/stt/*) +# ===================================================================== + + +def _resolve_stt_engine(engine: Optional[str]) -> str: + """Normalize the requested STT engine name; default is Transformers.""" + normalized = (engine or "transformers").strip().lower() + if normalized in ("", "transformers", "whisper"): + return "transformers" + if normalized in ("gguf", "ggml", "whisper_cpp", "whisper.cpp"): + return "gguf" + raise HTTPException( + status_code = 422, + detail = f"Unknown STT engine '{engine}'. Use 'transformers' or 'gguf'.", + ) + + +def _resolve_serving_stt_engine(engine: Optional[str]) -> str: + """Resolve the engine that will actually serve a model. + + whisper.cpp (gguf) only accepts curated ids, which Transformers serves too, + so when whisper-server is not installed (the common case: `unsloth studio + update` does not yet build it) fall back to Transformers instead of 501-ing + on every recording. Used for download/load/transcribe; unload targets a + specific engine via _resolve_stt_engine. + """ + resolved = _resolve_stt_engine(engine) + if resolved == "gguf": + from core.inference import stt_ggml_sidecar + if not stt_ggml_sidecar.is_available(): + return "transformers" + return resolved + + +def _stt_sidecar_for(engine: str): + if engine == "gguf": + from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar + return get_ggml_stt_sidecar() + from core.inference.stt_sidecar import get_stt_sidecar + return get_stt_sidecar() + + +@studio_router.get("/audio/stt/status") +async def stt_status( + model: Optional[str] = None, current_subject: str = Depends(get_current_subject) +): + """Report STT availability and which model, if any, is resident. + + ``model`` extends the Transformers ``downloaded_models`` check to a + custom Hugging Face repository beyond the curated defaults. + """ + from core.inference import stt_ggml_sidecar, stt_sidecar + from core.inference.stt_sidecar import ( + DEFAULT_STT_MODEL, + STT_MODELS, + get_stt_sidecar, + is_available, + ) + + sidecar = get_stt_sidecar() + ggml = stt_ggml_sidecar.get_ggml_stt_sidecar() + transformers_downloaded = [ + model_id for model_id in STT_MODELS if stt_sidecar.is_model_downloaded(model_id) + ] + if model and model not in STT_MODELS and stt_sidecar.is_model_downloaded(model): + transformers_downloaded.append(model) + return JSONResponse( + content = { + "available": is_available(), + "loaded_model": sidecar.loaded_model, + "loading": sidecar.is_loading(), + "device": sidecar.device, + "keep_alive_seconds": sidecar.keep_alive_seconds, + "default_model": DEFAULT_STT_MODEL, + "models": list(STT_MODELS.keys()), + # Transformers engine, same shape as "gguf" below so clients read + # either generically. Top-level fields above kept for old clients. + "transformers": { + "available": is_available(), + "loaded_model": sidecar.loaded_model, + "loading": sidecar.is_loading(), + "device": sidecar.device, + "keep_alive_seconds": sidecar.keep_alive_seconds, + "default_model": DEFAULT_STT_MODEL, + "models": list(STT_MODELS.keys()), + "downloaded_models": transformers_downloaded, + "download": stt_sidecar.download_status(), + }, + # whisper.cpp (GGUF) engine. + "gguf": { + "available": stt_ggml_sidecar.is_available(), + "loaded_model": ggml.loaded_model, + "loading": ggml.is_loading(), + "device": ggml.device, + "keep_alive_seconds": ggml.keep_alive_seconds, + "default_model": stt_ggml_sidecar.DEFAULT_GGML_STT_MODEL, + "models": list(stt_ggml_sidecar.GGML_STT_MODELS.keys()), + "downloaded_models": [ + model_id + for model_id in stt_ggml_sidecar.GGML_STT_MODELS + if stt_ggml_sidecar._cached_model_path(model_id) is not None + ], + "download": stt_ggml_sidecar.download_status(), + }, + } + ) + + +@studio_router.post("/audio/stt/download") +async def stt_download( + payload: SttLoadRequest, + current_subject: str = Depends(get_current_subject), + hf_token: Optional[str] = Depends(get_hf_token), +): + """Start a background download of a dictation model. + + Both engines download directly (a GGML checkpoint is a single file the Model + Hub's GGUF variant planner cannot express; a Transformers checkpoint is a + whole snapshot). Progress is reported by /audio/stt/status. + """ + from core.inference import stt_ggml_sidecar, stt_sidecar + from core.inference.stt_sidecar import ( + SttModelCompatibilityError, + SttModelIdError, + validate_remote_model, + ) + + engine = _resolve_serving_stt_engine(payload.engine) + module = stt_ggml_sidecar if engine == "gguf" else stt_sidecar + try: + # Transformers accepts custom `owner/model` repos, so confirm the repo is + # a Whisper checkpoint (metadata-only) before snapshot_download pulls a + # possibly-large non-STT repo into the shared cache. Curated ids + # short-circuit; GGUF only accepts curated ids, so it needs no check. + if engine != "gguf": + validated = await asyncio.to_thread(validate_remote_model, payload.model, hf_token) + # Pin the download to the commit that was just validated so the + # repo cannot be swapped between validation and snapshot_download. + await asyncio.to_thread( + module.start_model_download, + payload.model, + hf_token, + validated.get("revision"), + ) + else: + await asyncio.to_thread(module.start_model_download, payload.model, hf_token) + except SttModelIdError as e: + raise HTTPException(status_code = 422, detail = str(e)) + except SttModelCompatibilityError as e: + raise HTTPException(status_code = 422, detail = str(e)) + return JSONResponse(content = module.download_status()) + + +@studio_router.post("/audio/stt/load") +async def stt_load(payload: SttLoadRequest, current_subject: str = Depends(get_current_subject)): + """Load the selected STT model after the user starts local dictation.""" + from core.inference.stt_sidecar import ( + SttLoadCancelledError, + SttModelCompatibilityError, + SttModelIdError, + SttModelNotDownloadedError, + SttUnavailableError, + get_stt_sidecar, + ) + + sidecar = _stt_sidecar_for(_resolve_serving_stt_engine(payload.engine)) + try: + await asyncio.to_thread(sidecar.load, payload.model) + except SttModelNotDownloadedError as e: + raise HTTPException(status_code = 409, detail = str(e)) + except SttUnavailableError as e: + raise HTTPException(status_code = 501, detail = str(e)) + except SttLoadCancelledError as e: + raise HTTPException(status_code = 409, detail = str(e)) + except SttModelIdError as e: + raise HTTPException(status_code = 422, detail = str(e)) + except SttModelCompatibilityError as e: + raise HTTPException(status_code = 422, detail = str(e)) + except Exception as e: + logger.error(f"STT load error: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = safe_error_detail(e)) + return JSONResponse(content = {"loaded_model": sidecar.loaded_model, "device": sidecar.device}) + + +@studio_router.post("/audio/stt/validate") +async def stt_validate( + payload: SttLoadRequest, + current_subject: str = Depends(get_current_subject), + hf_token: Optional[str] = Depends(get_hf_token), +): + """Verify a Hub repository is a Whisper checkpoint before downloading it.""" + from core.inference.stt_sidecar import ( + SttModelCompatibilityError, + SttModelIdError, + validate_remote_model, + ) + + try: + result = await asyncio.to_thread(validate_remote_model, payload.model, hf_token) + except (SttModelIdError, SttModelCompatibilityError) as e: + raise HTTPException(status_code = 422, detail = str(e)) + return JSONResponse(content = result) + + +@studio_router.post("/audio/stt/unload") +async def stt_unload( + engine: Optional[str] = None, current_subject: str = Depends(get_current_subject) +): + """Release the local STT model when dictation is idle. + + Without an engine, both sidecars unload so an engine switch in Voice + settings always frees whichever backend was resident. + """ + if engine is None: + engines = ["transformers", "gguf"] + else: + # Use the serving resolver: a "gguf" pick without whisper-server is + # actually served by the Transformers fallback, so unload must target + # that same engine or the resident model is never freed. + engines = [_resolve_serving_stt_engine(engine)] + # Attempt every engine even if one raises, so failing to unload one never + # skips freeing the other (both can be resident after a switch). + failed: list[str] = [] + for name in engines: + try: + await asyncio.to_thread(_stt_sidecar_for(name).unload) + except Exception as exc: # noqa: BLE001 - report after attempting all engines + logger.warning("Failed to unload STT engine '%s': %s", name, exc) + failed.append(name) + if failed: + raise HTTPException( + status_code = 500, + detail = f"Failed to unload STT engine(s): {', '.join(failed)}", + ) + return JSONResponse(content = {"loaded_model": None, "device": None}) + + +async def _transcribe_audio_bytes( + raw: bytes, + model: Optional[str], + language: Optional[str], + fast: bool, + engine: Optional[str] = None, +) -> JSONResponse: + """Run STT for already-decoded request bytes.""" + from core.inference.stt_sidecar import ( + SttAudioDecodeError, + SttAudioTooLongError, + SttLanguageError, + SttLoadCancelledError, + SttModelCompatibilityError, + SttModelIdError, + SttModelNotDownloadedError, + SttUnavailableError, + ) + + if not raw: + raise HTTPException(status_code = 400, detail = "Audio is empty.") + if len(raw) > _MAX_AUDIO_RAW_BYTES: + raise HTTPException(status_code = 413, detail = "Audio is too large.") + + sidecar = _stt_sidecar_for(_resolve_serving_stt_engine(engine)) + try: + result = await asyncio.to_thread( + sidecar.transcribe, + raw, + model, + language, + fast, + ) + except SttUnavailableError as e: + raise HTTPException(status_code = 501, detail = str(e)) + except SttLoadCancelledError as e: + raise HTTPException(status_code = 409, detail = str(e)) + except SttModelNotDownloadedError as e: + raise HTTPException(status_code = 409, detail = str(e)) + except SttModelIdError as e: + raise HTTPException(status_code = 422, detail = str(e)) + except SttModelCompatibilityError as e: + raise HTTPException(status_code = 422, detail = str(e)) + except SttLanguageError as e: + raise HTTPException(status_code = 422, detail = str(e)) + except SttAudioTooLongError as e: + raise HTTPException(status_code = 413, detail = str(e)) + except SttAudioDecodeError as e: + raise HTTPException(status_code = 400, detail = str(e)) + except Exception as e: + logger.error(f"Transcription error: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = safe_error_detail(e)) + return JSONResponse(content = result) + + +@studio_router.post("/audio/transcribe") +async def transcribe_audio( + payload: TranscribeRequest, current_subject: str = Depends(get_current_subject) +): + """Transcribe dictation audio to text via the STT sidecar. + + Runs alongside the chat model without evicting it, so any model (including + text-only ones) can be driven by voice. + """ + b64 = payload.audio or "" + if not b64: + raise HTTPException(status_code = 400, detail = "No audio provided.") + if len(b64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException(status_code = 413, detail = "Audio is too large.") + try: + raw = base64.b64decode(b64, validate = True) + except Exception: + raise HTTPException(status_code = 400, detail = "Audio is not valid base64.") + return await _transcribe_audio_bytes( + raw, payload.model, payload.language, payload.fast, payload.engine + ) + + +@studio_router.post("/audio/transcribe/raw") +async def transcribe_audio_raw( + request: Request, + model: Optional[str] = None, + language: Optional[str] = None, + fast: bool = False, + engine: Optional[str] = None, + current_subject: str = Depends(get_current_subject), +): + """Transcribe a raw audio body without base64 or JSON conversion overhead.""" + chunks: list[bytes] = [] + size = 0 + async for chunk in request.stream(): + size += len(chunk) + if size > _MAX_AUDIO_RAW_BYTES: + raise HTTPException(status_code = 413, detail = "Audio is too large.") + chunks.append(chunk) + return await _transcribe_audio_bytes(b"".join(chunks), model, language, fast, engine) + + # ===================================================================== # OpenAI-Compatible Chat Completions (/chat/completions) # ===================================================================== @@ -6148,8 +6529,8 @@ def _decode_audio_base64(b64: str) -> np.ndarray: # cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally # bounds the *decoded* length, since a small compressed file (opus/flac/etc.) # can expand to a far larger PCM array than the encoded-size cap implies. -_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024 -_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3 +_MAX_AUDIO_RAW_BYTES = STT_AUDIO_RAW_MAX_BYTES +_MAX_AUDIO_B64_CHARS = STT_AUDIO_B64_MAX_CHARS _MAX_AUDIO_SECONDS = 30 * 60 _WAV_HEADER_BYTES = 44 _MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000 @@ -7035,6 +7416,51 @@ async def delete_openai_container( await client.close() +def _fill_recommended_sampling_openai(payload, model_id) -> None: + """Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to a + ChatCompletionRequest in place. + + Only the sampling fields the client did NOT explicitly send (tracked via + ``model_fields_set``) are overwritten, so a client that sets a field stays byte-identical + unless an operator pins it. Fields with neither a recommendation nor a pin keep their + existing (schema-default) value. + """ + from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES + + explicit = { + f: (getattr(payload, f) if f in payload.model_fields_set else None) + for f in SAMPLING_FIELD_NAMES + } + effective = resolve_effective_sampling(model_id, explicit) + for field, value in effective.items(): + setattr(payload, field, value) + + +# /v1/completions is proxied to llama-server verbatim; its repetition knob is "repeat_penalty", +# and every other sampling field keeps its name (mirrors _build_passthrough_payload). +_COMPLETIONS_SAMPLING_BODY_KEY = {"repetition_penalty": "repeat_penalty"} + + +def _fill_recommended_sampling_completions(body: dict, model_id) -> None: + """Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to a raw + ``/v1/completions`` body in place, so the legacy (non-chat) endpoint honors the same pins as + ``/v1/chat/completions``. + + Unlike :func:`_fill_recommended_sampling_openai`, which fills a ChatCompletionRequest whose + schema already carries per-field defaults, this body is proxied to llama-server as-is. A field + with no operator pin, client value, or per-model recommendation is therefore left untouched + (``fill_defaults = False``) so llama-server keeps its own default rather than being forced onto + this schema's value. llama-server names the repetition knob ``repeat_penalty``, so read and + write that alias for the client-sent value and any pin. + """ + from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES + + explicit = {f: body.get(_COMPLETIONS_SAMPLING_BODY_KEY.get(f, f)) for f in SAMPLING_FIELD_NAMES} + effective = resolve_effective_sampling(model_id, explicit, fill_defaults = False) + for field, value in effective.items(): + body[_COMPLETIONS_SAMPLING_BODY_KEY.get(field, field)] = value + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, @@ -7078,7 +7504,7 @@ async def openai_chat_completions( if payload.provider_id or payload.provider_type: # External provider: this request won't touch the local GGUF, so drop it # from the keep-warm count or its in-flight stream would falsely block a - # concurrent local auto-switch with model_switch_busy. + # concurrent local model switch from proceeding. from core.inference.llama_keepwarm import untrack_current_request untrack_current_request(request.scope) @@ -7364,6 +7790,13 @@ async def openai_chat_completions( completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) + # Apply recommended sampling + operator pins to the omitted fields before generating, + # so audio-input (non-whisper) generation honors `unsloth run --temperature` and + # per-model recommendations like chat does. Whisper (ASR) ignores these fields. + _fill_recommended_sampling_openai( + payload, getattr(backend, "active_model_name", None) or model_name + ) + def audio_input_generate(): if model_info.get("audio_type") == "whisper": return backend.generate_whisper_response( @@ -7507,6 +7940,18 @@ async def openai_chat_completions( ), ) + # Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to the + # fields the client omitted, so agents and API clients get the model's tuned defaults + # unless they set the field explicitly. Placed after external-provider routing (which + # returned above) so only local llama-server / transformers requests are touched, and it + # covers both the passthrough and non-passthrough branches below since both read payload.*. + _reco_model_id = ( + getattr(llama_backend, "model_identifier", None) + if using_gguf + else getattr(backend, "active_model_name", None) + ) or model_name + _fill_recommended_sampling_openai(payload, _reco_model_id) + # โ”€โ”€ Standard OpenAI function-calling pass-through (GGUF only) โ”€โ”€โ”€โ”€ # When a client (opencode / Claude Code via OpenAI compat / Cursor / # Continue / ...) sends standard OpenAI `tools` without Unsloth's @@ -10284,6 +10729,10 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if _resolved_max_tokens is not None else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) + # Apply per-model recommended sampling and any operator UNSLOTH_SAMPLING_* pin to the raw + # body so /v1/completions honors the same pins as /v1/chat/completions; it is otherwise a + # verbatim proxy that would keep llama-server's defaults for every omitted sampling field. + _fill_recommended_sampling_completions(body, getattr(llama_backend, "model_identifier", None)) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) @@ -11242,6 +11691,9 @@ async def _responses_stream( detail = "Image provided but current GGUF model does not support vision.", ) + # Streaming /v1/responses builds the passthrough body directly (bypassing + # openai_chat_completions), so apply recommended sampling here too. + _fill_recommended_sampling_openai(chat_req, getattr(llama_backend, "model_identifier", None)) body = _build_openai_passthrough_body( chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) @@ -12673,14 +13125,28 @@ async def anthropic_messages( # endpoint matches /v1/chat/completions. _has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision) - temperature = payload.temperature if payload.temperature is not None else 0.6 - top_p = payload.top_p if payload.top_p is not None else 0.95 - top_k = payload.top_k if payload.top_k is not None else 20 - min_p = payload.min_p if payload.min_p is not None else 0.01 - repetition_penalty = ( - payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 + # Fill omitted sampling fields with the per-model recommendation (or an operator + # UNSLOTH_SAMPLING_* pin); an explicit client value wins unless the operator pinned it. + # Anthropic sampling fields are Optional, so None already marks "client omitted". + from utils.inference.inference_config import resolve_effective_sampling + + _anthropic_sampling = resolve_effective_sampling( + getattr(llama_backend, "model_identifier", None) or model_name, + { + "temperature": payload.temperature, + "top_p": payload.top_p, + "top_k": payload.top_k, + "min_p": payload.min_p, + "repetition_penalty": payload.repetition_penalty, + "presence_penalty": payload.presence_penalty, + }, ) - presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0 + temperature = _anthropic_sampling["temperature"] + top_p = _anthropic_sampling["top_p"] + top_k = _anthropic_sampling["top_k"] + min_p = _anthropic_sampling["min_p"] + repetition_penalty = _anthropic_sampling["repetition_penalty"] + presence_penalty = _anthropic_sampling["presence_penalty"] stop = payload.stop_sequences or None # Translate Anthropic tool_choice to OpenAI format for llama-server. Falls diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 540647e3bc..84b89ad7d5 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""llama.cpp prebuilt update endpoints. +"""llama.cpp prebuilt update endpoints -- the single main update item. GET /api/llama/update-status -> is a newer prebuilt available + job state POST /api/llama/update -> download + atomically swap to the latest @@ -9,13 +9,19 @@ POST /api/llama/update -> download + atomically swap to the latest Detection reuses utils.llama_cpp_freshness; the swap reuses install_llama_prebuilt.py via utils.llama_cpp_update. Both fail open so the UI never blocks on a missing marker / offline GitHub. + +whisper.cpp updates piggyback here: the status payload carries a whisper +sub-status (update_available is the llama OR whisper union) and the apply job +chains a whisper phase after the llama phase when whisper is behind, with a +per-phase breakdown in job.phases. All pre-existing top-level fields keep +their shape, so older clients keep working unchanged. """ from __future__ import annotations import asyncio import threading -from typing import Optional +from typing import Literal, Optional from fastapi import APIRouter, Depends, Query from pydantic import BaseModel, Field @@ -38,6 +44,31 @@ class LlamaUpdateJob(BaseModel): progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.") started_at: Optional[str] = None finished_at: Optional[str] = None + phases: Optional[dict] = Field( + None, + description = ( + "Per-phase breakdown of a chained llama+whisper job " + "(name -> state/progress/to_tag/...); None for pre-chaining jobs." + ), + ) + + +class WhisperSubStatus(BaseModel): + """The whisper piggyback inside the llama update item.""" + + update_available: bool = Field( + False, description = "True when the chained apply would run a whisper phase." + ) + installed_tag: Optional[str] = None + latest_tag: Optional[str] = None + update_size_bytes: Optional[int] = None + skip_reason: Optional[str] = Field( + None, + description = ( + "Why the whisper phase would be skipped " + "(up_to_date | local_link | source_build | not_installed | ...)." + ), + ) class LlamaUpdateStatusResponse(BaseModel): @@ -46,7 +77,18 @@ class LlamaUpdateStatusResponse(BaseModel): description = "True when the install came from an Unsloth prebuilt (has a marker).", ) update_available: bool = Field( - False, description = "True when the latest release is genuinely newer than the install." + False, + description = ( + "True when an update would do something: llama.cpp is behind OR the " + "whisper piggyback is behind." + ), + ) + llama_update_available: bool = Field( + False, description = "True when the latest llama.cpp release is newer than the install." + ) + update_component: Optional[Literal["llama", "whisper"]] = Field( + None, + description = "Component whose versions the combined update banner should display.", ) stale: bool = Field( False, description = "Update available AND install older than the staleness threshold." @@ -62,6 +104,9 @@ class LlamaUpdateStatusResponse(BaseModel): update_size_bytes: Optional[int] = Field( None, description = "Download size of the prebuilt Update would fetch, in bytes." ) + whisper: Optional[WhisperSubStatus] = Field( + None, description = "Whisper piggyback sub-status; None when the probe is unavailable." + ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 0806c2f513..ed83a12f48 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -60,13 +60,52 @@ def _safe_is_dir(path) -> bool: # Shared with the hub inventory scans; keep the private aliases so existing -# importers (core.inference.local_model_resolver, tests) stay valid. +# importers stay valid. ``_HF_REPO_ID_RE`` is the Hub repo id shape ("owner/name"); +# anything else is treated as a local filesystem path. from utils.hidden_models import ( + _HF_REPO_ID_RE, + _existing_resolved_path, _safe_resolve, is_hidden_model as _is_hidden_model, ) +def hidden_model_matchers() -> tuple[list[str], list[str], list[str]]: + """Substring needles, exact repo ids, and exact resolved paths identifying + infra models (the RAG embedder and the llama.cpp install validation probe) + that pickers hide. Served by the ``/api/hub/hidden-models`` endpoint. A + configured HF-repo embedder is published as its exact lowercased repo id + (mirroring ``utils.hidden_models.is_hidden_model``) and a local-path + embedder as its exact resolved path only: a generic basename like "model" + must not substring-hide unrelated chat models.""" + from core.rag import config as rag_config + + needles = [ + # The validation probe's repo and its exact filename. The filename carries + # .gguf so it won't hide unrelated repos like ``user/stories260K-finetune-GGUF``. + "ggml-org/models", + "stories260k.gguf", + ] + exact_ids: list[str] = [] + exact_paths: list[str] = [] + for model in ( + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + ): + # Resolve an existing local path before the repo-id regex: a local embedder + # shaped like "models/embedder" is an exact path, not a Hub repo id. + existing_path = _existing_resolved_path(model) + if existing_path: + exact_paths.append(existing_path.lower()) + elif _HF_REPO_ID_RE.match(model): + exact_ids.append(model.lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + return needles, exact_ids, exact_paths + + backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) @@ -91,6 +130,7 @@ try: _pick_best_gguf, _extract_quant_label, _is_big_endian_gguf_path, + _is_mtp_drafter, is_audio_input_type, ) from core.inference import get_inference_backend @@ -123,6 +163,7 @@ except ImportError: _pick_best_gguf, _extract_quant_label, _is_big_endian_gguf_path, + _is_mtp_drafter, is_audio_input_type, ) from core.inference import get_inference_backend @@ -183,11 +224,8 @@ def derive_model_type( def _resolve_hf_cache_dir() -> Path: """Resolve local HF cache root used by hub downloads.""" - try: - from huggingface_hub.constants import HF_HUB_CACHE - return Path(HF_HUB_CACHE) - except Exception: - return Path.home() / ".cache" / "huggingface" / "hub" + from utils.hf_cache_settings import get_hf_cache_paths + return get_hf_cache_paths().hub_cache def _is_model_directory(d: Path) -> bool: @@ -329,10 +367,12 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca return found -def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: +def _scan_hf_cache(cache_dir: Path, *, active_cache: bool = True) -> List[LocalModelInfo]: if not cache_dir.exists() or not cache_dir.is_dir(): return [] + from hub.utils import inventory_scan as hf_cache_scan + found: List[LocalModelInfo] = [] for repo_dir in cache_dir.glob("models--*"): if not repo_dir.is_dir(): @@ -348,13 +388,21 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: except OSError: updated_at = None + partial = hf_cache_scan.is_snapshot_partial("model", model_id, repo_dir) + partial = partial or hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir) + + load_id = model_id + if not active_cache: + load_id = _resolve_hf_cache_realpath(repo_dir) or str(repo_dir.resolve()) found.append( LocalModelInfo( - id = model_id, + id = load_id, model_id = model_id, display_name = model_id.split("/")[-1], - path = str(repo_dir), + path = load_id if not active_cache else str(repo_dir), source = "hf_cache", + active_cache = active_cache, + partial = partial, updated_at = updated_at, ), ) @@ -735,26 +783,34 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: legacy_hf_cache_dir, lmstudio_model_dirs, ) + from utils.hf_cache_settings import known_hf_hub_caches hf_cache_dir = _resolve_hf_cache_dir() legacy_hf = legacy_hf_cache_dir() hf_default = hf_default_cache_dir() lm_dirs = lmstudio_model_dirs() - local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) - - # Resolve once; an inaccessible aux cache must skip that scan, not 500. - hf_cache_real = _safe_resolve(hf_cache_dir) - legacy_real = _safe_resolve(legacy_hf) - default_real = _safe_resolve(hf_default) - - # Scan legacy Unsloth HF cache for backward compatibility. - if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: - local_models += _scan_hf_cache(legacy_hf) - - # Scan HF system default cache (may differ under env overrides). - if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real: - local_models += _scan_hf_cache(hf_default) + local_models = _scan_models_dir(models_root) + active_cache_real = _safe_resolve(hf_cache_dir) + active_cache_key = os.path.normcase(active_cache_real) if active_cache_real else None + seen_hf: set[str] = set() + for cache_dir in ( + hf_cache_dir, + *known_hf_hub_caches(), + legacy_hf, + hf_default, + ): + cache_real = _safe_resolve(cache_dir) + if cache_real is None: + continue + cache_key = os.path.normcase(str(cache_real)) + if cache_key in seen_hf: + continue + seen_hf.add(cache_key) + local_models += _scan_hf_cache( + cache_dir, + active_cache = cache_key == active_cache_key, + ) # Scan LM Studio directories. for lm_dir in lm_dirs: @@ -776,7 +832,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: m for m in ( _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) - + _scan_hf_cache(folder_path) + + _scan_hf_cache(folder_path, active_cache = False) + _scan_lmstudio_dir(folder_path) ) if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) @@ -797,13 +853,23 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: # even when the model is also in the HF cache. deduped: dict[str, LocalModelInfo] = {} for model in local_models: - key = f"{model.id}\x00custom" if model.source == "custom" else model.id - if key not in deduped: + semantic_id = model.model_id if model.source == "hf_cache" and model.model_id else model.id + key = f"{semantic_id}\x00custom" if model.source == "custom" else semantic_id + existing = deduped.get(key) + prefer_model = existing is None + if existing is not None and model.source == existing.source == "hf_cache": + if model.partial != existing.partial: + prefer_model = not model.partial + elif bool(model.active_cache) != bool(existing.active_cache): + prefer_model = bool(model.active_cache) + else: + prefer_model = (model.updated_at or 0) > (existing.updated_at or 0) + if prefer_model: deduped[key] = model models = sorted( deduped.values(), - key = lambda item: (item.updated_at or 0), + key = lambda item: item.updated_at or 0, reverse = True, ) return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)] @@ -1161,10 +1227,7 @@ def _build_browse_allowlist( legacy_hf_cache_dir, well_known_model_dirs, ) - from utils.paths.external_media import ( - linux_run_media_mount_roots, - windows_drive_roots, - ) + from utils.paths import external_media from storage.studio_db import list_scan_folders candidates: list[Path] = [] @@ -1181,9 +1244,12 @@ def _build_browse_allowlist( _add(Path.home()) if media_roots is None: - media_roots = linux_run_media_mount_roots() + media_roots = [ + *external_media.linux_run_media_mount_roots(), + *external_media.macos_volume_roots(), + ] if drive_roots is None: - drive_roots = windows_drive_roots() + drive_roots = external_media.windows_drive_roots() for p in media_roots: _add(p) for p in drive_roots: @@ -1461,10 +1527,7 @@ def browse_folders( then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs - from utils.paths.external_media import ( - linux_run_media_mount_roots, - windows_drive_roots, - ) + from utils.paths import external_media from storage.studio_db import ( contains_sensitive_path_component, is_denied_system_path, @@ -1473,8 +1536,11 @@ def browse_folders( # Probe removable-media and Windows drive roots once; the allowlist and # chips reuse the result so a disconnected mapped drive isn't scanned twice. - media_roots = linux_run_media_mount_roots() - drive_roots = windows_drive_roots() + media_roots = [ + *external_media.linux_run_media_mount_roots(), + *external_media.macos_volume_roots(), + ] + drive_roots = external_media.windows_drive_roots() # Build once; the sandbox check and suggestion chips share it. allowed_roots = _build_browse_allowlist(media_roots, drive_roots) @@ -1750,9 +1816,11 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op async def get_model_config( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """Get configuration for a specific model (wraps load_model_defaults).""" + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) @@ -1991,19 +2059,19 @@ async def discard_remote_code_download( # Never delete a model that is loaded for inference. try: + from hub.services.models.deletion import _loaded_id_matches_repo from routes.inference import get_llama_cpp_backend + llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and llama_backend.model_identifier: - loaded = llama_backend.model_identifier.lower() - if loaded == model_name.lower() or loaded.startswith(model_name.lower()): + if _loaded_id_matches_repo(llama_backend.model_identifier, model_name): return {"deleted": False, "reason": "loaded"} except Exception: pass try: inference_backend = get_inference_backend() if inference_backend.active_model_name: - active = inference_backend.active_model_name.lower() - if active == model_name.lower() or active.startswith(model_name.lower()): + if _loaded_id_matches_repo(inference_backend.active_model_name, model_name): return {"deleted": False, "reason": "loaded"} except Exception: pass @@ -2471,6 +2539,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get async def check_vision_model( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2478,6 +2547,7 @@ async def check_vision_model( This endpoint wraps the backend is_vision_model function. """ + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: logger.info(f"Checking if vision model: {model_name}") # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). @@ -2503,6 +2573,7 @@ async def check_vision_model( async def check_embedding_model( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2510,6 +2581,7 @@ async def check_embedding_model( This endpoint wraps the backend is_embedding_model function. """ + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) @@ -2541,13 +2613,10 @@ def _read_native_context_length(repo_id: str, is_local: bool) -> Optional[int]: if is_local: roots = [Path(repo_id)] else: - from huggingface_hub import constants as hf_constants - + from hub.utils.hf_cache_state import iter_repo_cache_dirs if not _is_valid_repo_id(repo_id): return None - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"models--{repo_id.replace('/', '--')}".lower() - roots = [e for e in cache_dir.iterdir() if e.name.lower() == target] + roots = list(iter_repo_cache_dirs("model", repo_id)) for root in roots: for f in _iter_gguf_paths(root): @@ -2573,47 +2642,32 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio Q8_0 weights). Never raises. """ try: - from utils.models.model_config import ( - _extract_quant_label, - _is_big_endian_gguf_path, - _is_mtp_drafter, - ) - if is_local: roots = [Path(repo_id)] else: - from huggingface_hub import constants as hf_constants + from hub.utils.hf_cache_state import iter_repo_cache_dirs if not _is_valid_repo_id(repo_id): return None, 0 - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"models--{repo_id.replace('/', '--')}".lower() roots = [] - for entry in cache_dir.iterdir(): - if entry.name.lower() == target: - snaps = entry / "snapshots" - if snaps.is_dir(): - roots.extend(s for s in snaps.iterdir() if s.is_dir()) + for entry in iter_repo_cache_dirs("model", repo_id): + snaps = entry / "snapshots" + if snaps.is_dir(): + roots.extend(s for s in snaps.iterdir() if s.is_dir()) - want = quant.lower().replace("-", "").replace("_", "") + want = _normalized_quant_label(quant) best_total = 0 best_first: Optional[str] = None for root in roots: matches: list[tuple[str, Path]] = [] total = 0 for f in _iter_gguf_paths(root): - if _is_mmproj_filename(f.name): - continue try: rel = f.relative_to(root).as_posix() except ValueError: rel = f.name - if _is_mtp_drafter(rel): - continue - q = _extract_quant_label(rel) - if _is_big_endian_gguf_path(rel, q): - continue - if q.lower().replace("-", "").replace("_", "") != want: + q = _main_variant_gguf_label(rel) + if q is None or _normalized_quant_label(q) != want: continue try: total += f.stat().st_size @@ -2638,7 +2692,10 @@ async def get_kv_cache_estimate( repo_id: str = Query(..., description = "HF repo ID or local path"), quant: str = Query(..., description = "Quantization label (e.g. Q4_K_M)"), n_ctx: int = Query(..., ge = 1, description = "Context length to size the KV cache for"), - cache_type_kv: Optional[str] = Query(None, description = "KV cache dtype (e.g. q8_0)"), + cache_type_kv: Optional[str] = Query( + None, + description = "KV cache dtype (e.g. q8_0, q4_0, q5_0, iq4_nl, f32)", + ), current_subject: str = Depends(get_current_subject), ): """Estimate KV cache + weight bytes for a downloaded GGUF at n_ctx. @@ -2699,6 +2756,8 @@ async def get_gguf_variants( repo_id: str = Query( ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')" ), + prefer_local_cache: bool = False, + local_path: Optional[str] = None, hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"), hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), @@ -2710,9 +2769,16 @@ async def get_gguf_variants( response = await hub_gguf_variants.get_gguf_variants_response( repo_id, + prefer_local_cache = prefer_local_cache, + local_path = local_path, hf_token = hf_token, ) - local = is_local_path(repo_id) + context_model = ( + local_path + if prefer_local_cache and local_path and is_local_path(local_path) + else repo_id + ) + local = is_local_path(context_model) return GgufVariantsResponse( repo_id = response.repo_id, @@ -2734,7 +2800,7 @@ async def get_gguf_variants( # The header walk reads tokenizer arrays on dense models (tens of # ms per uncached file); keep it off the event loop. context_length = await asyncio.to_thread( - _read_native_context_length, repo_id, is_local = local + _read_native_context_length, context_model, is_local = local ), ) except HTTPException: @@ -2752,69 +2818,17 @@ async def get_gguf_download_progress( repo_id: str = Query(..., description = "HuggingFace repo ID"), variant: str = Query("", description = "Quantization variant (e.g. UD-TQ1_0)"), expected_bytes: int = Query(0, description = "Expected total download size in bytes"), + hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - """Download progress from cached GGUF files for a specific variant. - - Tracks completed shards in snapshots and in-progress (.incomplete) - downloads in the blobs directory. - """ - try: - if not _is_valid_repo_id(repo_id): - return { - "downloaded_bytes": 0, - "expected_bytes": expected_bytes, - "progress": 0, - } - - from huggingface_hub import constants as hf_constants - - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"models--{repo_id.replace('/', '--')}".lower() - variant_lower = variant.lower().replace("-", "").replace("_", "") - downloaded_bytes = 0 - in_progress_bytes = 0 - for entry in cache_dir.iterdir(): - if entry.name.lower() == target: - # Completed .gguf files for this variant in snapshots. - # Exclude mmproj so a vision adapter can't satisfy a same-label - # main variant (e.g. mmproj-F16 vs an F16 weight). - for f in _iter_gguf_paths(entry): - if _is_mmproj_filename(f.name): - continue - rel = f.relative_to(entry).as_posix() - quant = _extract_quant_label(rel) - if _is_big_endian_gguf_path(rel, quant): - continue - rel_key = rel.lower().replace("-", "").replace("_", "") - if not variant_lower or variant_lower in rel_key: - try: - downloaded_bytes += f.stat().st_size - except OSError: - continue # broken symlink / unreadable: skip - # In-progress (.incomplete) downloads in blobs. - blobs_dir = entry / "blobs" - if blobs_dir.is_dir(): - for f in blobs_dir.iterdir(): - if f.is_file() and f.name.endswith(".incomplete"): - try: - in_progress_bytes += f.stat().st_size - except OSError: - continue - break - - total_progress_bytes = downloaded_bytes + in_progress_bytes - progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0 - # Report 1.0 only when all bytes are in completed files. - if expected_bytes > 0 and downloaded_bytes >= expected_bytes: - progress = 1.0 - return { - "downloaded_bytes": total_progress_bytes, - "expected_bytes": expected_bytes, - "progress": round(progress, 3), - } - except Exception: - return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0} + """Compatibility route backed by the shared multi-cache progress service.""" + from hub.services.models import downloads + return await downloads.get_gguf_download_progress_response( + repo_id, + variant = variant, + expected_bytes = expected_bytes, + hf_token = hf_token, + ) def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: @@ -2839,98 +2853,12 @@ def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: @router.get("/download-progress") async def get_download_progress( repo_id: str = Query(..., description = "HuggingFace repo ID"), + hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - """Return download progress for any HuggingFace model repo. - - Checks the local HF cache for completed blobs and in-progress - (.incomplete) downloads. Gets the expected total size from the HF API - on the first call, then caches it for later polls. Also returns - ``cache_path``: the realpath of the snapshot dir (or cache repo root - if no snapshot yet) so the UI can show where weights live on disk. - """ - _empty = { - "downloaded_bytes": 0, - "expected_bytes": 0, - "progress": 0, - "cache_path": None, - } - try: - if not _is_valid_repo_id(repo_id): - return _empty - - from huggingface_hub import constants as hf_constants - - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"models--{repo_id.replace('/', '--')}".lower() - completed_bytes = 0 - in_progress_bytes = 0 - cache_path: Optional[str] = None - - for entry in cache_dir.iterdir(): - if entry.name.lower() != target: - continue - cache_path = _resolve_hf_cache_realpath(entry) - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - break - for f in blobs_dir.iterdir(): - if not f.is_file(): - continue - if f.name.endswith(".incomplete"): - in_progress_bytes += f.stat().st_size - else: - completed_bytes += f.stat().st_size - break - - downloaded_bytes = completed_bytes + in_progress_bytes - if downloaded_bytes == 0: - return {**_empty, "cache_path": cache_path} - - expected_bytes = _get_repo_size_cached(repo_id) - if expected_bytes <= 0: - # Total unknown; report bytes only, no percentage. - return { - "downloaded_bytes": downloaded_bytes, - "expected_bytes": 0, - "progress": 0, - "cache_path": cache_path, - } - - # 95% threshold (blob dedup can skew completed_bytes). Do NOT - # treat "no .incomplete files" as done: HF downloads sequentially, - # so none exist between files even when far from finished. - if completed_bytes >= expected_bytes * 0.95: - progress = 1.0 - else: - progress = min(downloaded_bytes / expected_bytes, 0.99) - return { - "downloaded_bytes": downloaded_bytes, - "expected_bytes": expected_bytes, - "progress": round(progress, 3), - "cache_path": cache_path, - } - except Exception as e: - logger.warning(f"Error checking download progress for {repo_id}: {e}") - return _empty - - -_repo_size_cache: dict[str, int] = {} - - -def _get_repo_size_cached(repo_id: str) -> int: - if repo_id in _repo_size_cache: - return _repo_size_cache[repo_id] - try: - from huggingface_hub import model_info as hf_model_info - - info = hf_model_info(repo_id, token = None, files_metadata = True) - total = sum(s.size for s in info.siblings if s.size) - _repo_size_cache[repo_id] = total - return total - except Exception as e: - logger.warning(f"Failed to get repo size for {repo_id}: {e}") - return 0 + """Compatibility route backed by the shared multi-cache progress service.""" + from hub.services.models import downloads + return await downloads.get_download_progress_response(repo_id, hf_token = hf_token) def _repo_in_any_hf_cache(model_name: str) -> bool: @@ -2943,25 +2871,13 @@ def _repo_in_any_hf_cache(model_name: str) -> bool: would delete a model they did not download via the scan. Mirrors the cache set in ``_all_hf_cache_scans`` but only probes for the one repo dir (cheap, no full scan). """ - from utils.paths import ( - hf_default_cache_dir, - legacy_hf_cache_dir, - resolve_cached_repo_id_case, - ) + from utils.paths import resolve_cached_repo_id_case dirname = f"models--{resolve_cached_repo_id_case(model_name).replace('/', '--')}" dirname_lower = dirname.lower() - candidates = [] - try: - from huggingface_hub.constants import HF_HUB_CACHE - candidates.append(Path(HF_HUB_CACHE)) - except Exception: - pass - for fn in (legacy_hf_cache_dir, hf_default_cache_dir): - try: - candidates.append(fn()) - except Exception: - continue + from hub.utils.hf_cache_state import hf_cache_roots + + candidates = hf_cache_roots() # resolve_cached_repo_id_case only normalizes the ACTIVE cache, but discard deletes # case-insensitively across all caches, so detect case-insensitively too -- else a # pre-existing case-variant repo is misreported as scan-created and deleted on decline. @@ -2985,38 +2901,8 @@ def _all_hf_cache_scans(): broken symlink, OS-redirected ~/.cache) is skipped, not fatal, so the Downloaded list never blanks out and downloads never leak into Recommended. """ - from huggingface_hub import scan_cache_dir - from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir - - scans = [] - # Guard the active cache too: degrade to "no downloads" instead of raising. - try: - scans.append(scan_cache_dir()) - except Exception as exc: - logger.warning("Could not scan active HF cache: %s", exc) - - seen: set[str] = set() - try: - # Resolve the active cache dir for dedup. - from huggingface_hub.constants import HF_HUB_CACHE - seen.add(str(Path(HF_HUB_CACHE).resolve())) - except Exception: - pass - - for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): - try: - extra = extra_fn() - # is_dir()/resolve() can raise on an inaccessible path; skip it. - if not extra.is_dir(): - continue - resolved = str(extra.resolve()) - if resolved in seen: - continue - seen.add(resolved) - scans.append(scan_cache_dir(cache_dir = str(extra))) - except Exception as exc: - logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc) - return scans + from hub.utils.inventory_scan import all_hf_cache_scans + return all_hf_cache_scans() def _is_gguf_filename(name: str) -> bool: @@ -3035,6 +2921,22 @@ def _is_main_gguf_filename(name: str) -> bool: return _is_gguf_filename(name) and not _is_mmproj_filename(name) +def _main_variant_gguf_label(rel_path: str) -> Optional[str]: + name = rel_path.rsplit("/", 1)[-1] + if not _is_main_gguf_filename(name): + return None + if _is_mtp_drafter(rel_path): + return None + label = _extract_quant_label(rel_path) + if _is_big_endian_gguf_path(rel_path, label): + return None + return label + + +def _normalized_quant_label(label: str) -> str: + return label.lower().replace("-", "").replace("_", "") + + def _repo_has_mmproj(repo_info) -> bool: """True if the repo ships a GGUF vision adapter (mmproj), so it can take image inputs. Cheap: scans already-listed file names only.""" @@ -3127,7 +3029,9 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): if repo_info.repo_type != "model": continue repo_id = repo_info.repo_id - if _is_hidden_model(repo_id): + # Pass the snapshot path too so the config check also hides + # custom Whisper checkpoints, not just curated repo ids. + if _is_hidden_model(repo_id, str(repo_info.repo_path)): continue total_size = _repo_gguf_size_bytes(repo_info) if total_size == 0: @@ -3184,7 +3088,9 @@ async def list_cached_models( if repo_info.repo_type != "model": continue repo_id = repo_info.repo_id - if _is_hidden_model(repo_id): + # Pass the snapshot path too so the config check also hides + # custom Whisper checkpoints, not just curated repo ids. + if _is_hidden_model(repo_id, str(repo_info.repo_path)): continue if _repo_has_gguf_files(repo_info): continue @@ -3242,124 +3148,177 @@ async def list_cached_models( async def delete_cached_model( repo_id: str = Body(...), variant: Optional[str] = Body(None), + cache_path: Optional[str] = Body(None), + hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - """Delete a cached model repo (or a specific GGUF variant) from the HF cache. + """Compatibility route backed by the shared multi-cache deletion service.""" + from hub.services.models import deletion + return await deletion.delete_cached_model_response(repo_id, variant, hf_token, cache_path) - With *variant*, only GGUF files matching that quant label are removed - (e.g. ``UD-Q4_K_XL``); otherwise the whole repo is deleted. Refuses - if the model is currently loaded for inference. - """ + +def _resolve_cached_model_path(repo_id: str, variant: Optional[str]) -> Path: + """Absolute path of a cached repo (newest snapshot dir) or, with *variant*, + that quant's main GGUF file (first split of a sharded quant). Paths come + from the HF cache scan only, so callers can't probe arbitrary paths.""" + cache_scans = _all_hf_cache_scans() + + matching_repos = [] + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + matching_repos.append(repo_info) + if not matching_repos: + raise HTTPException(status_code = 404, detail = "Model not found in cache") + + if variant: + want = _normalized_quant_label(variant) + candidate_revisions = sorted( + (rev for repo_info in matching_repos for rev in repo_info.revisions), + key = lambda rev: getattr(rev, "last_modified", 0) or 0, + reverse = True, + ) + for rev in candidate_revisions: + snapshot = getattr(rev, "snapshot_path", None) + matches = [] + for f in rev.files: + p = Path(f.file_path) + rel = f.file_name + if snapshot: + try: + rel = p.relative_to(snapshot).as_posix() + except ValueError: + pass + label = _main_variant_gguf_label(rel) + if label is None or _normalized_quant_label(label) != want: + continue + if p.exists() or p.is_symlink(): + matches.append((rel, p)) + if matches: + # Path-sorted so a sharded quant deterministically yields its first split. + return sorted(matches, key = lambda m: m[0].lower())[0][1] + raise HTTPException( + status_code = 404, + detail = f"Variant {variant} not found in cache for {repo_id}", + ) + + def repo_size(repo_info) -> int: + gguf_size = _repo_gguf_size_bytes(repo_info) + if gguf_size > 0: + return gguf_size + return sum( + (getattr(f, "size_on_disk", None) or 0) + for rev in repo_info.revisions + for f in rev.files + ) + + def repo_last_modified(repo_info) -> float: + return max( + (getattr(rev, "last_modified", 0) or 0 for rev in repo_info.revisions), + default = 0, + ) + + target_repo = max( + matching_repos, + key = lambda repo_info: (repo_size(repo_info), repo_last_modified(repo_info)), + ) + + # Whole repo: the newest revision's snapshot dir holds the visible files. + revisions = sorted( + (rev for rev in target_repo.revisions if getattr(rev, "snapshot_path", None)), + key = lambda rev: getattr(rev, "last_modified", 0) or 0, + reverse = True, + ) + for rev in revisions: + p = Path(rev.snapshot_path) + if p.exists(): + return p + p = Path(target_repo.repo_path) + if p.exists(): + return p + raise HTTPException(status_code = 404, detail = "Cached model path not found") + + +def _wsl_reveal_in_explorer(path: Path) -> bool: + import subprocess + + from utils.paths.path_utils import _IS_WSL + + if not _IS_WSL: + return False + try: + windows_path = subprocess.run( + ["wslpath", "-w", str(path)], + capture_output = True, + text = True, + check = True, + timeout = 10, + ).stdout.strip() + if not windows_path: + return False + argument = f"/select,{windows_path}" if path.is_file() else windows_path + subprocess.Popen(["explorer.exe", argument]) + return True + except (OSError, subprocess.SubprocessError): + return False + + +def _reveal_in_file_manager(path: Path) -> None: + """Open the OS file manager with *path* selected (best effort per platform).""" + import subprocess + + target = str(path) + if sys.platform == "darwin": + cmd = ["open", "-R", target] if path.is_file() else ["open", target] + subprocess.Popen(cmd) + elif os.name == "nt": + if path.is_file(): + subprocess.Popen(["explorer", f"/select,{target}"]) + else: + os.startfile(target) # noqa: S606 - local user's own file manager + elif not _wsl_reveal_in_explorer(path): + # No cross-desktop "select file" standard on Linux; open the directory. + directory = target if path.is_dir() else str(path.parent) + subprocess.Popen(["xdg-open", directory]) + + +class CachedModelPathResponse(BaseModel): + path: str + is_dir: bool + + +@router.get("/cached-model-path", response_model = CachedModelPathResponse) +async def get_cached_model_path( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + variant: str = Query("", description = "Quantization variant (empty for whole repo)"), + current_subject: str = Depends(get_current_subject), +): + """Absolute on-disk path of a cached repo or one of its GGUF variants.""" if not _is_valid_repo_id(repo_id): raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant.strip() or None) + return {"path": str(path), "is_dir": path.is_dir()} - # Refuse if the model is currently loaded. + +@router.post("/reveal-cached-model") +async def reveal_cached_model( + repo_id: str = Body(...), + variant: Optional[str] = Body(None), + current_subject: str = Depends(get_current_subject), +): + """Reveal a cached repo (or one GGUF variant's file) in the OS file manager.""" + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + variant = (variant or "").strip() or None + path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant) try: - from routes.inference import get_llama_cpp_backend - llama_backend = get_llama_cpp_backend() - if llama_backend.is_loaded and llama_backend.model_identifier: - loaded_id = llama_backend.model_identifier.lower() - if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()): - raise HTTPException( - status_code = 400, - detail = "Unload the model before deleting", - ) - except HTTPException: - raise - except Exception: - pass - - try: - inference_backend = get_inference_backend() - if inference_backend.active_model_name: - active = inference_backend.active_model_name.lower() - if active == repo_id.lower() or active.startswith(repo_id.lower()): - raise HTTPException( - status_code = 400, - detail = "Unload the model before deleting", - ) - except HTTPException: - raise - except Exception: - pass - - try: - cache_scans = _all_hf_cache_scans() - - target_repo = None - for hf_cache in cache_scans: - for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": - continue - if repo_info.repo_id.lower() == repo_id.lower(): - target_repo = repo_info - break - if target_repo is not None: - break - - if target_repo is None: - raise HTTPException(status_code = 404, detail = "Model not found in cache") - - # โ”€โ”€ Per-variant GGUF deletion โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - if variant: - deleted_bytes = 0 - deleted_count = 0 - for rev in target_repo.revisions: - for f in rev.files: - if not _is_gguf_filename(f.file_name): - continue - quant = _extract_quant_label(f.file_name) - if quant.lower() != variant.lower(): - continue - # Delete the blob (data) and the snapshot symlink. - try: - blob = Path(f.blob_path) - snap = Path(f.file_path) - size = blob.stat().st_size if blob.exists() else 0 - if snap.exists() or snap.is_symlink(): - snap.unlink() - if blob.exists(): - blob.unlink() - deleted_bytes += size - deleted_count += 1 - except Exception as e: - logger.warning(f"Failed to delete {f.file_name}: {e}") - - if deleted_count == 0: - raise HTTPException( - status_code = 404, - detail = f"Variant {variant} not found in cache for {repo_id}", - ) - - freed_mb = deleted_bytes / (1024 * 1024) - logger.info( - f"Deleted {deleted_count} file(s) for {repo_id} variant {variant}: " - f"{freed_mb:.1f} MB freed" - ) - return {"status": "deleted", "repo_id": repo_id, "variant": variant} - - # โ”€โ”€ Full repo deletion โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - revision_hashes = [rev.commit_hash for rev in target_repo.revisions] - if not revision_hashes: - raise HTTPException(status_code = 404, detail = "No revisions found for model") - - delete_strategy = hf_cache.delete_revisions(*revision_hashes) - logger.info( - f"Deleting cached model {repo_id}: " - f"{delete_strategy.expected_freed_size_str} will be freed" - ) - delete_strategy.execute() - - return {"status": "deleted", "repo_id": repo_id} - - except HTTPException: - raise + await asyncio.to_thread(_reveal_in_file_manager, path) except Exception as e: - logger.error(f"Error deleting cached model {repo_id}: {e}", exc_info = True) - raise HTTPException( - status_code = 500, - detail = "Failed to delete cached model", - ) + logger.error(f"Failed to reveal {path}: {e}") + raise HTTPException(status_code = 500, detail = "Failed to open file manager") + return {"status": "ok", "path": str(path)} @router.get("/checkpoints", response_model = CheckpointListResponse) diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 5a55c9b0bb..4e7e53f2f0 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -47,6 +47,20 @@ logger = structlog.get_logger(__name__) router = APIRouter() +def _provider_response(row: dict) -> ProviderResponse: + return ProviderResponse( + id = row["id"], + provider_type = row["provider_type"], + display_name = row["display_name"], + base_url = row["base_url"], + is_enabled = bool(row["is_enabled"]), + models = row.get("models") or [], + available_models = row.get("available_models") or [], + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + + # โ”€โ”€ Public key for API key encryption โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ @@ -89,18 +103,7 @@ async def get_pricing_snapshot(current_subject: str = Depends(get_current_subjec async def list_provider_configs(current_subject: str = Depends(get_current_subject)): """List all saved provider configurations.""" rows = providers_db.list_providers() - return [ - ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) - for row in rows - ] + return [_provider_response(row) for row in rows] @router.post("/", response_model = ProviderResponse, status_code = 201) @@ -124,18 +127,12 @@ async def create_provider_config( provider_type = payload.provider_type, display_name = payload.display_name, base_url = base_url, + models = payload.models, + available_models = payload.available_models, ) row = providers_db.get_provider(provider_id) - return ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) + return _provider_response(row) @router.put("/{provider_id}", response_model = ProviderResponse) @@ -154,20 +151,14 @@ async def update_provider_config( display_name = payload.display_name, base_url = payload.base_url, is_enabled = payload.is_enabled, + models = payload.models, + available_models = payload.available_models, ) if not updated: raise HTTPException(status_code = 400, detail = "No fields to update") row = providers_db.get_provider(provider_id) - return ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) + return _provider_response(row) @router.delete("/{provider_id}", status_code = 204) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 17e64df918..fef18a9145 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -60,6 +60,7 @@ from utils.embedding_model_settings import ( set_rag_embedding_model, validate_embedding_model, ) +from utils.hf_cache_settings import cache_status, get_hf_cache_paths, set_hf_cache_home router = APIRouter() @@ -89,6 +90,23 @@ class HelperPrecacheResponse(BaseModel): disabled_by_env: bool +class HuggingFaceCachePayload(BaseModel): + cache_home: Optional[str] = Field(default = None, max_length = 4096) + + +class HuggingFaceCacheResponse(BaseModel): + cache_home: str + hub_cache: str + xet_cache: str + source: Literal["default", "studio", "environment"] + editable: bool + is_custom: bool + available: bool + writable: bool + free_bytes: Optional[int] = None + environment_variable: Optional[str] = None + + class OpenAIAutoSwitchPayload(BaseModel): enabled: bool # None leaves the stored value untouched (partial updates can't clobber it). @@ -135,6 +153,30 @@ def _helper_precache_response(enabled: bool | None = None) -> HelperPrecacheResp ) +def _hugging_face_cache_response() -> HuggingFaceCacheResponse: + return HuggingFaceCacheResponse(**cache_status(get_hf_cache_paths())) + + +@router.get("/hugging-face-cache", response_model = HuggingFaceCacheResponse) +def get_hugging_face_cache( + current_subject: str = Depends(get_current_subject), +) -> HuggingFaceCacheResponse: + return _hugging_face_cache_response() + + +@router.put("/hugging-face-cache", response_model = HuggingFaceCacheResponse) +def update_hugging_face_cache( + payload: HuggingFaceCachePayload, current_subject: str = Depends(get_current_subject) +) -> HuggingFaceCacheResponse: + try: + set_hf_cache_home(payload.cache_home) + except RuntimeError as exc: + raise HTTPException(status_code = 409, detail = str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc + return _hugging_face_cache_response() + + @router.get("/upload-limit", response_model = UploadLimitResponse) def get_upload_limit(current_subject: str = Depends(get_current_subject)) -> UploadLimitResponse: return _upload_limit_response(get_upload_limit_mb()) @@ -416,6 +458,11 @@ def update_embedding_model( log = logger, ) from exc hf_token = (payload.hf_token or "").strip() or None + from utils.utils import hf_env_offline + + # Offline, both the Hub malware scan and the is-embedding check are unreachable and degrade + # to the local cache below; capture the state once. + local_only_load = hf_env_offline() # The env/default model needs no verification; saving it is a no-op override. # A local GGUF on the llama-server backend is accepted as-is: it is exactly # what the backend loads, and HF metadata cannot verify a local path. @@ -439,26 +486,41 @@ def update_embedding_model( # Fall back to the loader's own token so a gated/private repo is actually scanned # (a token-less scan fails open for exactly the repo that would still load). scan_token = hf_token or _ambient_hf_token() - # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under - # one blocks instead of passing as an unreferenced nested shard. - load_subdirs = tuple( - dict.fromkeys( - ( - *security_load_subdirs(model, scan_token), - *_st_module_subdirs(model, scan_token), + # Offline: subdir probes would hit the network and hang; the offline gate walks the + # whole cached snapshot, so no load-subdir hints are needed. + if local_only_load: + load_subdirs = () + else: + # Include ST module dirs (0_Transformer/) so a flagged pickle directly under one + # blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + ( + *security_load_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token), + ) ) ) - ) - if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + if evaluate_file_security( + model, + hf_token = scan_token, + load_subdirs = load_subdirs, + local_only_load = local_only_load, + ).blocked: # 403, not 409: the client routes every 409 into the forceable "save anyway" # flow, but this block is a hard, non-forceable security refusal. - raise HTTPException( - status_code = 403, + if local_only_load: + detail = ( + f"{model!r} has cached pickle weights that cannot be security-scanned " + "offline and no safetensors alternative, so it cannot be used as the " + "embedding model. Re-download it with safetensors weights while online." + ) + else: detail = ( f"{model!r} is flagged as unsafe by Hugging Face's security scan and " "cannot be used as the embedding model." - ), - ) + ) + raise HTTPException(status_code = 403, detail = detail) if model != default_embedding_model() and not payload.force and not is_local_gguf: from core.rag import config as rag_config @@ -468,15 +530,28 @@ def update_embedding_model( # which would wrongly 409 a valid online GGUF embedder. gguf_named = _llama_backend_active() and rag_config._names_gguf(model) if not gguf_named and not is_embedding_model(model, hf_token = hf_token): - raise HTTPException( - status_code = 409, - detail = ( - f"Could not verify {model!r} as an embedding model on " - "Hugging Face (it may be the wrong model type, gated, or " - "you may be offline)." - ), - ) - gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) + # Offline, is_embedding_model can only confirm the ST layout (modules.json); a + # transformers-native embedder (e.g. gte-modernbert) is unverifiable without Hub + # metadata. If already cached and loadable, accept it rather than raising a 409 that + # online would not (ST can load any cached encoder). Uncached -> 409. + from utils.utils import hf_cache_snapshot_is_loadable + + # Require a genuinely loadable cache (config + weights), not just a resolved refs/main, + # so a metadata-only partial cache still gets the forceable 409. + offline_cached = local_only_load and hf_cache_snapshot_is_loadable(model) + if not offline_cached: + raise HTTPException( + status_code = 409, + detail = ( + f"Could not verify {model!r} as an embedding model on " + "Hugging Face (it may be the wrong model type, gated, or " + "you may be offline)." + ), + ) + # The Hub GGUF probe (list_repo_files) can hang offline; skip it. Local check stays. + gguf_error = _local_gguf_backend_error(model) + if gguf_error is None and not local_only_load: + gguf_error = _hf_gguf_backend_error(model, hf_token) if gguf_error: raise HTTPException(status_code = 409, detail = gguf_error) set_rag_embedding_model(model) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 53b1c4d991..8be4283415 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -109,7 +109,9 @@ async def get_hardware_utilization(current_subject: str = Depends(get_current_su @router.get("/hardware/visible") async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)): from utils.hardware import get_visible_gpu_utilization - return get_visible_gpu_utilization() + + # Off the event loop: the ROCm fallbacks shell out (Windows perf counters, sysfs) and the System view polls this route. + return await asyncio.to_thread(get_visible_gpu_utilization) @router.post("/start") @@ -196,6 +198,7 @@ async def start_training( request.local_eval_datasets, "Local eval dataset" ) resume_output_dir: Optional[str] = None + resume_run: Optional[dict] = None if request.resume_from_checkpoint: try: resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint) @@ -208,7 +211,7 @@ async def start_training( if not resume_run or not can_resume_run(resume_run): raise HTTPException( status_code = 400, - detail = "Resume checkpoint must belong to a stopped run with saved trainer state.", + detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.", ) resume_checkpoint = get_resume_checkpoint_path(resume_output_dir) if not resume_checkpoint: @@ -329,6 +332,7 @@ async def start_training( else "unsloth", "use_rslora": request.use_rslora, "use_loftq": request.use_loftq, + "use_dora": request.use_dora, "train_on_completions": request.train_on_completions, "finetune_vision_layers": request.finetune_vision_layers, "finetune_language_layers": request.finetune_language_layers, @@ -412,53 +416,39 @@ async def start_training( try: from routes.training_vram import ( can_keep_chat_during_training, - free_chat_models_for_training, - summarize_resident_chat, + coordinate_models_for_training, ) - resident = summarize_resident_chat() - if not resident["any"]: - return - if resident.get("loading"): - # In-flight load can't be sized -> free rather than risk OOM. - freed = free_chat_models_for_training(reason = "chat model still loading") - logger.info("Freed in-flight chat load for training: %s", freed) - return - keep, info = can_keep_chat_during_training( - model_name = training_kwargs["model_name"], - hf_token = training_kwargs["hf_token"], - training_type = training_kwargs["training_type"], - load_in_4bit = training_kwargs["load_in_4bit"], - batch_size = training_kwargs["batch_size"], - max_seq_length = training_kwargs["max_seq_length"], - lora_rank = training_kwargs["lora_r"], - target_modules = training_kwargs["target_modules"], - gradient_checkpointing = training_kwargs["gradient_checkpointing"], - optimizer = training_kwargs["optim"], - gpu_ids = training_kwargs["gpu_ids"], - ) - if keep: - logger.info( - "Keeping chat model(s) loaded during training " - "(free ~%s GB, needs ~%s GB): %s", - info.get("usable_gb"), - info.get("required_gb"), - resident, + def _can_keep_resident_models(): + return can_keep_chat_during_training( + model_name = training_kwargs["model_name"], + hf_token = training_kwargs["hf_token"], + training_type = training_kwargs["training_type"], + load_in_4bit = training_kwargs["load_in_4bit"], + batch_size = training_kwargs["batch_size"], + max_seq_length = training_kwargs["max_seq_length"], + lora_rank = training_kwargs["lora_r"], + target_modules = training_kwargs["target_modules"], + gradient_checkpointing = training_kwargs["gradient_checkpointing"], + optimizer = training_kwargs["optim"], + gpu_ids = training_kwargs["gpu_ids"], ) - else: - freed = free_chat_models_for_training( - reason = "insufficient VRAM to run training alongside chat", - ) - logger.info("Freed chat model(s) for training: %s", freed) + + freed = coordinate_models_for_training(_can_keep_resident_models) + if freed: + logger.info("Freed models for training: %s", freed) except Exception as e: - logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e) + logger.warning("Inference/training memory coordination failed; proceeding: %s", e) # The hook runs only once start guards pass -> VRAM freed iff training starts. from utils.transformers_version import SidecarSwapInProgress try: success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + job_id = job_id, + before_spawn = _free_vram_for_training, + resume_source_run_id = resume_run["id"] if resume_run else None, + **training_kwargs, ) except SidecarSwapInProgress as exc: # Expected loss of the race against a sidecar install: a retryable @@ -521,7 +511,10 @@ async def stop_training( status = "idle", message = "No training job is currently running" ) - backend.stop_training(save = body.save) + if not backend.stop_training(save = body.save): + return TrainingStopResponse( + status = "idle", message = "No training job is currently running" + ) return TrainingStopResponse( status = "stopped", @@ -637,9 +630,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject "loss": getattr(progress, "loss", None), "learning_rate": getattr(progress, "learning_rate", None), } - output_dir = getattr(backend, "_output_dir", None) - if output_dir: - details["output_dir"] = output_dir + # Always present: an explicit null tells the client to drop a cached + # path (stop without save clears the run's output_dir). + details["output_dir"] = getattr(backend, "_output_dir", None) or None # Metric history for chart recovery after SSE reconnection. metric_history = None diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index fd96fe2175..8ddda11b1e 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -1,15 +1,13 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""VRAM coordination between chat/inference and training. +"""Memory coordination between inference and training. -Decides, from live free VRAM, whether a resident chat model can stay loaded -during training or must be unloaded, and unloads it across all backends -(HF/MLX orchestrator + llama.cpp GGUF server). In the route layer because the -GGUF accessor lives in routes/inference.py; backends are imported lazily. +Uses live free VRAM to keep resident chat and STT models when they fit. STT is +evicted before chat when training needs memory. """ -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from loggers import get_logger @@ -77,6 +75,37 @@ def summarize_resident_chat() -> Dict[str, Any]: } +def summarize_resident_stt() -> Dict[str, Any]: + """Report the resident dictation model (either engine). Never raises.""" + try: + from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar + from core.inference.stt_sidecar import get_stt_sidecar + + sidecar = get_stt_sidecar() + model = sidecar.loaded_model + device = sidecar.device + loading = sidecar.is_loading() + # whisper.cpp holds GPU memory via its subprocess, and both engines can be + # live at once (engine switch or direct /audio/stt/load). Always fold the + # GGUF sidecar in: a resident Transformers model must not mask a GGUF + # server still binding its backend, or admission lets training launch into + # that startup and OOM. + ggml = get_ggml_stt_sidecar() + if not model: + model = ggml.loaded_model + device = device or ggml.device + loading = loading or ggml.is_loading() + return { + "model": model, + "device": device, + "loading": loading, + "any": bool(model or loading), + } + except Exception as e: + logger.warning("Could not inspect STT sidecar: %s", e) + return {"model": None, "device": None, "loading": False, "any": False} + + def can_keep_chat_during_training( *, model_name: str, @@ -106,8 +135,8 @@ def can_keep_chat_during_training( resolve_requested_gpu_ids, ) - if get_device() != DeviceType.CUDA: - return False, {"mode": "non_cuda", "reason": "non_cuda"} + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + return False, {"mode": "non_accelerator", "reason": "non_accelerator"} # 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 @@ -196,6 +225,7 @@ def can_load_chat_during_training( max_seq_length: int, requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, + is_vulkan: bool = False, required_override_gb: Optional[float] = None, single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: @@ -204,11 +234,15 @@ def can_load_chat_during_training( 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. ``single_device_gpu`` is the - exact physical device token selected by a single-device runner. - `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.""" + required_override_gb over the visible pool. A Vulkan GGUF selection picks by ggml + Vulkan ordinal (separate index space from CUDA ids), so its requested_gpu_ids is + NOT resolved against the CUDA set (which would raise -> invalid_gpu_ids -> bypass + the OOM check); conservatively size an N-device request against the least-free + N visible GPUs instead. + ``single_device_gpu`` is the exact physical device token selected by a + single-device runner. `load_in_4bit` must be effective (LoRA can flip 4-bit + -> 16-bit). CPU/MLX allows the load; default-deny on any CUDA/XPU case it + can't size, so a load never OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -219,8 +253,8 @@ def can_load_chat_during_training( resolve_requested_gpu_ids, ) - if get_device() != DeviceType.CUDA: - return True, {"mode": "non_cuda", "reason": "non_cuda"} + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + return True, {"mode": "non_accelerator", "reason": "non_accelerator"} est_kwargs = dict( hf_token = hf_token or None, @@ -229,6 +263,11 @@ def can_load_chat_during_training( max_seq_length = max_seq_length or 2048, ) + # A Vulkan GGUF selection uses ggml Vulkan ordinals, not CUDA physical ids; + # size it against the full visible pool (GGUF self-placement) rather than + # resolving ordinals against the CUDA parent-visible set. + vulkan_gguf = is_gguf and is_vulkan + # 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) @@ -254,7 +293,9 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. - if single_device_gpu is not None: + if requested_gpu_ids and vulkan_gguf: + mode = "gguf_vulkan" + elif single_device_gpu is not None: mode = "single_device" elif is_gguf: mode = "gguf" @@ -267,7 +308,17 @@ def can_load_chat_during_training( return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if single_device_gpu is not None: + if requested_gpu_ids and vulkan_gguf: + # Vulkan ordinals cannot be mapped to CUDA physical indices. Budget + # the least-free N visible cards for an N-device request. If that + # conservative subset fits, any physical mapping of the ordinals + # fits, without collapsing a multi-GPU request to one card. + visible_free = list(free_by_index.values()) + if not visible_free: + return False, {"mode": "gguf_vulkan", "reason": "no_visible_gpus"} + n_pins = min(len(requested_gpu_ids), len(visible_free)) + free_vals = sorted(visible_free)[:n_pins] + elif single_device_gpu is not None: token = str(single_device_gpu).strip() if not token: # Empty token = a CPU-only single-device runner (e.g. a CPU @@ -295,7 +346,8 @@ def can_load_chat_during_training( return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] else: - # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. + # GGUF self-placement / auto Vulkan (no requested ids): llama.cpp picks + # the GPU(s), so any visible GPU is a candidate -> size the whole pool. free_vals = list(free_by_index.values()) if not free_vals: @@ -366,3 +418,110 @@ def free_chat_models_for_training(reason: str) -> List[str]: logger.warning("Could not unload GGUF chat model: %s", e) return freed + + +def free_stt_model_for_training(reason: str) -> List[str]: + """Unload the dictation model(s) before training. Never raises. + + The Transformers and GGUF sidecars are freed under independent exception + boundaries so a failure unloading one backend never skips freeing the other + (both can hold accelerator memory at once after an engine switch). + """ + freed: List[str] = [] + try: + from core.inference.stt_sidecar import get_stt_sidecar + sidecar = get_stt_sidecar() + if sidecar.is_loading() and sidecar.cancel_pending_load(): + logger.info("Cancelling STT model load for training (%s)", reason) + # The loader may still be in from_pretrained()/.to(device) holding + # VRAM; wait for it to observe the cancel and release first. + sidecar.wait_for_load_to_settle() + # A load that finished before seeing the cancel leaves a resident + # model; unload it so training gets the memory back. + if sidecar.loaded_model: + sidecar.unload() + freed.append("stt:loading") + else: + model = sidecar.loaded_model + if model: + logger.info("Unloading STT model '%s' for training (%s)", model, reason) + sidecar.unload() + freed.append(f"stt:{model}") + except Exception as e: + logger.warning("Could not unload Transformers STT model: %s", e) + + # Check the GGUF sidecar even after a cancelled/failed Transformers unload; + # both engines can hold memory at once (engine switch or direct load). + try: + from core.inference.stt_ggml_sidecar import get_ggml_stt_sidecar + ggml = get_ggml_stt_sidecar() + if ggml.is_loading() and ggml.cancel_pending_load(): + logger.info("Cancelling GGUF STT model load for training (%s)", reason) + # whisper-server may still be binding its backend; wait for the + # cancelled startup to be killed and reaped before training claims + # the memory (loaded_model stays unset until it is ready). + ggml.wait_for_load_to_settle() + if ggml.loaded_model: + ggml.unload() + freed.append("stt:gguf-loading") + else: + ggml_model = ggml.loaded_model + if ggml_model: + logger.info("Unloading GGUF STT model '%s' for training (%s)", ggml_model, reason) + ggml.unload() + freed.append(f"stt:{ggml_model}") + except Exception as e: + logger.warning("Could not unload GGUF STT model: %s", e) + + return freed + + +def coordinate_models_for_training( + can_keep: Callable[[], Tuple[bool, Dict[str, Any]]], +) -> List[str]: + """Keep resident models when they fit, evicting STT before chat.""" + resident_chat = summarize_resident_chat() + resident_stt = summarize_resident_stt() + if not resident_chat["any"] and not resident_stt["any"]: + return [] + + if resident_chat.get("loading"): + freed = free_stt_model_for_training(reason = "chat model still loading") + freed += free_chat_models_for_training(reason = "chat model still loading") + return freed + + freed: List[str] = [] + if resident_stt.get("loading"): + released_stt = free_stt_model_for_training(reason = "STT model still loading") + freed += released_stt + resident_stt = ( + {"model": None, "device": None, "loading": False, "any": False} + if released_stt + else summarize_resident_stt() + ) + if not resident_chat["any"] and not resident_stt["any"]: + return freed + + keep, info = can_keep() + if keep: + logger.info( + "Keeping resident models loaded during training (free ~%s GB, needs ~%s GB): %s", + info.get("usable_gb"), + info.get("required_gb"), + {"chat": resident_chat, "stt": resident_stt}, + ) + return freed + + if resident_stt["any"]: + freed += free_stt_model_for_training(reason = "insufficient training memory") + if not resident_chat["any"]: + return freed + keep, _info = can_keep() + if keep: + logger.info("Keeping chat model loaded after freeing STT: %s", resident_chat) + return freed + + freed += free_chat_models_for_training( + reason = "insufficient VRAM to run training alongside chat", + ) + return freed diff --git a/studio/backend/routes/whisper.py b/studio/backend/routes/whisper.py new file mode 100644 index 0000000000..08a8f269ec --- /dev/null +++ b/studio/backend/routes/whisper.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""whisper.cpp prebuilt status endpoint. + +GET /api/whisper/update-status -> is a newer prebuilt available + job state + +Detection reuses utils.whisper_cpp_freshness and fails open so the UI never +blocks on a missing marker / offline GitHub. There is no whisper-only update +trigger: whisper updates piggyback on the single main update item +(POST /api/llama/update chains a whisper phase when whisper is behind). +""" + +from __future__ import annotations + +import asyncio +from typing import Optional + +from fastapi import APIRouter, Depends, Query +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject +from utils.whisper_cpp_update import get_update_status + +router = APIRouter() + + +class WhisperUpdateJob(BaseModel): + state: str = Field("idle", description = "idle | running | success | error") + message: str = "" + from_tag: Optional[str] = None + to_tag: Optional[str] = None + reload_required: Optional[bool] = None + error: Optional[str] = None + progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.") + started_at: Optional[str] = None + finished_at: Optional[str] = None + + +class WhisperUpdateStatusResponse(BaseModel): + supported: bool = Field( + False, + description = "True when the install came from an Unsloth prebuilt (has a marker).", + ) + update_available: bool = Field( + False, description = "True when the latest release is genuinely newer than the install." + ) + stale: bool = Field( + False, description = "Update available AND install older than the staleness threshold." + ) + installed_tag: Optional[str] = None + latest_tag: Optional[str] = None + published_repo: Optional[str] = None + installed_at_utc: Optional[str] = None + age_days: Optional[int] = None + source_build: bool = Field( + False, description = "True when there is no marker (source build) but a prebuilt is offered." + ) + update_size_bytes: Optional[int] = Field( + None, description = "Download size of the prebuilt an update would fetch, in bytes." + ) + job: WhisperUpdateJob = Field(default_factory = WhisperUpdateJob) + + +@router.get("/update-status", response_model = WhisperUpdateStatusResponse) +async def whisper_update_status( + force_refresh: bool = Query( + False, description = "Bypass the 24h release cache for an explicit check." + ), + current_subject: str = Depends(get_current_subject), +) -> WhisperUpdateStatusResponse: + # Off the event loop: detection may probe the host and read GitHub. + status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh) + return WhisperUpdateStatusResponse(**status) diff --git a/studio/backend/run.py b/studio/backend/run.py index 398943cc2c..d9569c46f6 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1244,6 +1244,13 @@ def _apply_supplied_password(password_value: "Optional[str]") -> None: flush = True, ) sys.exit(1) + if any(ch.isspace() for ch in supplied): + print( + "Error: password cannot contain spaces; not starting.", + file = sys.stderr, + flush = True, + ) + sys.exit(1) if _is_current_password(supplied): print( "Error: the new password must differ from the current bootstrap " diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py index 07165cbe70..e6f40c5030 100644 --- a/studio/backend/storage/providers_db.py +++ b/studio/backend/storage/providers_db.py @@ -6,8 +6,12 @@ Same pattern as studio_db.py (module-level functions, raw sqlite3, WAL, per-function connections). API keys are NOT stored here: they live only in the browser (localStorage) and are sent encrypted per-request. + +Enabled model selections and discovered catalog IDs are stored server-side so +remote Studio clients see the same connection state (#7281). """ +import json import logging import sqlite3 import threading @@ -22,6 +26,33 @@ _schema_lock = threading.Lock() _schema_ready = False +def _encode_models_json(models: Optional[list[str]]) -> str: + if not models: + return "[]" + return json.dumps([str(model).strip() for model in models if str(model).strip()]) + + +def _decode_models_json(raw: Optional[str]) -> list[str]: + if not raw: + return [] + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + return [] + if not isinstance(parsed, list): + return [] + return [str(model).strip() for model in parsed if str(model).strip()] + + +def _row_models(row: sqlite3.Row) -> tuple[list[str], list[str]]: + return ( + _decode_models_json(row["models_json"] if "models_json" in row.keys() else None), + _decode_models_json( + row["available_models_json"] if "available_models_json" in row.keys() else None + ), + ) + + def _ensure_schema(conn: sqlite3.Connection) -> None: """Create the llm_providers table if absent. Called once per process.""" conn.execute("PRAGMA journal_mode=WAL") @@ -38,6 +69,13 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(llm_providers)").fetchall()} + if "models_json" not in existing_cols: + conn.execute("ALTER TABLE llm_providers ADD COLUMN models_json TEXT NOT NULL DEFAULT '[]'") + if "available_models_json" not in existing_cols: + conn.execute( + "ALTER TABLE llm_providers ADD COLUMN available_models_json TEXT NOT NULL DEFAULT '[]'" + ) def get_connection() -> sqlite3.Connection: @@ -59,17 +97,37 @@ def get_connection() -> sqlite3.Connection: return conn -def create_provider(id: str, provider_type: str, display_name: str, base_url: str) -> None: +def create_provider( + id: str, + provider_type: str, + display_name: str, + base_url: str, + models: Optional[list[str]] = None, + available_models: Optional[list[str]] = None, +) -> None: """Insert a new provider configuration.""" now = datetime.now(timezone.utc).isoformat() conn = get_connection() try: conn.execute( """ - INSERT INTO llm_providers (id, provider_type, display_name, base_url, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO llm_providers ( + id, provider_type, display_name, base_url, + models_json, available_models_json, + created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, provider_type, display_name, base_url, now, now), + ( + id, + provider_type, + display_name, + base_url, + _encode_models_json(models), + _encode_models_json(available_models), + now, + now, + ), ) conn.commit() finally: @@ -81,6 +139,8 @@ def update_provider( display_name: Optional[str] = None, base_url: Optional[str] = None, is_enabled: Optional[bool] = None, + models: Optional[list[str]] = None, + available_models: Optional[list[str]] = None, ) -> bool: """Update fields on an existing provider. Returns True if a row was updated.""" updates = [] @@ -94,6 +154,12 @@ def update_provider( if is_enabled is not None: updates.append("is_enabled = ?") params.append(1 if is_enabled else 0) + if models is not None: + updates.append("models_json = ?") + params.append(_encode_models_json(models)) + if available_models is not None: + updates.append("available_models_json = ?") + params.append(_encode_models_json(available_models)) if not updates: return False updates.append("updated_at = ?") @@ -128,7 +194,13 @@ def get_provider(id: str) -> Optional[dict]: conn = get_connection() try: row = conn.execute("SELECT * FROM llm_providers WHERE id = ?", (id,)).fetchone() - return dict(row) if row else None + if not row: + return None + data = dict(row) + models, available_models = _row_models(row) + data["models"] = models + data["available_models"] = available_models + return data finally: conn.close() @@ -138,6 +210,13 @@ def list_providers() -> list[dict]: conn = get_connection() try: rows = conn.execute("SELECT * FROM llm_providers ORDER BY created_at").fetchall() - return [dict(row) for row in rows] + providers: list[dict] = [] + for row in rows: + data = dict(row) + models, available_models = _row_models(row) + data["models"] = models + data["available_models"] = available_models + providers.append(data) + return providers finally: conn.close() diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index d889894d04..6972e7b7ff 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -192,13 +192,18 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error_message TEXT, duration_seconds REAL, loss_sparkline TEXT, - display_name TEXT + display_name TEXT, + resume_blocked INTEGER NOT NULL DEFAULT 0 ) """ ) existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()} if "display_name" not in existing_cols: conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT") + if "resume_blocked" not in existing_cols: + conn.execute( + "ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS training_metrics ( @@ -734,16 +739,43 @@ def create_run( config_json: str, started_at: str, total_steps: Optional[int], + *, + output_dir: Optional[str] = None, + cancel_requested: bool = False, + resumed_from_run_id: Optional[str] = None, ) -> None: conn = get_connection() try: conn.execute( """ - INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO training_runs ( + id, model_name, dataset_name, config_json, started_at, total_steps, + output_dir, resume_blocked + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, model_name, dataset_name, config_json, started_at, total_steps), + ( + id, + model_name, + dataset_name, + config_json, + started_at, + total_steps, + None if cancel_requested else output_dir, + int(cancel_requested), + ), ) + if resumed_from_run_id: + claimed = conn.execute( + """ + UPDATE training_runs SET resume_blocked = 1 + WHERE id = ? AND status IN ('stopped', 'error') + AND output_dir = ? AND resume_blocked = 0 + """, + (resumed_from_run_id, output_dir), + ) + if claimed.rowcount != 1: + raise RuntimeError("Resume source is no longer available") conn.commit() finally: conn.close() @@ -786,6 +818,8 @@ def finish_run( loss_sparkline: Optional[str] = None, output_dir: Optional[str] = None, error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: conn = get_connection() try: @@ -793,9 +827,16 @@ def finish_run( """ UPDATE training_runs SET status = ?, ended_at = ?, final_step = ?, final_loss = ?, - duration_seconds = ?, loss_sparkline = ?, output_dir = ?, - error_message = ? - WHERE id = ? + duration_seconds = ?, loss_sparkline = ?, + output_dir = CASE + WHEN resume_blocked = 1 OR ? = 1 THEN NULL + WHEN ? IS NOT NULL THEN ? + WHEN ? IN ('error', 'stopped') THEN output_dir + ELSE NULL + END, + error_message = ?, + resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END + WHERE id = ? AND status = 'running' """, ( status, @@ -804,8 +845,13 @@ def finish_run( final_loss, duration_seconds, loss_sparkline, + int(clear_output_dir), output_dir, + output_dir, + status, error_message, + int(clear_output_dir), + int(resume_blocked), id, ), ) @@ -865,6 +911,38 @@ def update_run_display_name(id: str, display_name: Optional[str]) -> None: conn.close() +def update_run_output_dir(id: str, output_dir: Optional[str]) -> None: + conn = get_connection() + try: + conn.execute( + """ + UPDATE training_runs SET output_dir = ? + WHERE id = ? AND status = 'running' AND resume_blocked = 0 + """, + (output_dir, id), + ) + conn.commit() + finally: + conn.close() + + +def mark_run_cancel_requested(id: str) -> bool: + """Clear resume/export state only while the exact run is still active.""" + conn = get_connection() + try: + cursor = conn.execute( + """ + UPDATE training_runs SET output_dir = NULL, resume_blocked = 1 + WHERE id = ? AND status = 'running' + """, + (id,), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + def list_runs(limit: int = 50, offset: int = 0) -> dict: conn = get_connection() try: @@ -874,15 +952,15 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at, r.ended_at, r.total_steps, r.final_step, r.final_loss, r.output_dir, r.duration_seconds, r.error_message, - r.loss_sparkline, r.display_name, r.config_json, + r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -917,13 +995,13 @@ def get_run(id: str) -> Optional[dict]: """ SELECT r.*, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -958,12 +1036,12 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: 0 AS resumed_later FROM training_runs r WHERE r.output_dir = ? - AND r.status = 'stopped' + AND r.status IN ('stopped', 'error') AND NOT EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) ORDER BY r.started_at DESC @@ -1066,8 +1144,12 @@ def cleanup_orphaned_runs() -> None: conn.execute( """ UPDATE training_runs - SET status = 'error', - error_message = 'Server restarted during training', + SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END, + error_message = CASE + WHEN resume_blocked = 1 THEN NULL + ELSE 'Server restarted during training' + END, + output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END, ended_at = ? WHERE status = 'running' """, diff --git a/studio/backend/tests/test_audio_sampling_fill.py b/studio/backend/tests/test_audio_sampling_fill.py new file mode 100644 index 0000000000..efea18b83e --- /dev/null +++ b/studio/backend/tests/test_audio_sampling_fill.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Audio (TTS) generation applies recommended sampling + operator pins, like chat. + +Regression guard for the fix that moved the sampling fill ahead of the audio generators: a +prior version resolved sampling only after the audio branches returned, so `unsloth run +--temperature` (UNSLOTH_SAMPLING_*) and per-model recommendations never reached audio +generation. These exercise the transformers TTS path of ``generate_audio`` (the direct +``/audio/generate`` route, which the chat-completions audio branches also delegate to). +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import ChatCompletionRequest +from utils.inference import inference_config as ic + + +class _FakeLlama: + # is_loaded False forces the transformers (non-GGUF) TTS branch in generate_audio. + is_loaded = False + _is_audio = False + + +class _FakeTransformersBackend: + def __init__(self): + self.active_model_name = "some/custom-tts" + self.models = {"some/custom-tts": {"is_audio": True}} + self.captured = {} + + def generate_audio_response(self, **kwargs): + self.captured.update(kwargs) + return (b"RIFFfake", 24000) + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + ic._recommended_sampling.cache_clear() + for field in ic.SAMPLING_FIELD_NAMES: + monkeypatch.delenv(ic._SAMPLING_FIELDS[field][0], raising = False) + yield + ic._recommended_sampling.cache_clear() + + +def _run_generate_audio( + monkeypatch, + *, + recommended = None, + temperature = None, +): + backend = _FakeTransformersBackend() + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: backend) + + async def _noop_switch(*a, **k): + return None + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _noop_switch) + + # Recommendation source == the Chat UI's .inference block. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(recommended or {})) + ic._recommended_sampling.cache_clear() + + kwargs = {"model": "some/custom-tts", "messages": [{"role": "user", "content": "hi"}]} + if temperature is not None: + kwargs["temperature"] = temperature + payload = ChatCompletionRequest(**kwargs) + + asyncio.run(inference_route.generate_audio(payload, request = None, current_subject = "t")) + return backend.captured + + +def test_audio_uses_recommended_sampling_when_omitted(monkeypatch): + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0, "top_k": 64}) + assert captured["temperature"] == 1.0 + assert captured["top_k"] == 64 + + +def test_audio_operator_pin_overrides_client(monkeypatch): + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0}, temperature = 0.2) + assert captured["temperature"] == 0.9 # operator pin wins even over an explicit client value + + +def test_audio_client_explicit_preserved(monkeypatch): + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0}, temperature = 0.2) + assert captured["temperature"] == 0.2 # explicit client value preserved over recommendation diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index b3e6255d55..6f2c672002 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -66,6 +66,66 @@ def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path): assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"] +def test_legacy_hf_scan_uses_snapshot_path_for_inactive_cache(tmp_path): + repo = tmp_path / "models--Org--Model" + snapshot = repo / "snapshots" / "revision" + snapshot.mkdir(parents = True) + + [row] = models_route._scan_hf_cache(tmp_path, active_cache = False) + + assert row.model_id == "Org/Model" + assert row.id == str(snapshot.resolve()) + assert row.path == str(snapshot.resolve()) + + +def test_collect_local_models_scans_previous_cache(monkeypatch, tmp_path): + active = tmp_path / "active" + previous = tmp_path / "previous" + active.mkdir() + snapshot = previous / "models--Org--Previous" / "snapshots" / "revision" + snapshot.mkdir(parents = True) + + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active) + monkeypatch.setattr("utils.paths.legacy_hf_cache_dir", lambda: tmp_path / "legacy") + monkeypatch.setattr("utils.paths.hf_default_cache_dir", lambda: tmp_path / "default") + monkeypatch.setattr("utils.paths.lmstudio_model_dirs", lambda: []) + monkeypatch.setattr("utils.hf_cache_settings.known_hf_hub_caches", lambda: [active, previous]) + monkeypatch.setattr("storage.studio_db.list_scan_folders", lambda: []) + + rows = models_route.collect_local_models(tmp_path / "models") + + previous_row = next(row for row in rows if row.model_id == "Org/Previous") + assert previous_row.id == str(snapshot.resolve()) + + +def test_collect_local_models_prefers_complete_previous_copy(monkeypatch, tmp_path): + active = tmp_path / "active" + previous = tmp_path / "previous" + active_partial = active / "models--Org--Model" / "blobs" / "abc.incomplete" + active_partial.parent.mkdir(parents = True) + active_partial.write_bytes(b"partial") + snapshot = previous / "models--Org--Model" / "snapshots" / "revision" + snapshot.mkdir(parents = True) + (snapshot / "model.safetensors").write_bytes(b"complete") + + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active) + monkeypatch.setattr("utils.paths.legacy_hf_cache_dir", lambda: tmp_path / "legacy") + monkeypatch.setattr("utils.paths.hf_default_cache_dir", lambda: tmp_path / "default") + monkeypatch.setattr("utils.paths.lmstudio_model_dirs", lambda: []) + monkeypatch.setattr( + "utils.hf_cache_settings.known_hf_hub_caches", + lambda: [active, previous], + ) + monkeypatch.setattr("storage.studio_db.list_scan_folders", lambda: []) + + rows = models_route.collect_local_models(tmp_path / "models") + + [row] = [row for row in rows if row.model_id == "Org/Model"] + assert row.id == str(snapshot.resolve()) + assert row.partial is False + assert row.active_cache is False + + def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monkeypatch, tmp_path): repo = _repo( "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive", @@ -131,6 +191,72 @@ def test_is_hidden_model_hides_validation_probe_everywhere(): assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF") +def test_is_hidden_model_hides_dictation_models(tmp_path): + assert models_route._is_hidden_model("unsloth/whisper-tiny") + assert models_route._is_hidden_model("unsloth/whisper-base") + assert models_route._is_hidden_model("unsloth/whisper-small") + assert models_route._is_hidden_model("unsloth/whisper-large-v3-turbo") + assert models_route._is_hidden_model( + "/hf/models--unsloth--whisper-large-v3/snapshots/abc/model.safetensors" + ) + assert not models_route._is_hidden_model("user/whisper-finetune") + assert not models_route._is_hidden_model( + "C:\\cache\\models--unsloth--whisper-small-finetune\\model.safetensors" + ) + custom = tmp_path / "custom-whisper" + custom.mkdir() + (custom / "config.json").write_text( + '{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}' + ) + (custom / "model.safetensors").write_bytes(b"weights") + assert models_route._is_hidden_model( + "user/custom-checkpoint", + str(custom / "model.safetensors"), + ) + named_only = tmp_path / "whisper-finetune" + named_only.mkdir() + (named_only / "config.json").write_text('{"model_type": "llama"}') + assert not models_route._is_hidden_model("user/whisper-finetune", str(named_only)) + + +def test_list_cached_models_hides_custom_whisper_by_config(monkeypatch, tmp_path): + # Regression: the legacy /cached-models picker must pass the snapshot path so + # the config check hides a custom (non-curated) Whisper checkpoint; a bare + # repo id cannot ("user/whisper-finetune" is not in the curated set). + repo_path = tmp_path / "models--user--whisper-finetune" + snap = repo_path / "snapshots" / "abc" + snap.mkdir(parents = True) + (snap / "config.json").write_text( + '{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}' + ) + (snap / "model.safetensors").write_bytes(b"weights") + + captured: list = [] + real_hidden = models_route._is_hidden_model + + def spy(*values): + captured.append(values) + return real_hidden(*values) + + monkeypatch.setattr(models_route, "_is_hidden_model", spy) + repo = _repo( + "user/whisper-finetune", + [SimpleNamespace(file_name = "model.safetensors", size_on_disk = 10)], + repo_path, + ) + monkeypatch.setattr( + models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])] + ) + + result = asyncio.run( + models_route.list_cached_models(current_subject = "test-user", hf_token = None) + ) + # The route passed the snapshot path (not just the repo id) ... + assert any(str(repo_path) in values for values in captured) + # ... so the custom Whisper checkpoint is hidden from the chat picker. + assert result["cached"] == [] + + def test_is_hidden_model_matches_repo_ids_exactly(monkeypatch): """A custom embedder with a generic basename is hidden by EXACT repo-id match only, so unrelated cached repos that merely contain the basename stay @@ -573,33 +699,14 @@ def _gfile(name: str, size: int, mtime: float) -> SimpleNamespace: ) -def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_path): - """An unreadable auxiliary cache (e.g. an inaccessible - ``~/.cache/huggingface/hub``) must be skipped, not abort the scan. - Regression guard for ``extra.is_dir()`` raising and wiping the response. - """ - import huggingface_hub - import utils.paths as paths_mod +def test_all_hf_cache_scans_uses_shared_inventory(monkeypatch, tmp_path): + from hub.utils import inventory_scan active = SimpleNamespace( repos = [_repo("Org/Active", [_file("Q4_K_M.gguf", 5_000)], tmp_path / "active")] ) - def _fake_scan(cache_dir = None): - if cache_dir is None: - return active - raise AssertionError("auxiliary scan should have been skipped") - - class _Boom: - def is_dir(self): - raise PermissionError(13, "Permission denied") - - def resolve(self): - raise PermissionError(13, "Permission denied") - - monkeypatch.setattr(huggingface_hub, "scan_cache_dir", _fake_scan) - monkeypatch.setattr(paths_mod, "legacy_hf_cache_dir", lambda: _Boom()) - monkeypatch.setattr(paths_mod, "hf_default_cache_dir", lambda: _Boom()) + monkeypatch.setattr(inventory_scan, "all_hf_cache_scans", lambda: [active]) scans = models_route._all_hf_cache_scans() assert scans == [active] @@ -686,13 +793,17 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True, []), ) - monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) + monkeypatch.setattr( + GV, + "_local_main_gguf_blobs_by_quant", + lambda _repo_id, repo_cache_dir = None: {}, + ) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16" - monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id, root = None: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -705,6 +816,52 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa assert flags["F16"] is False +def test_gguf_variants_route_scopes_local_probe_to_selected_cache(monkeypatch, tmp_path): + snapshot = tmp_path / "inactive" / "models--org--repo" / "snapshots" / "rev" + snapshot.mkdir(parents = True) + calls = [] + + async def scoped_variants(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + return SimpleNamespace( + repo_id = repo_id, + variants = [], + has_vision = False, + default_variant = None, + ) + + context_calls = [] + monkeypatch.setattr(GV, "get_gguf_variants_response", scoped_variants) + monkeypatch.setattr( + models_route, + "_read_native_context_length", + lambda model, *, is_local: context_calls.append((model, is_local)) or 8192, + ) + + result = asyncio.run( + models_route.get_gguf_variants( + repo_id = "org/repo", + prefer_local_cache = True, + local_path = str(snapshot), + hf_token = None, + current_subject = "test-user", + ) + ) + + assert calls == [ + ( + "org/repo", + { + "prefer_local_cache": True, + "local_path": str(snapshot), + "hf_token": None, + }, + ) + ] + assert context_calls == [(str(snapshot), True)] + assert result.context_length == 8192 + + def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): siblings = [ SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100), @@ -726,12 +883,16 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): siblings, ), ) - monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) + monkeypatch.setattr( + GV, + "_local_main_gguf_blobs_by_quant", + lambda _repo_id, repo_cache_dir = None: {}, + ) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10) - monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id, root = None: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -758,12 +919,16 @@ def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, t "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False, []), ) - monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) + monkeypatch.setattr( + GV, + "_local_main_gguf_blobs_by_quant", + lambda _repo_id, repo_cache_dir = None: {}, + ) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10) - monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id, root = None: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -774,66 +939,82 @@ def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, t assert result.variants[0].downloaded is False -def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path): - """A cached mmproj adapter must not count toward a same-label main - variant's download progress (mmproj-F16 vs an F16 weight).""" - import huggingface_hub.constants as hf_constants +def test_legacy_gguf_progress_delegates_to_shared_service(monkeypatch): + calls = [] - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - snap = tmp_path / "models--org--repo" / "snapshots" / "rev" - snap.mkdir(parents = True) - (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # only the adapter on disk + async def shared(repo_id, *, variant, expected_bytes, hf_token): + calls.append((repo_id, variant, expected_bytes, hf_token)) + return {"downloaded_bytes": 10, "expected_bytes": 20, "progress": 0.5} - result = asyncio.run( - models_route.get_gguf_download_progress( - repo_id = "org/repo", - variant = "F16", - expected_bytes = 20_000, - current_subject = "test-user", - ) + monkeypatch.setattr( + "hub.services.models.downloads.get_gguf_download_progress_response", + shared, ) - assert result["downloaded_bytes"] == 0 - assert result["progress"] == 0 - - -def test_gguf_download_progress_excludes_big_endian_sibling(monkeypatch, tmp_path): - import huggingface_hub.constants as hf_constants - - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - snap = tmp_path / "models--org--repo" / "snapshots" / "rev" - snap.mkdir(parents = True) - (snap / "model-Q4_K_M-be.gguf").write_bytes(b"y" * 20_000) - result = asyncio.run( models_route.get_gguf_download_progress( repo_id = "org/repo", variant = "Q4_K_M", - expected_bytes = 20_000, + expected_bytes = 20, + hf_token = "token", current_subject = "test-user", ) ) - assert result["downloaded_bytes"] == 0 - assert result["progress"] == 0 + assert result["progress"] == 0.5 + assert calls == [("org/repo", "Q4_K_M", 20, "token")] -def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path): - import huggingface_hub.constants as hf_constants +def test_legacy_model_progress_delegates_to_shared_service(monkeypatch): + calls = [] - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - snap = tmp_path / "models--org--repo" / "snapshots" / "rev" / "Q4_K_M" - snap.mkdir(parents = True) - (snap / "foo.gguf").write_bytes(b"x" * 20_000) + async def shared(repo_id, *, hf_token): + calls.append((repo_id, hf_token)) + return {"downloaded_bytes": 10, "expected_bytes": 20, "progress": 0.5} + + monkeypatch.setattr( + "hub.services.models.downloads.get_download_progress_response", + shared, + ) result = asyncio.run( - models_route.get_gguf_download_progress( + models_route.get_download_progress( repo_id = "org/repo", - variant = "Q4_K_M", - expected_bytes = 20_000, + hf_token = "token", current_subject = "test-user", ) ) - assert result["downloaded_bytes"] == 20_000 - assert result["progress"] == 1.0 + assert result["progress"] == 0.5 + assert calls == [("org/repo", "token")] + + +def test_legacy_delete_delegates_to_shared_service(monkeypatch): + calls = [] + + async def shared( + repo_id, + variant, + hf_token, + cache_path = None, + ): + calls.append((repo_id, variant, hf_token, cache_path)) + return {"status": "deleted", "repo_id": repo_id} + + monkeypatch.setattr( + "hub.services.models.deletion.delete_cached_model_response", + shared, + ) + + result = asyncio.run( + models_route.delete_cached_model( + repo_id = "org/repo", + variant = None, + cache_path = "/data/hf/hub", + hf_token = "token", + current_subject = "test-user", + ) + ) + + assert result == {"status": "deleted", "repo_id": "org/repo"} + assert calls == [("org/repo", None, "token", "/data/hf/hub")] diff --git a/studio/backend/tests/test_change_password_policy.py b/studio/backend/tests/test_change_password_policy.py new file mode 100644 index 0000000000..c73e9ed839 --- /dev/null +++ b/studio/backend/tests/test_change_password_policy.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import importlib.util +import sys +from pathlib import Path + +import pytest +from fastapi import HTTPException + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from models.auth import ChangePasswordRequest # noqa: E402 + +# Load routes/auth.py directly so collection does not execute routes/__init__.py, +# which pulls in the heavy training/models/inference routers. +_route_path = _BACKEND_ROOT / "routes" / "auth.py" +_spec = importlib.util.spec_from_file_location("_change_password_route", _route_path) +assert _spec is not None and _spec.loader is not None +auth_routes = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(auth_routes) + + +@pytest.fixture +def _user(monkeypatch): + monkeypatch.setattr( + auth_routes.storage, + "get_user_and_secret", + lambda username: ("salt", "hash", "jwt-secret", False), + ) + monkeypatch.setattr( + auth_routes.hashing, + "verify_password", + lambda password, salt, pwd_hash: password == "bootstrap-pw", + ) + + +def _change(new_password): + payload = ChangePasswordRequest( + current_password = "bootstrap-pw", + new_password = new_password, + ) + return asyncio.run(auth_routes.change_password(payload, None, "unsloth")) + + +def test_rejects_whitespace_only_password(_user): + with pytest.raises(HTTPException) as excinfo: + _change(" " * 8) + assert excinfo.value.status_code == 400 + assert "spaces" in excinfo.value.detail + + +def test_rejects_tabs_and_spaces_password(_user): + with pytest.raises(HTTPException) as excinfo: + _change(" \t \t \t \t ") + assert excinfo.value.status_code == 400 + + +def test_rejects_password_containing_spaces(_user): + with pytest.raises(HTTPException) as excinfo: + _change("correct horse battery") + assert excinfo.value.status_code == 400 + assert "spaces" in excinfo.value.detail + + +def test_allows_password_without_spaces(_user, monkeypatch): + monkeypatch.setattr(auth_routes.storage, "update_password", lambda *args, **kwargs: True) + monkeypatch.setattr(auth_routes, "create_access_token", lambda subject: "at") + monkeypatch.setattr(auth_routes, "create_refresh_token", lambda subject: "rt") + token = _change("correct-horse-battery") + assert token.access_token == "at" + assert token.must_change_password is False diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index a60ac700bf..896bf1a6cd 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -91,6 +91,28 @@ def test_chat_settings_payload_accepts_fast_mode_presets(): assert dumped["customPresets"][0]["params"]["fastMode"] is True +def test_chat_settings_payload_accepts_preset_load_config(): + payload = chat_history.ChatSettingsPayload.model_validate( + { + "customPresets": [ + { + "name": "GGUF preset", + "params": {"temperature": 0.7, "maxTokens": 512}, + "loadConfig": { + "customContextLength": 256, + "kvCacheDtype": "q8_0", + "tensorParallel": False, + }, + }, + ], + } + ) + + dumped = payload.model_dump(exclude_unset = True) + assert dumped["customPresets"][0]["loadConfig"]["customContextLength"] == 256 + assert dumped["customPresets"][0]["loadConfig"]["kvCacheDtype"] == "q8_0" + + def test_chat_settings_payload_accepts_nudge_tool_calls(): # extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the # frontend's persisted nudgeToolCalls needs a payload field (like diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 7daa4224aa..f1d973f004 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -170,6 +170,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): estimate = None, single_device_gpu = None, gpu_ids = None, + is_vulkan = False, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), @@ -185,6 +186,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): max_seq_length = 0, requested_gpu_ids = gpu_ids, is_gguf = True, + is_vulkan = is_vulkan, required_override_gb = required_override, single_device_gpu = single_device_gpu, ) @@ -234,6 +236,35 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): self.assertFalse(blocked) self.assertEqual(blocked_info["usable_gb"], 10.0) + def test_vulkan_pin_takes_precedence_over_unknown_diffusion_fallback(self): + # An uncached GGUF can carry a speculative single-device fallback while + # its explicit pin is actually a ggml Vulkan ordinal. Never interpret + # that ordinal as the same-numbered CUDA physical device. + ok, info, _ = self._run( + devices = _devices((0, 80, 0), (1, 80, 78)), + required_override = 20.0, + single_device_gpu = "0", + gpu_ids = [0], + is_vulkan = True, + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "gguf_vulkan") + self.assertEqual(info["usable_gb"], 2.0) + + def test_vulkan_multi_gpu_guard_counts_requested_devices(self): + # The ordinal mapping is unknown, so use the least-free two visible + # cards for a two-device request. Their aggregate capacity is still + # available instead of collapsing the request to one card. + ok, info, _ = self._run( + devices = _devices((0, 80, 70), (1, 80, 70), (2, 80, 0)), + required_override = 10.0, + gpu_ids = [0, 1], + is_vulkan = True, + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf_vulkan") + self.assertEqual(info["usable_gb"], 18.5) + def test_single_device_unresolved_token_sizes_against_worst_device(self): # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a # free-VRAM index. The runner still drives ONE device, so size against the @@ -295,7 +326,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): - def test_non_cuda_allows(self): + def test_non_accelerator_allows(self): with patch("utils.hardware.get_device", return_value = DeviceType.MLX): ok, info = tv.can_load_chat_during_training( model_name = "m", @@ -305,7 +336,30 @@ class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): requested_gpu_ids = None, ) self.assertTrue(ok) - self.assertEqual(info["mode"], "non_cuda") + self.assertEqual(info["mode"], "non_accelerator") + + def test_xpu_overcommit_is_refused(self): + # XPU must NOT get the blanket non-accelerator allow: an oversized + # chat model during resident training is refused, like CUDA. + with ( + patch("utils.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.auto_select_gpu_ids", + return_value = ( + None, + {"selection_mode": "auto", "required_gb": 50.0, "usable_gb": 4.0}, + ), + ), + ): + 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.assertNotEqual(info.get("mode"), "non_accelerator") def test_no_visible_gpus_refuses(self): # GGUF with an empty device list -> no candidate GPU -> default-deny. @@ -478,58 +532,19 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_manual_known_normal_gguf_bypasses_training_estimate(self): captured = [] config = SimpleNamespace(is_gguf = True) - with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False) as classify: self._guard( config = config, captured = captured, training_active = True, decision = (False, {"reason": "must not run"}), gpu_memory_mode = "manual", + requested_gpu_ids = [1, 3], ) + classify.assert_called_once_with(config) self.assertEqual(captured, []) - def test_manual_unknown_gguf_keeps_single_device_training_guard(self): - captured = [] - config = SimpleNamespace(is_gguf = True) - with ( - patch.object(self.route, "_classify_diffusion_gguf", return_value = None), - patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), - patch.object( - self.route.LlamaCppBackend, - "_diffusion_gpu_arg", - return_value = "2", - ), - ): - self._guard( - config = config, - captured = captured, - training_active = True, - decision = (True, {"mode": "single_device"}), - gpu_memory_mode = "manual", - ) - self.assertEqual(len(captured), 1) - self.assertEqual(captured[0]["single_device_gpu"], "2") - - def test_manual_diffusion_uses_single_device_guard(self): - captured = [] - config = SimpleNamespace(is_gguf = True) - with ( - patch.object(self.route, "_classify_diffusion_gguf", return_value = True), - patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), - ): - self._guard( - config = config, - captured = captured, - training_active = True, - decision = (True, {"mode": "gguf"}), - gpu_memory_mode = "manual", - requested_gpu_ids = [3, 1], - ) - self.assertEqual(len(captured), 1) - self.assertEqual(captured[0]["single_device_gpu"], "1") - self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) - - def test_unpinned_diffusion_uses_runner_default_gpu(self): + def test_manual_diffusion_keeps_single_device_training_guard(self): captured = [] config = SimpleNamespace(is_gguf = True) with ( @@ -540,11 +555,6 @@ class TestChatLoadGuardRoute(unittest.TestCase): "_effective_gpu_count", return_value = 2, ), - patch.object( - self.route.LlamaCppBackend, - "_diffusion_gpu_arg", - return_value = "3", - ) as gpu_arg, ): self._guard( config = config, @@ -552,9 +562,11 @@ class TestChatLoadGuardRoute(unittest.TestCase): training_active = True, decision = (True, {"mode": "single_device"}), gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], ) - gpu_arg.assert_called_once_with(None, cpu_only = False) - self.assertEqual(captured[0]["single_device_gpu"], "3") + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} @@ -801,6 +813,80 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): asyncio.run(self.route.validate_model(request, current_subject = "u")) self.assertEqual(guard_called, []) + def _validate_gguf_template( + self, + *, + template, + canonical_path = "/picked/model.gguf", + ): + # Drive validate_model for a native lease-backed GGUF template probe and + # capture what the embedded-template reader was called with. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "model.gguf", + gguf_variant = "Q4_K_M", + native_path_lease = "signed-lease", + include_chat_template = True, + ) + cfg = SimpleNamespace( + identifier = canonical_path, + display_name = "model.gguf", + is_gguf = True, + is_lora = False, + is_vision = False, + gguf_file = canonical_path, + path = None, + base_model = None, + ) + import utils.models.gguf_metadata as gguf_meta + + seen = {} + + def _fake_read(path): + seen["path"] = path + return template + + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = (canonical_path, "model.gguf", True), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(gguf_meta, "read_gguf_chat_template", _fake_read), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + resp = asyncio.run(self.route.validate_model(request, current_subject = "u")) + return resp, seen, guard_called + + def test_include_chat_template_reads_leased_gguf_embedded_template(self): + # The picker chat-template GET has no lease plumbing, so a native picked + # GGUF surfaces its default template through this lease-aware probe: the + # embedded template is read from the granted canonical path and returned. + resp, seen, _ = self._validate_gguf_template(template = "{{ messages }}") + self.assertEqual(resp.chat_template, "{{ messages }}") + # Read strictly the leased file's own embedded template, never a sibling + # sidecar: the grant authorizes just this one path. + self.assertEqual(seen["path"], "/picked/model.gguf") + + def test_include_chat_template_skips_training_guard(self): + # A template-only probe allocates no VRAM, so like include_context_length + # it must not be refused by the training guard. + _, _, guard_called = self._validate_gguf_template(template = "{{ messages }}") + self.assertEqual(guard_called, []) + + def test_include_chat_template_over_cap_is_dropped(self): + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + resp, _, _ = self._validate_gguf_template(template = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1)) + self.assertIsNone(resp.chat_template) + # โ”€โ”€ _estimate_gguf_required_gb (sizes the same weights the loader loads) โ”€โ”€โ”€โ”€โ”€โ”€ diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index bb51cabf76..2094d15066 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -403,11 +403,234 @@ def test_reader_ignores_api_endpoint_failure_line(): assert t.error == "cloudflared exited before emitting a tunnel URL" +# โ”€โ”€ public reachability probe โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +class _FakeResponse: + def __init__(self, body): + self._body = body + + def read(self, size = -1): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _patch_urlopen(monkeypatch, handler): + import urllib.request + monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: handler(req)) + + +@pytest.fixture(autouse = True) +def _stub_dns_wait(monkeypatch, request): + if request.node.name.startswith("test_verify_public_url"): + monkeypatch.setattr(ct, "_wait_for_dns", lambda *a, **kw: None) + + +def test_wait_for_dns_polls_until_answer(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + if len(calls) < 3: + return _FakeResponse(b'{"Status":3}') + return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}') + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5) + assert len(calls) == 3 + assert "name=words.trycloudflare.com" in calls[0] + + +def test_wait_for_dns_gives_up_at_deadline(monkeypatch): + _patch_urlopen(monkeypatch, lambda req: _FakeResponse(b'{"Status":3}')) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 0.05) + + +def test_wait_for_dns_retries_transient_doh_error(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + if len(calls) < 3: + raise OSError("transient") + return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}') + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5) + assert len(calls) == 3 + + +def test_wait_for_dns_bails_on_persistent_doh_errors(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + raise OSError("blocked") + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5) + assert len(calls) == ct._DNS_MAX_DOH_ERRORS + + +def test_verify_public_url_accepts_studio_marker(monkeypatch): + seen = {} + + def handler(req): + seen["url"] = req.full_url + return _FakeResponse(b'{"status":"healthy","service":"Unsloth UI Backend"}') + + _patch_urlopen(monkeypatch, handler) + assert ct.verify_public_url("https://words.trycloudflare.com") is True + assert seen["url"] == "https://words.trycloudflare.com/api/health" + + +def test_verify_public_url_waits_for_dns_first(monkeypatch): + order = [] + monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: order.append(("dns", host))) + + def handler(req): + order.append(("probe", req.full_url)) + return _FakeResponse(b'{"service":"Unsloth UI Backend"}') + + _patch_urlopen(monkeypatch, handler) + assert ct.verify_public_url("https://words.trycloudflare.com") is True + assert order[0] == ("dns", "words.trycloudflare.com") + assert order[1][0] == "probe" + + +def test_verify_public_url_dns_wait_and_probe_share_deadline(monkeypatch): + # An exhausted DNS wait leaves the probe a single attempt, not a fresh window. + calls = [] + monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: None) + + def handler(req): + calls.append(req.full_url) + raise OSError("unreachable") + + _patch_urlopen(monkeypatch, handler) + assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0) is False + assert len(calls) == 1 + + +def test_verify_public_url_retries_then_succeeds(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + if len(calls) < 3: + raise OSError("Name or service not known") + return _FakeResponse(b'{"service":"Unsloth UI Backend"}') + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + assert ct.verify_public_url("https://words.trycloudflare.com") is True + assert len(calls) == 3 + + +def test_verify_public_url_rejects_unreachable_host(monkeypatch): + def handler(req): + raise OSError("Name or service not known") + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False + + +def test_verify_public_url_rejects_foreign_responder(monkeypatch): + # e.g. a Cloudflare error page: no service marker in the body. + _patch_urlopen(monkeypatch, lambda req: _FakeResponse(b"error 1033")) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False + + +@pytest.fixture(autouse = True) +def _stub_public_probe(monkeypatch, request): + # start_studio_tunnel tests use fake hostnames; keep them off the network. + if not request.node.name.startswith("test_start_studio_tunnel"): + return + monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: True) + + def test_start_studio_tunnel_no_binary(monkeypatch): monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None) assert ct.start_studio_tunnel(8080) is None +def test_start_studio_tunnel_drops_url_that_is_not_publicly_reachable(monkeypatch): + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" + + def wait_for_ready(self, timeout): + return self.url + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: False) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_returns_url_once_probe_passes(monkeypatch): + probed = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + self.protocol = protocol + + def start(self): + self.url = "https://words.trycloudflare.com" + + def wait_for_ready(self, timeout): + return self.url + + def stop(self): + pass + + def _probe(url, **kw): + probed.append(url) + return True + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + monkeypatch.setattr(ct, "verify_public_url", _probe) + try: + assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com" + assert probed == ["https://words.trycloudflare.com"] + finally: + ct.stop_studio_tunnel() + + def test_start_studio_tunnel_registers_before_wait(monkeypatch): # The tunnel must be visible to stop_studio_tunnel() during the readiness # wait, else a shutdown in that window orphans cloudflared. diff --git a/studio/backend/tests/test_colab_embed.py b/studio/backend/tests/test_colab_embed.py new file mode 100644 index 0000000000..dae0c7dae0 --- /dev/null +++ b/studio/backend/tests/test_colab_embed.py @@ -0,0 +1,479 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression coverage for Colab iframe embedding (#7344).""" + +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import colab + + +def _mock_google_colab_modules(colab_mod): + """Mock ``google`` and ``google.colab`` for environments without Google packages.""" + google_mod = types.ModuleType("google") + google_mod.colab = colab_mod + return {"google": google_mod, "google.colab": colab_mod} + + +def test_short_colab_url_truncates_proxy_host(): + url = "https://8888-gpu-a100-s-kkb-usc1f0-9hzedjcxrlu8-f.us-central1-0.prod.colab.dev/" + assert colab._short_colab_url(url, 8888) == "https://8888-gpu-..." + + +def test_short_colab_url_falls_back_on_unexpected_shape(): + assert colab._short_colab_url("https://example.com", 8888) == "https://example.com" + + +def test_is_colab_proxy_url_requires_https_proxy(): + assert colab._is_colab_proxy_url("https://8888-test.prod.colab.dev/", 8888) is True + assert colab._is_colab_proxy_url("http://localhost:8888", 8888) is False + assert colab._is_colab_proxy_url("http://127.0.0.1:8888", 8888) is False + + +def test_ready_card_html_does_not_open_colab_proxy_in_new_tab(): + """Colab proxy hosts 404 as top-level tabs (#7349 reporter); never window.open them.""" + html = colab._ready_card_html("https://8888-test.prod.colab.dev/", 8888) + assert "window.open" not in html + assert 'href="https://8888-test.prod.colab.dev/"' not in html + assert "start(cloudflare=True)" in html + + +def test_ready_card_html_points_to_cloudflare_when_link_ready(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html( + "https://8888-test.prod.colab.dev/", + 8888, + has_cloudflare_link = True, + ) + assert "Cloudflare link above" in html + + +def test_ready_card_html_warns_when_cloudflare_tunnel_missing(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html( + "https://8888-test.prod.colab.dev/", + 8888, + cloudflare_requested = True, + ) + assert "Could not open a Cloudflare tunnel" in html + + +def test_warn_colab_cloudflare_missing_logs_on_colab_without_tunnel(monkeypatch): + warnings: list[str] = [] + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr(colab.logger, "warning", lambda msg, **kwargs: warnings.append(msg)) + colab._warn_colab_cloudflare_missing(use_cloudflare = True, cloudflare_url = None) + assert warnings + assert "Cloudflare tunnel unavailable" in warnings[0] + + +def test_warn_colab_cloudflare_missing_skips_when_tunnel_ready(monkeypatch, caplog): + import logging + + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + with caplog.at_level(logging.WARNING): + colab._warn_colab_cloudflare_missing( + use_cloudflare = True, + cloudflare_url = "https://share.trycloudflare.com", + ) + assert "Cloudflare tunnel unavailable" not in caplog.text + + +def test_is_colab_runtime_uses_backend_colab_detector(monkeypatch): + fake_main = types.ModuleType("main") + fake_main._IS_COLAB = True + monkeypatch.setitem(sys.modules, "main", fake_main) + assert colab._is_colab_runtime() is True + fake_main._IS_COLAB = False + assert colab._is_colab_runtime() is False + + +def test_ready_card_html_uses_cloudflare_hint_on_colab_runtime_localhost(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + html = colab._ready_card_html("http://localhost:8888", 8888) + assert "window.open" not in html + assert "start(cloudflare=True)" in html + + +def test_ready_card_html_keeps_open_button_for_localhost_outside_colab(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + html = colab._ready_card_html("http://localhost:8888", 8888) + assert "window.open" in html + assert 'href="http://localhost:8888"' in html + assert "Open Unsloth Studio" in html + + +def test_embed_kernel_port_iframe_uses_colab_helper(monkeypatch): + colab_output = MagicMock() + google_colab = SimpleNamespace(output = colab_output) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + with patch.dict("sys.modules", _mock_google_colab_modules(google_colab)): + assert colab._embed_kernel_port_iframe(8888) is True + colab_output.serve_kernel_port_as_iframe.assert_called_once_with( + 8888, + height = colab._COLAB_IFRAME_HEIGHT, + width = "100%", + ) + + +def test_embed_kernel_port_iframe_returns_false_without_colab(): + with patch.dict("sys.modules", _mock_google_colab_modules(None)): + assert colab._embed_kernel_port_iframe(8888) is False + + +def test_embed_kernel_port_iframe_skips_colabtools_without_runtime(monkeypatch): + """colabtools can queue JS without appending an iframe; only trust the helper on Colab.""" + colab_output = MagicMock() + google_colab = SimpleNamespace(output = colab_output) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + with patch.dict("sys.modules", _mock_google_colab_modules(google_colab)): + assert colab._embed_kernel_port_iframe(8888) is False + colab_output.serve_kernel_port_as_iframe.assert_not_called() + + +def test_show_and_embed_prefers_kernel_port_iframe(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "kernel_iframe"] + + +def test_show_and_embed_falls_back_to_html_iframe(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: False) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append((url, port)) or True, + ) + + colab._show_and_embed(8888) + + assert calls == [("https://8888-test.prod.colab.dev/", 8888)] + + +def test_colab_wants_cloudflare_auto_enables_on_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + assert colab._colab_wants_cloudflare(None) is True + assert colab._colab_wants_cloudflare(True) is True + assert colab._colab_wants_cloudflare(False) is False + + +def test_colab_wants_cloudflare_defaults_off_outside_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + assert colab._colab_wants_cloudflare(None) is False + assert colab._colab_wants_cloudflare(True) is True + + +def test_finalize_colab_admin_password_skips_outside_runtime(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + assert colab._finalize_colab_admin_password() is None + + +def test_finalize_colab_admin_password_clears_bootstrap_gate(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr(colab, "_load_colab_login_credentials", lambda: None) + stored: list[tuple[str, str]] = [] + monkeypatch.setattr( + colab, + "_store_colab_login_credentials", + lambda username, password: stored.append((username, password)), + ) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(return_value = "alpha-beta-gamma"), + generate_bootstrap_password = MagicMock(return_value = "alpha-beta-gamma"), + requires_password_change = MagicMock(return_value = True), + update_password = MagicMock(return_value = True), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result == ("unsloth", "alpha-beta-gamma") + storage.ensure_default_admin.assert_called_once() + storage.update_password.assert_called_once_with("unsloth", "alpha-beta-gamma") + assert stored == [("unsloth", "alpha-beta-gamma")] + + +def test_start_skips_finalize_when_cloudflare_disabled(monkeypatch): + import time + + finalize_calls: list[str] = [] + monkeypatch.setattr(colab, "_is_studio_healthy", lambda port: True) + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_finalize_colab_admin_password", + lambda: finalize_calls.append("finalize") or ("unsloth", "secret"), + ) + monkeypatch.setattr( + colab, "start_cloudflare_tunnel", lambda port: "https://share.trycloudflare.com" + ) + monkeypatch.setattr(colab, "_publish_cloudflare_url", lambda url: None) + monkeypatch.setattr(colab, "_show_and_embed", lambda port, **kwargs: None) + monkeypatch.setattr(colab, "_stop_cloudflare_tunnel", lambda: None) + monkeypatch.setattr(time, "sleep", lambda _: (_ for _ in ()).throw(KeyboardInterrupt)) + + colab.start(cloudflare = False) + + assert finalize_calls == [] + + +def test_finalize_colab_admin_password_redisplay_on_rerun(monkeypatch): + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_load_colab_login_credentials", + lambda: ("unsloth", "saved-pass"), + ) + monkeypatch.setattr(colab, "_colab_credentials_still_valid", lambda username, password: True) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(), + generate_bootstrap_password = MagicMock(), + requires_password_change = MagicMock(return_value = False), + update_password = MagicMock(), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result == ("unsloth", "saved-pass") + storage.update_password.assert_not_called() + + +def test_finalize_colab_admin_password_drops_stale_cached_credentials(monkeypatch): + """After an in-app password change the cached first-run password no longer + authenticates, so it must not be redisplayed (#7349 Codex review).""" + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_load_colab_login_credentials", + lambda: ("unsloth", "stale-pass"), + ) + monkeypatch.setattr(colab, "_colab_credentials_still_valid", lambda username, password: False) + cleared: list[bool] = [] + monkeypatch.setattr(colab, "_clear_colab_login_credentials", lambda: cleared.append(True)) + + storage = SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + ensure_default_admin = MagicMock(), + get_bootstrap_password = MagicMock(), + generate_bootstrap_password = MagicMock(), + requires_password_change = MagicMock(return_value = False), + update_password = MagicMock(), + ) + auth_pkg = types.ModuleType("auth") + auth_pkg.storage = storage + with patch.dict("sys.modules", {"auth": auth_pkg, "auth.storage": storage}): + result = colab._finalize_colab_admin_password() + + assert result is None + assert cleared == [True] + storage.update_password.assert_not_called() + + +def test_colab_credentials_still_valid_matches_stored_hash(monkeypatch): + from auth.hashing import hash_password + + salt, pwd_hash = hash_password("right-pass") + storage = SimpleNamespace( + get_user_and_secret = MagicMock(return_value = (salt, pwd_hash, "jwt", False)), + ) + with patch.dict("sys.modules", {"auth.storage": storage}): + assert colab._colab_credentials_still_valid("unsloth", "right-pass") is True + assert colab._colab_credentials_still_valid("unsloth", "wrong-pass") is False + + +def test_colab_credentials_still_valid_false_when_user_missing(monkeypatch): + storage = SimpleNamespace(get_user_and_secret = MagicMock(return_value = None)) + with patch.dict("sys.modules", {"auth.storage": storage}): + assert colab._colab_credentials_still_valid("unsloth", "any") is False + + +def test_colab_login_html_includes_credentials(): + html = colab._colab_login_html("unsloth", "alpha-beta-gamma-delta") + assert "unsloth" in html + assert "alpha-beta-gamma-delta" in html + + +def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch): + displayed: list[str] = [] + ipython_display = SimpleNamespace( + HTML = lambda html: SimpleNamespace(html = html), + display = lambda html: displayed.append(html.html), + ) + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + with patch.dict("sys.modules", {"IPython.display": ipython_display}): + colab._show_and_embed( + 8888, + cloudflare_url = "https://share.trycloudflare.com", + colab_login = ("unsloth", "secret-pass"), + ) + + assert len(displayed) == 2 + assert "share.trycloudflare.com" in displayed[0] + assert "secret-pass" in displayed[1] + + +def test_show_and_embed_skips_iframe_on_colab_when_cloudflare_ready(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com") + + assert calls == [] + + +def test_show_and_embed_uses_kernel_helper_on_colab_runtime_despite_localhost(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"http://localhost:{port}") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "kernel_iframe"] + + +def test_show_and_embed_skips_kernel_helper_for_localhost_outside_colab(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"http://localhost:{port}") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: False) + + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["show_link", "html_iframe"] + + +def test_show_and_embed_still_embeds_when_show_link_fails(monkeypatch): + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: f"https://{port}-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None: (_ for _ in ()).throw(RuntimeError("no display")), + ) + monkeypatch.setattr( + colab, + "_embed_kernel_port_iframe", + lambda port: calls.append("kernel_iframe") or True, + ) + monkeypatch.setattr( + colab, + "_embed_html_iframe", + lambda url, port: calls.append("html_iframe") or True, + ) + + colab._show_and_embed(8888) + + assert calls == ["kernel_iframe"] diff --git a/studio/backend/tests/test_combined_update.py b/studio/backend/tests/test_combined_update.py new file mode 100644 index 0000000000..b96d3d030c --- /dev/null +++ b/studio/backend/tests/test_combined_update.py @@ -0,0 +1,735 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic tests for the combined llama+whisper update item. + +llama.cpp is the single main update item; whisper.cpp piggybacks on it. These +pin the union status (update_available = llama behind OR whisper behind), the +chained apply (llama phase first, whisper phase only when behind), the failure +policy (llama failure aborts; whisper failure keeps the llama partial success), +the silent whisper skips, and the backward-compatible payload shape. +""" + +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import utils.llama_cpp_freshness as freshness # noqa: E402 +import utils.llama_cpp_update as upd # noqa: E402 +import utils.whisper_cpp_freshness as wfresh # noqa: E402 +import utils.whisper_cpp_update as wupd # noqa: E402 + +MARKER = "UNSLOTH_PREBUILT_INFO.json" +WHISPER_MARKER = "UNSLOTH_WHISPER_PREBUILT_INFO.json" + +# The top-level status and job fields that predate the whisper piggyback; the +# combined payload must stay an exact superset so current UI code keeps working. +LEGACY_STATUS_FIELDS = { + "supported", + "update_available", + "stale", + "installed_tag", + "latest_tag", + "published_repo", + "installed_at_utc", + "age_days", + "source_build", + "update_size_bytes", + "job", +} +LEGACY_JOB_FIELDS = { + "state", + "message", + "from_tag", + "to_tag", + "reload_required", + "error", + "progress", + "started_at", + "finished_at", +} + + +class _FakeInstallerPopen: + """Stands in for the streamed llama installer process.""" + + def __init__( + self, + cmd, + *, + returncode = 0, + lines = None, + on_start = None, + **kwargs, + ): + if on_start is not None: + on_start(list(cmd)) + self.returncode = returncode + self.stdout = iter(lines or []) + + def wait(self): + return self.returncode + + def kill(self): + pass + + +def _patch_llama_installer( + monkeypatch, + *, + returncode = 0, + lines = None, + on_start = None, +): + # Only intercept the installer invocation: importing routes.inference inside + # the worker can Popen unrelated host probes (ldconfig etc). + def _popen(cmd, **kw): + is_installer = any("install_llama_prebuilt" in str(part) for part in cmd) + return _FakeInstallerPopen( + cmd, + returncode = returncode if is_installer else 0, + lines = lines if is_installer else None, + on_start = on_start if is_installer else None, + ) + + monkeypatch.setattr(upd.subprocess, "Popen", _popen) + + +def _write_llama_install(dir_: Path, tag: str) -> str: + """Create a fake llama prebuilt install and return the llama-server path.""" + bin_dir = dir_ / "build" / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + binary = bin_dir / "llama-server" + binary.write_text("stub") + (dir_ / MARKER).write_text( + json.dumps( + { + "tag": tag, + "release_tag": tag, + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": "2020-01-01T00:00:00Z", + } + ) + ) + return str(binary) + + +def _write_whisper_install( + dir_: Path, + tag: str, + backend: str = "cpu", +) -> str: + """Create a fake whisper prebuilt install and return the whisper-server path.""" + bin_dir = dir_ / "build" / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + binary = bin_dir / "whisper-server" + binary.write_text("stub") + (dir_ / WHISPER_MARKER).write_text( + json.dumps( + { + "release_tag": tag, + "upstream_tag": tag.split("-")[0], + "published_repo": "unslothai/whisper.cpp", + "backend": backend, + "installed_at_utc": "2020-01-01T00:00:00Z", + } + ) + ) + return str(binary) + + +@pytest.fixture(autouse = True) +def _clean_state(monkeypatch, tmp_path): + freshness.reset_caches() + wfresh.reset_caches() + upd._reset_job_for_tests() + upd._resolve_memo.clear() + wupd._resolve_memo.clear() + monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".llama_cache") + monkeypatch.setattr(wfresh, "_cache_dir", lambda: tmp_path / ".whisper_cache") + for var in ( + "LLAMA_SERVER_PATH", + "UNSLOTH_LLAMA_CPP_PATH", + "WHISPER_SERVER_PATH", + "UNSLOTH_WHISPER_CPP_PATH", + ): + monkeypatch.delenv(var, raising = False) + # Never hit the network in these tests. + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + monkeypatch.setattr(wfresh, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + yield + freshness.reset_caches() + wfresh.reset_caches() + upd._reset_job_for_tests() + upd._resolve_memo.clear() + wupd._resolve_memo.clear() + + +def _setup_llama( + monkeypatch, + tmp_path, + *, + installed = "b9493", + latest = "b9518", +): + """Marker-managed llama install; behind when installed != latest.""" + install_dir = tmp_path / "llama.cpp" + binary = _write_llama_install(install_dir, installed) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest) + return install_dir + + +def _setup_whisper( + monkeypatch, + tmp_path, + *, + installed = "v1.9.1-unsloth.1", + latest = "v1.9.2-unsloth.1", +): + """Marker-managed whisper install; behind when latest is newer.""" + install_dir = tmp_path / "whisper.cpp" + binary = _write_whisper_install(install_dir, installed) + monkeypatch.setattr(wupd, "_find_binary", lambda: binary) + monkeypatch.setattr(wupd, "_installer_script", lambda: tmp_path / "install_whisper_prebuilt.py") + monkeypatch.setattr(wfresh, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest) + return install_dir + + +def _patch_whisper_phase( + monkeypatch, + events, + *, + to_tag = "v1.9.2-unsloth.1", + error = None, +): + """Record whisper phase runs without touching a real installer.""" + + def _run(phase, set_progress): + events.append("whisper") + if error is not None: + raise RuntimeError(error) + set_progress(0.5) + return { + "to_tag": to_tag, + "reload_required": False, + "message": f"Updated whisper.cpp to {to_tag}.", + } + + monkeypatch.setattr(wupd, "run_chained_phase", _run) + + +def _wait_for_job(): + deadline = time.time() + 10 + while time.time() < deadline: + with upd._job_lock: + job = dict(upd._job) + if job["state"] in ("success", "error"): + return job + time.sleep(0.05) + with upd._job_lock: + return dict(upd._job) + + +# --- status: the single item folds whisper in --- + + +def test_status_payload_is_exact_superset_of_legacy_fields(monkeypatch, tmp_path): + _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path) + st = upd.get_update_status(force_refresh = True) + assert LEGACY_STATUS_FIELDS <= set(st) + assert LEGACY_JOB_FIELDS <= set(st["job"]) + # The new fields ride alongside, never replacing the legacy ones. + assert st["llama_update_available"] is True + assert st["whisper"]["update_available"] is True + assert st["whisper"]["latest_tag"] == "v1.9.2-unsloth.1" + assert st["update_component"] == "llama" + + +def test_status_union_whisper_only_surfaces_update(monkeypatch, tmp_path): + # llama current, whisper behind: the single item still shows an update. + _setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518") + _setup_whisper(monkeypatch, tmp_path) + st = upd.get_update_status(force_refresh = True) + assert st["llama_update_available"] is False + assert st["whisper"]["update_available"] is True + assert st["update_available"] is True + assert st["update_component"] == "whisper" + assert st["installed_tag"] == "b9518" + assert st["latest_tag"] == "b9518" + assert st["whisper"]["installed_tag"] == "v1.9.1-unsloth.1" + assert st["whisper"]["latest_tag"] == "v1.9.2-unsloth.1" + + +def test_status_whisper_current_does_not_flip_union(monkeypatch, tmp_path): + _setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518") + _setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1") + st = upd.get_update_status(force_refresh = True) + assert st["update_available"] is False + assert st["whisper"]["skip_reason"] == "up_to_date" + assert st["update_component"] is None + + +def test_status_survives_whisper_probe_failure(monkeypatch, tmp_path): + # The piggyback fails open: llama status still works without a whisper probe. + _setup_llama(monkeypatch, tmp_path) + + def _boom(*, force_refresh = False): + raise RuntimeError("probe exploded") + + monkeypatch.setattr(wupd, "chained_phase_plan", _boom) + st = upd.get_update_status(force_refresh = True) + assert st["update_available"] is True + assert st["whisper"] is None + + +# --- whisper chained_phase_plan: silent skips --- + + +def test_whisper_plan_skips_local_link(monkeypatch, tmp_path): + monkeypatch.setattr(wupd, "_find_binary", lambda: str(tmp_path / "whisper-server")) + monkeypatch.setattr(wupd, "_active_install_is_local_link", lambda b: True) + plan = wupd.chained_phase_plan() + assert plan["update_available"] is False + assert plan["skip_reason"] == "local_link" + assert plan["phase"] is None + + +def test_whisper_plan_skips_source_build(monkeypatch, tmp_path): + binary = tmp_path / "whisper.cpp" / "build" / "bin" / "whisper-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.setattr(wupd, "_find_binary", lambda: str(binary)) + plan = wupd.chained_phase_plan() + assert plan["skip_reason"] == "source_build" + assert plan["phase"] is None + + +def test_whisper_update_targets_canonical_root_when_inner_marker_exists(tmp_path): + install_dir = tmp_path / "whisper.cpp" + binary = install_dir / "build" / "bin" / "whisper-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + (install_dir / WHISPER_MARKER).write_text("{}") + (binary.parent / WHISPER_MARKER).write_text("{}") + assert wupd._install_dir_for(str(binary)) == install_dir + + +def test_whisper_plan_skips_when_not_installed(monkeypatch): + monkeypatch.setattr(wupd, "_find_binary", lambda: None) + plan = wupd.chained_phase_plan() + assert plan["skip_reason"] == "not_installed" + assert plan["phase"] is None + + +def test_whisper_plan_eligible_when_behind(monkeypatch, tmp_path): + install_dir = _setup_whisper(monkeypatch, tmp_path) + script = tmp_path / "install_whisper_prebuilt.py" + script.write_text("stub") + plan = wupd.chained_phase_plan(force_refresh = True) + assert plan["update_available"] is True + assert plan["skip_reason"] is None + assert plan["phase"]["install_dir"] == install_dir + assert plan["phase"]["repo"] == "unslothai/whisper.cpp" + assert plan["phase"]["backend"] == "cpu" + # Pin to the exact release the freshness check offered: unpinned, the + # installer's download-host /releases/latest pointer can lag published_at + # and reinstall an older build in a loop. + assert plan["phase"]["pin_release_tag"] == "v1.9.2-unsloth.1" + + +def test_whisper_plan_requires_a_repairable_pair_for_slim_installs(monkeypatch, tmp_path): + install_dir = _setup_whisper(monkeypatch, tmp_path) + marker_path = install_dir / WHISPER_MARKER + marker = json.loads(marker_path.read_text()) + marker["install_kind"] = "slim" + marker_path.write_text(json.dumps(marker)) + wfresh.reset_caches() + monkeypatch.setattr( + wupd, + "_resolve_prebuilt_for_host", + lambda **kwargs: {"prebuilt_available": False}, + ) + + plan = wupd.chained_phase_plan(force_refresh = True) + assert plan["update_available"] is False + assert plan["skip_reason"] == "paired_llama_unavailable" + + repaired = wupd.chained_phase_plan( + force_refresh = True, + paired_llama_will_update = True, + ) + assert repaired["update_available"] is True + assert repaired["phase"] is not None + + +def test_whisper_phase_pins_installer_to_checked_release(monkeypatch, tmp_path): + calls = [] + monkeypatch.setattr( + wupd._flow, + "stream_installer", + lambda cmd, env, **kw: calls.append(cmd), + ) + monkeypatch.setattr(wupd, "reset_caches", lambda **kw: None) + monkeypatch.setattr(wupd, "latest_published_release", lambda repo, **kw: "v9") + install_dir = tmp_path / "whisper.cpp" + binary = _write_whisper_install(install_dir, "v9") + monkeypatch.setattr(wupd, "_find_binary", lambda: binary) + wupd.run_chained_phase( + { + "install_dir": install_dir, + "repo": "unslothai/whisper.cpp", + "asset": None, + "backend": "cpu", + "script": tmp_path / "install_whisper_prebuilt.py", + "pin_release_tag": "v9", + }, + lambda f: None, + ) + cmd = calls[0] + assert "--published-release-tag" in cmd + assert cmd[cmd.index("--published-release-tag") + 1] == "v9" + + +def test_whisper_phase_exit_2_is_a_failed_phase(monkeypatch, tmp_path): + # No install occurred, so incompatibility must remain an actionable job + # error instead of producing a false success toast and hiding the banner. + def _raise_exit_2(cmd, env, **kw): + raise wupd._flow.InstallerExit(2, "installer exited 2: incompatible release") + + monkeypatch.setattr(wupd._flow, "stream_installer", _raise_exit_2) + install_dir = tmp_path / "whisper.cpp" + binary = _write_whisper_install(install_dir, "v1") + monkeypatch.setattr(wupd, "_find_binary", lambda: binary) + with pytest.raises(wupd._flow.InstallerExit) as exc_info: + wupd.run_chained_phase( + { + "install_dir": install_dir, + "repo": "unslothai/whisper.cpp", + "asset": None, + "backend": "cpu", + "script": tmp_path / "install_whisper_prebuilt.py", + "pin_release_tag": None, + }, + lambda f: None, + ) + assert exc_info.value.returncode == 2 + + +def test_llama_update_survives_unavailable_whisper_module(monkeypatch, tmp_path): + import builtins + + llama_dir = _setup_llama(monkeypatch, tmp_path) + monkeypatch.setattr(upd, "_whisper_chain_status", lambda **kw: None) + _patch_llama_installer( + monkeypatch, + on_start = lambda cmd: _write_llama_install(llama_dir, "b9518"), + ) + real_import = builtins.__import__ + + def guarded_import( + name, + globals = None, + locals = None, + fromlist = (), + level = 0, + ): + if name == "utils" and "whisper_cpp_update" in fromlist: + raise AssertionError("whisper module was re-imported after its failed probe") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", guarded_import) + + # A failed optional whisper probe must not be followed by an unconditional + # import. The valid llama phase still starts and completes. + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "success", job + assert job["phases"]["llama"]["state"] == "success" + assert job["phases"]["whisper"]["state"] == "skipped" + assert job["phases"]["whisper"]["reason"] == "unavailable" + + +def test_macos_status_uses_compatible_resolver_release(monkeypatch, tmp_path): + _setup_whisper( + monkeypatch, + tmp_path, + installed = "v1.9.1-unsloth.1", + latest = "v1.9.2-unsloth.1", + ) + monkeypatch.setattr(wupd.sys, "platform", "darwin") + monkeypatch.setattr( + wupd, + "_resolve_prebuilt_for_host", + lambda **kw: { + "prebuilt_available": True, + "release_tag": "v1.9.1-unsloth.1", + }, + ) + + status = wupd.get_update_status(force_refresh = True) + assert status["latest_tag"] == "v1.9.1-unsloth.1" + assert status["update_available"] is False + assert status["stale"] is False + + +def test_whisper_phase_integrity_failure_is_not_swallowed(monkeypatch, tmp_path): + def _raise_exit_1(cmd, env, **kw): + raise wupd._flow.InstallerExit(1, "installer exited 1: checksum mismatch") + + monkeypatch.setattr(wupd._flow, "stream_installer", _raise_exit_1) + install_dir = tmp_path / "whisper.cpp" + binary = _write_whisper_install(install_dir, "v1") + monkeypatch.setattr(wupd, "_find_binary", lambda: binary) + with pytest.raises(wupd._flow.InstallerExit, match = "checksum mismatch"): + wupd.run_chained_phase( + { + "install_dir": install_dir, + "repo": "unslothai/whisper.cpp", + "asset": None, + "backend": "cpu", + "script": tmp_path / "install_whisper_prebuilt.py", + "pin_release_tag": None, + }, + lambda f: None, + ) + + +# --- apply: the chained job --- + + +def test_apply_runs_llama_then_whisper(monkeypatch, tmp_path): + llama_dir = _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path) + (tmp_path / "install_whisper_prebuilt.py").write_text("stub") + + events = [] + _patch_llama_installer( + monkeypatch, + on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")), + ) + _patch_whisper_phase(monkeypatch, events) + + res = upd.start_update() + assert res["started"] is True, res + job = _wait_for_job() + assert job["state"] == "success", job + assert events == ["llama", "whisper"] # llama phase strictly first + assert job["phases"]["llama"]["state"] == "success" + assert job["phases"]["llama"]["to_tag"] == "b9518" + assert job["phases"]["whisper"]["state"] == "success" + assert job["phases"]["whisper"]["to_tag"] == "v1.9.2-unsloth.1" + # Legacy top-level fields keep their llama meaning. + assert job["from_tag"] == "b9493" + assert job["to_tag"] == "b9518" + assert "Updated llama.cpp to b9518." in job["message"] + assert "Updated whisper.cpp to v1.9.2-unsloth.1." in job["message"] + assert job["progress"] == 1.0 + assert LEGACY_JOB_FIELDS <= set(job) + + +def test_apply_llama_only_when_whisper_current(monkeypatch, tmp_path): + llama_dir = _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1") + + events = [] + _patch_llama_installer( + monkeypatch, + on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")), + ) + _patch_whisper_phase(monkeypatch, events) + + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "success", job + assert events == ["llama"] + assert job["phases"]["whisper"]["state"] == "skipped" + assert job["phases"]["whisper"]["reason"] == "up_to_date" + + +def test_apply_whisper_only_noops_llama(monkeypatch, tmp_path): + # llama current + whisper behind: the same single apply runs, with the llama + # phase a cheap already-matches no-op and the whisper phase doing the work. + _setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518") + _setup_whisper(monkeypatch, tmp_path) + (tmp_path / "install_whisper_prebuilt.py").write_text("stub") + + events = [] + _patch_llama_installer(monkeypatch, on_start = lambda cmd: events.append("llama")) + _patch_whisper_phase(monkeypatch, events) + + res = upd.start_update() + assert res["started"] is True, res + job = _wait_for_job() + assert job["state"] == "success", job + assert events == ["whisper"] # the llama installer never ran + # The legacy job-level to_tag means "llama tag"; a whisper-only round + # leaves it unset so the UI never reports a llama update that never ran. + assert job["to_tag"] is None + assert job["phases"]["llama"]["state"] == "skipped" + assert job["phases"]["llama"]["reason"] == "up_to_date" + assert job["phases"]["whisper"]["state"] == "success" + assert "Updated whisper.cpp to v1.9.2-unsloth.1." in job["message"] + + +def test_whisper_reload_never_raises_job_reload_flag(monkeypatch, tmp_path): + # A whisper-only update that had to unload a warm sidecar reports + # reload_required on its phase, but the JOB flag stays down: the chat + # frontend resyncs (and clears the local checkpoint) off the job flag, + # which must mean "the llama server changed", not "the sidecar restarted". + _setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518") + _setup_whisper(monkeypatch, tmp_path) + (tmp_path / "install_whisper_prebuilt.py").write_text("stub") + + def _whisper_phase(phase, set_progress): + return { + "to_tag": "v1.9.2-unsloth.1", + "reload_required": True, + "message": "Updated whisper.cpp to v1.9.2-unsloth.1.", + } + + monkeypatch.setattr(wupd, "run_chained_phase", _whisper_phase) + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "success", job + assert job["phases"]["whisper"]["reload_required"] is True + assert not job["reload_required"] + + +def test_apply_refuses_when_both_current(monkeypatch, tmp_path): + _setup_llama(monkeypatch, tmp_path, installed = "b9518", latest = "b9518") + _setup_whisper(monkeypatch, tmp_path, installed = "v1.9.2-unsloth.1", latest = "v1.9.2-unsloth.1") + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "up_to_date" + + +def test_apply_llama_failure_aborts_before_whisper(monkeypatch, tmp_path): + _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path) + (tmp_path / "install_whisper_prebuilt.py").write_text("stub") + + events = [] + _patch_llama_installer(monkeypatch, returncode = 2, lines = ["boom: disk full\n"]) + _patch_whisper_phase(monkeypatch, events) + + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "error", job + assert "boom" in (job["error"] or "") + assert events == [] # whisper never attempted + assert job["phases"]["llama"]["state"] == "error" + assert job["phases"]["whisper"]["state"] == "skipped" + assert job["phases"]["whisper"]["reason"] == "aborted" + assert job["message"] == "llama.cpp update failed." + + +def test_apply_whisper_failure_keeps_llama_partial_success(monkeypatch, tmp_path): + llama_dir = _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path) + (tmp_path / "install_whisper_prebuilt.py").write_text("stub") + + # An active model makes the llama phase report reload_required. + import threading + from types import ModuleType + + class _FakeBackend: + def __init__(self): + self._serial_load_lock = threading.Lock() + self._llama_update_in_progress = False + self.is_active = True + + def unload_model(self): + self.is_active = False + + backend = _FakeBackend() + routes_pkg = ModuleType("routes") + routes_pkg.__path__ = [] + inference_mod = ModuleType("routes.inference") + inference_mod.get_llama_cpp_backend = lambda: backend + monkeypatch.setitem(sys.modules, "routes", routes_pkg) + monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) + + events = [] + _patch_llama_installer( + monkeypatch, + on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")), + ) + _patch_whisper_phase(monkeypatch, events, error = "whisper installer exploded") + + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "error", job + assert events == ["llama", "whisper"] + # The message says both halves: llama landed, whisper did not. + assert "Updated llama.cpp to b9518." in job["message"] + assert "whisper.cpp update failed." in job["message"] + assert "whisper installer exploded" in (job["error"] or "") + # The llama phase's reload_required survives the whisper failure. + assert job["reload_required"] is True + assert job["to_tag"] == "b9518" + assert job["phases"]["llama"]["state"] == "success" + assert job["phases"]["whisper"]["state"] == "error" + + +def test_apply_skips_whisper_local_link_silently(monkeypatch, tmp_path): + llama_dir = _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path) + monkeypatch.setattr(wupd, "_active_install_is_local_link", lambda b: True) + + events = [] + _patch_llama_installer( + monkeypatch, + on_start = lambda cmd: (events.append("llama"), _write_llama_install(llama_dir, "b9518")), + ) + _patch_whisper_phase(monkeypatch, events) + + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "success", job + assert events == ["llama"] + assert job["phases"]["whisper"]["state"] == "skipped" + assert job["phases"]["whisper"]["reason"] == "local_link" + assert job["message"] == "Updated llama.cpp to b9518." + + +def test_chained_progress_windows(monkeypatch, tmp_path): + # The llama phase fills roughly the first 0.7 slice and whisper the rest. + llama_dir = _setup_llama(monkeypatch, tmp_path) + _setup_whisper(monkeypatch, tmp_path) + (tmp_path / "install_whisper_prebuilt.py").write_text("stub") + + seen = {} + + def _whisper_phase(phase, set_progress): + with upd._job_lock: + seen["at_whisper_start"] = upd._job["progress"] + set_progress(0.5) + with upd._job_lock: + seen["mid_whisper"] = upd._job["progress"] + return {"to_tag": "v1.9.2-unsloth.1", "reload_required": False, "message": "ok"} + + monkeypatch.setattr(wupd, "run_chained_phase", _whisper_phase) + _patch_llama_installer( + monkeypatch, + lines = ["Downloading app.tar.gz: 100.0% (35.0 MiB/35.0 MiB) at 9.0 MiB/s\n"], + on_start = lambda cmd: _write_llama_install(llama_dir, "b9518"), + ) + + assert upd.start_update()["started"] is True + job = _wait_for_job() + assert job["state"] == "success", job + assert seen["at_whisper_start"] == pytest.approx(0.7) + assert seen["mid_whisper"] == pytest.approx(0.7 + 0.5 * 0.3) + assert job["progress"] == 1.0 diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py index 804221ec7e..181e0c9fad 100644 --- a/studio/backend/tests/test_consent_gate.py +++ b/studio/backend/tests/test_consent_gate.py @@ -873,6 +873,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): if fn == "config.json": import json @@ -899,6 +900,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -932,6 +934,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -972,6 +975,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -1008,6 +1012,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -1037,6 +1042,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -1079,6 +1085,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -1120,6 +1127,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile @@ -1182,6 +1190,7 @@ class TestScannerCoversAllExecutableCode: repo, fn, token = None, + cache_dir = None, ): import json import tempfile diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index b3fa98b604..a6c18bd8de 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -106,6 +106,56 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch): assert unverified.status_code == 409 +def test_offline_cached_non_st_model_is_accepted(client, monkeypatch): + # Offline, a cached transformers-native embedder (no modules.json) is unverifiable via HF + # metadata, but ST can load any cached encoder, so accept it (no 409). + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + import utils.models as _models + import utils.utils as _uu + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: True) + r = c.put("/embedding-model", json = {"embedding_model": "acme/gte-modernbert"}) + assert r.status_code == 200 + assert saved.get("model") == "acme/gte-modernbert" + + +def test_offline_partial_or_uncached_model_still_409(client, monkeypatch): + # Offline but not loadable (uncached or metadata-only partial cache): keep the forceable + # 409, since the cache-only load would fail anyway. + c, _saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + import utils.models as _models + import utils.utils as _uu + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: False) + r = c.put("/embedding-model", json = {"embedding_model": "acme/uncached-embedder"}) + assert r.status_code == 409 + + +def test_offline_skips_remote_gguf_probe(client, monkeypatch): + # Offline + llama backend: the remote GGUF probe (list_repo_files) must be skipped so a + # dead-DNS session cannot hang. + c, _saved = client + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_local_gguf_backend_error", lambda model: None) + + def _boom(*a, **k): + raise AssertionError("hit the network for the GGUF probe") + + monkeypatch.setattr(settings, "_hf_gguf_backend_error", _boom) + import utils.models as _models + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: True) + r = c.put("/embedding-model", json = {"embedding_model": "acme/embedder"}) + assert r.status_code == 200 + + def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py index 761ea08e3f..5097f9f53a 100644 --- a/studio/backend/tests/test_export_absolute_paths.py +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -158,6 +158,7 @@ def _install_lightweight_backend_stubs(monkeypatch): utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None utils_model_config._extract_quant_label = lambda value: value utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False + utils_model_config._is_mtp_drafter = lambda *args, **kwargs: False utils_model_config.is_audio_input_type = lambda *args, **kwargs: None monkeypatch.setitem( sys.modules, diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py index b4c8f5d242..e02c33a0f1 100644 --- a/studio/backend/tests/test_file_security.py +++ b/studio/backend/tests/test_file_security.py @@ -165,6 +165,23 @@ def test_skips_local_path(): assert "local" in d.reason +def test_scans_inactive_hf_cache_snapshot_path(tmp_path): + # An inactive HF cache loads by snapshot path; the gate must recover the repo id + + # commit from models--org--repo/snapshots/ and scan that exact commit, not exempt + # it and not fall back to the default branch (an older commit may hold a dropped pickle). + snapshot = tmp_path / "models--evil--repo" / "snapshots" / "deadbeef" + snapshot.mkdir(parents = True) + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status) as model_info: + d = evaluate_file_security(str(snapshot)) + assert d.blocked is True + assert model_info.call_args.args[0] == "evil/repo" + assert model_info.call_args.kwargs["revision"] == "deadbeef" + + 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. diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 62596fcc8a..6d1fac980b 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -103,6 +103,10 @@ def _build_cache( @pytest.fixture def hf_cache(tmp_path, monkeypatch): monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: _types.SimpleNamespace(hub_cache = tmp_path), + ) monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) return tmp_path @@ -117,6 +121,61 @@ def _fail_get_paths_info(*_args, **_kwargs): class TestLoadReusesCachedCopy: + def test_download_uses_selected_cache_for_lookup_preflight_and_write( + self, tmp_path, monkeypatch + ): + backend = LlamaCppBackend() + selected = tmp_path / "selected" / "hub" + startup = tmp_path / "startup" / "hub" + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(startup)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: _types.SimpleNamespace(hub_cache = selected), + ) + seen = {"lookups": [], "disk": [], "downloads": []} + + def cached_lookup( + repo_id, + filename, + *, + cache_dir = None, + **_kwargs, + ): + seen["lookups"].append((repo_id, filename, cache_dir)) + return None + + def disk_usage(path): + seen["disk"].append(str(path)) + return _types.SimpleNamespace(free = 1024) + + def download(repo_id, filename, _token, **kwargs): + seen["downloads"].append((repo_id, filename, kwargs.get("cache_dir"))) + return str(selected / filename) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch( + "huggingface_hub.get_paths_info", + lambda _repo, paths, **_kwargs: [ + _types.SimpleNamespace(path = path, size = 4) for path in paths + ], + ), + patch("huggingface_hub.try_to_load_from_cache", cached_lookup), + patch("core.inference.llama_cpp.shutil.disk_usage", disk_usage), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + download, + ), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(selected / MAIN) + assert seen == { + "lookups": [(REPO, MAIN, str(selected))], + "disk": [str(selected)], + "downloads": [(REPO, MAIN, str(selected))], + } + def test_online_reuse_after_revision_bump(self, hf_cache): """A new repo revision does not replace a complete cached model.""" backend = LlamaCppBackend() diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index b17274197f..271a882b11 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -591,10 +591,23 @@ def test_load_request_accepts_gpu_ids(): @pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) def test_response_models_emit_gpu_ids(model_cls): if model_cls is LoadResponse: - obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + obj = model_cls( + status = "loaded", + model = "m", + display_name = "m", + inference = {}, + gpu_ids = [1], + requested_gpu_ids = [1, 2], + ) else: - obj = model_cls(gpu_ids = [1]) + obj = model_cls(gpu_ids = [1], requested_gpu_ids = [1, 2]) assert obj.model_dump()["gpu_ids"] == [1] + assert obj.model_dump()["requested_gpu_ids"] == [1, 2] + + +def test_gguf_load_and_status_responses_include_requested_gpu_pool(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert route_src.count("requested_gpu_ids = llama_backend.requested_gpu_ids") == 3 def test_gpu_ids_property_default_and_reset(): @@ -625,6 +638,10 @@ def _target_state_gpu_ids(backend, gpu_ids): def test_gpu_ids_reload_detection_is_order_insensitive(): backend = _loaded_backend("auto") backend._gpu_ids = [0, 1] + # A real non-narrowed load records the raw request too; the non-diffusion + # dedupe now compares that raw pin (#7239). Set it to match the effective pin + # (no narrowing) so this exercises the order-insensitive comparison. + backend._requested_gpu_ids = [0, 1] # Same set, different order -> no reload. assert _target_state_gpu_ids(backend, [1, 0]) is True # Different set -> reload. @@ -633,6 +650,26 @@ def test_gpu_ids_reload_detection_is_order_insensitive(): assert _target_state_gpu_ids(backend, None) is False +def test_gpu_ids_reload_detection_accepts_raw_and_effective_pin(): + backend = _loaded_backend("auto") + backend._requested_gpu_ids = [0, 1] + backend._gpu_ids = [0] + backend._last_load_kwargs = {"gpu_ids": [0, 1], "model_identifier": "owner/repo"} + + # The original request still matches after the fitter narrows it. + assert _target_state_gpu_ids(backend, [1, 0]) is True + assert backend.requested_gpu_ids == [0, 1] + # The status response echoes the effective pin, which must also round-trip. + # Treat the incoming subset as the latest intent so status and a future + # reload do not restore GPU 1 after the user removed it. + assert _target_state_gpu_ids(backend, [0]) is True + assert backend.requested_gpu_ids == [0] + assert backend._last_load_kwargs == {"gpu_ids": [0], "model_identifier": "owner/repo"} + # A genuinely different placement pool still reloads. + assert _target_state_gpu_ids(backend, [1]) is False + assert _target_state_gpu_ids(backend, None) is False + + def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): # The diffusion runner drives only its single lowest device, so the backend # records [lowest]. A later multi-GPU request that still resolves to that @@ -642,6 +679,7 @@ def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): backend._is_diffusion = True backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick assert _target_state_gpu_ids(backend, [3, 1]) is True + assert backend.requested_gpu_ids == [1] assert _target_state_gpu_ids(backend, [1]) is True # Lowest device changes (2, not 1) -> reload. assert _target_state_gpu_ids(backend, [3, 2]) is False @@ -649,6 +687,56 @@ def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): assert _target_state_gpu_ids(backend, None) is False +def test_remote_vulkan_diffusion_preflight_runs_before_teardown(monkeypatch): + def _mark_diffusion(probe, path): + assert path == "/cache/model.gguf" + probe._is_diffusion = True + + monkeypatch.setattr(LlamaCppBackend, "_read_gguf_metadata", _mark_diffusion) + assert LlamaCppBackend._gguf_path_is_diffusion("/cache/model.gguf", "owner/model") is True + + src = inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + preflight = src.index("_preflight_model_path = self._download_gguf(") + teardown = src.index("# โ”€โ”€ Phase 1: kill old process") + assert preflight < teardown + assert "model_path = _preflight_model_path or self._download_gguf(" in src + + +def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch): + backend = LlamaCppBackend() + killed = [] + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)]) + monkeypatch.setattr( + backend, + "_download_gguf", + lambda **_kwargs: "/cache/diffusion.gguf", + ) + monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: True) + monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True)) + monkeypatch.setattr( + llama_cpp_module, + "_resolve_repo_id_casing", + lambda repo: repo, + ) + monkeypatch.setattr( + llama_cpp_module, + "_hf_offline_if_dns_dead", + lambda: __import__("contextlib").nullcontext(), + ) + + with pytest.raises(ValueError, match = "DiffusionGemma"): + backend.load_model( + hf_repo = "owner/model", + hf_variant = "Q4_K_M", + model_identifier = "owner/model", + gpu_ids = [0], + ) + + assert killed == [] + + def test_start_diffusion_server_resets_tensor_parallel(): # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model # phase 1 only kills the process, it skips the unload reset). Diffusion is never @@ -656,18 +744,16 @@ def test_start_diffusion_server_resets_tensor_parallel(): # diffusion re-Apply reloads against stale tensor-parallel state. src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) assert "self._tensor_parallel = False" in src + assert "self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None" in src -def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): - # The route-level reload dedupe mirrors the backend: for a loaded diffusion - # model it compares the request against the single recorded device, not the - # full requested list, or a same-device multi-GPU pick reloads needlessly. +def test_route_matches_loaded_settings_uses_shared_gpu_pin_matcher(): + # Route-level and backend race dedupe must share one normalization path so + # raw, effective, and diffusion pins cannot drift apart. route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] - guard = match_impl.index("if llama_backend.is_diffusion:") - collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") - compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") - assert guard < collapse < compare + assert "if not llama_backend.matches_gpu_ids(request.gpu_ids):" in match_impl + assert "llama_backend._record_matching_gpu_request(request.gpu_ids)" in match_impl # โ”€โ”€ Manual tensor split: child enumeration pinned to the picker's order โ”€โ”€โ”€โ”€โ”€โ”€ @@ -733,20 +819,211 @@ def test_split_pin_without_mask_only_sets_pci_order(monkeypatch): def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch): - # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR - # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP - # would index into the already-reduced set). + # ROCm with the mask sourced from HIP: the pin must land in + # HIP_VISIBLE_DEVICES too, and an inherited ROCR mask is cleared so the + # mask can't apply twice (ROCR re-indexes, then HIP would index into the + # already-reduced set). _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) - torch_stub = _types.ModuleType("torch") - torch_stub.version = _types.SimpleNamespace(hip = "6.0") - monkeypatch.setitem(sys.modules, "torch", torch_stub) - env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"} + _rocm_torch_stub(monkeypatch) + env = { + "CUDA_VISIBLE_DEVICES": "3,1", + "HIP_VISIBLE_DEVICES": "3,1", + "ROCR_VISIBLE_DEVICES": "3,1", + } LlamaCppBackend._pin_visible_gpu_order_for_split(env) assert env["CUDA_VISIBLE_DEVICES"] == "1,3" assert env["HIP_VISIBLE_DEVICES"] == "1,3" assert "ROCR_VISIBLE_DEVICES" not in env +def test_split_pin_preserves_inherited_rocr_mask(monkeypatch): + # Mask sourced from ROCR alone (e.g. an AMD SDK parent): the pin must + # re-emit at the ROCr layer, not swap to HIP -- clearing ROCR re-exposes + # every agent to HSA enumeration, which can segfault at startup on an + # unsupported GPU the parent mask was hiding (#7272 review). CUDA carries + # the post-ROCR ordinals, mirroring the prefer_rocr emission. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + _rocm_torch_stub(monkeypatch) + env = {"ROCR_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["ROCR_VISIBLE_DEVICES"] == "1,3" + assert env["CUDA_VISIBLE_DEVICES"] == "0,1" + assert "HIP_VISIBLE_DEVICES" not in env + + +def test_split_pin_keeps_hip_on_windows_despite_stray_rocr(monkeypatch): + # On Windows the ROCR var is dead (no ROCr layer) and the resolver never + # reads it, so a stray value must not flip the pin to the ROCR emission: + # the HIP mask is the only effective selector there. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + monkeypatch.setattr(sys, "platform", "win32") + env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "9"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + assert env["HIP_VISIBLE_DEVICES"] == "1,3" + assert "ROCR_VISIBLE_DEVICES" not in env + + +def _rocm_torch_stub(monkeypatch): + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + # prefer_rocr is Linux-only (ROCR is an ROCr variable); pin the platform so + # these Linux-behaviour tests also pass on a Windows dev box. + monkeypatch.setattr(sys, "platform", "linux") + + +def test_subset_pin_masks_via_rocr_on_rocm(monkeypatch): + # A GPU-subset pin must exclude the rest at the ROCr/HSA layer: HIP masking + # still enumerates every agent first, which segfaults the build on an + # unsupported deselected GPU (e.g. a gfx1103 iGPU under a gfx110X prebuilt). + # ROCR drops it at the driver layer; only one mask is set (HIP cleared). + _rocm_torch_stub(monkeypatch) + env = {"HIP_VISIBLE_DEVICES": "9"} # stale/inherited HIP mask must not survive + LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True) + assert env["ROCR_VISIBLE_DEVICES"] == "0" + assert env["CUDA_VISIBLE_DEVICES"] == "0" + assert "HIP_VISIBLE_DEVICES" not in env + + +def test_prefer_rocr_remaps_cuda_to_post_rocr_ordinals(monkeypatch): + # ROCR re-indexes the visible agents from 0, and HIP (cleared here) falls back + # to CUDA_VISIBLE_DEVICES -- so on the prefer_rocr path CUDA must carry the + # post-ROCR ordinals, not the physical ids, else a non-zero pick indexes out + # of range and the child sees no GPU and drops to CPU (#7272 review). + _rocm_torch_stub(monkeypatch) + # Single non-zero GPU: ROCR keeps the physical id, CUDA becomes ordinal 0. + env = {} + LlamaCppBackend._emit_child_gpu_visibility(env, "1", prefer_rocr = True) + assert env["ROCR_VISIBLE_DEVICES"] == "1" + assert env["CUDA_VISIBLE_DEVICES"] == "0" + assert "HIP_VISIBLE_DEVICES" not in env + # Multi-GPU subset: ROCR keeps the physical ids, CUDA is the 0-based ordinals. + env = {} + LlamaCppBackend._emit_child_gpu_visibility(env, "1,3", prefer_rocr = True) + assert env["ROCR_VISIBLE_DEVICES"] == "1,3" + assert env["CUDA_VISIBLE_DEVICES"] == "0,1" + assert "HIP_VISIBLE_DEVICES" not in env + + +def test_subset_pin_default_still_uses_hip_and_clears_rocr(monkeypatch): + # Without prefer_rocr the masking is unchanged: HIP narrows, inherited ROCR + # is cleared so the two can't double-mask. + _rocm_torch_stub(monkeypatch) + env = {"ROCR_VISIBLE_DEVICES": "0,1"} + LlamaCppBackend._emit_child_gpu_visibility(env, "1") + assert env["HIP_VISIBLE_DEVICES"] == "1" + assert "ROCR_VISIBLE_DEVICES" not in env + + +def test_cpu_only_pin_keeps_hip_even_with_prefer_rocr(monkeypatch): + # The CPU-only sentinel never routes through ROCR (no portable "hide all" + # spelling); it hides every GPU via HIP. + _rocm_torch_stub(monkeypatch) + env = {} + LlamaCppBackend._emit_child_gpu_visibility(env, "-1", prefer_rocr = True) + assert env["HIP_VISIBLE_DEVICES"] == "-1" + assert "ROCR_VISIBLE_DEVICES" not in env + + +def _amd_sdk_torch_stub(monkeypatch): + # AMD SDK wheel: torch.version.hip is None but __version__ encodes rocm. + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = None) + torch_stub.__version__ = "2.9.1+rocm7.2.1" + monkeypatch.setitem(sys.modules, "torch", torch_stub) + monkeypatch.setattr(sys, "platform", "linux") + + +def test_prefer_rocr_falls_back_to_hip_on_windows(monkeypatch): + # ROCR_VISIBLE_DEVICES is a Linux ROCr variable (Windows HIP has no ROCr + # layer), so on Windows ROCm prefer_rocr must keep the HIP mask or a nonzero + # pick loses its only effective selector (#7272 review). + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + monkeypatch.setattr(sys, "platform", "win32") + env = {"ROCR_VISIBLE_DEVICES": "9"} + LlamaCppBackend._emit_child_gpu_visibility(env, "1", prefer_rocr = True) + assert env["HIP_VISIBLE_DEVICES"] == "1" + assert env["CUDA_VISIBLE_DEVICES"] == "1" + assert "ROCR_VISIBLE_DEVICES" not in env + + +def test_amd_sdk_wheel_hip_none_still_masks_rocr(monkeypatch): + # An AMD SDK wheel leaves torch.version.hip unset but has "rocm" in __version__. + # It must still get the ROCR mask, else only CUDA_VISIBLE_DEVICES is set and an + # unsupported iGPU keeps enumerating and can crash llama-server. + _amd_sdk_torch_stub(monkeypatch) + env = {"HIP_VISIBLE_DEVICES": "9"} + LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True) + assert env["ROCR_VISIBLE_DEVICES"] == "0" + assert "HIP_VISIBLE_DEVICES" not in env + + +def test_cuda_wheel_hip_none_gets_no_rocm_mask(monkeypatch): + # A CUDA wheel (hip=None, no "rocm" in __version__) must NOT get a HIP/ROCR mask + # -- only CUDA_VISIBLE_DEVICES -- so the version-string check can't false-positive. + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = None) + torch_stub.__version__ = "2.9.1+cu124" + monkeypatch.setitem(sys.modules, "torch", torch_stub) + env = {} + LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True) + assert env["CUDA_VISIBLE_DEVICES"] == "0" + assert "ROCR_VISIBLE_DEVICES" not in env + assert "HIP_VISIBLE_DEVICES" not in env + + +def test_resolve_physical_ids_reads_rocr_on_amd_sdk_wheel(monkeypatch): + # _resolve_visible_physical_ids must use the same ROCm detection as + # _emit_child_gpu_visibility: on an AMD SDK wheel (hip=None, rocm in + # __version__) an inherited ROCR mask IS the ordinal->physical mapping. + # Reading it as "no mask" labels ordinal 0 as physical 0 and the child's + # ROCR pin then re-exposes the GPU the mask was hiding (#7272 review). + _amd_sdk_torch_stub(monkeypatch) + for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + monkeypatch.delenv(var, raising = False) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + assert LlamaCppBackend._resolve_visible_physical_ids() == [1] + + +def test_resolve_physical_ids_ignores_rocr_on_cuda_wheel(monkeypatch): + # A CUDA wheel (hip=None, no "rocm") keeps CUDA-only semantics: a stray + # ROCR var must not be read as the mask. + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = None) + torch_stub.__version__ = "2.9.1+cu124" + monkeypatch.setitem(sys.modules, "torch", torch_stub) + for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + monkeypatch.delenv(var, raising = False) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + assert LlamaCppBackend._resolve_visible_physical_ids() is None + + +def test_resolve_physical_ids_ignores_rocr_on_windows(monkeypatch): + # ROCR_VISIBLE_DEVICES is a Linux ROCr variable: Windows HIP has no ROCr + # layer, so a stray ROCR var there does not mask the runtime. Reading it as + # the ordinal->physical mapping would label ordinal 0 with a stale ROCR id + # while the runtime still enumerates every adapter, so auto-selection could + # budget one card and pin another (#7272 review). HIP must still be honoured. + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = None) + torch_stub.__version__ = "2.9.1+rocm7.2.1" # AMD SDK wheel + monkeypatch.setitem(sys.modules, "torch", torch_stub) + monkeypatch.setattr(sys, "platform", "win32") + for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + monkeypatch.delenv(var, raising = False) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + assert LlamaCppBackend._resolve_visible_physical_ids() is None + # HIP precedence is unchanged on Windows. + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1") + assert LlamaCppBackend._resolve_visible_physical_ids() == [1] + + # โ”€โ”€ Diffusion single-device selection โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d4f2fbe993..3dab7ef368 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -119,7 +119,8 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaisesRegex( - ValueError, "unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG" + ValueError, + "unsupported when CUDA_VISIBLE_DEVICES uses non-numeric or subdevice", ): resolve_requested_gpu_ids([1]) @@ -130,6 +131,26 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): ): self.assertEqual(resolve_requested_gpu_ids([]), [1, 3]) + def test_vulkan_ordinals_bypass_cuda_parent_visible_validation(self): + # Vulkan build on a CPU-only torch host: no CUDA parent-visible set and a + # zero physical count, yet a valid Vulkan ordinal must not be rejected as + # a CUDA physical id (issue #7239). + with ( + patch.dict(os.environ, {}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 0), + ): + # As a CUDA physical id, [0] is outside the empty parent-visible set. + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([0]) + # As Vulkan ordinals, [0] and [0, 1] pass through unchanged. + self.assertEqual(resolve_requested_gpu_ids([0], is_vulkan = True), [0]) + self.assertEqual(resolve_requested_gpu_ids([0, 1], is_vulkan = True), [0, 1]) + # Malformed ordinals are still rejected. + with self.assertRaisesRegex(ValueError, "duplicate GPU IDs"): + resolve_requested_gpu_ids([0, 0], is_vulkan = True) + with self.assertRaisesRegex(ValueError, "non-negative"): + resolve_requested_gpu_ids([-1], is_vulkan = True) + def test_apply_gpu_ids_only_updates_cuda_visible_devices(self): with patch.dict( os.environ, @@ -846,12 +867,177 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): class TestRouteErrors(unittest.TestCase): - def test_prepare_gpu_selection_rejects_gpu_ids_on_non_cuda_backend(self): + def test_prepare_gpu_selection_rejects_gpu_ids_on_non_accelerator_backend(self): with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): with self.assertRaises(ValueError) as exc_info: prepare_gpu_selection([0], model_name = "unsloth/test") - self.assertIn("only supported on CUDA devices", str(exc_info.exception)) + self.assertIn("only supported on CUDA and Intel XPU", str(exc_info.exception)) + + def test_inference_route_resolves_gguf_gpu_ids(self): + # GGUF gpu_ids are now supported: /load routes them through the same + # resolution as non-GGUF loads (rejecting only genuinely invalid ids with + # the resolver's actionable message) rather than a blanket "not supported" + # reject, so /validate can stay consistent with /load (#7239). + import utils.hardware.hardware as hardware_mod + + inference_route = _load_route_module( + "inference_route_module_for_gguf_gpu_ids_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + def _fake_resolve(ids, is_vulkan = False): + raise ValueError("SENTINEL requested GPUs are outside the parent-visible set") + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + # Patch both the package re-export and the defining module so the stub + # fires no matter which import path the route uses. + patch("utils.hardware.resolve_requested_gpu_ids", _fake_resolve), + patch.object(hardware_mod, "resolve_requested_gpu_ids", _fake_resolve), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route._load_model_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + + # The selection was routed through resolution (not the old blanket reject). + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("SENTINEL", exc_info.exception.detail) + self.assertNotIn("not supported for GGUF", exc_info.exception.detail) + + def test_load_rejects_unavailable_vulkan_ordinal_before_training_guard(self): + inference_route = _load_route_module( + "inference_route_module_for_vulkan_preflight_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [99]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch("utils.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(inference_route, "_classify_diffusion_gguf", return_value = None), + patch.object( + inference_route.LlamaCppBackend, + "_is_vulkan_backend", + return_value = True, + ), + patch.object( + inference_route.LlamaCppBackend, + "_find_llama_server_binary", + return_value = "/tmp/llama-server", + ), + patch.object( + inference_route.LlamaCppBackend, + "_get_gpu_memory", + return_value = [(0, 8 * 1024**3, 16 * 1024**3)], + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ) as training_guard, + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route._load_model_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("Vulkan GPU ordinal(s) [99]", exc_info.exception.detail) + training_guard.assert_not_called() + + def test_vulkan_ordinals_are_allowed_on_xpu_hosts(self): + import utils.hardware.hardware as hardware_mod + + inference_route = _load_route_module( + "inference_route_module_for_xpu_vulkan_test", + "routes/inference.py", + ) + config = SimpleNamespace(is_gguf = True) + + with ( + patch("utils.hardware.get_device", return_value = DeviceType.XPU), + patch.object( + inference_route.LlamaCppBackend, + "_is_vulkan_backend", + return_value = True, + ), + patch.object(inference_route, "_classify_diffusion_gguf", return_value = False), + patch.object(hardware_mod, "resolve_requested_gpu_ids", return_value = [0, 1]), + patch.object( + inference_route.LlamaCppBackend, + "_find_llama_server_binary", + return_value = None, + ), + ): + resolved = asyncio.run( + inference_route._resolve_gguf_gpu_ids_for_request(config, [1, 0]) + ) + + self.assertEqual(resolved, [0, 1]) def test_inference_route_validates_gpu_ids_for_gguf(self): # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still @@ -861,7 +1047,7 @@ class TestRouteErrors(unittest.TestCase): import utils.hardware.hardware as hardware_mod inference_route = _load_route_module( - "inference_route_module_for_gguf_gpu_ids_test", + "inference_route_module_for_gguf_gpu_ids_test2", "routes/inference.py", ) request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) @@ -886,6 +1072,17 @@ class TestRouteErrors(unittest.TestCase): "ModelConfig", SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), + # Patch both the package re-export and the defining module so the stub + # fires no matter which import path the route uses. + patch( + "utils.hardware.resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), patch.object( inference_route, "_guard_chat_load_against_training", @@ -893,11 +1090,6 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), - patch.object( - hardware_mod, - "resolve_requested_gpu_ids", - side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), - ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -1439,18 +1631,61 @@ class TestAutoSelectWithNoneRequired(_GpuCacheResetMixin, unittest.TestCase): self.assertEqual(metadata["selection_mode"], "fallback_all") -class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase): - def test_auto_select_returns_non_cuda_for_xpu(self): - with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): +class TestXpuSelection(_GpuCacheResetMixin, unittest.TestCase): + def test_auto_select_supports_xpu(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = (1.0, {}), + ), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "devices": [ + {"index": 0, "vram_total_gb": 8, "vram_used_gb": 1}, + ] + }, + ), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": None, + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0], + ), + ): selected, metadata = auto_select_gpu_ids("unsloth/test") - self.assertIsNone(selected) - self.assertEqual(metadata["selection_mode"], "non_cuda") + self.assertEqual(selected, [0]) + self.assertEqual(metadata["selection_mode"], "auto") - def test_prepare_gpu_selection_rejects_explicit_ids_on_xpu(self): - with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): - with self.assertRaisesRegex(ValueError, "only supported on CUDA"): - prepare_gpu_selection([0], model_name = "unsloth/test") + def test_prepare_gpu_selection_accepts_explicit_ids_on_xpu(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": "0", + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0], + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 1), + ): + selected, metadata = prepare_gpu_selection([0], model_name = "unsloth/test") + + self.assertEqual(selected, [0]) + self.assertEqual(metadata["selection_mode"], "explicit") class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase): diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py index 733933271b..ba6d057123 100644 --- a/studio/backend/tests/test_gpu_selection_sandbox.py +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -294,13 +294,13 @@ class TestAutoSelectGpuIds(unittest.TestCase): # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB self.assertEqual(len(selected), 2) - def test_non_cuda_returns_none(self): + def test_non_accelerator_returns_none(self): from utils.hardware.hardware import auto_select_gpu_ids import utils.hardware.hardware as hw with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): selected, meta = auto_select_gpu_ids("test/model") self.assertIsNone(selected) - self.assertEqual(meta["selection_mode"], "non_cuda") + self.assertEqual(meta["selection_mode"], "non_accelerator") class TestGetDeviceMap(unittest.TestCase): diff --git a/studio/backend/tests/test_hf_cache_settings.py b/studio/backend/tests/test_hf_cache_settings.py new file mode 100644 index 0000000000..1875d61809 --- /dev/null +++ b/studio/backend/tests/test_hf_cache_settings.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import os +import sys +import threading +import time +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from hub.services.models.common import _local_model_info +from utils import hf_cache_settings +from utils import native_path_leases + + +@pytest.fixture() +def settings_store(monkeypatch, tmp_path): + store = {} + monkeypatch.setattr(hf_cache_settings, "_EXPLICIT_CACHE_ENV", {}) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + monkeypatch.setattr( + "storage.studio_db.get_app_setting", + lambda key, fallback = None: store.get(key, fallback), + ) + monkeypatch.setattr( + "storage.studio_db.upsert_app_settings", + lambda values: store.update(values) or values, + ) + return store + + +def test_studio_cache_switch_is_live_and_keeps_history(settings_store, tmp_path): + first = tmp_path / "external-a" / "huggingface" + second = tmp_path / "external-b" / "huggingface" + first.parent.mkdir() + second.parent.mkdir() + + selected = hf_cache_settings.set_hf_cache_home(str(first)) + assert selected.hub_cache == first / "hub" + assert selected.xet_cache == first / "xet" + assert selected.child_env({}) == { + "HF_HUB_CACHE": str(first / "hub"), + "HF_XET_CACHE": str(first / "xet"), + } + + hf_cache_settings.set_hf_cache_home(str(second)) + assert settings_store[hf_cache_settings.CACHE_HISTORY_SETTING_KEY] == [str(first)] + assert first / "hub" in hf_cache_settings.known_hf_hub_caches() + + reset = hf_cache_settings.set_hf_cache_home(None) + assert reset.source == "default" + assert second in hf_cache_settings.known_hf_cache_homes() + + +def test_environment_cache_is_read_only(monkeypatch, tmp_path): + custom = tmp_path / "managed" + monkeypatch.setattr( + hf_cache_settings, + "_EXPLICIT_CACHE_ENV", + {"HF_HOME": str(custom)}, + ) + paths = hf_cache_settings.get_hf_cache_paths() + assert paths.source == "environment" + assert paths.editable is False + assert paths.hub_cache == custom / "hub" + with pytest.raises(RuntimeError, match = "environment variable"): + hf_cache_settings.set_hf_cache_home(str(tmp_path / "other")) + + +def test_explicit_hub_cache_is_the_displayed_location(monkeypatch, tmp_path): + custom_hub = tmp_path / "models-cache" + custom_hub.mkdir() + monkeypatch.setattr( + hf_cache_settings, + "_EXPLICIT_CACHE_ENV", + {"HF_HUB_CACHE": str(custom_hub)}, + ) + + paths = hf_cache_settings.get_hf_cache_paths() + status = hf_cache_settings.cache_status(paths) + + assert paths.cache_home == custom_hub + assert paths.hub_cache == custom_hub + assert status["cache_home"] == str(custom_hub) + assert status["available"] is True + assert custom_hub / "hub" not in hf_cache_settings.known_hf_hub_caches() + + +def test_explicit_hub_cache_display_wins_over_hf_home(monkeypatch, tmp_path): + hf_home = tmp_path / "hf-home" + custom_hub = tmp_path / "other-disk" / "models-cache" + hf_home.mkdir() + custom_hub.mkdir(parents = True) + monkeypatch.setattr( + hf_cache_settings, + "_EXPLICIT_CACHE_ENV", + {"HF_HOME": str(hf_home), "HF_HUB_CACHE": str(custom_hub)}, + ) + + paths = hf_cache_settings.get_hf_cache_paths() + + assert paths.cache_home == custom_hub + assert paths.hub_cache == custom_hub + assert paths.xet_cache == hf_home / "xet" + assert custom_hub / "hub" not in hf_cache_settings.known_hf_hub_caches() + assert hf_home / "hub" in hf_cache_settings.known_hf_hub_caches() + + +def test_xet_only_override_keeps_model_cache_editable(settings_store, monkeypatch, tmp_path): + xet_cache = tmp_path / "chunks" + stored = tmp_path / "stored-cache" + settings_store[hf_cache_settings.CACHE_HOME_SETTING_KEY] = str(stored) + monkeypatch.setattr( + hf_cache_settings, + "_EXPLICIT_CACHE_ENV", + {"HF_XET_CACHE": str(xet_cache)}, + ) + + paths = hf_cache_settings.get_hf_cache_paths() + + assert paths.cache_home == stored + assert paths.hub_cache == stored / "hub" + assert paths.xet_cache == xet_cache + assert paths.editable is True + + selected = tmp_path / "selected-cache" + selected.parent.mkdir(exist_ok = True) + updated = hf_cache_settings.set_hf_cache_home(str(selected)) + assert updated.hub_cache == selected / "hub" + assert updated.xet_cache == xet_cache + + +def test_worker_environment_is_applied_before_import(monkeypatch, tmp_path): + hub = str(tmp_path / "hub") + xet = str(tmp_path / "xet") + observed = {} + + class Module: + @staticmethod + def run(): + import os + return os.environ["HF_HUB_CACHE"], os.environ["HF_XET_CACHE"] + + def fake_import(name): + import os + + observed["name"] = name + observed["hub"] = os.environ.get("HF_HUB_CACHE") + return Module + + monkeypatch.setattr(native_path_leases.importlib, "import_module", fake_import) + result = native_path_leases.run_without_native_path_secret( + "fake.worker", + "run", + {"HF_HUB_CACHE": hub, "HF_XET_CACHE": xet}, + ) + assert observed == {"name": "fake.worker", "hub": hub} + assert result == (hub, xet) + + +def test_spawn_environment_is_applied_then_restored(monkeypatch, tmp_path): + hub = str(tmp_path / "hub") + xet = str(tmp_path / "xet") + monkeypatch.setenv("HF_HUB_CACHE", "parent-hub") + monkeypatch.delenv("HF_XET_CACHE", raising = False) + + with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": hub, "HF_XET_CACHE": xet}): + import os + assert os.environ["HF_HUB_CACHE"] == hub + assert os.environ["HF_XET_CACHE"] == xet + + assert os.environ["HF_HUB_CACHE"] == "parent-hub" + assert "HF_XET_CACHE" not in os.environ + + +def test_spawn_environment_supports_nested_contexts(monkeypatch): + monkeypatch.setenv("HF_HUB_CACHE", "parent") + + with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "outer"}): + assert os.environ["HF_HUB_CACHE"] == "outer" + with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "inner"}): + assert os.environ["HF_HUB_CACHE"] == "inner" + assert os.environ["HF_HUB_CACHE"] == "outer" + + assert os.environ["HF_HUB_CACHE"] == "parent" + + +def test_spawn_environment_serializes_threads(monkeypatch): + monkeypatch.setenv("HF_HUB_CACHE", "parent") + first_entered = threading.Event() + release_first = threading.Event() + observations: list[tuple[str, str]] = [] + + def first(): + with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "first"}): + observations.append(("first", os.environ["HF_HUB_CACHE"])) + first_entered.set() + assert release_first.wait(timeout = 2) + + def second(): + assert first_entered.wait(timeout = 2) + with hf_cache_settings.child_environment_for_spawn({"HF_HUB_CACHE": "second"}): + observations.append(("second", os.environ["HF_HUB_CACHE"])) + + first_thread = threading.Thread(target = first) + second_thread = threading.Thread(target = second) + first_thread.start() + second_thread.start() + assert first_entered.wait(timeout = 2) + time.sleep(0.02) + assert observations == [("first", "first")] + release_first.set() + first_thread.join(timeout = 2) + second_thread.join(timeout = 2) + + assert observations == [("first", "first"), ("second", "second")] + assert os.environ["HF_HUB_CACHE"] == "parent" + + +def test_cache_switch_invalidates_inventory(settings_store, tmp_path, monkeypatch): + invalidations = [] + monkeypatch.setattr( + "hub.utils.inventory_scan.invalidate_hf_cache_scans", + lambda: invalidations.append(True), + ) + selected = tmp_path / "external" / "huggingface" + selected.parent.mkdir() + + hf_cache_settings.set_hf_cache_home(str(selected)) + + assert invalidations == [True] + + +def test_cache_validation_write_tests_hub_and_xet(settings_store, tmp_path, monkeypatch): + selected = tmp_path / "external" / "huggingface" + selected.parent.mkdir() + tested = [] + real_named_temporary_file = hf_cache_settings.tempfile.NamedTemporaryFile + + def recording_write_test(*args, **kwargs): + tested.append(Path(kwargs["dir"])) + return real_named_temporary_file(*args, **kwargs) + + monkeypatch.setattr( + hf_cache_settings.tempfile, + "NamedTemporaryFile", + recording_write_test, + ) + + hf_cache_settings.set_hf_cache_home(str(selected)) + + assert tested == [selected / "hub", selected / "xet"] + + +def test_cache_validation_rejects_unwritable_child(settings_store, tmp_path, monkeypatch): + selected = tmp_path / "external" / "huggingface" + selected.parent.mkdir() + + def reject_hub(*args, **kwargs): + if Path(kwargs["dir"]).name == "hub": + raise PermissionError("read-only") + raise AssertionError("xet should not be tested after hub fails") + + monkeypatch.setattr(hf_cache_settings.tempfile, "NamedTemporaryFile", reject_hub) + + with pytest.raises(ValueError, match = "permission"): + hf_cache_settings.set_hf_cache_home(str(selected)) + + +def test_inactive_cache_model_loads_from_snapshot_path(tmp_path): + snapshot = tmp_path / "snapshots" / "revision" + snapshot.mkdir(parents = True) + row = _local_model_info( + scan_path = snapshot, + load_path = snapshot, + source = "hf_cache", + model_format = "safetensors", + model_id = "org/model", + active_cache = False, + ) + assert row.model_id == "org/model" + assert row.active_cache is False + assert row.load_id == str(snapshot) diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 48aff29659..a037ea2579 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -101,13 +101,23 @@ def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", - lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), + lambda repo_type, repo_id, mode, *a, **k: prepared.append( + (repo_type, repo_id, mode, k.get("root")) + ), ) - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + selected_cache = "/captured/hub" + out = xf.hf_hub_download_with_xet_fallback( + DL_REPO, + FILE, + None, + cache_dir = selected_cache, + ) assert out == "/cache/model.gguf" assert seen_disable_xet == [False, True] # Xet first, then HTTP - assert prepared == [("model", DL_REPO, "http")], "shim must run Unsloth's marker-aware prep" + assert prepared == [ + ("model", DL_REPO, "http", Path(selected_cache)) + ], "shim must prepare the cache captured by the download" def test_shim_snapshot_injects_studio_prepare(monkeypatch): @@ -120,10 +130,22 @@ def test_shim_snapshot_injects_studio_prepare(monkeypatch): return "/tmp/snap-dir" monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) - out = xf.snapshot_download_with_xet_fallback("org/model") + selected_cache = "/captured/hub" + out = xf.snapshot_download_with_xet_fallback( + "org/model", + cache_dir = selected_cache, + ) assert out == "/tmp/snap-dir" assert captured["repo_id"] == "org/model" - assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http + prepared = [] + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", + lambda repo_type, repo_id, mode, *a, **k: prepared.append( + (repo_type, repo_id, mode, k.get("root")) + ), + ) + captured["prepare_for_http_fn"]("model", "org/model") + assert prepared == [("model", "org/model", "http", Path(selected_cache))] def test_degrades_gracefully_without_shared_helper(monkeypatch): diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index 3ebad861ad..02ccc68b11 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -200,14 +200,9 @@ def _gpu_linux_host(caps): ) -def test_host_is_blackwell_includes_datacenter_parts(): - assert ilp._host_is_blackwell(_gpu_linux_host(["10.0"])) is True # B200 sm_100 - assert ilp._host_is_blackwell(_gpu_linux_host(["10.3"])) is True # B300 sm_103 - assert ilp._host_is_blackwell(_gpu_linux_host(["12.0"])) is True # RTX 50 sm_120 - assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121 - assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper - assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere - assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins +# _host_is_blackwell / _blackwell_min_toolkit_for_host are prebuilt_core +# re-exports; their value tables moved verbatim to +# tests/studio/install/test_prebuilt_core.py. def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile): @@ -285,16 +280,6 @@ def test_drop_blackwell_incapable_windows_cuda_applies_to_datacenter(): assert [a.name for a in kept] == [cuda13.name] -def test_blackwell_min_toolkit_is_sm_aware(): - # Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it. - f = ilp._blackwell_min_toolkit_for_host - assert f(_gpu_linux_host(["10.0"])) == (12, 8) # B200 - assert f(_gpu_linux_host(["12.0"])) == (12, 8) # RTX 50 - assert f(_gpu_linux_host(["10.3"])) == (12, 9) # B300 - assert f(_gpu_linux_host(["12.1"])) == (12, 9) # DGX Spark - assert f(_gpu_linux_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins - - def test_sm103_host_drops_cuda128_windows_build(): # B300 (sm_103) needs cuda-12.9: a legacy win-cuda-12.8 build must be dropped. host = _host( diff --git a/studio/backend/tests/test_install_whisper_prebuilt_checksums.py b/studio/backend/tests/test_install_whisper_prebuilt_checksums.py new file mode 100644 index 0000000000..19bece9d0c --- /dev/null +++ b/studio/backend/tests/test_install_whisper_prebuilt_checksums.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Trust-anchor tests for install_whisper_prebuilt.py. + +Whisper verifies each download against the release's own +whisper-prebuilt-sha256.json checksum index (the same model as +install_llama_prebuilt.py), not a committed pins file. These pin the index +parser, the fail-closed behaviour when an asset is not covered, the +tampered-manifest guard, and the newest-release resolution. +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path + +import pytest + +_studio = Path(__file__).resolve().parent.parent.parent +if str(_studio) not in sys.path: + sys.path.insert(0, str(_studio)) + +iwp = importlib.import_module("install_whisper_prebuilt") + +if not hasattr(iwp, "parse_release_checksums"): + pytest.skip("checksum-model symbols not present - check branch", allow_module_level = True) + +_A = "0" * 64 +_B = "1" * 64 +_TAG = "v1.9.1-unsloth.1" +_REPO = "unslothai/whisper.cpp" + + +def _index(**overrides) -> dict: + payload = { + "schema_version": 1, + "component": "whisper.cpp", + "release_tag": _TAG, + "upstream_tag": "v1.9.1", + "artifacts": { + "whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz": {"sha256": _A}, + "whisper-v1.9.1-unsloth.1-linux-x64-cuda12-portable.tar.gz": {"sha256": _B}, + }, + } + payload.update(overrides) + return payload + + +# parse_release_checksums / expected_sha256_for are prebuilt_core re-exports; +# their valid/fail-closed matrix is asserted against the real whisper +# descriptor in tests/studio/install/test_prebuilt_core.py. The download-host +# fast-path tests below still route through this module's parse wrapper. + +# release tag resolution. + + +def test_resolve_release_tag_explicit_override_passthrough(): + assert iwp.resolve_release_tag(_REPO, published_release_tag = "v1.9.1-unsloth.2") == ( + "v1.9.1-unsloth.2" + ) + + +def test_resolve_release_tag_resolves_newest_when_no_override(monkeypatch): + monkeypatch.setattr(iwp, "resolve_newest_release_tag", lambda repo: "v9.9.9-unsloth.9") + assert iwp.resolve_release_tag(_REPO, published_release_tag = None) == "v9.9.9-unsloth.9" + + +def test_resolve_newest_release_tag_picks_latest_published(monkeypatch): + releases = [ + {"tag_name": "v1.9.1-unsloth.1", "published_at": "2026-01-01T00:00:00Z"}, + {"tag_name": "v1.9.1-unsloth.3", "published_at": "2026-03-01T00:00:00Z"}, + {"tag_name": "v1.9.1-unsloth.2", "published_at": "2026-02-01T00:00:00Z"}, + {"tag_name": "draft", "published_at": "2026-09-01T00:00:00Z", "draft": True}, + {"tag_name": "pre", "published_at": "2026-09-01T00:00:00Z", "prerelease": True}, + ] + monkeypatch.setattr(iwp, "fetch_json", lambda url: releases) + assert iwp.resolve_newest_release_tag(_REPO) == "v1.9.1-unsloth.3" + + +def test_resolve_newest_release_tag_none_published_fails_closed(monkeypatch): + monkeypatch.setattr(iwp, "fetch_json", lambda url: [{"tag_name": "d", "draft": True}]) + with pytest.raises(iwp.PrebuiltFallback): + iwp.resolve_newest_release_tag(_REPO) + + +def test_pins_symbols_are_gone(): + # The committed-pins trust model was removed in favour of llama's runtime index. + for gone in ("load_pins", "pins_path", "resolve_expected_sha256", "PINS_FILENAME"): + assert not hasattr(iwp, gone), f"{gone} should have been removed" + + +# Download-host fast path (resolve + fetch the JSON assets with no GitHub API). + +_CPU_ASSET = "whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz" + + +def _manifest() -> dict: + return { + "schema_version": 1, + "component": "whisper.cpp", + "upstream_tag": "v1.9.1", + "artifacts": [{"asset": _CPU_ASSET, "os": "linux", "arch": "x64", "backend": "cpu"}], + } + + +def _no_api(monkeypatch): + """Fail loudly if any code path touches api.github.com.""" + + def _boom(*a, **k): + raise AssertionError("api.github.com was used on the fast path") + + monkeypatch.setattr(iwp, "fetch_json", _boom) + monkeypatch.setattr(iwp, "github_release", _boom) + monkeypatch.setattr(iwp, "fetch_release_bundle", _boom) + + +def test_fetch_release_for_install_prefers_download_host(monkeypatch): + _no_api(monkeypatch) + monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG) + + def _dhj(url): + if url.endswith(iwp.SHA256_ASSET_NAME): + return _index() + if url.endswith(iwp.MANIFEST_ASSET_NAME): + return _manifest() + raise AssertionError(f"unexpected url {url}") + + monkeypatch.setattr(iwp, "_download_host_json", _dhj) + bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = None) + assert bundle.release_tag == _TAG + assert checks[_CPU_ASSET] == _A + # asset_urls point at the download host (github.com), not the API. + assert bundle.asset_urls[iwp.SHA256_ASSET_NAME].startswith( + f"https://github.com/{_REPO}/releases/" + ) + assert bundle.asset_urls[_CPU_ASSET].startswith( + f"https://github.com/{_REPO}/releases/download/" + ) + walked = iwp._fetch_release_candidate(_REPO, _TAG) + assert iwp.SHA256_ASSET_NAME in walked.asset_urls + assert _CPU_ASSET in walked.asset_urls + + +def test_fetch_release_for_install_explicit_tag_skips_the_head(monkeypatch): + # An explicit tag needs no /releases/latest HEAD: resolving it must not call it. + monkeypatch.setattr( + iwp, + "_download_host_latest_release_tag", + lambda repo: (_ for _ in ()).throw(AssertionError("HEAD used for an explicit tag")), + ) + monkeypatch.setattr( + iwp, + "_download_host_json", + lambda url: _index() if url.endswith(iwp.SHA256_ASSET_NAME) else _manifest(), + ) + _no_api(monkeypatch) + bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = _TAG) + assert bundle.release_tag == _TAG + + +def test_fetch_release_for_install_falls_back_to_api(monkeypatch): + # Fast path returns None (e.g. a 404) -> the API path resolves the release. + monkeypatch.setattr(iwp, "_resolve_release_via_download_host", lambda repo, tag: None) + sentinel = iwp.ReleaseBundle(repo = _REPO, release_tag = _TAG, manifest = _manifest(), asset_urls = {}) + monkeypatch.setattr(iwp, "resolve_release_tag", lambda repo, *, published_release_tag: _TAG) + monkeypatch.setattr(iwp, "fetch_release_bundle", lambda repo, tag: sentinel) + monkeypatch.setattr(iwp, "fetch_release_checksums", lambda bundle: {_CPU_ASSET: _A}) + bundle, checks = iwp.fetch_release_for_install(_REPO, published_release_tag = None) + assert bundle is sentinel + assert checks == {_CPU_ASSET: _A} + + +def test_resolve_via_download_host_sha_404_returns_none(monkeypatch): + import urllib.error + + monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG) + + def _dhj(url): + raise urllib.error.HTTPError(url, 404, "not found", {}, None) + + monkeypatch.setattr(iwp, "_download_host_json", _dhj) + assert iwp._resolve_release_via_download_host(_REPO, None) is None + + +def test_resolve_via_download_host_tag_mismatch_returns_none(monkeypatch): + # A checksum index whose self-reported release_tag disagrees is rejected (None). + monkeypatch.setattr(iwp, "_download_host_latest_release_tag", lambda repo: _TAG) + monkeypatch.setattr( + iwp, "_download_host_json", lambda url: _index(release_tag = "v1.9.1-unsloth.2") + ) + assert iwp._resolve_release_via_download_host(_REPO, None) is None + + +def test_download_host_latest_release_tag_parses_redirect(monkeypatch): + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def geturl(self): + return f"https://github.com/{_REPO}/releases/tag/{_TAG}" + + class _Opener: + def open( + self, + req, + timeout = None, + ): + return _Resp() + + monkeypatch.setattr(iwp, "_URL_OPENER", _Opener()) + assert iwp._download_host_latest_release_tag(_REPO) == _TAG + + +def test_download_host_latest_release_tag_404_returns_none(monkeypatch): + import urllib.error + + class _Opener: + def open( + self, + req, + timeout = None, + ): + raise urllib.error.HTTPError(req.full_url, 404, "nf", {}, None) + + monkeypatch.setattr(iwp, "_URL_OPENER", _Opener()) + assert iwp._download_host_latest_release_tag(_REPO) is None diff --git a/studio/backend/tests/test_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py index b735bd1132..8373cdd6bb 100644 --- a/studio/backend/tests/test_linux_external_media_paths.py +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -254,8 +254,10 @@ def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tm ) fake_external_media = SimpleNamespace( linux_run_media_mount_roots = lambda: [media_root], + macos_volume_roots = lambda: [], windows_drive_roots = lambda: [], ) + fake_paths.external_media = fake_external_media fake_studio_db = SimpleNamespace( list_scan_folders = lambda: [], contains_sensitive_path_component = studio_db.contains_sensitive_path_component, diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 049058e511..4332a440a5 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -335,3 +335,24 @@ class TestRetryContract: 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 + + +class TestMmprojRetryFailureMessage: + """#7302: bare mmproj crashes must not be reported as projector-format.""" + + def test_confirmed_projector_keeps_historical_wording(self): + msg = LlamaCppBackend._mmproj_retry_failure_message( + projector_confirmed = True, + detail = "llama-server failed to start", + ) + assert msg.startswith("Vision projector incompatible with this llama.cpp") + assert "llama-server failed to start" in msg + + def test_bare_crash_does_not_claim_projector_incompatibility(self): + msg = LlamaCppBackend._mmproj_retry_failure_message( + projector_confirmed = False, + detail = "llama-server failed to start. Check that the GGUF file is valid", + ) + assert "Vision projector incompatible" not in msg + assert "crashed with --mmproj" in msg + assert "GGUF file is valid" in msg diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 1d15647967..27c1b17a85 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -637,7 +637,9 @@ def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch 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 = "") + return _types.SimpleNamespace( + stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "", returncode = 0 + ) monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run) @@ -678,6 +680,95 @@ def test_probe_server_capabilities_reports_outdated_binary(tmp_path): assert caps["found"] is True assert caps["mtp_token"] is None assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_reads_mtp_from_multiline_help(tmp_path): + # Enum on the indented line: first-line-only probing falsely reported + # "lacks MTP" (#7302). + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--spec-type TYPE\n" + " speculative decoding type\n" + " (none,draft-simple,draft-mtp,ngram-mod)\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["mtp_token"] == "draft-mtp" + assert caps["supports_mtp"] is True + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_empty_help_fails_open(tmp_path): + # --help prints nothing: must not claim the prebuilt lacks MTP (#7302). + fake = tmp_path / "llama-server" + fake.write_text("#!/usr/bin/env bash\nexit 0\n") + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +@_NEEDS_BASH +def test_probe_server_capabilities_no_spec_type_is_definitive(tmp_path): + # Nonempty --help without --spec-type: pre-spec binary, not inconclusive. + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--gpu-layers N\n GPU layers to offload\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_failed_help_with_output_is_inconclusive(tmp_path): + fake = tmp_path / "llama-server" + fake.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "--help" ]; then\n' + " echo 'illegal instruction'\n" + " exit 1\n" + "fi\n" + ) + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +@_NEEDS_BASH +def test_probe_server_capabilities_crash_on_help_fails_open(tmp_path): + fake = tmp_path / "llama-server" + fake.write_text("#!/usr/bin/env bash\nkill -SEGV $$\n") + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +def test_mtp_token_from_spec_help_prefers_draft_mtp(): + assert ( + LlamaCppBackend._mtp_token_from_spec_help("--spec-type none,draft-mtp,mtp,ngram-mod") + == "draft-mtp" + ) + assert LlamaCppBackend._mtp_token_from_spec_help("--spec-type [none|mtp|ngram-cache]") == "mtp" + assert LlamaCppBackend._mtp_token_from_spec_help("--spec-type none,ngram-mod") is None + # No incidental substring matches. + assert LlamaCppBackend._mtp_token_from_spec_help("prompt cache") is None def test_probe_server_capabilities_handles_missing_binary(): @@ -685,6 +776,7 @@ 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["mtp_probe_inconclusive"] is True assert caps["supports_cache_ram"] is False assert caps["supports_ctx_checkpoints"] is False assert caps["supports_no_cache_prompt"] is False @@ -1176,12 +1268,14 @@ def _resolver_backend( *, ngram_supported = True, mtp_token = "draft-mtp", + mtp_probe_inconclusive = False, ): """Backend with a deterministic probe so the resolver is hermetic.""" fake = { "found": True, "mtp_token": mtp_token, "supports_mtp": bool(mtp_token), + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": "new" if ngram_supported else None, "supports_ngram_mod": bool(ngram_supported), "spec_draft_n_max_flag": "--spec-draft-n-max", @@ -1879,6 +1973,24 @@ def test_spec_fallback_reason_set_when_binary_lacks_mtp(monkeypatch): assert backend.spec_fallback_reason == "binary_no_mtp" +def test_spec_fallback_reason_none_when_mtp_probe_inconclusive(monkeypatch): + backend = _resolver_backend( + monkeypatch, + mtp_token = None, + mtp_probe_inconclusive = True, + ) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason is None + + def test_spec_fallback_reason_none_when_mtp_engages(monkeypatch): backend = _resolver_backend(monkeypatch) backend._build_speculative_flags( diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index f12384231f..9e23242b97 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -119,6 +119,9 @@ def _clean_state(monkeypatch, tmp_path): monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) # Never hit the network in these tests. monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + # Keep the whisper piggyback out of the llama-only tests: no host probe, no + # whisper phase (test_combined_update.py covers the chained flow). + monkeypatch.setattr(upd, "_whisper_chain_status", lambda **kwargs: None) yield freshness.reset_caches() upd._reset_job_for_tests() diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py index 0ecfeee018..cc450d55cc 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -100,6 +100,21 @@ def test_status_response_exposes_update_size_bytes(): assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None +def test_status_response_exposes_update_component(): + model = rl.LlamaUpdateStatusResponse( + supported = True, + update_available = True, + llama_update_available = False, + update_component = "whisper", + whisper = { + "update_available": True, + "installed_tag": "v1", + "latest_tag": "v2", + }, + ) + assert model.model_dump()["update_component"] == "whisper" + + def test_status_handler_runs_off_event_loop(monkeypatch): seen = {} diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py index 6b44f61972..79c9977c84 100644 --- a/studio/backend/tests/test_local_llama_cpp_link.py +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -21,6 +21,13 @@ from utils import llama_cpp_update as u from core.inference.llama_cpp import LlamaCppBackend +@pytest.fixture(autouse = True) +def _no_whisper_piggyback(monkeypatch): + # Keep the whisper piggyback probe off the host: these tests exercise the + # llama local-link contract only. + monkeypatch.setattr(u, "_whisper_chain_status", lambda **kwargs: None) + + def _make_link(link: Path, target: Path) -> None: """Create a directory junction (Windows) / symlink (POSIX); neither needs elevation.""" diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 11aeee6d77..36061b5375 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -14,6 +14,7 @@ import pytest from fastapi import FastAPI, HTTPException, Request from fastapi.responses import Response from fastapi.testclient import TestClient +from starlette.middleware.gzip import GZipMiddleware _BACKEND_ROOT = Path(__file__).resolve().parents[1] @@ -33,6 +34,7 @@ def main_module(): def _make_protected_app( max_bytes: int, main_module, + request_max_bytes_getter = None, upload_passthrough_prefixes: tuple = (), upload_passthrough_max_bytes_getter = None, ): @@ -40,7 +42,13 @@ def _make_protected_app( app.add_middleware( main_module.MaxBodyMiddleware, max_bytes_getter = lambda: max_bytes, - protected_prefixes = ("/v1/chat/completions", "/api/settings", "/api/train"), + protected_prefixes = ( + "/v1/chat/completions", + "/api/inference", + "/api/settings", + "/api/train", + ), + request_max_bytes_getter = request_max_bytes_getter, upload_passthrough_prefixes = upload_passthrough_prefixes, upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter, ) @@ -67,6 +75,10 @@ def _make_protected_app( total += len(chunk) return {"ok": True, "chunks": chunks, "total": total} + @app.post("/api/inference/audio/transcribe/raw") + async def transcribe_raw(request: Request): + return {"ok": True, "total": len(await request.body())} + @app.get("/api/train/status") async def status_get(): return {"ok": True, "get": True} @@ -96,6 +108,43 @@ class TestMaxBodyMiddleware: assert r.status_code == 200 assert r.json()["unprotected"] is True + def test_route_specific_cap_overrides_default(self, main_module): + app = _make_protected_app( + 4096, + main_module, + request_max_bytes_getter = lambda path: ( + 128 if path.endswith("/transcribe/raw") else 4096 + ), + ) + c = TestClient(app) + + rejected = c.post( + "/api/inference/audio/transcribe/raw", + content = b"x" * 129, + ) + accepted = c.post( + "/api/inference/audio/transcribe/raw", + content = b"x" * 128, + ) + + assert rejected.status_code == 413 + assert accepted.status_code == 200 + assert accepted.json()["total"] == 128 + + def test_stt_routes_use_audio_specific_caps(self, main_module): + from utils.upload_limits import ( + STT_AUDIO_JSON_MAX_BYTES, + STT_AUDIO_RAW_MAX_BYTES, + ) + assert ( + main_module._get_request_body_max_bytes("/api/inference/audio/transcribe/raw") + == STT_AUDIO_RAW_MAX_BYTES + ) + assert ( + main_module._get_request_body_max_bytes("/api/inference/audio/transcribe") + == STT_AUDIO_JSON_MAX_BYTES + ) + def test_settings_put_body_over_cap_rejected(self, main_module): app = _make_protected_app(1024, main_module) c = TestClient(app) @@ -471,6 +520,71 @@ class TestSecurityHeadersMiddleware: assert b"server" in names +class TestFrontendAssets: + def test_hashed_assets_are_compressed_and_cached(self, tmp_path, main_module): + content = b"export const value = 'responsive';\n" * 200 + (tmp_path / "page-abc123.js").write_bytes(content) + app = FastAPI() + assets_app = GZipMiddleware( + main_module.ImmutableStaticFiles(directory = tmp_path), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") + + response = TestClient(app).get( + "/assets/page-abc123.js", + headers = {"Accept-Encoding": "gzip"}, + ) + + assert response.status_code == 200 + assert response.content == content + assert response.headers["content-encoding"] == "gzip" + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + assert "accept-encoding" in response.headers["vary"].lower() + + def test_asset_revalidation_keeps_immutable_cache_header(self, tmp_path, main_module): + (tmp_path / "page-abc123.js").write_text("export {};", encoding = "utf-8") + app = FastAPI() + app.mount( + "/assets", + main_module.ImmutableStaticFiles(directory = tmp_path), + name = "assets", + ) + client = TestClient(app) + first = client.get("/assets/page-abc123.js") + + response = client.get( + "/assets/page-abc123.js", + headers = {"If-None-Match": first.headers["etag"]}, + ) + + assert response.status_code == 304 + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + + def test_range_request_is_not_compressed(self, tmp_path, main_module): + content = b"export const value = 'responsive';\n" * 200 + (tmp_path / "page-abc123.js").write_bytes(content) + app = FastAPI() + assets_app = main_module._AssetGZipMiddleware( + main_module.ImmutableStaticFiles(directory = tmp_path), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") + + response = TestClient(app).get( + "/assets/page-abc123.js", + headers = {"Accept-Encoding": "gzip", "Range": "bytes=0-99"}, + ) + + assert response.status_code == 206 + assert response.headers.get("content-encoding") != "gzip" + assert response.headers["content-range"] == f"bytes 0-99/{len(content)}" + assert response.content == content[:100] + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + + # /api/health auth gate diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index fafaea0043..d49a2281a0 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -922,3 +922,413 @@ def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch): "vision", "vision answer", ] + + +class _FakeLRUPromptCache: + def __init__( + self, + max_size = 10, + max_bytes = 1 << 63, + ): + self.max_size = max_size + self.max_bytes = max_bytes + self.entries = {} + + def fetch_nearest_cache(self, key, tokens): + import copy + + stored = self.entries.get(key, {}) + exact = stored.get(tuple(tokens)) + if exact is not None: + return copy.deepcopy(exact), [] + best = None + for candidate, cache in stored.items(): + if len(candidate) < len(tokens) and tuple(tokens[: len(candidate)]) == candidate: + if best is None or len(candidate) > len(best[0]): + best = (candidate, cache) + if best is not None: + return copy.deepcopy(best[1]), list(tokens[len(best[0]) :]) + return None, list(tokens) + + def insert_cache( + self, + key, + tokens, + prompt_cache, + *, + cache_type = "assistant", + ): + import copy + self.entries.setdefault(key, {})[tuple(tokens)] = copy.deepcopy(prompt_cache) + + +class _FakeCacheEntry: + def __init__( + self, + offset = 0, + nbytes = 1, + ): + self.offset = offset + self.nbytes = nbytes + + +def _install_fake_prompt_cache_api(monkeypatch, trimmable = True): + from core.inference import mlx_inference + + def _make_prompt_cache(_model): + return [_FakeCacheEntry()] + + def _can_trim_prompt_cache(_cache): + return trimmable + + def _trim_prompt_cache(cache, num): + cache[0].offset = max(cache[0].offset - num, 0) + return num + + monkeypatch.setattr( + mlx_inference, + "_mlx_prompt_cache_api", + lambda: ( + _FakeLRUPromptCache, + _make_prompt_cache, + _can_trim_prompt_cache, + _trim_prompt_cache, + ), + ) + + +def test_mlx_prompt_cache_max_bytes_budget(monkeypatch): + from core.inference.mlx_inference import ( + PROMPT_CACHE_FALLBACK_BYTES, + PROMPT_CACHE_MEMORY_FRACTION, + _prompt_cache_max_bytes, + ) + + monkeypatch.delenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", raising = False) + assert _prompt_cache_max_bytes(None) == PROMPT_CACHE_FALLBACK_BYTES + assert _prompt_cache_max_bytes(20.0) == int(20.0 * 1e9 * PROMPT_CACHE_MEMORY_FRACTION) + + monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "4096") + assert _prompt_cache_max_bytes(20.0) == 4096 + monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "0") + assert _prompt_cache_max_bytes(20.0) == 0 + monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "not-a-number") + assert _prompt_cache_max_bytes(20.0) == int(20.0 * 1e9 * PROMPT_CACHE_MEMORY_FRACTION) + + +def test_mlx_prompt_cache_never_returns_empty_remainder(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + from core.inference.mlx_inference import _MLXPromptCacheHistory + + history = _MLXPromptCacheHistory(6, 1 << 30) + tokens = list(range(10)) + cache, rest = history.fetch(object(), "key", tokens) + assert len(rest) == 10 + cache[0].offset = len(tokens) + history.insert("key", tokens, cache) + + _cache, rest = history.fetch(object(), "key", tokens) + assert rest == tokens[-1:] + + longer = tokens + [99, 100] + _cache, rest = history.fetch(object(), "key", longer) + assert rest == [99, 100] + + _install_fake_prompt_cache_api(monkeypatch, trimmable = False) + history = _MLXPromptCacheHistory(6, 1 << 30) + cache, _rest = history.fetch(object(), "key", tokens) + cache[0].offset = len(tokens) + history.insert("key", tokens, cache) + _cache, rest = history.fetch(object(), "key", tokens) + assert rest == tokens, "untrimmable entry must not be reused" + + +def test_mlx_prompt_cache_key_isolates_adapter_state(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + class _Tok: + bos_token = None + + def encode( + self, + text, + add_special_tokens = True, + ): + return [ord(c) for c in text] + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = _Tok() + backend.active_model_name = "model-a" + + prompt = "shared prefix" + _rest, cache, key, tokens, cached = backend._prepare_prompt_cache(prompt, True) + assert cached == 0 + cache[0].offset = len(tokens) + backend._prompt_cache_history.insert(key, tokens, cache) + + _rest, _cache, _key, _tokens, cached_same = backend._prepare_prompt_cache(prompt, True) + assert cached_same > 0 + _rest, _cache, _key, _tokens, cached_flipped = backend._prepare_prompt_cache(prompt, False) + assert cached_flipped == 0 + + +def _install_fake_text_stack( + monkeypatch, + token_map, + captured, + markers = None, +): + import types as _types + + from core.inference import mlx_inference + + _install_fake_mlx(monkeypatch) + monkeypatch.setattr( + mlx_inference, + "_temporary_mlx_adapter_state", + lambda _model, _state: __import__("contextlib").nullcontext(), + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda _tok, messages, **_kw: messages[-1]["content"], + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.render_with_native_template_fallback", + lambda formatted_prompt, **_kw: SimpleNamespace( + prompt = formatted_prompt, + reasoning_channel_markers = markers, + ), + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_a, **_kw: "", + ) + + class _Resp: + def __init__(self, token, processed): + self.token = token + self.text = f"<{token}>" + self.prompt_tokens = processed + self.prompt_tps = 10.0 + self.generation_tokens = 1 + self.generation_tps = 5.0 + + def _stream_generate(_model, _tokenizer, **kwargs): + captured.append(kwargs) + processed = len(kwargs["prompt"]) + cache = kwargs.get("prompt_cache") + if cache is not None: + cache[0].offset += processed + for token in token_map["generated"]: + if cache is not None: + cache[0].offset += 1 + yield _Resp(token, processed) + + mlx_lm_pkg = _types.ModuleType("mlx_lm") + mlx_lm_pkg.stream_generate = _stream_generate + mlx_lm_sample = _types.ModuleType("mlx_lm.sample_utils") + mlx_lm_sample.make_sampler = lambda **_kw: object() + mlx_lm_sample.make_logits_processors = lambda **_kw: [] + monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample) + + class _Tok: + bos_token = None + chat_template = "x" + + def encode( + self, + text, + add_special_tokens = True, + ): + return list(token_map[text]) + + def decode( + self, + ids, + skip_special_tokens = False, + ): + return "".join(str(i) for i in ids) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = _Tok() + backend._is_vlm = False + backend.active_model_name = "model-a" + return backend + + +def _run_turn(backend, prompt): + list( + backend.generate_chat_response( + messages = [{"role": "user", "content": prompt}], + max_new_tokens = 4, + ) + ) + + +def test_mlx_text_reuses_prompt_cache_on_the_next_turn(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + captured = [] + token_map = { + "P1": [1, 2, 3], + "P2": [1, 2, 3, 7, 8, 9, 10], + "generated": [7, 8], + } + backend = _install_fake_text_stack(monkeypatch, token_map, captured) + + _run_turn(backend, "P1") + assert captured[0]["prompt"] == [1, 2, 3] + assert "prompt_cache" in captured[0] + assert backend.last_generation_stats["timings"]["cache_n"] == 0 + + _run_turn(backend, "P2") + assert captured[1]["prompt"] == [9, 10], "turn two should prefill only the new tail" + + stats = backend.last_generation_stats + assert stats["timings"]["cache_n"] == 5 + assert stats["timings"]["prompt_n"] == 2 + assert stats["usage"]["prompt_tokens"] == 7 + + +def test_mlx_text_without_lru_prompt_cache_prefills_the_full_prompt(monkeypatch): + from core.inference import mlx_inference + + monkeypatch.setattr(mlx_inference, "_mlx_prompt_cache_api", lambda: None) + captured = [] + token_map = {"P1": [1, 2, 3], "generated": [7]} + backend = _install_fake_text_stack(monkeypatch, token_map, captured) + + _run_turn(backend, "P1") + assert captured[0]["prompt"] == "P1" + assert "prompt_cache" not in captured[0] + assert backend.last_generation_stats["timings"]["cache_n"] == 0 + + +def test_mlx_text_tracks_tokens_on_the_native_reasoning_path(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + captured = [] + token_map = {"P1": [1, 2, 3], "P2": [1, 2, 3, 7, 8, 9], "generated": [7, 8]} + backend = _install_fake_text_stack(monkeypatch, token_map, captured, markers = ("", "")) + + _run_turn(backend, "P1") + _run_turn(backend, "P2") + assert captured[1]["prompt"] == [9] + + +def test_mlx_presence_penalty_latches_the_first_decode_step(): + mx = pytest.importorskip("mlx.core") + import numpy as np + + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + processor = _make_mlx_presence_penalty_processor(2.0) + logits = mx.zeros((1, 5)) + out = processor(mx.array([3]), logits) + assert np.array_equal(np.array(out), np.zeros((1, 5))), "prompt must not be penalized" + out = processor(mx.array([3, 1]), mx.zeros((1, 5))) + penalized = np.array(out)[0] + assert penalized[1] == -2.0 + assert penalized[3] == 0.0 + + +def test_mlx_prompt_cache_survives_reset_but_not_unload(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + _install_fake_mlx(monkeypatch) + sys.modules["mlx.core"].clear_cache = lambda: None + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + backend.active_model_name = "model-a" + history = backend._prompt_cache() + assert history is not None + + backend.reset_generation_state() + assert backend._prompt_cache_history is history + + backend.unload_model("model-a") + assert backend._prompt_cache_history is None + + +def test_mlx_prompt_cache_skips_entries_over_budget(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + from core.inference.mlx_inference import _MLXPromptCacheHistory + + history = _MLXPromptCacheHistory(6, 1000) + history.insert("key", [1, 2, 3], [_FakeCacheEntry(offset = 3, nbytes = 400)]) + assert len(history._lru.entries.get("key", {})) == 1 + + history.insert("key", list(range(50)), [_FakeCacheEntry(offset = 50, nbytes = 5000)]) + stored = history._lru.entries.get("key", {}) + assert tuple([1, 2, 3]) in stored + assert tuple(range(50)) not in stored + + +def test_mlx_prompt_cache_keys_on_what_the_kv_covers(monkeypatch): + _install_fake_prompt_cache_api(monkeypatch) + from core.inference.mlx_inference import _MLXPromptCacheHistory + + class _Entry: + def __init__( + self, + offset, + nbytes = 1, + ): + self.offset = offset + self.nbytes = nbytes + + history = _MLXPromptCacheHistory(6, 1 << 30) + + history.insert("key", list(range(10)), [_Entry(offset = 8)]) + assert tuple(range(8)) in history._lru.entries["key"] + assert tuple(range(10)) not in history._lru.entries["key"] + + history.insert("other", list(range(4)), [_Entry(offset = 9)]) + assert "other" not in history._lru.entries + + +def test_mlx_prompt_cache_only_stores_verifiable_prefix_coverage(monkeypatch): + mx = pytest.importorskip("mlx.core") + from mlx_lm.models.cache import CacheList, ChunkedKVCache, KVCache, RotatingKVCache + + _install_fake_prompt_cache_api(monkeypatch) + from core.inference.mlx_inference import _kv_prefix_coverage, _MLXPromptCacheHistory + + def feed(entry, n): + for _ in range(n): + block = mx.zeros((1, 2, 1, 4), dtype = mx.float16) + entry.update_and_fetch(block, block) + mx.eval(entry.state) + return entry + + plain = feed(KVCache(), 30) + unwrapped = feed(RotatingKVCache(max_size = 100, keep = 2), 30) + wrapped = feed(RotatingKVCache(max_size = 10, keep = 2), 30) + chunked = feed(ChunkedKVCache(chunk_size = 8), 30) + slid = feed(ChunkedKVCache(chunk_size = 8), 30) + slid.maybe_trim_front() + + assert _kv_prefix_coverage([plain]) == 30 + assert _kv_prefix_coverage([unwrapped]) == 30 + assert _kv_prefix_coverage([chunked]) == 30 + assert wrapped.offset == 30 and wrapped.state[0].shape[2] == 10 + assert _kv_prefix_coverage([wrapped]) is None + assert slid.start_position > 0 + assert _kv_prefix_coverage([slid]) is None + assert _kv_prefix_coverage([CacheList(feed(KVCache(), 30), feed(KVCache(), 30))]) == 30 + assert _kv_prefix_coverage([CacheList(feed(KVCache(), 30), wrapped)]) is None + assert _kv_prefix_coverage([feed(KVCache(), 30), feed(KVCache(), 29)]) is None + assert _kv_prefix_coverage([]) is None + + history = _MLXPromptCacheHistory(6, 1 << 40) + for unsafe in (wrapped, slid): + history.insert("key", list(range(30)), [unsafe]) + assert "key" not in history._lru.entries + + history.insert("key", list(range(30)), [plain]) + assert tuple(range(30)) in history._lru.entries["key"] diff --git a/studio/backend/tests/test_mlx_stop_checkpoint.py b/studio/backend/tests/test_mlx_stop_checkpoint.py new file mode 100644 index 0000000000..d4a00cc6c8 --- /dev/null +++ b/studio/backend/tests/test_mlx_stop_checkpoint.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for MLX stop-and-save checkpoint handling.""" + +import importlib.util +import json +import sys +import types +from pathlib import Path + +import numpy as np +from safetensors.numpy import save_file + + +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load_worker_module(): + spec = importlib.util.spec_from_file_location( + "training_worker_under_test", + _BACKEND / "core" / "training" / "worker.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +worker = _load_worker_module() + + +class _FakeTrainer: + def __init__(self, step: int): + self._global_step = step + self._train_loss_history = [] + self.model = object() + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors") + save_file( + {"state": np.ones(1, dtype = np.float32)}, + checkpoint / "optimizer_state.safetensors", + ) + return checkpoint + + +def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._mlx_has_checkpoint_at_step(out, 5) is True + + +def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True + + +def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists( + tmp_path, monkeypatch +): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + saved_steps: list[int] = [] + + def _save_state(_value, path, name): + save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name)) + + def _save_trainer_state(state, ckpt_dir, **_kwargs): + Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8") + saved_steps.append(int(state["global_step"])) + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = lambda model, path: _save_state( + model, path, "adapters.safetensors" + ), + save_optimizer_state = lambda optimizer, path: _save_state( + optimizer, path, "optimizer_state.safetensors" + ), + save_trainer_state = _save_trainer_state, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True + assert saved_steps == [10] + assert (out / "checkpoint-10" / "trainer_state.json").is_file() + + +def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path): + out = tmp_path / "outputs" / "run_x" + ckpt = out / "checkpoint-5" + ckpt.mkdir(parents = True) + (ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8") + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path): + # An older checkpoint does not cover the current step, so this still fails. + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False + + +def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + def _boom(*_args, **_kwargs): + raise RuntimeError("save failed") + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = _boom, + save_optimizer_state = lambda *_a, **_k: None, + save_trainer_state = lambda *_a, **_k: None, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False diff --git a/studio/backend/tests/test_model_picker_regression.py b/studio/backend/tests/test_model_picker_regression.py new file mode 100644 index 0000000000..f38a4d0b8d --- /dev/null +++ b/studio/backend/tests/test_model_picker_regression.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression guards for the model-picker per-model-config feature (the set of +bugs that got the predecessor PR reverted). Pure-function / validation checks +only, so they run on CPU in the backend pytest job with no model download. + +Covers, at the backend layer: + - infra-model hiding: the RAG embedder (bge-small-en-v1.5) and the llama.cpp + install-validation probe (ggml-org/models / stories260K) stay hidden, while + normal chat repos are not hidden; + - the HF token is honored from the dedicated header with the query string as a + fallback, never the other way around; + - the chat-template byte caps reject oversized overrides (both the char-count + fast path and the UTF-8 byte path) and the sidecar reader is size-bounded. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +# Keep this test runnable without the optional structlog dependency (mirrors +# tests/test_cached_gguf_routes.py), since importing routes.models pulls it in. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route +from core.rag import config as rag_config +from hub.dependencies import get_hf_token +from models.inference import LoadRequest +from picker.schemas import MAX_CHAT_TEMPLATE_BYTES +from picker.service import _read_bounded_text +from utils.hidden_models import is_hidden_model + + +@pytest.fixture(autouse = True) +def _pin_default_embedder(monkeypatch): + """Pin the effective embedder to Studio's static default so hiding is + deterministic and cannot depend on ambient RAG config / env.""" + default = "unsloth/bge-small-en-v1.5" + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", default, raising = False) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: default) + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: default) + monkeypatch.setattr(rag_config, "default_gguf_repo", lambda: default) + + +# --------------------------------------------------------------------------- # +# Infra-model hiding (the "infra models resurfaced in the picker" regression) # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "value", + [ + "ggml-org/models", # the probe repo id + "unsloth/bge-small-en-v1.5", # the RAG embedder repo + "unsloth/bge-small-en-v1.5-GGUF", # its GGUF companion + "/root/.cache/huggingface/hub/x/stories260K.gguf", # probe on disk + "/root/.cache/x/Stories260K.GGUF", # case-insensitive + r"C:\\models\\stories260K.gguf", # windows-style path + "/opt/models/bge-small-en-v1.5", # embedder basename folder + "/opt/models/bge-small-en-v1.5-Q8_0.gguf", # suffixed local weight + ], +) +def test_infra_models_are_hidden(value): + assert is_hidden_model(value) is True + + +@pytest.mark.parametrize( + "value", + [ + "unsloth/gemma-3-270m-it-GGUF", # a normal small chat GGUF + "unsloth/Qwen3-0.6B", # a normal non-GGUF chat model + "user/stories260K-finetune-GGUF", # repo id merely contains "stories260k" + "user/model-chat", # generic repo must not be hidden + "meta-llama/Llama-3.1-8B-Instruct", + ], +) +def test_normal_models_are_not_hidden(value): + assert is_hidden_model(value) is False + + +def test_is_hidden_model_ignores_empty_values(): + assert is_hidden_model(None) is False + assert is_hidden_model("") is False + assert is_hidden_model(None, "", "unsloth/gemma-3-270m-it-GGUF") is False + + +def test_hidden_model_matchers_expose_probe_needles(): + needles, exact_ids, _exact_paths = models_route.hidden_model_matchers() + lowered = [n.lower() for n in needles] + assert "ggml-org/models" in lowered + assert "stories260k.gguf" in lowered + # The configured embedder is exposed as an exact repo id, never as a + # basename needle that would substring-hide unrelated chat models. + assert "bge-small-en-v1.5" not in lowered + assert "unsloth/bge-small-en-v1.5" in exact_ids + + +def test_hidden_model_matchers_custom_repo_publishes_exact_ids(monkeypatch): + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + needles, exact_ids, exact_paths = models_route.hidden_model_matchers() + assert needles == ["ggml-org/models", "stories260k.gguf"] + assert "org/model" in exact_ids + assert "org/model-gguf" in exact_ids + assert exact_paths == [] + + +def test_hidden_model_matchers_local_owner_name_path_is_exact_path(monkeypatch, tmp_path): + # A local embedder shaped like owner/name that exists on disk must be an + # exact resolved path, not a Hub repo id (mirroring is_hidden_model), so the + # local row stays hidden instead of showing as a chat model. + (tmp_path / "models" / "embedder").mkdir(parents = True) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "ggml-org/models") + _needles, exact_ids, exact_paths = models_route.hidden_model_matchers() + resolved = str((tmp_path / "models" / "embedder").resolve()).lower() + assert resolved in exact_paths + assert "models/embedder" not in exact_ids + + +# --------------------------------------------------------------------------- # +# HF token via header, query string only as a fallback (the token-leak fix) # +# --------------------------------------------------------------------------- # + + +def test_get_hf_token_strips_and_returns(): + assert get_hf_token(" hf_abc ") == "hf_abc" + + +@pytest.mark.parametrize("value", [None, "", " ", "\n\t"]) +def test_get_hf_token_blank_is_none(value): + assert get_hf_token(value) is None + + +@pytest.mark.parametrize( + "value,expected", + [(" hf_x ", "hf_x"), ("", None), (" ", None), (None, None), (1234, None)], +) +def test_normalize_hf_token(value, expected): + assert models_route._normalize_hf_token(value) == expected + + +def test_header_token_wins_over_query(): + header, query = "hf_header", "hf_query" + resolved = models_route._normalize_hf_token(header) or models_route._normalize_hf_token(query) + assert resolved == "hf_header" + + +def test_query_token_is_fallback_when_header_absent(): + resolved = models_route._normalize_hf_token(None) or models_route._normalize_hf_token( + "hf_query" + ) + assert resolved == "hf_query" + + +# --------------------------------------------------------------------------- # +# Chat-template byte caps (the unbounded-template hardening) # +# --------------------------------------------------------------------------- # + + +def _load_request(**overrides): + data = {"model_path": "unsloth/test-model-GGUF", "gguf_variant": "Q4_K_M"} + data.update(overrides) + return LoadRequest.model_validate(data) + + +def test_blank_chat_template_override_normalizes_to_none(): + assert _load_request(chat_template_override = " \n\t").chat_template_override is None + + +def test_nonblank_chat_template_override_preserved_verbatim(): + template = " {{ messages }} " + assert _load_request(chat_template_override = template).chat_template_override == template + + +def test_chat_template_at_byte_limit_is_accepted(): + template = "a" * MAX_CHAT_TEMPLATE_BYTES # exactly the limit, 1 byte/char + assert ( + len(_load_request(chat_template_override = template).chat_template_override) + == MAX_CHAT_TEMPLATE_BYTES + ) + + +def test_chat_template_over_char_limit_is_rejected(): + with pytest.raises(Exception): # pydantic ValidationError wrapping ValueError + _load_request(chat_template_override = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1)) + + +def test_chat_template_over_byte_limit_is_rejected(): + # Char count stays under the limit but UTF-8 bytes exceed it (3 bytes/char), + # so only the byte-count branch can catch this. + multibyte = "โ‚ฌ" * (MAX_CHAT_TEMPLATE_BYTES // 2) # euro sign, 3 bytes each + assert len(multibyte) <= MAX_CHAT_TEMPLATE_BYTES + assert len(multibyte.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES + with pytest.raises(Exception): + _load_request(chat_template_override = multibyte) + + +def test_read_bounded_text_reads_within_limit(tmp_path): + p = tmp_path / "t.json" + p.write_text("hello", encoding = "utf-8") + assert _read_bounded_text(p, 16) == "hello" + + +def test_read_bounded_text_rejects_over_limit(tmp_path): + p = tmp_path / "big.json" + p.write_bytes(b"x" * 100) + assert _read_bounded_text(p, 50) is None + + +def test_read_bounded_text_at_limit_is_read(tmp_path): + p = tmp_path / "exact.json" + p.write_bytes(b"x" * 50) + assert _read_bounded_text(p, 50) == "x" * 50 + + +def test_read_bounded_text_missing_file_is_none(tmp_path): + assert _read_bounded_text(tmp_path / "nope.json", 50) is None diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index edf55812e2..d84f8c94a7 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -112,13 +112,21 @@ def patch_hub_gguf(monkeypatch): blob_ids = [local_blob], gguf_files = {"model-Q4_K_M.gguf": 1000}, ) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = tmp_path), + ) monkeypatch.setattr( GV, "list_gguf_variants", lambda r, hf_token = None: (_variants(), False, [remote_sibling]), raising = True, ) - monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr( + GV, + "iter_hf_cache_snapshots", + lambda _repo_id, root = None: [snap], + ) monkeypatch.setattr( CI, "all_hf_cache_scans", @@ -217,6 +225,10 @@ def test_variant_update_check_detects_companion_only_update( companion_path: 100, }, ) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = tmp_path), + ) siblings = [ patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "mainsha"), patch_hub_gguf.sibling(companion_path, 100, "new-companion"), @@ -227,7 +239,11 @@ def test_variant_update_check_detects_companion_only_update( lambda r, hf_token = None: (_variants(), has_vision, siblings), raising = True, ) - monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr( + GV, + "iter_hf_cache_snapshots", + lambda _repo_id, root = None: [snap], + ) monkeypatch.setattr( CI, "all_hf_cache_scans", @@ -314,6 +330,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): file_name = "model.safetensors", size_on_disk = 100, blob_path = str(repo_path / "blobs" / "modelsha"), + blob_last_modified = 3_000.0, ), ] ) @@ -336,6 +353,98 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): assert rows[0]["repo_id"] == "Org/SafeTensorRepo" assert rows[0]["model_format"] == "safetensors" assert rows[0]["size_bytes"] == 100 + assert rows[0]["last_modified"] == 3_000.0 + + +def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path): + repo_path = tmp_path / "models--Org--GgufRepo" + repo = SimpleNamespace( + repo_id = "Org/GgufRepo", + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + size_on_disk = 100, + blob_path = None, + blob_last_modified = 5_000.0, + ), + ] + ) + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + monkeypatch.setattr( + CI.hf_cache_scan, + "is_gguf_repo_partial", + lambda *args, **kwargs: False, + ) + monkeypatch.setattr( + CI, + "_gguf_variant_state_summary", + lambda _repo_id, **_kwargs: (False, 0), + ) + + rows = CI._scan_cached_gguf() + + assert len(rows) == 1 + assert rows[0]["repo_id"] == "Org/GgufRepo" + assert rows[0]["model_format"] == "gguf" + assert rows[0]["size_bytes"] == 100 + assert rows[0]["last_modified"] == 5_000.0 + + +def test_cached_model_scan_hides_custom_whisper_repo(monkeypatch, tmp_path): + repo_path = tmp_path / "models--Org--CustomWhisper" + snapshot = repo_path / "snapshots" / ("a" * 40) + snapshot.mkdir(parents = True) + (snapshot / "config.json").write_text( + '{"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}' + ) + repo = SimpleNamespace( + repo_id = "Org/CustomWhisper", + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "config.json", + size_on_disk = 10, + blob_path = None, + ), + SimpleNamespace( + file_name = "model.safetensors", + size_on_disk = 100, + blob_path = str(repo_path / "blobs" / "modelsha"), + ), + ] + ) + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + monkeypatch.setattr( + CI, + "_cached_model_snapshot_path", + lambda _repo_path: snapshot, + ) + monkeypatch.setattr( + CI.hf_cache_scan, + "is_snapshot_partial", + lambda *args, **kwargs: False, + ) + + assert CI._scan_cached_models() == [] # โ”€โ”€ hf_hub_download_with_xet_fallback force_download bypass (X2/F2) โ”€โ”€โ”€ @@ -584,7 +693,12 @@ def test_reclaim_replaced_gguf_variant_prunes_old_revision_only(monkeypatch, tmp invalidated = [] monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: invalidated.append(True)) - result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"NEWsha"})) + result = D.reclaim_replaced_gguf_variant( + repo_id, + "Q4_K_M", + frozenset({"NEWsha"}), + hub_cache = tmp_path, + ) assert result["removed_snapshots"] == 1 assert result["deleted_blobs"] == 1 @@ -631,8 +745,85 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch monkeypatch.setattr(CI, "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo_info])]) monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: None) - result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"REMOTEsha256"})) + result = D.reclaim_replaced_gguf_variant( + repo_id, + "Q4_K_M", + frozenset({"REMOTEsha256"}), + hub_cache = tmp_path, + ) assert snap.exists() is True # the current file must survive assert result["removed_snapshots"] == 0 assert result["deleted_blobs"] == 0 + + +def test_reclaim_replaced_gguf_variant_only_mutates_worker_cache(monkeypatch, tmp_path): + repo_id = "org/repo-GGUF" + cache_a = tmp_path / "cache-a" + cache_b = tmp_path / "cache-b" + + def cached_repo(cache_dir, revision): + repo_path = cache_dir / "models--org--repo-GGUF" + snap = repo_path / "snapshots" / revision / "model-Q4_K_M.gguf" + blob = repo_path / "blobs" / "OLDsha" + snap.parent.mkdir(parents = True, exist_ok = True) + blob.parent.mkdir(parents = True, exist_ok = True) + blob.write_bytes(b"old") + snap.symlink_to(blob) + return ( + SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = snap.name, + file_path = str(snap), + blob_path = str(blob), + ) + ] + ) + ], + ), + snap, + blob, + ) + + repo_a, snap_a, blob_a = cached_repo(cache_a, "a" * 40) + repo_b, snap_b, blob_b = cached_repo(cache_b, "b" * 40) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo_a]), SimpleNamespace(repos = [repo_b])], + ) + monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: None) + + result = D.reclaim_replaced_gguf_variant( + repo_id, + "Q4_K_M", + frozenset({"NEWsha"}), + hub_cache = cache_b, + ) + + assert result["removed_snapshots"] == 1 + assert snap_b.exists() is False + assert blob_b.exists() is False + assert snap_a.exists() is True + assert blob_a.exists() is True + + +def _mmproj_repo(*file_names: str): + return SimpleNamespace( + revisions = [SimpleNamespace(files = [SimpleNamespace(file_name = n) for n in file_names])] + ) + + +def test_repo_has_mmproj_requires_gguf_projector(): + # A non-GGUF sidecar whose name merely contains "mmproj" must NOT mark the + # repo vision-capable; the runtime's projector detection is GGUF-only. + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj_config.json")) is False + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "README-mmproj.md")) is False + # A real GGUF projector still marks the repo vision-capable. + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj-F16.gguf")) is True diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py index 12f6c497ab..a50765898b 100644 --- a/studio/backend/tests/test_models_get_model_config_case_resolution.py +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -84,7 +84,7 @@ def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, mon # covers the active cache; discard deletes case-insensitively, so detection must too, # else a decline deletes a pre-existing user repo). import utils.paths as paths_pkg - import huggingface_hub.constants as hf_constants + import hub.utils.paths as hub_paths active = tmp_path / "active" legacy = tmp_path / "legacy" @@ -96,9 +96,12 @@ def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, mon # No active-cache variant; case resolution is a no-op here. monkeypatch.setattr(paths_pkg, "resolve_cached_repo_id_case", lambda name: name) - monkeypatch.setattr(paths_pkg, "legacy_hf_cache_dir", lambda: legacy) - monkeypatch.setattr(paths_pkg, "hf_default_cache_dir", lambda: default) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(active)) + monkeypatch.setattr(hub_paths, "legacy_hf_cache_dir", lambda: legacy) + monkeypatch.setattr(hub_paths, "hf_default_cache_dir", lambda: default) + monkeypatch.setattr( + "utils.hf_cache_settings.known_hf_hub_caches", + lambda: [active], + ) assert models_route._repo_in_any_hf_cache("unsloth/foo") is True # Absent from every cache -> reported absent. diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py new file mode 100644 index 0000000000..ccc6b5f76a --- /dev/null +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -0,0 +1,942 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Offline RAG embedding-model handling (issue #6817). + +Offline the studio must never call the Hub (a DNS-dead session hangs on retries). Using a fake +HF cache under a temp HF_HUB_CACHE, assert that offline: is_embedding_model classifies from the +cached modules.json without the Hub; the file-security gate fails CLOSED on an unscanned pickle +weight with no safetensors alternative and allows an inert cache; the embedder threads +local_files_only into the load. Online behavior is unchanged (bounded timeout + cache fallback). +""" + +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from utils.security import evaluate_file_security +from utils.utils import ( + hf_cache_snapshot_dir, + hf_cache_snapshot_is_loadable, + hf_env_offline, + st_repo_id_candidates, +) + +# Minimal sentence-transformers modules.json (the marker the gate keys on). +MODULES_JSON = ( + '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}]' +) + + +def _modules_json(*paths): + """modules.json listing one Transformer module per path (a load root).""" + import json + return json.dumps( + [ + { + "idx": i, + "name": str(i), + "path": p, + "type": "sentence_transformers.models.Transformer", + } + for i, p in enumerate(paths) + ] + ) + + +_COMMIT = "0123456789abcdef0123456789abcdef01234567" + + +def _fs_case_sensitive(root): + """Whether root's filesystem is case-sensitive (Linux yes; macOS/Windows usually no). The gate + mirrors the loader, whose file lookups follow the same rule, so some cases only exist on one.""" + probe = Path(root) / "_case_probe" + probe.write_text("x") + try: + return not (Path(root) / "_CASE_PROBE").exists() + finally: + probe.unlink() + + +def _requires_case_sensitive_fs(root): + if not _fs_case_sensitive(root): + pytest.skip("requires a case-sensitive filesystem") + + +def _requires_case_insensitive_fs(root): + if _fs_case_sensitive(root): + pytest.skip("requires a case-insensitive filesystem") + + +def _make_cache( + root, + repo_id, + files, + commit = _COMMIT, +): + """Build a canonical HF-cache snapshot (refs/main + snapshots//) for repo_id under + root from {relpath: contents}; returns the snapshot dir.""" + from huggingface_hub.file_download import repo_folder_name + + repo_dir = Path(root) / repo_folder_name(repo_id = repo_id, repo_type = "model") + (repo_dir / "refs").mkdir(parents = True, exist_ok = True) + (repo_dir / "refs" / "main").write_text(commit) + snapshot = repo_dir / "snapshots" / commit + snapshot.mkdir(parents = True, exist_ok = True) + for rel, contents in files.items(): + path = snapshot / rel + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(contents) + return snapshot + + +def _no_network(): + """Patch model_info to fail loudly if any offline path reaches the network.""" + return patch("huggingface_hub.model_info", side_effect = AssertionError("hit the network")) + + +def _is_embedding_model(*args, **kwargs): + from utils.models.model_config import is_embedding_model + return is_embedding_model(*args, **kwargs) + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + """Point the HF cache at a fresh temp dir. + + get_hf_cache_paths() reads an import-time env snapshot, not live os.environ, + so point it (and thus active_hf_hub_cache + the snapshot lookup's selected + root) at this temp cache too.""" + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = root), + ) + return root + + +@pytest.fixture(autouse = True) +def _clean_env(monkeypatch): + """Start each test online with an empty detection cache; offline tests opt in.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + from utils.models import model_config as mc + + mc._embedding_detection_cache.clear() + yield + mc._embedding_detection_cache.clear() + + +# โ”€โ”€ hf_env_offline โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "]) +def test_hf_env_offline_true(monkeypatch, value): + monkeypatch.setenv("HF_HUB_OFFLINE", value) + assert hf_env_offline() is True + + +@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) +def test_hf_env_offline_false(monkeypatch, value): + monkeypatch.setenv("HF_HUB_OFFLINE", value) + assert hf_env_offline() is False + + +def test_hf_env_offline_honors_transformers_flag(monkeypatch): + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + assert hf_env_offline() is True + + +def test_hf_env_offline_default_false(): + assert hf_env_offline() is False + + +# โ”€โ”€ st_repo_id_candidates โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def test_candidates_slashless_adds_st_alias(): + assert st_repo_id_candidates("all-MiniLM-L6-v2") == [ + "all-MiniLM-L6-v2", + "sentence-transformers/all-MiniLM-L6-v2", + ] + + +def test_candidates_with_org_is_verbatim(): + assert st_repo_id_candidates("org/model") == ["org/model"] + + +def test_candidates_empty_name(): + assert st_repo_id_candidates(" ") == [] + + +# โ”€โ”€ hf_cache_snapshot_dir โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def test_snapshot_dir_resolves_active_commit(hf_cache): + snapshot = _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_dir_none_when_uncached(hf_cache): + assert hf_cache_snapshot_dir("org/missing") is None + + +def test_snapshot_dir_uses_st_alias_for_slashless(hf_cache): + snapshot = _make_cache( + hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON} + ) + assert hf_cache_snapshot_dir("all-MiniLM-L6-v2") == snapshot + + +def test_snapshot_dir_none_when_snapshot_missing(hf_cache): + from huggingface_hub.file_download import repo_folder_name + + repo_dir = hf_cache / repo_folder_name(repo_id = "org/broken", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text("deadbeef") # no snapshots/deadbeef dir + assert hf_cache_snapshot_dir("org/broken") is None + + +def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch): + # An unexpanded $VAR in HF_HUB_CACHE must resolve where the loader looks. + real = tmp_path / "hub" + real.mkdir() + monkeypatch.setenv("MY_HF_CACHE", str(real)) + monkeypatch.setenv("HF_HUB_CACHE", "$MY_HF_CACHE") + monkeypatch.delenv("HF_HOME", raising = False) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + snapshot = _make_cache(real, "org/emb", {"modules.json": MODULES_JSON}) + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch): + # ST uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must inspect it too. + st_home = tmp_path / "st_home" + st_home.mkdir() + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) + monkeypatch.delenv("HF_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HOME", raising = False) + snapshot = _make_cache(st_home, "org/emb", {"modules.json": MODULES_JSON}) + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_dir_prefers_selected_cache_over_st_home(tmp_path, monkeypatch): + # The RAG loader passes cache_folder=active_hf_hub_cache(), which overrides + # SENTENCE_TRANSFORMERS_HOME, so the snapshot + offline security lookup must + # search the selected cache even when ST_HOME points elsewhere. Otherwise the + # gate scans a cache the model never loads from and a pickle weight in the + # selected cache slips through. + st_home = tmp_path / "st_home" + st_home.mkdir() + selected = tmp_path / "hub" + selected.mkdir() + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) + monkeypatch.delenv("HF_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HOME", raising = False) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = selected), + ) + snapshot = _make_cache(selected, "org/emb", {"modules.json": MODULES_JSON}) # only in selected + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_is_loadable_with_config_and_weights(hf_cache): + _make_cache(hf_cache, "org/emb", {"config.json": "{}", "model.safetensors": "x"}) + assert hf_cache_snapshot_is_loadable("org/emb") is True + + +def test_snapshot_is_not_loadable_when_metadata_only(hf_cache): + # A partial cache (refs/main resolves but no weights) is not loadable. + _make_cache(hf_cache, "org/partial", {"config.json": "{}", "modules.json": MODULES_JSON}) + assert hf_cache_snapshot_is_loadable("org/partial") is False + + +def test_snapshot_is_not_loadable_when_uncached(hf_cache): + assert hf_cache_snapshot_is_loadable("org/missing") is False + + +def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch): + # A pickle under SENTENCE_TRANSFORMERS_HOME must still fail closed offline. + st_home = tmp_path / "st_home" + st_home.mkdir() + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) + monkeypatch.delenv("HF_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HOME", raising = False) + _make_cache(st_home, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) + with _no_network(): + assert evaluate_file_security("org/pk", local_only_load = True).blocked is True + + +# โ”€โ”€ is_embedding_model: offline (no network) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def test_offline_true_for_cached_st_model(hf_cache, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON, "config.json": "{}"}) + with _no_network(): + assert _is_embedding_model("org/emb") is True + + +def test_offline_false_for_cached_non_st_model(hf_cache, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _make_cache(hf_cache, "org/plain", {"config.json": "{}", "model.safetensors": "x"}) + with _no_network(): + assert _is_embedding_model("org/plain") is False + + +def test_offline_false_when_uncached(hf_cache, monkeypatch): + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/missing") is False + + +def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _make_cache(hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON}) + with _no_network(): + assert _is_embedding_model("all-MiniLM-L6-v2") is True + + +def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch): + # An online lookup memoizes True for an UNCACHED repo (tags say embedding, no weights). Once + # offline, is_embedding_model must reclassify from the empty cache and return False, not the + # stale online True that would make settings accept a repo _get() cannot load. + with patch( + "huggingface_hub.model_info", + side_effect = lambda *a, **k: SimpleNamespace( + tags = ["sentence-transformers"], pipeline_tag = None + ), + ): + assert _is_embedding_model("org/uncached-emb") is True # memoized True online + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache + + +def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch): + # Because the offline branch never records a memo, once an uncached repo's snapshot + # materializes (another process populates the cache) the next call re-reports True. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/later") is False # uncached + _make_cache(hf_cache, "org/later", {"modules.json": MODULES_JSON}) + assert _is_embedding_model("org/later") is True # cache now present, no stale negative + + +# โ”€โ”€ is_embedding_model: online (bounded + fallback) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def test_online_passes_bounded_timeout(hf_cache): + seen = {} + + def _mi( + name, + token = None, + timeout = None, + **kw, + ): + seen["timeout"] = timeout + return SimpleNamespace(tags = ["sentence-transformers"], pipeline_tag = None) + + with patch("huggingface_hub.model_info", side_effect = _mi): + assert _is_embedding_model("org/emb") is True + assert seen["timeout"] == 15.0 + + +def test_online_error_falls_back_to_cache_marker(hf_cache): + _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) + with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")): + assert _is_embedding_model("org/emb") is True + + +def test_online_error_without_cache_returns_false(hf_cache): + with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")): + assert _is_embedding_model("org/missing") is False + + +# โ”€โ”€ evaluate_file_security: offline fail-closed gate โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def _offline_decision(name): + return evaluate_file_security(name, local_only_load = True) + + +def test_gate_allows_safetensors_only(hf_cache): + _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}) + with _no_network(): + assert _offline_decision("org/st").blocked is False + + +def test_gate_blocks_pickle_without_safetensors(hf_cache): + _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) + with _no_network(): + decision = _offline_decision("org/pk") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_allows_pickle_with_safetensors_sibling(hf_cache): + _make_cache(hf_cache, "org/both", {"pytorch_model.bin": "x", "model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/both").blocked is False + + +def test_gate_blocks_sharded_pickle(hf_cache): + _make_cache( + hf_cache, + "org/shard", + { + "pytorch_model-00001-of-00002.bin": "a", + "pytorch_model-00002-of-00002.bin": "b", + }, + ) + with _no_network(): + assert _offline_decision("org/shard").blocked is True + + +def test_gate_blocks_indexed_pickle_shard_in_subdirectory(hf_cache): + # from_pretrained follows weight_map paths relative to the root index, so these nested shards + # are deserialized even though they are not direct children of the load root (iterdir misses + # them). The online gate blocks index-referenced subdir pickles; the offline gate must too. + _make_cache( + hf_cache, + "org/indexed-shard", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"layer.weight": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-shard") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_pickle_shard_with_nonstandard_stem(hf_cache): + # The index tells the loader to deserialize this file, so a pickle EXTENSION is enough -- the + # shard's stem need not match the on-disk weight-name heuristic (which only guesses bare files). + _make_cache( + hf_cache, + "org/indexed-odd", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/evil-00001-of-00001.bin"}}', + "shards/evil-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-odd") + assert decision.blocked is True + assert any(u["path"] == "shards/evil-00001-of-00001.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_safetensors_index_pointing_to_pickle_shard(hf_cache): + # load_state_dict picks safetensors vs torch.load by each shard's own suffix, so a + # model.safetensors.index.json that maps a weight to a .bin shard still deserializes it. The + # index's own existence must not suppress the shard it names. + _make_cache( + hf_cache, + "org/st-index-pickle", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/st-index-pickle") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_shard_with_no_pickle_extension(hf_cache): + # Transformers torch.loads any indexed shard not ending in .safetensors, so an unconventional + # extensionless name is still a deserialization target. + _make_cache( + hf_cache, + "org/indexed-noext", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload"}}', + "shards/payload": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-noext") + assert decision.blocked is True + assert any(u["path"] == "shards/payload" for u in decision.unsafe_files) + + +_UPPER_INDEX_FILES = { + "PYTORCH_MODEL.BIN.INDEX.JSON": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", +} + + +def test_gate_blocks_uppercase_index_on_case_insensitive_fs(hf_cache): + # On a case-insensitive volume (Windows/macOS) from_pretrained opens an oddly-cased index when it + # requests the canonical lowercase name, so the loader-mirror lookup resolves it and blocks. + _requires_case_insensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + decision = _offline_decision("org/upper-index") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_allows_uppercase_index_on_case_sensitive_fs(hf_cache): + # On a case-sensitive FS from_pretrained's os.path.isfile of the canonical lowercase name misses + # the uppercase artifact and never loads its shard, so the gate must not over-block it. + _requires_case_sensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + assert _offline_decision("org/upper-index").blocked is False + + +def test_gate_blocks_indexed_shard_named_with_backslash(hf_cache): + # On POSIX a backslash is a literal filename char, so from_pretrained joins the raw weight_map + # value and deserializes a file actually named "dir\payload.bin"; the gate must probe it verbatim. + import os + + if os.sep != "/": + pytest.skip("backslash is a path separator off POSIX") + _make_cache( + hf_cache, + "org/backslash", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "dir\\\\payload.bin"}}', + "dir\\payload.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/backslash") + assert decision.blocked is True + assert any(u["path"] == "dir\\payload.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_shard_with_uppercase_safetensors_suffix(hf_cache): + # load_state_dict's endswith(".safetensors") is case-sensitive, so a shard named payload.SAFETENSORS + # falls to torch.load. The gate must classify shard suffixes case-sensitively to match it. + _make_cache( + hf_cache, + "org/upper-suffix", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload.SAFETENSORS"}}', + "shards/payload.SAFETENSORS": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-suffix") + assert decision.blocked is True + assert any(u["path"] == "shards/payload.SAFETENSORS" for u in decision.unsafe_files) + + +def test_gate_allows_stale_safetensors_index_beside_direct_safetensors(hf_cache): + # A complete direct model.safetensors is selected before either index, so a stale + # model.safetensors.index.json referencing a .bin shard never deserializes -> must not block. + _make_cache( + hf_cache, + "org/direct-plus-stale-index", + { + "model.safetensors": "tensors", + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + assert _offline_decision("org/direct-plus-stale-index").blocked is False + + +def test_gate_blocks_pytorch_index_with_uppercase_safetensors_decoy(hf_cache): + # On a case-sensitive FS, from_pretrained asks for the canonical lowercase model.safetensors, does + # not find an uppercase decoy, and selects the pytorch index instead. The decoy must not suppress. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy", + { + "MODEL.SAFETENSORS": "decoy", + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_direct_pickle_with_uppercase_safetensors_decoy(hf_cache): + # Same decoy against a direct pytorch_model.bin: the loader selects the pickle, so the uppercase + # safetensors must not suppress it on a case-sensitive FS. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy-direct", + {"MODEL.SAFETENSORS": "decoy", "pytorch_model.bin": "pickle"}, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy-direct") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_pickle_shard_in_module_subdir(hf_cache): + # A weight index inside a sentence-transformers module load root points at a nested pickle shard. + _make_cache( + hf_cache, + "org/mod-indexed", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "0_Transformer/shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/mod-indexed") + assert decision.blocked is True + assert any( + u["path"] == "0_Transformer/shards/pytorch_model-00001-of-00001.bin" + for u in decision.unsafe_files + ) + + +def test_gate_allows_indexed_pickle_shard_with_safetensors_sibling(hf_cache): + # A base model.safetensors makes the loader ignore the pickle index entirely, so it must not + # block (mirrors the direct-file safetensors-sibling suppression). + _make_cache( + hf_cache, + "org/indexed-both", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + "model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/indexed-both").blocked is False + + +def test_gate_allows_indexed_safetensors_shard_in_subdirectory(hf_cache): + # A safetensors index lists inert shards -- following it must never block (guards against a + # scanner that flags every indexed shard regardless of format). + _make_cache( + hf_cache, + "org/st-indexed", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}' + ), + "shards/model-00001-of-00001.safetensors": "tensors", + }, + ) + with _no_network(): + assert _offline_decision("org/st-indexed").blocked is False + + +def test_gate_blocks_on_index_path_traversal(hf_cache): + # A weight_map entry escaping the snapshot via ".." is abnormal/hostile -> fail closed. + _make_cache( + hf_cache, + "org/escape", + {"pytorch_model.bin.index.json": '{"weight_map": {"w": "../../../../etc/evil.bin"}}'}, + ) + with _no_network(): + assert _offline_decision("org/escape").blocked is True + + +def test_gate_allows_symlinked_sharded_safetensors(tmp_path, monkeypatch): + # Real HF caches store snapshot files as symlinks into blobs/. A resolve()-based containment + # check would escape the snapshot and false-block every sharded model; the lexical gate must not. + import hashlib + import os + + from huggingface_hub.file_download import repo_folder_name + + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = root), + ) + repo_dir = root / repo_folder_name(repo_id = "org/sym", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text(_COMMIT) + blobs = repo_dir / "blobs" + blobs.mkdir() + snapshot = repo_dir / "snapshots" / _COMMIT + (snapshot / "shards").mkdir(parents = True) + + def _blobbed(rel, content): + digest = hashlib.sha256(content.encode()).hexdigest() + (blobs / digest).write_text(content) + target = snapshot / rel + target.parent.mkdir(parents = True, exist_ok = True) + target.symlink_to(os.path.relpath(blobs / digest, target.parent)) + + _blobbed("config.json", "{}") + _blobbed( + "model.safetensors.index.json", + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}', + ) + _blobbed("shards/model-00001-of-00001.safetensors", "tensors") + with _no_network(): + assert _offline_decision("org/sym").blocked is False + + +def test_gate_allows_index_without_weight_map(hf_cache): + # An index whose top-level JSON has no dict weight_map lets the loader resolve no shards, so it + # must not crash or block on its own (only inert safetensors are cached here). + _make_cache( + hf_cache, + "org/no-wm", + {"model.safetensors.index.json": "[]", "model.safetensors": "x"}, + ) + with _no_network(): + assert _offline_decision("org/no-wm").blocked is False + + +def test_gate_allows_nothing_cached(hf_cache): + with _no_network(): + assert _offline_decision("org/missing").blocked is False + + +def test_gate_allows_gguf_only(hf_cache): + _make_cache(hf_cache, "org/gg", {"model.gguf": "x"}) + with _no_network(): + assert _offline_decision("org/gg").blocked is False + + +def test_gate_blocks_pickle_in_module_subdir(hf_cache): + # 0_Transformer is a module load root (listed in modules.json), so its pickle blocks. + _make_cache( + hf_cache, + "org/mod", + {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"}, + ) + with _no_network(): + assert _offline_decision("org/mod").blocked is True + + +def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache): + _make_cache( + hf_cache, + "org/mod2", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin": "x", + "0_Transformer/model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/mod2").blocked is False + + +def test_gate_allows_unreferenced_nested_pickle(hf_cache): + # A pickle in a dir NOT referenced by modules.json (e.g. nemo/) is never deserialized, so it + # must not block the offline load (matches the online gate). + _make_cache( + hf_cache, + "org/aux", + { + "modules.json": MODULES_JSON, # Transformer at the root only + "model.safetensors": "w", + "nemo/pytorch_model.bin": "x", + }, + ) + with _no_network(): + assert _offline_decision("org/aux").blocked is False + + +def test_gate_blocks_adapter_pickle_without_safetensors(hf_cache): + _make_cache(hf_cache, "org/ad", {"config.json": "{}", "adapter_model.bin": "x"}) + with _no_network(): + decision = _offline_decision("org/ad") + assert decision.blocked is True + assert any(u["path"] == "adapter_model.bin" for u in decision.unsafe_files) + + +def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache): + _make_cache(hf_cache, "org/ad2", {"adapter_model.bin": "x", "adapter_model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/ad2").blocked is False + + +def test_gate_blocks_base_pickle_with_only_adapter_safetensors_decoy(hf_cache): + # A decoy adapter_model.safetensors must NOT suppress a base pytorch_model.bin (the base + # loader would still deserialize the unscanned pickle). + _make_cache(hf_cache, "org/decoy", {"pytorch_model.bin": "x", "adapter_model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/decoy").blocked is True + + +def test_gate_blocks_adapter_pickle_with_only_base_safetensors_decoy(hf_cache): + # Symmetric: a base model.safetensors must NOT suppress an adapter_model.bin. + _make_cache(hf_cache, "org/decoy2", {"adapter_model.bin": "x", "model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/decoy2").blocked is True + + +def test_gate_reports_snapshot_relative_path(hf_cache): + _make_cache( + hf_cache, + "org/mod3", + {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"}, + ) + with _no_network(): + decision = _offline_decision("org/mod3") + assert decision.blocked is True + assert any(u["path"] == "0_Transformer/pytorch_model.bin" for u in decision.unsafe_files) + + +# โ”€โ”€ evaluate_file_security: online path unchanged โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def test_online_default_blocks_unsafe(): + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with patch( + "huggingface_hub.model_info", + side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status), + ): + assert evaluate_file_security("org/x").blocked is True + + +def test_online_default_allows_clean(): + status = {"scansDone": True, "filesWithIssues": []} + with patch( + "huggingface_hub.model_info", + side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status), + ): + assert evaluate_file_security("org/x").blocked is False + + +# โ”€โ”€ embeddings guard + loader โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def test_guard_offline_blocks_pickle_only(hf_cache): + from core.rag.embeddings import UnsafeEmbeddingModelError, _guard_model_security + _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) + with _no_network(): + with pytest.raises(UnsafeEmbeddingModelError): + _guard_model_security("org/pk", local_only = True) + + +def test_guard_offline_allows_safetensors(hf_cache): + from core.rag.embeddings import _guard_model_security + _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}) + with _no_network(): + _guard_model_security("org/st", local_only = True) # must not raise + + +def _install_fake_sentence_transformers(monkeypatch, captured): + class FakeSentenceTransformer: + def __init__( + self, + name, + *, + device = None, + model_kwargs = None, + local_files_only = False, + **kw, + ): + captured["name"] = name + captured["device"] = device + captured["local_files_only"] = local_files_only + + module = types.ModuleType("sentence_transformers") + module.SentenceTransformer = FakeSentenceTransformer + monkeypatch.setitem(sys.modules, "sentence_transformers", module) + + +def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch): + from core.rag import embeddings + + snapshot = _make_cache( + hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"} + ) + # TRANSFORMERS_OFFLINE only: a cached model loads from its local snapshot dir (a local path, + # never the Hub), offline-safe on ANY sentence-transformers version. + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + captured = {} + _install_fake_sentence_transformers(monkeypatch, captured) + with _no_network(): + embeddings._get("org/st") + assert captured["name"] == str(snapshot) + + +def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch): + from core.rag import embeddings + + empty = tmp_path / "hub" + empty.mkdir() + monkeypatch.setenv("HF_HUB_CACHE", str(empty)) + monkeypatch.delenv("HF_HOME", raising = False) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + # No cache -> repo-id load forced cache-only (fails fast offline, not a hang). + monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None) + captured = {} + _install_fake_sentence_transformers(monkeypatch, captured) + embeddings._get("org/uncached-xyz") + assert captured["name"] == "org/uncached-xyz" + assert captured["local_files_only"] is True + + +def test_get_online_omits_local_files_only(monkeypatch): + from core.rag import embeddings + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + # Isolate the loader wiring from the online guard's network calls. + monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None) + captured = {} + _install_fake_sentence_transformers(monkeypatch, captured) + embeddings._get("org/online") + assert captured["local_files_only"] is False diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 295549c443..d1e61d0546 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -119,10 +119,21 @@ def _build_cache( return snap +def _symlink_or_skip(link: Path, target: Path) -> None: + try: + link.symlink_to(target) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + @pytest.fixture def hf_cache(tmp_path, monkeypatch): """Point ``huggingface_hub.constants.HF_HUB_CACHE`` at a temp dir.""" monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: _types.SimpleNamespace(hub_cache = tmp_path), + ) return tmp_path @@ -220,6 +231,10 @@ class TestGgufVariantFileResolution: return f"/fake/{repo_id}/{filename}" monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: _types.SimpleNamespace(hub_cache = tmp_path), + ) with ( patch( "huggingface_hub.list_repo_files", @@ -427,6 +442,40 @@ class TestGgufVariantFileResolution: assert out == str(snap / "mmproj-F16.gguf") + def test_download_companion_uses_selected_cache_not_import_time_default( + self, monkeypatch, tmp_path + ): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + import_time_cache = tmp_path / "import-time-cache" + selected_cache = tmp_path / "selected-cache" + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(import_time_cache)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: _types.SimpleNamespace(hub_cache = selected_cache), + ) + repo = "unsloth/vision-GGUF" + snap = _build_cache(selected_cache, repo, {"mmproj-F16.gguf": 4}) + backend = LlamaCppBackend() + + offline_error = type("OfflineModeIsEnabled", (Exception,), {}) + + def fail_list(*_args, **_kwargs): + raise offline_error("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("selected-cache companion must not download") + + with ( + patch("huggingface_hub.list_repo_files", fail_list), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fail_download, + ), + ): + out = backend._download_mmproj(hf_repo = repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -453,6 +502,10 @@ class TestGgufVariantFileResolution: return f"/fake/{repo_id}/{filename}" monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: _types.SimpleNamespace(hub_cache = tmp_path), + ) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -1084,7 +1137,7 @@ class TestListLocalGgufVariantsSubdir: target.write_bytes(b"\0" * 20) out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") - assert out == str(target.resolve()) + assert out == str(target.absolute()) def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path): from utils.models.model_config import _find_local_gguf_by_variant @@ -1094,6 +1147,57 @@ class TestListLocalGgufVariantsSubdir: assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None + def test_find_local_gguf_by_variant_keeps_split_symlink_name(self, tmp_path): + from utils.models.model_config import _find_local_gguf_by_variant + + blobs = tmp_path / "blobs" + blobs.mkdir() + snap = tmp_path / "snapshots" / "rev" / "BF16" + snap.mkdir(parents = True) + (tmp_path / "snapshots" / "rev" / "config.json").write_text("{}") + for i, sha in enumerate(("aa" * 32, "bb" * 32), start = 1): + (blobs / sha).write_bytes(b"\0" * 10) + _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha) + + out = _find_local_gguf_by_variant(str(tmp_path / "snapshots" / "rev"), "BF16") + assert out is not None + assert Path(out).name == "model-BF16-00001-of-00002.gguf" + + def test_detect_gguf_model_keeps_split_symlink_name(self, tmp_path): + from utils.models.model_config import detect_gguf_model + + blobs = tmp_path / "blobs" + blobs.mkdir() + snap = tmp_path / "snapshots" / "rev" + snap.mkdir(parents = True) + for i, (sha, size) in enumerate((("cc" * 32, 10), ("dd" * 32, 20)), start = 1): + (blobs / sha).write_bytes(b"\0" * size) + _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha) + + out = detect_gguf_model(str(snap)) + assert out is not None + assert Path(out).name == "model-BF16-00001-of-00002.gguf" + + def test_lone_split_symlink_uses_colocated_target_shards(self, tmp_path): + from utils.models.model_config import _find_local_gguf_by_variant, detect_gguf_model + + target_dir = tmp_path / "external" / "BF16" + target_dir.mkdir(parents = True) + target = target_dir / "model-BF16-00001-of-00002.gguf" + target.write_bytes(b"\0" * 10) + (target_dir / "model-BF16-00002-of-00002.gguf").write_bytes(b"\0" * 10) + + local = tmp_path / "local" + local.mkdir() + (local / "config.json").write_text("{}") + link = local / target.name + _symlink_or_skip(link, target) + + expected = str(target.absolute()) + assert _find_local_gguf_by_variant(str(local), "BF16") == expected + assert detect_gguf_model(str(local)) == expected + assert detect_gguf_model(str(link)) == expected + def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path): from utils.models.model_config import ModelConfig diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 1ee9ef36d3..9c6c20e6b6 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -68,7 +68,13 @@ class _LoadRecorder: request, fastapi_request, current_subject = None, + *, + current_request_counted = False, ): + # Mirror the production load boundary before recording any replacement. + await inference_route._wait_for_model_switch_idle( + current_request_counted = current_request_counted + ) self.calls.append(request) if self.fail: from fastapi import HTTPException @@ -94,7 +100,6 @@ def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder): # gate that auto-switch already owns, so it calls the impl directly). monkeypatch.setattr(inference_route, "_load_model_impl", recorder) monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) - monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) def _run_hook(model = "some/model"): @@ -1091,6 +1096,7 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch from pathlib import Path import routes.models as models_route from utils import paths as upaths + from utils import hf_cache_settings import storage.studio_db as studio_db scanned = [] @@ -1111,13 +1117,18 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch ) monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active") monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr( + hf_cache_settings, + "known_hf_hub_caches", + lambda: [tmp_path / "active", tmp_path / "previous"], + ) monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy") monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default") monkeypatch.setattr(upaths, "lmstudio_model_dirs", lambda: [tmp_path / "lmstudio"]) monkeypatch.setattr( studio_db, "list_scan_folders", lambda: [{"path": str(tmp_path / "custom")}] ) - for sub in ("active", "legacy", "default", "lmstudio", "custom"): + for sub in ("active", "previous", "legacy", "default", "lmstudio", "custom"): (tmp_path / sub).mkdir() resolver._build_index() @@ -1126,6 +1137,7 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch lm = {p for k, p in scanned if k == "lm"} assert str((tmp_path / "legacy").resolve()) in hf assert str((tmp_path / "default").resolve()) in hf + assert str((tmp_path / "previous").resolve()) in hf assert str((tmp_path / "custom").resolve()) in hf assert str((tmp_path / "lmstudio").resolve()) in lm @@ -1205,10 +1217,9 @@ def test_middleware_ignores_non_post(monkeypatch): # โ”€โ”€ review round 4: swap guard, idle variant identity, load-by-path, stash clear โ”€โ”€ -def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): - # A cross-model swap must 409 (not kill) while another inference request is in - # flight; the requesting call itself is excluded from the count. - from fastapi import HTTPException +def test_auto_switch_waits_for_another_inference_to_finish(monkeypatch): + # A cross-model swap queues while another request is generating, then loads + # after that request drains. The requesting call itself is excluded. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M") @@ -1222,10 +1233,18 @@ def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): ) monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one monkeypatch.setattr(kw, "_pending", 0) - with pytest.raises(HTTPException) as exc: - _run_hook("org/B-GGUF:Q8_0") - assert exc.value.status_code == 409 - assert rec.calls == [] + + async def _drive(): + task = asyncio.create_task( + inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester") + ) + await asyncio.sleep(0.05) + assert rec.calls == [] + kw._note_end() # the other generation finishes; this request remains counted + await asyncio.wait_for(task, timeout = 1) + + asyncio.run(_drive()) + assert len(rec.calls) == 1 def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch): @@ -1411,13 +1430,12 @@ def test_concurrent_same_target_requests_load_once(monkeypatch): monkeypatch.setattr(kw, "_pending", 0) inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) _run_hook("org/B-GGUF:Q8_0") - assert len(rec.calls) == 1 # loads once, no 409 + assert len(rec.calls) == 1 -def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch): - # A concurrent request heading to a different target still blocks the swap: the - # same-target exclusion must not swallow a genuinely conflicting request. - from fastapi import HTTPException +def test_queued_different_target_does_not_deadlock_current_swap(monkeypatch): + # A concurrent request already queued for another target is not generating, + # so it must not prevent the current serialized swap from proceeding. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF") @@ -1432,10 +1450,8 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat monkeypatch.setattr(kw, "_inflight", 2) monkeypatch.setattr(kw, "_pending", 0) inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) - with pytest.raises(HTTPException) as exc: - _run_hook("org/B-GGUF:Q8_0") - assert exc.value.status_code == 409 - assert rec.calls == [] + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): @@ -1481,6 +1497,37 @@ def test_load_route_holds_lifecycle_gate(monkeypatch): assert "_load_model_impl" in src +def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded(): + # Both replacement directions drain active inference, then recheck whether a + # sidecar install reserved the lifecycle gate during that wait. Exact-model + # reuse exits earlier, so an already-loaded model never waits on unrelated inference. + import inspect + + src = inspect.getsource(inference_route._load_model_impl) + gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:")) + gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait) + unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait) + standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1) + standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait) + unload_gguf = src.index("llama_backend.unload_model()", standard_wait) + already_loaded = src.index('status = "already_loaded"') + + assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth + assert standard_wait < standard_sidecar_check < unload_gguf + + +def test_switch_waiter_deregisters_before_swap_gate_release(): + # A waiter left registered after the swap gate is released would let a swap on + # another event loop count the finished request as still queued, pass the drain + # early, and unload the model that request is about to generate against. + import inspect + + src = inspect.getsource(inference_route._maybe_auto_switch_model) + deregister = src.index("_note_switch_waiter(key, -1)") + release = src.index("_auto_switch_process_lock.release()") + assert deregister < release + + def _anthropic_payload(max_tokens = None): from models.inference import AnthropicMessagesRequest, AnthropicMessage return AnthropicMessagesRequest( @@ -1519,9 +1566,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): # โ”€โ”€ review round 6: concurrency ordering, external untrack, unload gate, ids โ”€โ”€ -def test_pending_same_target_request_does_not_force_409(monkeypatch): +def test_pending_same_target_request_does_not_block_swap(monkeypatch): # A second same-target request blocked in the middleware (pending, not yet - # generating) must not make the first request 409: pending is excluded. + # generating) must not block the first request: pending is excluded. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF") @@ -1536,13 +1583,13 @@ def test_pending_same_target_request_does_not_force_409(monkeypatch): monkeypatch.setattr(kw, "_inflight", 1) # just the caller monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware _run_hook("org/B-GGUF:Q8_0") - assert len(rec.calls) == 1 # loads once, no 409 + assert len(rec.calls) == 1 -def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch): +def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch): # The real middleware counts a concurrent same-model request as in-flight - # before it resolves and registers a target waiter. The raw-request waiter, - # registered before resolve, must still exclude it so the first request loads. + # before it resolves and registers a target waiter. Treat it as active until + # its target is known, then recognize it as another queued switch request. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF") @@ -1556,10 +1603,20 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat ) monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin monkeypatch.setattr(kw, "_pending", 0) - # The twin has only registered its raw requested model (not yet a target waiter). - inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) - _run_hook("org/B-GGUF:Q8_0") - assert len(rec.calls) == 1 # loads once, no 409 + # The twin is still resolving, so it is counted in-flight but has not joined + # the concrete target queue yet. + + async def _drive(): + task = asyncio.create_task( + inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester") + ) + await asyncio.sleep(0.05) + assert rec.calls == [] + inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + await asyncio.wait_for(task, timeout = 1) + + asyncio.run(_drive()) + assert len(rec.calls) == 1 def test_external_untrack_decrements_inflight_and_is_idempotent(): @@ -1595,11 +1652,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): assert not backend.is_loaded # torn down despite the active request -def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): +def test_auto_switch_waits_when_unsloth_stream_active(monkeypatch): # The GGUF slot is empty but an Unsloth model is streaming (counted in-flight). - # _load_model_impl would unload it, so auto-switch must 409, not only when a - # GGUF is loaded. - from fastapi import HTTPException + # The replacement waits for it just as it does for a GGUF generation. from core.inference import llama_keepwarm as kw backend = _FakeBackend(None) # no GGUF loaded @@ -1613,10 +1668,18 @@ def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): ) monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request monkeypatch.setattr(kw, "_pending", 0) - with pytest.raises(HTTPException) as exc: - _run_hook("org/B-GGUF:Q8_0") - assert exc.value.status_code == 409 - assert rec.calls == [] # the active Unsloth model is not torn down + + async def _drive(): + task = asyncio.create_task( + inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester") + ) + await asyncio.sleep(0.05) + assert rec.calls == [] + kw._note_end() + await asyncio.wait_for(task, timeout = 1) + + asyncio.run(_drive()) + assert len(rec.calls) == 1 def test_public_model_id_prefers_advertised_over_path(): @@ -3097,6 +3160,8 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch): request, fastapi_request, current_subject = None, + *, + current_request_counted = False, ): with slock: state["cur"] += 1 @@ -3114,7 +3179,6 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch): monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load) monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) - monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) barrier = threading.Barrier(2) diff --git a/studio/backend/tests/test_password_prompt.py b/studio/backend/tests/test_password_prompt.py index 372d6a2aa4..1af8836065 100644 --- a/studio/backend/tests/test_password_prompt.py +++ b/studio/backend/tests/test_password_prompt.py @@ -183,6 +183,22 @@ def test_loop_short_password_reprompts(monkeypatch): assert "at least 8 characters" in out +def test_loop_whitespace_only_reprompts(monkeypatch): + ok, applied, out = _run_loop(monkeypatch, _keys(" " * 8, "long-enough-pw", "long-enough-pw")) + assert ok is True + assert applied == ["long-enough-pw"] + assert "contain spaces" in out + + +def test_loop_password_with_inner_space_reprompts(monkeypatch): + ok, applied, out = _run_loop( + monkeypatch, _keys("has space pw", "long-enough-pw", "long-enough-pw") + ) + assert ok is True + assert applied == ["long-enough-pw"] + assert "contain spaces" in out + + def test_loop_rejects_current_password(monkeypatch): ok, applied, out = _run_loop( monkeypatch, _keys("bootstrap-pw", "fresh-password", "fresh-password") diff --git a/studio/backend/tests/test_picker_service.py b/studio/backend/tests/test_picker_service.py new file mode 100644 index 0000000000..1bdfc135e3 --- /dev/null +++ b/studio/backend/tests/test_picker_service.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +from types import SimpleNamespace + +from picker.service import ( + MAX_TEMPLATE_METADATA_BYTES, + _chat_template_from_dir, + _chat_template_from_processor_json, + _chat_template_from_tokenizer_config, + _chat_template_from_tokenizer_dir, + _find_gguf_in_dir, + _iter_ggufs, + read_default_chat_template, + validate_chat_template, +) + + +def test_iter_ggufs_skips_gguf_companions(tmp_path): + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + main = tmp_path / "model-Q8_0.gguf" + main.write_bytes(b"") + (tmp_path / "mmproj-F16.gguf").write_bytes(b"") + (tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"") + (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"") + (tmp_path / "model-Q8_0-be.gguf").write_bytes(b"") + + assert _iter_ggufs(tmp_path) == [main] + + +def test_find_gguf_in_dir_matches_quant_label(tmp_path): + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + main = tmp_path / "model-Q8_0.gguf" + main.write_bytes(b"") + (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"") + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + + assert _find_gguf_in_dir(tmp_path, "Q8_0") == main + assert _find_gguf_in_dir(tmp_path, "Q4_K") is None + + +def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path): + smaller = tmp_path / "a-model-Q4_K_M.gguf" + larger = tmp_path / "z-model-Q8_0.gguf" + smaller.write_bytes(b"0") + larger.write_bytes(b"00") + + assert _find_gguf_in_dir(tmp_path, None) == larger + + +def test_find_gguf_in_dir_without_variant_prefers_first_split(tmp_path): + first = tmp_path / "model-Q4_K_M-00001-of-00003.gguf" + second = tmp_path / "model-Q4_K_M-00002-of-00003.gguf" + third = tmp_path / "model-Q4_K_M-00003-of-00003.gguf" + first.write_bytes(b"0") + second.write_bytes(b"000") + third.write_bytes(b"00") + + assert _find_gguf_in_dir(tmp_path, None) == first + + first.unlink() + assert _find_gguf_in_dir(tmp_path, None) == second + + +def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path): + target = tmp_path / "model-IQ4_XS-3.53bpw.gguf" + target.write_bytes(b"") + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + + assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target + assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target + assert _find_gguf_in_dir(tmp_path, "Q4_K") is None + + +def test_validate_chat_template_accepts_valid_and_empty(): + assert validate_chat_template("{{ messages[0].content }}").valid is True + assert validate_chat_template("").valid is True + assert validate_chat_template(" ").valid is True + + +def test_validate_chat_template_reports_syntax_error_with_line(): + result = validate_chat_template("{% if %}{% endif %}") + assert result.valid is False + assert result.error is not None + assert result.error.startswith("Line ") + + +def test_chat_template_from_tokenizer_config_reads_string(): + assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO" + assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None + assert _chat_template_from_tokenizer_config({}) is None + + +def test_chat_template_from_tokenizer_config_prefers_named_default(): + config = { + "chat_template": [ + {"name": "tool_use", "template": "TOOL"}, + {"name": "default", "template": "DEFAULT"}, + ] + } + assert _chat_template_from_tokenizer_config(config) == "DEFAULT" + + +def test_chat_template_from_tokenizer_config_falls_back_to_first_entry(): + config = { + "chat_template": [ + {"name": "tool_use", "template": "TOOL"}, + {"name": "other", "template": "OTHER"}, + ] + } + assert _chat_template_from_tokenizer_config(config) == "TOOL" + + +def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path): + (tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA" + + +def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG" + + +def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG" + + +def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # Selecting a variant must not flip precedence to the embedded GGUF template. + assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG" + + +def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch): + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # With no tokenizer sidecar, the embedded GGUF template is still the fallback. + assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF" + + +def test_chat_template_from_dir_returns_none_when_absent(tmp_path): + assert _chat_template_from_dir(tmp_path) is None + + +def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch): + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path]) + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # A directly selected .gguf must prefer a maintained sidecar over its embedded copy. + assert read_default_chat_template(str(gguf)) == "FROM_CONFIG" + + +def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch): + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path]) + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # With no sidecar next to the file, the embedded GGUF template is the fallback. + assert read_default_chat_template(str(gguf)) == "FROM_GGUF" + + +def test_tokenizer_config_over_size_limit_is_skipped_not_parsed(tmp_path): + # An oversized tokenizer_config.json must be skipped before json.loads so a + # hostile sidecar cannot exhaust memory. + padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "HELLO", "_pad": padding}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) is None + + +def test_processor_json_over_size_limit_is_skipped_not_parsed(tmp_path): + padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024) + (tmp_path / "chat_template.json").write_text( + json.dumps({"default": "HELLO", "_pad": padding}), encoding = "utf-8" + ) + assert _chat_template_from_processor_json(tmp_path) is None + + +def test_tokenizer_config_at_size_limit_is_still_read(tmp_path): + # A normal-sized config is unaffected by the bound (regression guard). + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG" + + +def test_remote_template_over_size_limit_is_skipped_before_download(monkeypatch): + # An uncached Hub repo whose template exceeds the cap must be skipped via the + # remote size pre-check, never downloaded. + import huggingface_hub + + monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: []) + + def _fail_download(*args, **kwargs): + raise AssertionError("oversized remote template must not be downloaded") + + def _fake_get_paths_info(self, repo_id, paths, **kwargs): + return [SimpleNamespace(path = p, size = MAX_TEMPLATE_METADATA_BYTES + 1) for p in paths] + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fail_download) + monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info) + + assert read_default_chat_template("org/oversized-model") is None + + +def test_remote_oversized_jinja_falls_through_to_tokenizer_template(tmp_path, monkeypatch): + # A raw chat_template.jinja between the response cap (MAX_CHAT_TEMPLATE_BYTES) + # and the download bound (MAX_TEMPLATE_METADATA_BYTES) must not be returned: the + # route drops it, so the remote path must skip the oversized Jinja and fall + # through to the smaller tokenizer_config.json. + import huggingface_hub + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + + big_jinja = tmp_path / "chat_template.jinja" + big_jinja.write_text("{{ x }}" * (MAX_CHAT_TEMPLATE_BYTES // 4), encoding = "utf-8") + assert MAX_CHAT_TEMPLATE_BYTES < big_jinja.stat().st_size < MAX_TEMPLATE_METADATA_BYTES + tokenizer_config = tmp_path / "tokenizer_config.json" + tokenizer_config.write_text(json.dumps({"chat_template": "SMALL_TEMPLATE"}), encoding = "utf-8") + files = { + "chat_template.jinja": big_jinja, + "tokenizer_config.json": tokenizer_config, + } + selected_cache = tmp_path / "selected-cache" / "hub" + observed_cache_dirs = [] + + monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: []) + monkeypatch.setattr("picker.service.active_hf_hub_cache", lambda: str(selected_cache)) + + def _fake_download(repo_id, rel, **kwargs): + observed_cache_dirs.append(kwargs.get("cache_dir")) + target = files.get(rel) + if target is None: + raise FileNotFoundError(rel) + return str(target) + + def _fake_get_paths_info(self, repo_id, paths, **kwargs): + return [ + SimpleNamespace( + path = p, + size = files[p].stat().st_size if p in files else 0, + ) + for p in paths + ] + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info) + + assert read_default_chat_template("org/big-jinja-model") == "SMALL_TEMPLATE" + assert observed_cache_dirs + assert set(observed_cache_dirs) == {str(selected_cache)} diff --git a/studio/backend/tests/test_providers_db_models.py b/studio/backend/tests/test_providers_db_models.py new file mode 100644 index 0000000000..ca9dffbd70 --- /dev/null +++ b/studio/backend/tests/test_providers_db_models.py @@ -0,0 +1,70 @@ +# 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 provider model persistence (unslothai/unsloth#7281).""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import storage.providers_db as providers_db + + +@pytest.fixture() +def isolated_providers_db(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + db_path = tmp_path / "studio.db" + monkeypatch.setattr(providers_db, "studio_db_path", lambda: db_path) + monkeypatch.setattr(providers_db, "ensure_dir", lambda _path: None) + providers_db._schema_ready = False + yield db_path + providers_db._schema_ready = False + + +def test_create_and_list_provider_models(isolated_providers_db: Path): + providers_db.create_provider( + id = "ollama1", + provider_type = "ollama", + display_name = "Home Ollama", + base_url = "http://127.0.0.1:11434", + models = ["llama3.2", "qwen2.5"], + available_models = ["llama3.2", "qwen2.5", "mistral"], + ) + + row = providers_db.get_provider("ollama1") + assert row is not None + assert row["models"] == ["llama3.2", "qwen2.5"] + assert row["available_models"] == ["llama3.2", "qwen2.5", "mistral"] + + listed = providers_db.list_providers() + assert len(listed) == 1 + assert listed[0]["models"] == ["llama3.2", "qwen2.5"] + + +def test_update_provider_models(isolated_providers_db: Path): + providers_db.create_provider( + id = "vllm1", + provider_type = "vllm", + display_name = "Remote vLLM", + base_url = "http://studio-host:8000/v1", + models = ["meta-llama/Llama-3.2-1B-Instruct"], + available_models = ["meta-llama/Llama-3.2-1B-Instruct"], + ) + + assert providers_db.update_provider( + id = "vllm1", + models = ["meta-llama/Llama-3.2-3B-Instruct"], + available_models = [ + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ], + ) + + row = providers_db.get_provider("vllm1") + assert row is not None + assert row["models"] == ["meta-llama/Llama-3.2-3B-Instruct"] + assert row["available_models"] == [ + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ] diff --git a/studio/backend/tests/test_rag_embeddings.py b/studio/backend/tests/test_rag_embeddings.py index 28a2f69426..197ae4c495 100644 --- a/studio/backend/tests/test_rag_embeddings.py +++ b/studio/backend/tests/test_rag_embeddings.py @@ -5,8 +5,10 @@ and token counting must be serialized (else threads panic "Already borrowed").""" import os +import sys import threading import time +from types import SimpleNamespace import numpy as np import pytest @@ -130,6 +132,35 @@ def test_token_counter_enables_parallelism_only_during_call(monkeypatch): assert os.environ.get("TOKENIZERS_PARALLELISM") == "false" # restored after +def test_sentence_transformer_load_uses_live_cache(monkeypatch, tmp_path): + observed = {} + + class FakeSentenceTransformer: + def __init__(self, name, **kwargs): + observed["name"] = name + observed.update(kwargs) + + monkeypatch.setitem( + sys.modules, + "sentence_transformers", + SimpleNamespace(SentenceTransformer = FakeSentenceTransformer), + ) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_guard_model_security", lambda *_a, **_k: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + monkeypatch.setattr( + "utils.hf_cache_settings.active_hf_hub_cache", + lambda: str(tmp_path / "selected-hub"), + ) + embeddings._model = None + embeddings._name = None + + embeddings._get("Org/Embedder") + + assert observed["name"] == "Org/Embedder" + assert observed["cache_folder"] == str(tmp_path / "selected-hub") + + class _SentinelLlamaBackend: """Stand-in for LlamaServerBackend; never spawns a real server.""" diff --git a/studio/backend/tests/test_resolve_quant_gguf.py b/studio/backend/tests/test_resolve_quant_gguf.py index 840c4d8d4c..a137237e80 100644 --- a/studio/backend/tests/test_resolve_quant_gguf.py +++ b/studio/backend/tests/test_resolve_quant_gguf.py @@ -68,8 +68,6 @@ def test_skips_mtp_drafter_for_main_weights(tmp_path): def test_prefers_the_complete_snapshot(tmp_path, monkeypatch): - from huggingface_hub import constants as hf_constants - cache = tmp_path / "hub" snaps = cache / "models--org--repo" / "snapshots" # Partial older snapshot: one small shard. @@ -78,7 +76,10 @@ def test_prefers_the_complete_snapshot(tmp_path, monkeypatch): complete_first = _write(snaps / "bbbb" / "model-00001-of-00002-Q4_K_M.gguf", 30) _write(snaps / "bbbb" / "model-00002-of-00002-Q4_K_M.gguf", 40) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(cache)) + monkeypatch.setattr( + "utils.hf_cache_settings.known_hf_hub_caches", + lambda: [cache], + ) path, total = models_route._resolve_quant_gguf("org/repo", "Q4_K_M", is_local = False) diff --git a/studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py b/studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py new file mode 100644 index 0000000000..bdafdeae9b --- /dev/null +++ b/studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py @@ -0,0 +1,554 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The System tab's multi-GPU view must show system-wide VRAM on ROCm (#7072). + +When amd-smi is unavailable, get_visible_gpu_utilization fell back to torch, +whose readings are process-local: a model held by the separate llama-server +process read as ~0 VRAM used even with the GPU full. These tests cover the +per-GPU system-wide overlay the multi-device endpoint now applies, matched by +physical device identity. +""" + +from __future__ import annotations + +import importlib +import sys +import types +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent +if str(_BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(_BACKEND_DIR)) + + +def _maybe_stub(name: str, builder): + # Stub only if the real module is missing, so we never shadow it for later tests. + try: + importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +def _build_structlog_stub(): + m = types.ModuleType("structlog") + m.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", _build_structlog_stub) + +import utils.hardware.hardware as hw # noqa: E402 + + +def _device( + index, + used, + total, + *, + ordinal = None, +): + return { + "index": index, + "index_kind": "physical", + "visible_ordinal": index if ordinal is None else ordinal, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": used, + "vram_total_gb": total, + "vram_utilization_pct": round((used / total) * 100, 1) if total > 0 else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + + +# โ”€โ”€ Linux per-card sysfs โ”€โ”€ + + +def _fake_drm(tmp_path, monkeypatch, cards): + """Fake /sys/class/drm tree; glob returns cards REVERSED so the PCI sort must order them. + + ``cards``: (card_no, pci_bdf, driver, vram) tuples; vram is (used_gb, total_gb) + or None for a device with no mem_info_vram_* files. + """ + drivers = tmp_path / "drivers" + card_paths = [] + for card_no, bdf, driver, vram in cards: + pci_dir = tmp_path / "pci" / bdf + pci_dir.mkdir(parents = True, exist_ok = True) + drv_dir = drivers / driver + drv_dir.mkdir(parents = True, exist_ok = True) + (pci_dir / "driver").symlink_to(drv_dir) + if vram is not None: + used, total = vram + (pci_dir / "mem_info_vram_used").write_text(str(int(used * 1024**3))) + (pci_dir / "mem_info_vram_total").write_text(str(int(total * 1024**3))) + card_dir = tmp_path / "drm" / f"card{card_no}" + card_dir.mkdir(parents = True, exist_ok = True) + (card_dir / "device").symlink_to(pci_dir) + card_paths.append(str(card_dir)) + monkeypatch.setattr(hw.glob, "glob", lambda pattern: list(reversed(card_paths))) + return card_paths + + +def test_linux_vram_keyed_by_pci_excludes_foreign_adapters(monkeypatch, tmp_path): + # Foreign (non-amdgpu) adapters contribute no entry, so they cannot shift ordinals. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_drm( + tmp_path, + monkeypatch, + [ + (0, "0000:00:02.0", "i915", (0.5, 2.0)), # foreign adapter: excluded + (1, "0000:03:00.0", "amdgpu", (40, 48)), # AMD device 0 + (2, "0000:41:00.0", "amdgpu", (1, 8)), # AMD device 1 + ], + ) + assert hw._rocm_linux_sysfs_vram_by_pci_gb() == { + "0000:03:00.0": (40.0, 48.0), + "0000:41:00.0": (1.0, 8.0), + } + + +def test_linux_vram_omits_bad_cards_without_shifting(monkeypatch, tmp_path): + # A zero-total card has no entry; identity keying means its absence renumbers nothing. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_drm( + tmp_path, + monkeypatch, + [ + (0, "0000:03:00.0", "amdgpu", (0, 0)), # zero total -> no entry + (1, "0000:41:00.0", "amdgpu", (2, 16)), + ], + ) + assert hw._rocm_linux_sysfs_vram_by_pci_gb() == {"0000:41:00.0": (2.0, 16.0)} + + +def test_linux_vram_omits_amd_card_without_vram_files(monkeypatch, tmp_path): + # An APU with no mem_info_vram_* files has no entry; the discrete card keeps its address. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_drm( + tmp_path, + monkeypatch, + [ + (0, "0000:03:00.0", "amdgpu", None), # APU: no VRAM sysfs files + (1, "0000:41:00.0", "amdgpu", (2, 16)), + ], + ) + assert hw._rocm_linux_sysfs_vram_by_pci_gb() == {"0000:41:00.0": (2.0, 16.0)} + + +# โ”€โ”€ KFD topology: the authoritative ROCm device order โ”€โ”€ + + +_AMD = 4098 # 0x1002 +_NVIDIA = 4318 # 0x10DE -- the open kernel module also registers KFD nodes + + +def _fake_kfd(tmp_path, monkeypatch, nodes): + """Fake KFD topology nodes tree, returned out of node order so the sort must order it. + + ``nodes``: (node_id, simd_count, location_id, domain, vendor_id); simd_count 0 + marks a CPU node, location_id None omits the property. + """ + node_paths = [] + for node_id, simd_count, location_id, domain, vendor_id in nodes: + d = tmp_path / "kfd" / str(node_id) + d.mkdir(parents = True, exist_ok = True) + lines = [f"cpu_cores_count {0 if simd_count else 8}", f"simd_count {simd_count}"] + if location_id is not None: + lines.append(f"location_id {location_id}") + lines.append(f"domain {domain}") + if vendor_id is not None: + lines.append(f"vendor_id {vendor_id}") + (d / "properties").write_text("\n".join(lines) + "\n") + node_paths.append(str(d)) + monkeypatch.setattr(hw.glob, "glob", lambda pattern: list(reversed(node_paths))) + return node_paths + + +def test_kfd_lists_gpu_nodes_in_device_order(monkeypatch, tmp_path): + # The CPU node (simd_count 0) takes no ordinal; GPU nodes in node-id order are HIP's order. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_kfd( + tmp_path, + monkeypatch, + [ + (0, 0, None, 0, None), # CPU node + (1, 304, (0x03 << 8) | (0x00 << 3) | 0, 0, _AMD), # 0000:03:00.0 -> dev 0 + (2, 304, (0x41 << 8) | (0x00 << 3) | 0, 0, _AMD), # 0000:41:00.0 -> dev 1 + ], + ) + assert hw._rocm_kfd_gpu_pci_ids() == ["0000:03:00.0", "0000:41:00.0"] + + +def test_kfd_decodes_domain_device_and_function(monkeypatch, tmp_path): + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_kfd(tmp_path, monkeypatch, [(1, 64, (0xC1 << 8) | (0x1F << 3) | 5, 0x1234, _AMD)]) + assert hw._rocm_kfd_gpu_pci_ids() == ["1234:c1:1f.5"] + + +def test_kfd_skips_non_amd_gpu_nodes(monkeypatch, tmp_path): + # An NVIDIA KFD node is not a HIP device: it must take no ordinal, else it + # shifts every AMD GPU and ROCm device 1 resolves to AMD GPU 0. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_kfd( + tmp_path, + monkeypatch, + [ + (0, 0, None, 0, None), # CPU + (1, 128, (0x01 << 8) | 0, 0, _NVIDIA), # NVIDIA: no ordinal + (2, 304, (0x03 << 8) | 0, 0, _AMD), # AMD device 0 + (3, 304, (0x41 << 8) | 0, 0, _AMD), # AMD device 1 + ], + ) + assert hw._rocm_kfd_gpu_pci_ids() == ["0000:03:00.0", "0000:41:00.0"] + + +def test_kfd_fails_closed_when_a_gpu_has_no_location(monkeypatch, tmp_path): + # Dropping an unplaceable AMD GPU shifts later ordinals; fail closed for the whole map. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _fake_kfd( + tmp_path, + monkeypatch, + [ + (1, 304, None, 0, _AMD), # AMD GPU with no location_id + (2, 304, (0x41 << 8) | 0, 0, _AMD), + ], + ) + assert hw._rocm_kfd_gpu_pci_ids() == [] + + +def test_kfd_fails_closed_when_a_node_is_unreadable(monkeypatch, tmp_path): + # An unreadable node could be a GPU; assuming otherwise would shift ordinals. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + paths = _fake_kfd( + tmp_path, + monkeypatch, + [ + (1, 304, (0x03 << 8) | 0, 0, _AMD), + (2, 304, (0x41 << 8) | 0, 0, _AMD), + ], + ) + (Path(paths[0]) / "properties").unlink() + assert hw._rocm_kfd_gpu_pci_ids() == [] + + +def test_kfd_absent_yields_no_device_order(monkeypatch): + monkeypatch.setattr(hw.glob, "glob", lambda pattern: []) + assert hw._rocm_kfd_gpu_pci_ids() == [] + + +# โ”€โ”€ overlay โ”€โ”€ + + +def _patch_pci_map(monkeypatch, bdfs): + """Declare the ROCm device order by PCI address (index N is device N) and clear + the visibility masks the overlay requires unset. + """ + for var in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", + ): + monkeypatch.delenv(var, raising = False) + monkeypatch.setattr(hw, "_rocm_kfd_gpu_pci_ids", lambda: list(bdfs)) + + +def _pci(n): + """A distinct, well-formed PCI address for card n.""" + return f"0000:{n:02x}:00.0" + + +def test_overlay_windows_is_noop_keeps_torch(monkeypatch): + # Windows is intentionally not overlaid (perf counters can't map to ROCm ordinals): keep torch. + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr( + hw, + "_rocm_linux_sysfs_vram_by_pci_gb", + lambda: (_ for _ in ()).throw(AssertionError("sysfs must not run on Windows")), + ) + devices = [_device(0, used = 0.02, total = 8.0)] + _patch_pci_map(monkeypatch, [_pci(0)]) + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 # untouched + + +def test_overlay_linux_matches_by_device_ordinal(monkeypatch): + # Devices arriving as [index 1, index 0] each get their own GPU's figures by ordinal. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr( + hw, + "_rocm_linux_sysfs_vram_by_pci_gb", + lambda: {_pci(0): (30.0, 45.0), _pci(1): (0.5, 8.0)}, # dev 0 big, dev 1 small + ) + devices = [_device(1, used = 0.01, total = 8.0), _device(0, used = 0.02, total = 45.0)] + _patch_pci_map(monkeypatch, [_pci(0), _pci(1)]) + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.5 # index 1 -> device 1 (small) + assert devices[0]["vram_total_gb"] == 8.0 + assert devices[1]["vram_used_gb"] == 30.0 # index 0 -> device 0 (big) + assert devices[1]["vram_total_gb"] == 45.0 + + +def test_overlay_linux_ordinal_hole_does_not_shift(monkeypatch): + # Device 0's card dropped: index 0 keeps torch, index 1 still maps to ordinal 1 (no compaction). + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(1): (0.5, 8.0)}) + devices = [_device(0, used = 0.02, total = 45.0), _device(1, used = 0.01, total = 8.0)] + _patch_pci_map(monkeypatch, [_pci(0), _pci(1)]) + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 # no ordinal 0 -> torch kept + assert devices[1]["vram_used_gb"] == 0.5 # ordinal 1 -> device 1, not device 0 + + +def test_overlay_linux_skips_unified_memory_card(monkeypatch): + # Unified-memory APU: the smaller sysfs total must not shrink torch's GTT-backed pool. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (0.4, 1.0)}) + devices = [_device(0, used = 12.0, total = 96.0)] # torch's unified pool + _patch_pci_map(monkeypatch, [_pci(0)]) + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 12.0 + assert devices[0]["vram_total_gb"] == 96.0 + + +def test_overlay_linux_skips_partitioned_device(monkeypatch): + # Partitioned MI300: the whole-card sysfs total dwarfs the partition, so the overlay must not overwrite it. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (40.0, 192.0)}) + devices = [_device(0, used = 1.0, total = 24.0)] # torch partition + _patch_pci_map(monkeypatch, [_pci(0)]) + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 1.0 # partition figures kept + assert devices[0]["vram_total_gb"] == 24.0 + + +def test_overlay_linux_out_of_range_index_untouched(monkeypatch): + # A masked host exposing physical index 5 with no card 5: keep torch data. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr( + hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (30.0, 45.0), _pci(1): (0.5, 8.0)} + ) + devices = [_device(5, used = 0.02, total = 45.0)] + _patch_pci_map(monkeypatch, [_pci(0), _pci(1)]) + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 + + +def test_overlay_ignores_adapters_rocm_cannot_enumerate(monkeypatch): + # A HIP-unenumerable amdgpu adapter has no KFD node, so device 0 resolves to + # the supported GPU's own address, never the display card's. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr( + hw, + "_rocm_linux_sysfs_vram_by_pci_gb", + # Both in DRM sysfs with similar capacity -- what the total-size guard can't separate. + lambda: {_pci(9): (30.0, 45.0), _pci(3): (12.0, 45.0)}, + ) + _patch_pci_map(monkeypatch, [_pci(3)]) # KFD lists only the supported GPU + devices = [_device(0, used = 0.02, total = 45.0)] # torch sees that one GPU + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 12.0 # the supported GPU's own figures + + +def test_overlay_skips_masked_subsets(monkeypatch): + # Under a mask the index is not verifiably a host ordinal, so keep torch's figures. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _patch_pci_map(monkeypatch, [_pci(0), _pci(1), _pci(2), _pci(3)]) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1,3") + monkeypatch.setattr( + hw, + "_rocm_linux_sysfs_vram_by_pci_gb", + lambda: {_pci(1): (30.0, 48.0), _pci(3): (12.0, 48.0)}, + ) + devices = [_device(1, used = 0.02, total = 48.0), _device(3, used = 0.01, total = 48.0)] + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 # torch kept + assert devices[1]["vram_used_gb"] == 0.01 + + +def test_overlay_skips_device_cgroup_filtered_container(monkeypatch): + # A device-cgroup container sets no env var yet compacts torch's indices from + # zero while KFD/DRM list every GPU, so the count mismatch must disable the overlay. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _patch_pci_map(monkeypatch, [_pci(0), _pci(1), _pci(2), _pci(3)]) # host has 4 + monkeypatch.setattr( + hw, + "_rocm_linux_sysfs_vram_by_pci_gb", + lambda: {_pci(0): (30.0, 48.0), _pci(2): (12.0, 48.0)}, + ) + devices = [_device(0, used = 0.02, total = 48.0)] # container sees 1, as index 0 + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 # torch kept, not host GPU 0's 30.0 + + +def test_overlay_skips_without_kfd_topology(monkeypatch): + # No KFD means no identity to join on; fall back to torch rather than guess. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_rocm_kfd_gpu_pci_ids", lambda: []) + monkeypatch.setattr( + hw, + "_rocm_linux_sysfs_vram_by_pci_gb", + lambda: (_ for _ in ()).throw(AssertionError("must not read sysfs without KFD")), + ) + devices = [_device(0, used = 0.02, total = 45.0)] + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 + + +def test_overlay_empty_devices_is_noop(monkeypatch): + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + hw._overlay_system_wide_vram([]) # must not raise + + +# โ”€โ”€ integration: the ROCm torch fallback applies the overlay โ”€โ”€ + + +def test_visible_utilization_rocm_fallback_overlays(monkeypatch): + for _var in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", + ): + monkeypatch.delenv(_var, raising = False) + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) # amd-smi unavailable + monkeypatch.setattr( + hw, + "_get_parent_visible_gpu_spec", + lambda: {"raw": None, "numeric_ids": [0, 1], "supports_explicit_gpu_ids": True}, + ) + monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [0, 1]) + monkeypatch.setattr( + hw, + "_torch_get_per_device_info", + lambda ids: [ + {"index": 0, "visible_ordinal": 0, "used_gb": 0.02, "total_gb": 45.0}, + {"index": 1, "visible_ordinal": 1, "used_gb": 0.01, "total_gb": 8.0}, + ], + ) + overlaid = [] + monkeypatch.setattr( + hw, "_overlay_system_wide_vram", lambda devices: overlaid.append(len(devices)) + ) + result = hw.get_visible_gpu_utilization() + assert result["available"] is True + assert overlaid == [2] + + +def test_visible_utilization_relative_index_skips_overlay(monkeypatch): + # UUID/MIG mask gives relative indices; the overlay matches physical index, so it must not run. + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) + monkeypatch.setattr( + hw, + "_get_parent_visible_gpu_spec", + lambda: {"raw": "GPU-uuid-a", "numeric_ids": None, "supports_explicit_gpu_ids": False}, + ) + monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: []) # UUID mask + monkeypatch.setattr(hw, "_torch_get_physical_gpu_count", lambda: 1) + monkeypatch.setattr( + hw, + "_torch_get_per_device_info", + lambda ids: [{"index": 0, "visible_ordinal": 0, "used_gb": 0.02, "total_gb": 8.0}], + ) + called = [] + monkeypatch.setattr(hw, "_overlay_system_wide_vram", lambda devices: called.append(1)) + result = hw.get_visible_gpu_utilization() + assert result["index_kind"] == "relative" + assert called == [] + + +def test_visible_utilization_nvidia_fallback_skips_overlay(monkeypatch): + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) + monkeypatch.setattr( + hw, + "_get_parent_visible_gpu_spec", + lambda: {"raw": None, "numeric_ids": [0], "supports_explicit_gpu_ids": True}, + ) + monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [0]) + monkeypatch.setattr( + hw, + "_torch_get_per_device_info", + lambda ids: [{"index": 0, "visible_ordinal": 0, "used_gb": 1.0, "total_gb": 24.0}], + ) + called = [] + monkeypatch.setattr(hw, "_overlay_system_wide_vram", lambda devices: called.append(1)) + result = hw.get_visible_gpu_utilization() + assert result["available"] is True + assert called == [] + + +def test_any_visibility_mask_is_detected(monkeypatch): + # Any of these makes the index not a host-physical ordinal, so each must disable the overlay. + for var in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", + ): + monkeypatch.delenv(var, raising = False) + assert hw._rocm_visibility_mask_active() is False + for var in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", + ): + monkeypatch.setenv(var, "1") + assert hw._rocm_visibility_mask_active() is True, var + monkeypatch.setenv(var, " ") # empty is not an active filter + assert hw._rocm_visibility_mask_active() is False, var + monkeypatch.delenv(var, raising = False) + + +def test_overlay_skips_under_gpu_device_ordinal(monkeypatch): + # GPU_DEVICE_ORDINAL=1 surfaces GPU 1 as torch ordinal 0, so index 0 is not GPU 0; overlay must not run. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + _patch_pci_map(monkeypatch, [_pci(0)]) + monkeypatch.setenv("GPU_DEVICE_ORDINAL", "1") + monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (30.0, 45.0)}) + devices = [_device(0, used = 0.02, total = 45.0)] + hw._overlay_system_wide_vram(devices) + assert devices[0]["vram_used_gb"] == 0.02 + + +def test_visible_utilization_delegates_gating_to_the_overlay(monkeypatch): + # The call site no longer pre-checks masks; the overlay gates itself, so a physical payload always reaches it. + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "2,3") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1") + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) + monkeypatch.setattr( + hw, + "_get_parent_visible_gpu_spec", + lambda: {"raw": "1", "numeric_ids": [1], "supports_explicit_gpu_ids": True}, + ) + monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [1]) + monkeypatch.setattr( + hw, + "_torch_get_per_device_info", + lambda ids: [{"index": 1, "visible_ordinal": 0, "used_gb": 0.02, "total_gb": 8.0}], + ) + # Real overlay + gating: the layered mask must leave torch's figures. + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_rocm_kfd_gpu_pci_ids", lambda: [_pci(0), _pci(1)]) + monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(1): (30.0, 8.0)}) + result = hw.get_visible_gpu_utilization() + assert result["index_kind"] == "physical" + assert result["devices"][0]["vram_used_gb"] == 0.02 # untouched diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index 699d0b74f5..ad46f6ee41 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -163,6 +163,9 @@ class TestDeviceNameFallback: "AMD Radeon 8060S", "Radeon 8050S Graphics", # cut-down Strix Halo SKU "AMD Radeon 8050S", + # gfx1151 Gorgon Halo (Ryzen AI Max 400 refresh) + "Radeon 8065S Graphics", # Ryzen AI Max+ 495 + "AMD Radeon 8065S", # case variants "RADEON 8060S GRAPHICS", "radeon 8050s", diff --git a/studio/backend/tests/test_sampling_resolution.py b/studio/backend/tests/test_sampling_resolution.py new file mode 100644 index 0000000000..1ebbae2502 --- /dev/null +++ b/studio/backend/tests/test_sampling_resolution.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Effective sampling resolution: per-model recommendation + operator pins. + +Precedence per field: operator UNSLOTH_SAMPLING_* pin -> client explicit value -> +per-model recommendation (load_inference_config) -> static schema default. +""" + +import pytest + +from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES +from utils.inference import inference_config as ic + +_SCHEMA_DEFAULTS = { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.01, + "repetition_penalty": 1.0, + "presence_penalty": 0.0, +} + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + # The recommended lookup is lru-cached; clear it so a patched config takes effect. + ic._recommended_sampling.cache_clear() + for field in SAMPLING_FIELD_NAMES: + monkeypatch.delenv(ic._SAMPLING_FIELDS[field][0], raising = False) + yield + ic._recommended_sampling.cache_clear() + + +def _all_omitted(): + return {f: None for f in SAMPLING_FIELD_NAMES} + + +def _set_recommended(monkeypatch, mapping): + # _recommended_sampling sources from load_inference_config -- the exact block the Chat UI + # seeds from -- so patch that directly. Fields absent from `mapping` fall to schema defaults. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(mapping)) + ic._recommended_sampling.cache_clear() + + +def test_recommended_applies_when_client_omits(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 1.0 + assert eff["top_k"] == 64 + assert eff["min_p"] == 0.0 + # A field with no recommendation keeps the static schema default. + assert eff["top_p"] == 0.95 + + +def test_client_explicit_beats_recommended(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0}) + eff = resolve_effective_sampling("some/model", {**_all_omitted(), "temperature": 0.2}) + assert eff["temperature"] == 0.2 + + +def test_operator_pin_beats_client_and_recommended(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0}) + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + eff = resolve_effective_sampling("some/model", {**_all_omitted(), "temperature": 0.2}) + assert eff["temperature"] == 0.9 + + +def test_unknown_model_matches_ui_inference_block(monkeypatch): + # An unknown model gets the same values the Chat UI would seed (load_inference_config's + # default.yaml fallback: temp 0.7 / top_k -1), NOT the request schema defaults. + ui_block = { + "temperature": 0.7, + "top_p": 0.95, + "top_k": -1, + "min_p": 0.01, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, + } + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(ui_block)) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/unknown-model", _all_omitted()) + assert eff["temperature"] == 0.7 + assert eff["top_k"] == -1 + assert eff["min_p"] == 0.01 + + +def test_empty_recommendation_falls_back_to_schema_defaults(monkeypatch): + # If load_inference_config yields nothing usable, the resolver falls back to the request + # schema defaults. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: {}) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff == _SCHEMA_DEFAULTS + + +@pytest.mark.parametrize( + "model", + ["unsloth/gemma-4-E4B", "unsloth/Qwen3-4B", "unsloth/Qwen3.5-9B", "someorg/unknown-xyz"], +) +def test_recommendation_matches_ui_source(model): + # Parity guard: what the server recommends for omitted fields equals the Chat UI's source + # (load_inference_config) for every field the UI adopts (mergeBackendRecommendedInference). + ic._recommended_sampling.cache_clear() + ui = ic.load_inference_config(model) + rec = ic._recommended_sampling(model) + for f in ic._UI_RECOMMENDED_FIELDS: + cleaned = ic._clean_sampling_value(f, ui.get(f)) + if cleaned is not None: + assert rec.get(f) == cleaned, f"{model}:{f} rec={rec.get(f)} ui={ui.get(f)}" + + +def test_repetition_penalty_not_auto_recommended(monkeypatch): + # The Chat UI's mergeBackendRecommendedInference never adopts a backend repetition_penalty + # (e.g. lfm2's family value 1.05), so the server must not auto-apply one either. It stays at + # the schema default unless the client sends it or an operator pins it. + monkeypatch.setattr( + ic, "load_inference_config", lambda mid: {"temperature": 0.7, "repetition_penalty": 1.05} + ) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/lfm2-model", _all_omitted()) + assert eff["temperature"] == 0.7 # a UI-adopted field is recommended + assert eff["repetition_penalty"] == 1.0 # rep is NOT auto-recommended (matches the UI) + # An operator can still pin it explicitly. + monkeypatch.setenv("UNSLOTH_SAMPLING_REPETITION_PENALTY", "1.05") + eff2 = resolve_effective_sampling("some/lfm2-model", _all_omitted()) + assert eff2["repetition_penalty"] == 1.05 + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("0.5", 0.5), + ("abc", None), # unparseable + ("9.0", None), # above temperature max (2.0) + ("-1", None), # below temperature min (0.0) + (" ", None), # blank + ("nan", None), # NaN would pass a naive range check + ("inf", None), # non-finite + ("-inf", None), # non-finite + ], +) +def test_operator_override_parsing(monkeypatch, raw, expected): + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", raw) + assert ic._operator_sampling_override("temperature") == expected + + +def test_out_of_range_recommendation_is_dropped(monkeypatch): + # A malformed model recommendation (out of range) is ignored, so the request keeps the + # schema default rather than forwarding a bad value to llama-server. + _set_recommended(monkeypatch, {"temperature": 5.0, "top_k": 64}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 0.6 # 5.0 is outside [0, 2] -> schema default + assert eff["top_k"] == 64 # a valid recommendation is still applied + + +def test_operator_override_top_k_int_and_range(monkeypatch): + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "40") + assert ic._operator_sampling_override("top_k") == 40 + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "200") # above max 100 + assert ic._operator_sampling_override("top_k") is None + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "-1") # min allowed + assert ic._operator_sampling_override("top_k") == -1 + + +@pytest.mark.parametrize( + "field, val", + [ + ("top_k", 10**400), # oversized int on an int field: int() ok, but math.isfinite raises + ("top_k", float("nan")), # NaN reaching an int field: int(nan) raises ValueError + ("top_k", float("inf")), # inf reaching an int field: int(inf) raises OverflowError + ( + "temperature", + 10**400, + ), # oversized int on a float field: float(huge_int) raises OverflowError + ], +) +def test_clean_sampling_value_rejects_unrepresentable(field, val): + # None of these may raise; each is unusable and must be dropped to None (regression: an + # oversized value used to raise OverflowError before the range check could drop it). + assert ic._clean_sampling_value(field, val) is None + + +def test_oversized_operator_override_ignored(monkeypatch): + # A huge integer string parses via int() but overflows float(); math.isfinite would raise + # OverflowError and 500 the request. It must be ignored like any other bad override and the + # field must fall back to the schema default -- no exception. + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "9" * 400) + assert ic._operator_sampling_override("top_k") is None + _set_recommended(monkeypatch, {}) # no per-model recommendation -> schema default applies + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["top_k"] == 20 # schema default, resolved without raising + + +def test_oversized_recommendation_ignored(monkeypatch): + # A malformed per-model recommendation carrying an oversized int must not raise while + # resolving either; the field simply falls back to the schema default. + _set_recommended(monkeypatch, {"temperature": 10**400, "top_k": 64}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 0.6 # oversized -> dropped -> schema default + assert eff["top_k"] == 64 # a valid recommendation is still applied + + +def test_fill_recommended_sampling_openai_payload(monkeypatch): + from models.inference import ChatCompletionRequest + from routes.inference import _fill_recommended_sampling_openai + + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + + # Client sent only temperature; top_k / min_p were omitted. + payload = ChatCompletionRequest( + model = "m", messages = [{"role": "user", "content": "hi"}], temperature = 0.2 + ) + _fill_recommended_sampling_openai(payload, "some/model") + assert payload.temperature == 0.2 # explicit client value preserved + assert payload.top_k == 64 # recommended fills the omitted field + assert payload.min_p == 0.0 + assert payload.top_p == 0.95 # no recommendation -> schema default unchanged + + +def test_fill_recommended_sampling_openai_operator_pin_overrides_client(monkeypatch): + from models.inference import ChatCompletionRequest + from routes.inference import _fill_recommended_sampling_openai + + monkeypatch.setattr(ic, "load_model_defaults", lambda mid: {}) + monkeypatch.setattr(ic, "get_family_inference_params", lambda mid: {}) + ic._recommended_sampling.cache_clear() + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + + payload = ChatCompletionRequest( + model = "m", messages = [{"role": "user", "content": "hi"}], temperature = 0.2 + ) + _fill_recommended_sampling_openai(payload, "some/model") + assert payload.temperature == 0.9 # operator pin wins even over an explicit client value + + +def test_fill_recommended_sampling_completions_body(monkeypatch): + # /v1/completions is a raw proxy: recommendations fill omitted fields, but a field with no + # recommendation and no pin is left absent so llama-server keeps its own default (unlike the + # chat schema, which carries per-field defaults). + from routes.inference import _fill_recommended_sampling_completions + + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + + body = {"prompt": "hi", "temperature": 0.2} + _fill_recommended_sampling_completions(body, "some/model") + assert body["temperature"] == 0.2 # explicit client value preserved + assert body["top_k"] == 64 # recommendation fills the omitted field + assert body["min_p"] == 0.0 + # No recommendation and no pin -> NOT injected (llama-server keeps its default). + assert "top_p" not in body + assert "presence_penalty" not in body + assert "repeat_penalty" not in body + + +def test_fill_recommended_sampling_completions_operator_pin(monkeypatch): + # An operator pin overrides the client's raw-body value, and the repetition pin is written + # under llama-server's "repeat_penalty" key (the schema field is repetition_penalty). + from routes.inference import _fill_recommended_sampling_completions + + monkeypatch.setattr(ic, "load_inference_config", lambda mid: {}) + ic._recommended_sampling.cache_clear() + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + monkeypatch.setenv("UNSLOTH_SAMPLING_REPETITION_PENALTY", "1.2") + + body = {"prompt": "hi", "temperature": 0.2, "repeat_penalty": 1.05} + _fill_recommended_sampling_completions(body, "some/model") + assert body["temperature"] == 0.9 # operator pin wins over the client's explicit value + assert body["repeat_penalty"] == 1.2 # repetition pin lands on llama-server's key + assert "repetition_penalty" not in body # never leak the schema field name into the body diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py index 2970b1a6bb..64201477e3 100644 --- a/studio/backend/tests/test_sandbox_tools.py +++ b/studio/backend/tests/test_sandbox_tools.py @@ -297,6 +297,8 @@ class TestSandboxEnvIsolation: "PYTHONPATH", "VIRTUAL_ENV", "SystemRoot", + "PATHEXT", # Windows only; minimal list so cwd scripts cannot hijack + "NoDefaultCurrentDirectoryInExePath", # Windows only; no cwd-first lookup } extras = set(env.keys()) - allowed assert not extras, f"sandbox env added unexpected keys: {extras}" @@ -305,6 +307,220 @@ class TestSandboxEnvIsolation: assert env["PYTHONPATH"].endswith("sandbox_site") assert "leak-me" not in env["PYTHONPATH"] + def test_host_git_dir_appended_after_curated(self, monkeypatch, tmp_path): + # #7317: Windows Git lives under Program Files, not System32. Sandbox + # PATH resolves bare `git` by appending the dir of the git the HOST + # shell resolves (shutil.which), after the curated prefix. + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + git_dir = prog / "Git" / "cmd" + git_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_dir / "git.exe")) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(git_dir) in parts + # Curated prefix stays ahead of host Git so Studio python/pip win. + assert parts.index(str(git_dir)) > 0 + + def test_host_path_dirs_not_inherited(self, monkeypatch, tmp_path): + """Host PATH dirs (user-writable, git-lookalike) are never inherited; + only the resolved git dir is. No git resolved -> nothing appended.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + venv_scripts = tmp_path / "venv" / "Scripts" + venv_scripts.mkdir(parents = True) + fake_git = tmp_path / "scratch" / "Git" / "cmd" + fake_git.mkdir(parents = True) + monkeypatch.setenv( + "PATH", + os.pathsep.join([str(venv_scripts), str(fake_git), os.environ.get("PATH", "")]), + ) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: None) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(venv_scripts) not in parts + # A git-suffixed but unresolved (user-writable) dir is NOT trusted. + assert str(fake_git) not in parts + + def test_git_cmd_shim_extension_added_to_pathext(self, monkeypatch, tmp_path): + """A host git resolved as a .cmd shim under a trusted root stays + resolvable under the restricted PATHEXT (cwd lookup disabled).""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + git_dir = prog / "Git" / "cmd" + git_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_dir / "git.cmd")) + env = _build_safe_env(str(tmp_path)) + assert str(git_dir) in env["PATH"].split(os.pathsep) + assert env["PATHEXT"] == ".EXE;.COM;.CMD" + + def test_user_writable_git_dir_refused(self, monkeypatch, tmp_path): + """Git resolved from a per-user manager (Scoop shims) is NOT trusted: + an attacker could drop rg.exe beside it and hit the auto-approve gate.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr( + tools_mod, "_windows_program_roots", lambda: [str(tmp_path / "Program Files")] + ) + shim_dir = tmp_path / "users" / "alice" / "scoop" / "shims" + shim_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(shim_dir / "git.exe")) + env = _build_safe_env(str(tmp_path)) + assert str(shim_dir) not in env["PATH"].split(os.pathsep) + # No trusted git launcher -> PATHEXT stays minimal. + assert env["PATHEXT"] == ".EXE;.COM" + + def test_trust_uses_known_folder_not_env_override(self, monkeypatch, tmp_path): + """Trust is driven by the resolved Program Files roots, so a git under + an attacker-overridden %ProgramFiles% env value is still refused.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "RealProgramFiles" + (real_prog).mkdir() + evil = tmp_path / "attacker" + (evil / "Git" / "cmd").mkdir(parents = True) + # Resolver returns the genuine root; env is overridden to the evil dir. + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + monkeypatch.setenv("ProgramFiles", str(evil)) + monkeypatch.setattr( + tools_mod.shutil, "which", lambda name: str(evil / "Git" / "cmd" / "git.exe") + ) + env = _build_safe_env(str(tmp_path)) + assert str(evil / "Git" / "cmd") not in env["PATH"].split(os.pathsep) + + def test_canonical_git_dir_appended(self, monkeypatch, tmp_path): + """The PATH entry is the realpath of the trusted dir, not a junction + alias, so it cannot be retargeted after the trust check.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "Program Files" + real_git = real_prog / "Git" / "cmd" + real_git.mkdir(parents = True) + link = tmp_path / "link" + try: + link.symlink_to(real_prog, target_is_directory = True) + except (OSError, NotImplementedError): + pytest.skip("symlink unsupported in this environment") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + monkeypatch.setattr( + tools_mod.shutil, + "which", + lambda name: str(link / "Git" / "cmd" / "git.exe"), + ) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(real_git) in parts # canonical, not the `link/...` alias + + def test_windows_temp_git_dir_refused(self, monkeypatch, tmp_path): + """A git under a world-writable %SystemRoot% subdir (Windows\\Temp) is + NOT trusted, even though it sits under the Windows root.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr( + tools_mod, "_windows_program_roots", lambda: [str(tmp_path / "Program Files")] + ) + temp_git = tmp_path / "Windows" / "Temp" / "Git" / "cmd" + temp_git.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(temp_git / "git.exe")) + env = _build_safe_env(str(tmp_path)) + assert str(temp_git) not in env["PATH"].split(os.pathsep) + + def test_trusted_program_dir_matches_via_realpath(self, monkeypatch, tmp_path): + """The trust check canonicalizes paths, so a symlinked/short alias of + Program Files still matches (stand-in for 8.3 PROGRA~1 on Windows).""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "Program Files" + (real_prog / "Git" / "cmd").mkdir(parents = True) + alias = tmp_path / "PROGRA~1" + try: + alias.symlink_to(real_prog, target_is_directory = True) + except (OSError, NotImplementedError): + pytest.skip("symlink unsupported in this environment") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + git_via_alias = alias / "Git" / "cmd" / "git.exe" + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_via_alias)) + env = _build_safe_env(str(tmp_path)) + parts = [os.path.normcase(os.path.realpath(p)) for p in env["PATH"].split(os.pathsep)] + assert os.path.normcase(str(real_prog / "Git" / "cmd")) in parts + + def test_scan_past_untrusted_git_shim(self, monkeypatch, tmp_path): + """When an untrusted shim sorts first on PATH, the scan still finds a + later trusted Program Files git.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + trusted_git = prog / "Git" / "cmd" + trusted_git.mkdir(parents = True) + (trusted_git / "git.EXE").write_text("") # match PATHEXT case on this FS + shim = tmp_path / "scoop" / "shims" + shim.mkdir(parents = True) + (shim / "git.EXE").write_text("") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + # shutil.which returns the untrusted shim first. + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(shim / "git.EXE")) + monkeypatch.setenv("PATH", os.pathsep.join([str(shim), str(trusted_git)])) + monkeypatch.setenv("PATHEXT", ".EXE") + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(trusted_git) in parts + assert str(shim) not in parts + + def test_program_roots_fails_closed_without_known_folder_api(self, monkeypatch): + """When the known-folder API is unavailable, no roots are trusted: env + vars (even %SystemDrive%) are caller-overrideable, so we never derive a + trusted root from them.""" + import core.inference.tools as tools_mod + + # ctypes fails on this Linux host, so the API path raises and we fail + # closed. Any attacker override of these env vars must be irrelevant. + monkeypatch.setenv("ProgramFiles", r"D:\attacker-writable") + monkeypatch.setenv("ProgramW6432", r"D:\attacker-writable") + monkeypatch.setenv("SystemDrive", "D:") + assert tools_mod._windows_program_roots() == [] + + def test_augment_native_program_roots_derives_native_sibling(self): + """A 32-bit process only sees the x86 root; the native sibling is + derived by stripping the ` (x86)` suffix.""" + import core.inference.tools as tools_mod + + roots = tools_mod._augment_native_program_roots([r"C:\Program Files (x86)"]) + lowered = [r.lower() for r in roots] + assert r"c:\program files (x86)" in lowered + assert r"c:\program files" in lowered + + def test_no_default_current_directory_in_exe_path_set_on_windows(self, monkeypatch, tmp_path): + """cmd/CreateProcess must not search cwd for bare names in the sandbox.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: None) + env = _build_safe_env(str(tmp_path)) + assert env["NoDefaultCurrentDirectoryInExePath"] == "1" + def test_home_points_at_sandbox_workdir(self, tmp_path): from core.inference.tools import _build_safe_env diff --git a/studio/backend/tests/test_setup_cache_env_hf_home.py b/studio/backend/tests/test_setup_cache_env_hf_home.py index 4520c93a51..c5d5f820a8 100644 --- a/studio/backend/tests/test_setup_cache_env_hf_home.py +++ b/studio/backend/tests/test_setup_cache_env_hf_home.py @@ -28,6 +28,9 @@ def _isolate_studio_home(monkeypatch, tmp_path): def _load_storage_roots(): + # Each test models a fresh backend process. The cache resolver intentionally + # snapshots explicit environment variables once per process. + sys.modules.pop("utils.hf_cache_settings", None) spec = importlib.util.spec_from_file_location("storage_roots_under_test", _STORAGE_ROOTS_PATH) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) @@ -40,10 +43,10 @@ def _clear_hf_env(monkeypatch): def test_custom_hf_home_seeds_hub_and_xet(monkeypatch, tmp_path): - sr = _load_storage_roots() _clear_hf_env(monkeypatch) custom = tmp_path / "shared" / "huggingface" monkeypatch.setenv("HF_HOME", str(custom)) + sr = _load_storage_roots() sr._setup_cache_env() @@ -54,9 +57,9 @@ def test_custom_hf_home_seeds_hub_and_xet(monkeypatch, tmp_path): def test_default_when_hf_home_unset(monkeypatch, tmp_path): - sr = _load_storage_roots() _clear_hf_env(monkeypatch) monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + sr = _load_storage_roots() sr._setup_cache_env() @@ -67,11 +70,11 @@ def test_default_when_hf_home_unset(monkeypatch, tmp_path): def test_explicit_hub_cache_is_not_overridden(monkeypatch, tmp_path): - sr = _load_storage_roots() _clear_hf_env(monkeypatch) monkeypatch.setenv("HF_HOME", str(tmp_path / "home")) explicit = tmp_path / "explicit" / "hub" monkeypatch.setenv("HF_HUB_CACHE", str(explicit)) + sr = _load_storage_roots() sr._setup_cache_env() @@ -81,11 +84,11 @@ def test_explicit_hub_cache_is_not_overridden(monkeypatch, tmp_path): def test_legacy_huggingface_hub_cache_alias_is_honored(monkeypatch, tmp_path): - sr = _load_storage_roots() _clear_hf_env(monkeypatch) monkeypatch.setenv("HF_HOME", str(tmp_path / "home")) legacy = tmp_path / "legacy" / "hub" monkeypatch.setenv("HUGGINGFACE_HUB_CACHE", str(legacy)) + sr = _load_storage_roots() sr._setup_cache_env() @@ -96,15 +99,16 @@ def test_legacy_huggingface_hub_cache_alias_is_honored(monkeypatch, tmp_path): def test_whitespace_hf_home_falls_back_to_default(monkeypatch, tmp_path): # A blank/whitespace HF_HOME must not become " /hub"; fall back to default. - sr = _load_storage_roots() _clear_hf_env(monkeypatch) monkeypatch.setenv("HF_HOME", " ") monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + sr = _load_storage_roots() sr._setup_cache_env() import os + assert os.environ["HF_HOME"] == str(tmp_path / "xdg" / "huggingface") assert os.environ["HF_HUB_CACHE"] == str(tmp_path / "xdg" / "huggingface" / "hub") @@ -114,9 +118,9 @@ def test_unwritable_hf_home_does_not_crash(monkeypatch, tmp_path): blocker = tmp_path / "blocker" blocker.write_text("not a dir") unwritable = blocker / "hf" - sr = _load_storage_roots() _clear_hf_env(monkeypatch) monkeypatch.setenv("HF_HOME", str(unwritable)) + sr = _load_storage_roots() sr._setup_cache_env() # must not raise diff --git a/studio/backend/tests/test_stt_download_validation.py b/studio/backend/tests/test_stt_download_validation.py new file mode 100644 index 0000000000..a612b14531 --- /dev/null +++ b/studio/backend/tests/test_stt_download_validation.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The /audio/stt/download route must validate a custom Transformers repo before +snapshot_download pulls it into the shared HF cache. + +Regression for a Codex finding: the Transformers engine accepts arbitrary +`owner/model` repos, so an authenticated caller could make Studio download a +large non-STT repository before load-time validation ever ran. Whisper- +compatibility is now enforced (metadata-only, no weights) before the background +download starts. The GGUF engine only accepts curated ids, so it is not gated. +""" + +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path + +import pytest +from fastapi import HTTPException + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +import core.inference.stt_ggml_sidecar as ggml_module # noqa: E402 +import core.inference.stt_sidecar as stt_module # noqa: E402 +import routes.inference as ri # noqa: E402 +from core.inference.stt_sidecar import SttModelCompatibilityError # noqa: E402 +from models.inference import SttLoadRequest # noqa: E402 + + +def _run(coro): + return asyncio.run(coro) + + +def test_custom_non_whisper_repo_is_rejected_before_download(monkeypatch): + started: list = [] + validated: list = [] + + def fake_validate(model, hf_token = None): + validated.append(model) + raise SttModelCompatibilityError( + f"STT model '{model}' is not a compatible Transformers Whisper model." + ) + + def fake_download(model, hf_token = None): + started.append(model) + + monkeypatch.setattr(stt_module, "validate_remote_model", fake_validate) + monkeypatch.setattr(stt_module, "start_model_download", fake_download) + + with pytest.raises(HTTPException) as excinfo: + _run( + ri.stt_download( + SttLoadRequest(model = "owner/chat-model", engine = "transformers"), + current_subject = "tester", + hf_token = None, + ) + ) + + assert excinfo.value.status_code == 422 + assert validated == ["owner/chat-model"] + # The download never starts for a repo that failed the Whisper check. + assert started == [] + + +def test_validated_transformers_repo_downloads(monkeypatch): + started: list = [] + revision = "a" * 40 + + monkeypatch.setattr( + stt_module, + "validate_remote_model", + lambda model, hf_token = None: {"model": model, "revision": revision}, + ) + monkeypatch.setattr( + stt_module, + "start_model_download", + lambda model, hf_token = None, revision = None: started.append((model, revision)), + ) + monkeypatch.setattr(stt_module, "download_status", lambda: {"downloading": True}) + + resp = _run( + ri.stt_download( + SttLoadRequest(model = "owner/real-whisper", engine = "transformers"), + current_subject = "tester", + hf_token = None, + ) + ) + + assert resp.status_code == 200 + assert started == [("owner/real-whisper", revision)] + + +def test_gguf_engine_skips_the_transformers_repo_check(monkeypatch): + started: list = [] + + def fail_if_called(model, hf_token = None): + raise AssertionError("GGUF downloads must not run the Transformers repo check") + + # whisper-server present, so the GGUF request stays on the GGUF engine. + monkeypatch.setattr(ggml_module, "is_available", lambda: True) + monkeypatch.setattr(stt_module, "validate_remote_model", fail_if_called) + monkeypatch.setattr( + ggml_module, "start_model_download", lambda model, hf_token = None: started.append(model) + ) + monkeypatch.setattr(ggml_module, "download_status", lambda: {"downloading": True}) + + resp = _run( + ri.stt_download( + SttLoadRequest(model = "small", engine = "gguf"), + current_subject = "tester", + hf_token = None, + ) + ) + + assert resp.status_code == 200 + assert started == ["small"] + + +def test_resolve_serving_stt_engine_falls_back_when_whisper_server_absent(monkeypatch): + # A curated GGUF request downgrades to Transformers when whisper-server is not + # installed (both engines serve curated ids), but stays GGUF when it is. + monkeypatch.setattr(ggml_module, "is_available", lambda: False) + assert ri._resolve_serving_stt_engine("gguf") == "transformers" + monkeypatch.setattr(ggml_module, "is_available", lambda: True) + assert ri._resolve_serving_stt_engine("gguf") == "gguf" + # Transformers is unaffected by whisper-server availability. + monkeypatch.setattr(ggml_module, "is_available", lambda: False) + assert ri._resolve_serving_stt_engine("transformers") == "transformers" + + +def test_gguf_download_falls_back_to_transformers_when_server_absent(monkeypatch): + """Selecting the default curated model on a host without whisper-server must + download through the Transformers engine, not 501/dead-end on GGUF.""" + gguf_started: list = [] + tf_started: list = [] + + monkeypatch.setattr(ggml_module, "is_available", lambda: False) # no whisper-server + # validate_remote_model no-ops curated ids in production; keep it a no-op here. + monkeypatch.setattr( + stt_module, "validate_remote_model", lambda model, hf_token = None: {"model": model} + ) + monkeypatch.setattr( + stt_module, + "start_model_download", + lambda model, hf_token = None, revision = None: tf_started.append(model), + ) + monkeypatch.setattr(stt_module, "download_status", lambda: {"downloading": True}) + monkeypatch.setattr( + ggml_module, + "start_model_download", + lambda model, hf_token = None: gguf_started.append(model), + ) + + resp = _run( + ri.stt_download( + SttLoadRequest(model = "small", engine = "gguf"), + current_subject = "tester", + hf_token = None, + ) + ) + + assert resp.status_code == 200 + assert tf_started == ["small"] # served by Transformers instead of dead-ending on GGUF + assert gguf_started == [] diff --git a/studio/backend/tests/test_stt_ggml_sidecar.py b/studio/backend/tests/test_stt_ggml_sidecar.py new file mode 100644 index 0000000000..686fd8f546 --- /dev/null +++ b/studio/backend/tests/test_stt_ggml_sidecar.py @@ -0,0 +1,780 @@ +# 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 http.server +import io +import json +import os +import sys +import threading +import time +import wave +from pathlib import Path + +import numpy as np +import pytest + +import core.inference.stt_ggml_sidecar as ggml_module +from core.inference.stt_ggml_sidecar import ( + DEFAULT_GGML_STT_MODEL, + GGML_STT_MODELS, + GGML_STT_REPOS, + GgmlSttSidecar, + SttEngineUnavailableError, + find_whisper_server_binary, + resolve_ggml_model_id, +) +from core.inference.stt_sidecar import ( + SttLanguageError, + SttLoadCancelledError, + SttModelIdError, + SttModelNotDownloadedError, + SttUnavailableError, +) + + +@pytest.fixture(autouse = True) +def isolate_runtime_and_stub_audio_decoder(monkeypatch, tmp_path): + """Unit tests exercise orchestration, not PyAV container parsing.""" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + monkeypatch.delenv("WHISPER_SERVER_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_WHISPER_CPP_PATH", raising = False) + monkeypatch.setenv("PATH", "") + monkeypatch.setattr( + ggml_module, + "_decode_audio_bounded", + lambda audio: np.zeros(16000, dtype = np.float32), + ) + + +# --------------------------------------------------------------------------- +# Model id resolution +# --------------------------------------------------------------------------- + + +def test_curated_ids_resolve(): + for model_id in GGML_STT_MODELS: + assert resolve_ggml_model_id(model_id) == model_id + + +def test_default_model_resolves_from_none_and_blank(): + assert resolve_ggml_model_id(None) == DEFAULT_GGML_STT_MODEL + assert resolve_ggml_model_id(" ") == DEFAULT_GGML_STT_MODEL + + +def test_custom_repo_ids_are_rejected(): + with pytest.raises(SttModelIdError): + resolve_ggml_model_id("owner/model") + with pytest.raises(SttModelIdError): + resolve_ggml_model_id("large-v2") + + +def test_curated_ids_mirror_transformers_sidecar(): + from core.inference.stt_sidecar import STT_MODELS + assert list(GGML_STT_MODELS.keys()) == list(STT_MODELS.keys()) + + +def test_curated_filenames_match_repo_naming(): + # unslothai/whisper--GGUF hosts whisper-.bin; keep the download + # filename in lockstep with the repo so it resolves instead of 404ing. + for model_id, repo in GGML_STT_REPOS.items(): + expected = repo.split("/", 1)[1].removesuffix("-GGUF") + ".bin" + assert GGML_STT_MODELS[model_id] == expected + + +# --------------------------------------------------------------------------- +# Binary discovery +# --------------------------------------------------------------------------- + + +def test_env_binary_override_wins(monkeypatch, tmp_path): + binary = tmp_path / "whisper-server" + binary.write_text("#!/bin/sh\n") + binary.chmod(0o755) # find_whisper_server_binary requires an executable + monkeypatch.setenv("WHISPER_SERVER_PATH", str(binary)) + assert find_whisper_server_binary() == str(binary) + + +def test_env_dir_override_scans_layouts(monkeypatch, tmp_path): + monkeypatch.delenv("WHISPER_SERVER_PATH", raising = False) + build_bin = tmp_path / "build" / "bin" + build_bin.mkdir(parents = True) + binary = build_bin / "whisper-server" + binary.write_text("#!/bin/sh\n") + binary.chmod(0o755) # find_whisper_server_binary requires an executable + monkeypatch.setenv("UNSLOTH_WHISPER_CPP_PATH", str(tmp_path)) + assert find_whisper_server_binary() == str(binary) + + +def test_missing_binary_reports_unavailable(monkeypatch, tmp_path): + monkeypatch.delenv("WHISPER_SERVER_PATH", raising = False) + monkeypatch.setenv("UNSLOTH_WHISPER_CPP_PATH", str(tmp_path / "nope")) + monkeypatch.setattr(ggml_module, "_managed_whisper_cpp_dir", lambda: tmp_path / "gone") + monkeypatch.setattr(ggml_module.shutil, "which", lambda name: None) + assert find_whisper_server_binary() is None + assert not ggml_module.is_available() + with pytest.raises(SttEngineUnavailableError): + ggml_module.ensure_engine_available() + + +def test_non_executable_binary_is_not_runnable(monkeypatch, tmp_path): + if sys.platform == "win32": + pytest.skip("X_OK is an existence check on Windows") + binary = tmp_path / "whisper-server" + binary.write_text("#!/bin/sh\n") # written but not chmod +x + monkeypatch.setenv("WHISPER_SERVER_PATH", str(binary)) + monkeypatch.setattr(ggml_module.shutil, "which", lambda name: None) + assert find_whisper_server_binary() is None + + +# --------------------------------------------------------------------------- +# Slim-install launch guard +# --------------------------------------------------------------------------- + + +def _slim_install( + tmp_path, + *, + install_kind = "slim", + with_ggml = True, + linked_libraries = None, + backend = "cpu", + linked_runtime_directories = None, + runtime_wiring_version = None, +) -> str: + """A managed-looking install tree: marker at the root, server in build/bin.""" + install_dir = tmp_path / "whisper.cpp" + bin_dir = install_dir / "build" / "bin" + bin_dir.mkdir(parents = True) + binary = bin_dir / "whisper-server" + binary.write_text("#!/bin/sh\n") + binary.chmod(0o755) + marker: dict = { + "schema_version": 1, + "component": "whisper.cpp", + "release_tag": "v1.9.1-unsloth.1", + "backend": backend, + "paired_llama_tag": "b10069-mix-fb3d4ca", + } + if install_kind is not None: + marker["install_kind"] = install_kind + if linked_libraries is not None: + marker["linked_libraries"] = linked_libraries + if linked_runtime_directories is not None: + marker["linked_runtime_directories"] = linked_runtime_directories + for name in linked_runtime_directories: + catalog = bin_dir / name + catalog.mkdir() + (catalog / "kernel.dat").write_bytes(b"kernel") + if runtime_wiring_version is not None: + marker["runtime_wiring_version"] = runtime_wiring_version + (install_dir / "UNSLOTH_WHISPER_PREBUILT_INFO.json").write_text(json.dumps(marker)) + if with_ggml: + names = ( + ("ggml.dll", "ggml-base.dll") + if sys.platform == "win32" + else ("libggml.so.0", "libggml-base.so.0") + ) + for name in names: + (bin_dir / name).write_bytes(b"ggml") + return str(binary) + + +def test_slim_guard_flags_missing_ggml_links(monkeypatch, tmp_path): + # A slim marker whose linked ggml runtime is gone must read as engine + # unavailable (reinstall), never crash into a server launch. + binary = _slim_install(tmp_path, with_ggml = False) + assert ggml_module.slim_runtime_intact(binary) is False + monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary) + assert not ggml_module.is_available() + with pytest.raises(SttEngineUnavailableError, match = "ggml"): + ggml_module.ensure_engine_available() + + +def test_slim_guard_passes_with_links_in_place(monkeypatch, tmp_path): + names = ["libggml.so.0", "libggml-base.so.0"] + binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names) + assert ggml_module.slim_runtime_intact(binary) is True + monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary) + assert ggml_module.ensure_engine_available() == binary + + +def test_slim_guard_verifies_the_marker_linked_libraries(monkeypatch, tmp_path): + # New markers record the exact wired filenames; one missing name flips the + # install to unavailable even when the legacy core ggml names are present. + names = ["libggml.dylib", "libggml-base.dylib", "libggml-metal.dylib"] + binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names) + bin_dir = Path(binary).parent + for name in names[:-1]: + (bin_dir / name).write_bytes(b"ggml") + assert ggml_module.slim_runtime_intact(binary) is False # metal dylib absent + (bin_dir / names[-1]).write_bytes(b"ggml") + assert ggml_module.slim_runtime_intact(binary) is True + monkeypatch.setattr(ggml_module, "find_whisper_server_binary", lambda: binary) + assert ggml_module.ensure_engine_available() == binary + + +def test_slim_guard_malformed_authoritative_marker_fails_closed(tmp_path): + for bad in ("not-a-list", [], [1, 2]): + root = tmp_path / f"case_{type(bad).__name__}_{len(str(bad))}" + root.mkdir() + binary = _slim_install(root, with_ggml = True, linked_libraries = bad) + assert ggml_module.slim_runtime_intact(binary) is False + + +def test_slim_guard_prefers_authoritative_root_marker(tmp_path): + names = ["libggml.so.0", "libggml-base.so.0"] + binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = names) + packaging_marker = Path(binary).parent / "UNSLOTH_WHISPER_PREBUILT_INFO.json" + packaging_marker.write_text(json.dumps({"backend": "slim", "release_tag": "packaging"})) + assert ggml_module._whisper_install_marker(binary)["install_kind"] == "slim" + assert ggml_module.slim_runtime_intact(binary) is True + + +def test_slim_guard_rejects_invalid_root_even_with_inner_marker(tmp_path): + binary = _slim_install(tmp_path, with_ggml = True, linked_libraries = ["libggml.so.0"]) + root_marker = Path(binary).parents[2] / "UNSLOTH_WHISPER_PREBUILT_INFO.json" + root_marker.write_text("not json") + (Path(binary).parent / root_marker.name).write_text(json.dumps({"backend": "slim"})) + assert ggml_module.slim_runtime_intact(binary) is False + + +def test_slim_guard_rejects_missing_rocm_catalog(tmp_path): + names = ["libggml.so.0", "libggml-base.so.0", "libggml-hip.so"] + binary = _slim_install( + tmp_path, + linked_libraries = names, + backend = "rocm", + linked_runtime_directories = ["hipblaslt", "rocblas"], + runtime_wiring_version = 2, + ) + bin_dir = Path(binary).parent + (bin_dir / "libggml-hip.so").write_bytes(b"ggml") + assert ggml_module.slim_runtime_intact(binary) is True + (bin_dir / "rocblas" / "kernel.dat").unlink() + assert ggml_module.slim_runtime_intact(binary) is False + + +def test_slim_guard_accepts_windows_rocm_dll_overlay(monkeypatch, tmp_path): + monkeypatch.setattr(ggml_module.sys, "platform", "win32") + names = ["ggml.dll", "ggml-base.dll", "ggml-hip.dll", "amdhip64.dll"] + binary = _slim_install( + tmp_path, + linked_libraries = names, + backend = "rocm", + linked_runtime_directories = [], + runtime_wiring_version = 2, + ) + for name in names: + (Path(binary).parent / name).write_bytes(b"dll") + assert ggml_module.slim_runtime_intact(binary) is True + + +def test_slim_guard_ignores_fat_and_markerless_installs(tmp_path): + # Fat installs carry their own ggml; no marker means source/custom build. + fat = _slim_install(tmp_path / "fat", install_kind = None, with_ggml = False) + assert ggml_module.slim_runtime_intact(fat) is True + bare = tmp_path / "bare" / "whisper-server" + bare.parent.mkdir(parents = True) + bare.write_text("#!/bin/sh\n") + assert ggml_module.slim_runtime_intact(str(bare)) is True + + +# --------------------------------------------------------------------------- +# whisper-server child-process environment +# --------------------------------------------------------------------------- + + +def _loader_path_var() -> str: + return {"win32": "PATH", "darwin": "DYLD_LIBRARY_PATH"}.get(sys.platform, "LD_LIBRARY_PATH") + + +def test_child_env_scrubs_secrets_and_adds_lib_dir(monkeypatch, tmp_path): + monkeypatch.setenv("HF_TOKEN", "secret-token") # exact name + monkeypatch.setenv("MY_API_KEY", "nope") # marker substring + monkeypatch.setenv("HTTPS_PROXY", "http://u:p@px:8080") # url-name + monkeypatch.setenv("SOME_REMOTE", "https://u:pw@host/repo") # url-userinfo value + monkeypatch.setenv("STT_KEEPME", "keep") # benign + binary = tmp_path / "whisper-server" + binary.write_text("#!/bin/sh\n") + env = ggml_module._whisper_server_child_env(str(binary)) + for scrubbed in ("HF_TOKEN", "MY_API_KEY", "HTTPS_PROXY", "SOME_REMOTE"): + assert scrubbed not in env + assert env.get("STT_KEEPME") == "keep" + assert str(tmp_path.resolve()) in env[_loader_path_var()].split(os.pathsep) + + +def test_child_env_isolates_home_and_cred_locations(monkeypatch, tmp_path): + # The downloaded server must not see the real home (token caches live + # there) nor explicit cred-store pointers like HF_HOME / NETRC. + monkeypatch.setenv("HOME", "/real/home") + monkeypatch.setenv("HF_HOME", "/real/hf") + monkeypatch.setenv("NETRC", "/real/.netrc") + monkeypatch.setattr(ggml_module, "_managed_whisper_cpp_dir", lambda: tmp_path / "managed") + binary = tmp_path / "whisper-server" + binary.write_text("#!/bin/sh\n") + env = ggml_module._whisper_server_child_env(str(binary)) + assert env["HOME"] == str(tmp_path / "managed" / ".child_home") + assert "HF_HOME" not in env + assert "NETRC" not in env + assert (tmp_path / "managed" / ".child_home").is_dir() + + +def test_child_env_wsl_rocm_prepends_system_hip(monkeypatch, tmp_path): + if sys.platform != "linux": + pytest.skip("WSL ROCm library precedence is Linux-only") + rocm = tmp_path / "rocm-lib" + rocm.mkdir() + bindir = tmp_path / "bin" + bindir.mkdir() + binary = bindir / "whisper-server" + binary.write_text("#!/bin/sh\n") + monkeypatch.setattr(ggml_module, "_wsl_system_rocm_lib_dirs", lambda: [str(rocm)]) + env = ggml_module._whisper_server_child_env(str(binary)) + parts = env["LD_LIBRARY_PATH"].split(os.pathsep) + assert parts[0] == str(rocm.resolve()) # system HIP wins + assert str(bindir.resolve()) in parts # bundle libs still present + assert env.get("HSA_ENABLE_DXG_DETECTION") == "1" + + +def test_child_env_adds_cuda_runtime_dirs_for_cuda_bundle(monkeypatch, tmp_path): + # Versioned CUDA backend modules are valid too. They still need the + # CUDA-from-PyTorch wheel dirs for libcudart/libcublas at launch. + if sys.platform == "darwin": + pytest.skip("no CUDA on macOS") + import utils.prebuilt.runtime_libs as rl + + bindir = tmp_path / "bin" + bindir.mkdir() + (bindir / "whisper-server").write_text("#!/bin/sh\n") + module_name = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so.0" + (bindir / module_name).write_text("") + cuda_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib" + cuda_dir.mkdir(parents = True) + monkeypatch.setattr(rl, "python_runtime_dirs", lambda: [str(cuda_dir)]) + env = ggml_module._whisper_server_child_env(str(bindir / "whisper-server")) + parts = env[_loader_path_var()].split(os.pathsep) + assert str(bindir.resolve()) in parts + assert str(cuda_dir.resolve()) in parts + assert parts.index(str(bindir.resolve())) < parts.index(str(cuda_dir.resolve())) + + +def test_child_env_omits_cuda_runtime_dirs_for_cpu_bundle(monkeypatch, tmp_path): + # No libggml-cuda.so beside the binary -> a static CPU/Metal bundle -> the CUDA + # wheel discovery must not run and must not touch the loader path. + if sys.platform == "darwin": + pytest.skip("no CUDA on macOS") + import utils.prebuilt.runtime_libs as rl + + bindir = tmp_path / "bin" + bindir.mkdir() + (bindir / "whisper-server").write_text("#!/bin/sh\n") + cuda_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib" + cuda_dir.mkdir(parents = True) + called = {"n": 0} + + def _fake_dirs(): + called["n"] += 1 + return [str(cuda_dir)] + + monkeypatch.setattr(rl, "python_runtime_dirs", _fake_dirs) + env = ggml_module._whisper_server_child_env(str(bindir / "whisper-server")) + parts = env[_loader_path_var()].split(os.pathsep) + assert str(cuda_dir.resolve()) not in parts + assert called["n"] == 0 + + +def test_engine_unavailable_is_stt_unavailable(): + # Routes map SttUnavailableError to HTTP 501; the engine error must share it. + assert issubclass(SttEngineUnavailableError, SttUnavailableError) + + +# --------------------------------------------------------------------------- +# WAV packaging +# --------------------------------------------------------------------------- + + +def test_pcm_to_wav_bytes_shape_and_rate(): + pcm = np.zeros(3200, dtype = np.float32) + data = ggml_module._pcm_to_wav_bytes(pcm) + with wave.open(io.BytesIO(data)) as w: + assert w.getnchannels() == 1 + assert w.getsampwidth() == 2 + assert w.getframerate() == 16000 + assert w.getnframes() == 3200 + + +def test_pcm_to_wav_bytes_clips_out_of_range(): + pcm = np.array([2.0, -2.0], dtype = np.float32) + data = ggml_module._pcm_to_wav_bytes(pcm) + with wave.open(io.BytesIO(data)) as w: + frames = np.frombuffer(w.readframes(2), dtype = "= {"downloading", "model", "error"} diff --git a/studio/backend/tests/test_stt_review_fixes.py b/studio/backend/tests/test_stt_review_fixes.py new file mode 100644 index 0000000000..e4495506a3 --- /dev/null +++ b/studio/backend/tests/test_stt_review_fixes.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regressions for a fresh review pass on the local STT dictation feature: + +1. Curated GGUF dictation repos (unslothai/whisper-*-GGUF) must be hidden from + chat pickers, not just their Transformers safetensors companions. +2. The GGUF sidecar's loaded_model/device status accessors must be lock-free so + they never block behind an in-flight transcription (which holds self._lock). +3. A "gguf" unload on a host without whisper-server must target the Transformers + fallback that actually served it, and unload-all must attempt both backends + even if one raises. +4. free_stt_model_for_training must free the GGUF sidecar even when the + Transformers unload raises (independent exception boundaries). +""" + +from __future__ import annotations + +import asyncio +import sys +import threading +from pathlib import Path + +import pytest +from fastapi import HTTPException + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + + +# 1. Hidden-model GGUF companions ------------------------------------------------ +def test_curated_gguf_dictation_repos_are_hidden(): + from utils.hidden_models import _HIDDEN_STT_REPO_IDS, is_hidden_model + for repo in ( + "unslothai/whisper-tiny-GGUF", + "unslothai/whisper-base-GGUF", + "unslothai/whisper-small-GGUF", + "unslothai/whisper-large-v3-turbo-GGUF", + "unslothai/whisper-large-v3-GGUF", + ): + assert repo in _HIDDEN_STT_REPO_IDS + assert is_hidden_model(repo) is True + # Case-insensitive, matching how the cache stores the repo id. + assert is_hidden_model(repo.lower()) is True + + # A same-prefix but genuinely different repo is NOT hidden. + assert is_hidden_model("unslothai/whisper-large-v3-GGUF-finetune") is False + + +# 2. GGUF status accessors are lock-free ---------------------------------------- +def test_gguf_status_accessors_do_not_block_on_the_inference_lock(): + from core.inference.stt_ggml_sidecar import GgmlSttSidecar + + sidecar = GgmlSttSidecar() + + class _AliveProc: + pid = 4321 + + def poll(self): + return None # still running + + sidecar._process = _AliveProc() + sidecar._model_id = "small" + + holder_has_lock = threading.Event() + release = threading.Event() + + def _hold_inference_lock(): + # Mimic transcribe() holding self._lock across the whole HTTP call. + with sidecar._lock: + holder_has_lock.set() + release.wait(timeout = 5) + + holder = threading.Thread(target = _hold_inference_lock) + holder.start() + assert holder_has_lock.wait(timeout = 5) + + result: dict = {} + + def _read_status(): + result["model"] = sidecar.loaded_model + result["device"] = sidecar.device + + reader = threading.Thread(target = _read_status) + reader.start() + reader.join(timeout = 2) + blocked = reader.is_alive() + + release.set() + holder.join(timeout = 5) + reader.join(timeout = 5) + + assert not blocked, "loaded_model/device blocked on self._lock (should be lock-free)" + assert result == {"model": "small", "device": "whisper.cpp"} + + +def test_process_alive_snapshots_process_against_concurrent_unload(): + # _process_alive() must read self._process exactly once. The lock-free + # readers (loaded_model/device) can run while unload() nulls self._process; + # the old `self._process is not None and self._process.poll() is None` read it + # twice, so a null landing between the two reads called None.poll(). A + # property that yields the live process on the first read and None afterwards + # reproduces that interleaving deterministically. + from core.inference.stt_ggml_sidecar import GgmlSttSidecar + + class _AliveProc: + def poll(self): + return None # still running + + live = _AliveProc() + reads = {"n": 0} + + class _RacingSidecar(GgmlSttSidecar): + @property + def _process(self): + reads["n"] += 1 + return live if reads["n"] == 1 else None + + @_process.setter + def _process(self, value): + pass # __init__ assigns None; the property drives the read + + sidecar = GgmlSttSidecar() + sidecar.__class__ = _RacingSidecar # data descriptor wins over the instance attr + + # Snapshot fix: exactly one read, no AttributeError from a second None read. + assert sidecar._process_alive() is True + assert reads["n"] == 1 + + +# 3. Unload resolves through the serving engine + attempts every backend --------- +def test_gguf_unload_targets_transformers_fallback_without_whisper_server(monkeypatch): + import core.inference.stt_ggml_sidecar as ggml_module + import routes.inference as ri + + monkeypatch.setattr(ggml_module, "is_available", lambda: False) # no whisper-server + + calls: list = [] + + class _Sidecar: + def __init__(self, name): + self.name = name + + def unload(self): + calls.append(self.name) + + monkeypatch.setattr(ri, "_stt_sidecar_for", lambda name: _Sidecar(name)) + + resp = asyncio.run(ri.stt_unload(engine = "gguf", current_subject = "tester")) + assert resp.status_code == 200 + # gguf is served by the Transformers fallback here, so that is what unloads. + assert calls == ["transformers"] + + +def test_unload_all_attempts_both_backends_even_when_one_fails(monkeypatch): + import routes.inference as ri + + attempted: list = [] + + class _Sidecar: + def __init__(self, name): + self.name = name + + def unload(self): + attempted.append(self.name) + if self.name == "transformers": + raise RuntimeError("boom") + + monkeypatch.setattr(ri, "_stt_sidecar_for", lambda name: _Sidecar(name)) + + with pytest.raises(HTTPException) as excinfo: + asyncio.run(ri.stt_unload(engine = None, current_subject = "tester")) + + assert excinfo.value.status_code == 500 + # gguf is still attempted after the transformers unload raised. + assert attempted == ["transformers", "gguf"] + + +# 4. free_stt_model_for_training isolates the two backends ----------------------- +def test_free_stt_frees_gguf_even_when_transformers_unload_raises(monkeypatch): + import routes.training_vram as tv + + class _TransformersSidecar: + def is_loading(self): + return False + + @property + def loaded_model(self): + return "whisper-small" + + def unload(self): + raise RuntimeError("transformers unload failed") + + class _GgmlSidecar: + def __init__(self): + self.unloaded = False + + def is_loading(self): + return False + + @property + def loaded_model(self): + return None if self.unloaded else "small" + + def unload(self): + self.unloaded = True + + ggml = _GgmlSidecar() + monkeypatch.setattr( + "core.inference.stt_sidecar.get_stt_sidecar", lambda: _TransformersSidecar() + ) + monkeypatch.setattr("core.inference.stt_ggml_sidecar.get_ggml_stt_sidecar", lambda: ggml) + + freed = tv.free_stt_model_for_training("test") + + # The Transformers failure must not skip GGUF eviction. + assert ggml.unloaded is True + assert any("small" in entry for entry in freed) diff --git a/studio/backend/tests/test_stt_review_fixes_2.py b/studio/backend/tests/test_stt_review_fixes_2.py new file mode 100644 index 0000000000..f0bdab42b5 --- /dev/null +++ b/studio/backend/tests/test_stt_review_fixes_2.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regressions for the second review pass on the local STT dictation feature: + +1. scripts/build_whisper_cpp.sh must not rm -rf a whisper.cpp/src tree under a + custom Studio home unless Studio itself created it (ownership marker), the + same policy studio/setup.sh applies before its destructive replacements. +2. _snapshot_is_complete must reject pickle (pytorch_model.bin) checkpoints + outright; only safetensors weights count as a usable snapshot. +3. _snapshot_is_complete must require tokenizer assets (tokenizer.json or + vocab.json + merges.txt); weights + config alone decode to blank text. +4. Custom-repo downloads must pin the revision validated beforehand and + restrict snapshot_download to the model/tokenizer/config/preprocessor file + classes (TOCTOU + unbounded-download hardening). +5. The GGML sidecar's readiness probe must not treat an arbitrary local HTTP + responder as whisper-server (mic audio would be posted to it), and the port + reservation must stay held until just before spawn. +""" + +from __future__ import annotations + +import http.server +import json +import os +import socket +import stat +import subprocess +import sys +import threading +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +import core.inference.stt_ggml_sidecar as ggml_module +import core.inference.stt_sidecar as stt_sidecar_module +from core.inference.stt_ggml_sidecar import GgmlSttSidecar, SttEngineUnavailableError +from core.inference.stt_sidecar import validate_remote_model + +_BUILD_SCRIPT = _BACKEND_ROOT.parents[1] / "scripts" / "build_whisper_cpp.sh" + + +# 1. build_whisper_cpp.sh ownership gate ---------------------------------------- + + +def _stub_tools(tmp_path: Path) -> dict: + """PATH with git/cmake stubs so the script never reaches a real build.""" + bin_dir = tmp_path / "stub-bin" + bin_dir.mkdir(exist_ok = True) + for tool in ("git", "cmake"): + stub = bin_dir / tool + stub.write_text("#!/bin/sh\necho stub-%s-invoked >&2\nexit 1\n" % tool) + stub.chmod(stub.stat().st_mode | stat.S_IEXEC) + env = dict(os.environ) + env["PATH"] = f"{bin_dir}:{env['PATH']}" + return env + + +def _run_build_script(env: dict) -> subprocess.CompletedProcess: + return subprocess.run( + ["sh", str(_BUILD_SCRIPT)], + env = env, + capture_output = True, + text = True, + timeout = 60, + ) + + +def test_build_script_refuses_unowned_dir_in_custom_studio_home(tmp_path): + home = tmp_path / "studio-home" + src = home / "whisper.cpp" / "src" + src.mkdir(parents = True) + user_file = src / "user-data.txt" + user_file.write_text("precious") + + env = _stub_tools(tmp_path) + env["UNSLOTH_STUDIO_HOME"] = str(home) + result = _run_build_script(env) + + assert result.returncode != 0 + assert "not marked as an Unsloth-owned" in result.stderr + # The unowned tree, and the user's file inside it, survived untouched. + assert user_file.read_text() == "precious" + + +def test_build_script_proceeds_when_marker_present(tmp_path): + home = tmp_path / "studio-home" + install = home / "whisper.cpp" + (install / "src").mkdir(parents = True) + (install / ".unsloth-studio-owned").write_text("") + + env = _stub_tools(tmp_path) + env["UNSLOTH_STUDIO_HOME"] = str(home) + result = _run_build_script(env) + + # Past the guard: it fails later at the stubbed git clone, not the gate. + assert "not marked as an Unsloth-owned" not in result.stderr + assert "stub-git-invoked" in result.stderr + + +def test_build_script_marks_fresh_custom_install_dir(tmp_path): + home = tmp_path / "studio-home" + home.mkdir() + + env = _stub_tools(tmp_path) + env["UNSLOTH_STUDIO_HOME"] = str(home) + _run_build_script(env) + + # A directory the script creates is marked so re-runs stay allowed. + assert (home / "whisper.cpp" / ".unsloth-studio-owned").is_file() + + +def test_build_script_keeps_legacy_home_behavior(tmp_path): + fake_home = tmp_path / "user-home" + src = fake_home / ".unsloth" / "whisper.cpp" / "src" + src.mkdir(parents = True) + + env = _stub_tools(tmp_path) + env.pop("UNSLOTH_STUDIO_HOME", None) + env.pop("STUDIO_HOME", None) + env["HOME"] = str(fake_home) + result = _run_build_script(env) + + # The legacy managed dir is always Studio-owned; no gate, straight to git. + assert "not marked as an Unsloth-owned" not in result.stderr + assert "stub-git-invoked" in result.stderr + + +# 2 + 3. _snapshot_is_complete -------------------------------------------------- + + +def _base_snapshot(tmp_path: Path) -> Path: + snap = tmp_path / "snap" + snap.mkdir() + (snap / "config.json").write_text("{}") + (snap / "preprocessor_config.json").write_text("{}") + (snap / "tokenizer.json").write_text("{}") + return snap + + +def test_pickle_checkpoint_snapshot_is_never_complete(tmp_path): + # A cached pytorch_model.bin is a pickle RCE load path; the snapshot must + # read as incomplete no matter how many shards are present, so update + # re-resolves and _select_snapshot_files fails it closed. + snap = _base_snapshot(tmp_path) + index = { + "weight_map": { + "a": "pytorch_model-00001-of-00002.bin", + "b": "pytorch_model-00002-of-00002.bin", + } + } + (snap / "pytorch_model.bin.index.json").write_text(json.dumps(index)) + (snap / "pytorch_model-00001-of-00002.bin").write_bytes(b"w" * 8) + (snap / "pytorch_model-00002-of-00002.bin").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is False + + # A single-file pickle checkpoint is likewise rejected; the safetensors + # equivalent in the same dir makes it complete. + (snap / "pytorch_model.bin").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is False + (snap / "model.safetensors").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is True + + +def test_safe_index_naming_pickle_shards_is_not_complete(tmp_path): + # A safetensors index that references .bin shards would still pickle-load + # via Transformers' per-shard dispatch; the cached snapshot must read as + # incomplete so it re-resolves and fails closed at selection. + snap = _base_snapshot(tmp_path) + (snap / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"a": "pytorch_model-00001-of-00001.bin"}}) + ) + (snap / "pytorch_model-00001-of-00001.bin").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is False + + +def test_snapshot_without_tokenizer_assets_is_incomplete(tmp_path): + snap = _base_snapshot(tmp_path) + (snap / "model.safetensors").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is True + + # Weights + config but no tokenizer decodes to blank text; not complete. + (snap / "tokenizer.json").unlink() + assert stt_sidecar_module._snapshot_is_complete(snap) is False + + # The slow vocab.json + merges.txt pair is an accepted alternative. + (snap / "vocab.json").write_text("{}") + assert stt_sidecar_module._snapshot_is_complete(snap) is False + (snap / "merges.txt").write_text("") + assert stt_sidecar_module._snapshot_is_complete(snap) is True + + +# 4. Revision pinning and allow_patterns ---------------------------------------- + + +def test_validate_remote_model_returns_the_validated_revision(monkeypatch): + revision = "a" * 40 + + class _FakeApi: + def __init__(self, token = None): + pass + + def model_info( + self, + repo, + expand = None, + timeout = None, + ): + return SimpleNamespace(config = {"model_type": "whisper"}, sha = revision) + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "HfApi", _FakeApi) + result = validate_remote_model("someone/custom-whisper") + assert result["revision"] == revision + + +def test_download_pins_revision_and_limits_patterns(monkeypatch): + captured = {} + validated_revision = "a" * 40 + head_revision = "b" * 40 + + def fake_snapshot_download(**kwargs): + captured.update(kwargs) + return "/cached" + + class _FakeApi: + def __init__(self, token = None): + pass + + def model_info( + self, + repo, + revision = None, + files_metadata = None, + timeout = None, + ): + names = ( + "config.json", + "preprocessor_config.json", + "tokenizer.json", + "model.safetensors", + ) + siblings = [ + SimpleNamespace(rfilename = name, size = 10, blob_id = name, lfs = None) for name in names + ] + return SimpleNamespace(siblings = siblings, sha = head_revision) + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "HfApi", _FakeApi) + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + + state = stt_sidecar_module._SnapshotDownloadState() + # The revision resolved at validation time wins over the current head. + state._run("someone/custom-whisper", None, revision = validated_revision) + assert captured["revision"] == validated_revision + patterns = captured["allow_patterns"] + assert "model.safetensors" in patterns and "tokenizer.json" in patterns + # No wildcard that would admit arbitrary repo contents. + assert "*" not in patterns + + # Without a validated revision (curated repos), pin to the metadata head. + captured.clear() + state._run("someone/custom-whisper", None) + assert captured["revision"] == head_revision + assert captured["allow_patterns"] + + +# 5. GGML readiness must identify whisper-server -------------------------------- + + +class _CannedHandler(http.server.BaseHTTPRequestHandler): + body = b"" + + def do_GET(self): # noqa: N802 + payload = type(self).body + self.send_response(200) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + pass + + +def _serve(body: bytes): + handler = type("Handler", (_CannedHandler,), {"body": body}) + server = http.server.HTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target = server.serve_forever, daemon = True) + thread.start() + return server, server.server_address[1] + + +def _fake_alive_process(): + return SimpleNamespace(poll = lambda: None, pid = 999999) + + +def test_wait_for_server_rejects_a_foreign_http_responder(monkeypatch): + server, port = _serve(b"hello from some other local app") + try: + monkeypatch.setattr(ggml_module, "_SERVER_START_TIMEOUT_SECONDS", 1.0) + with pytest.raises(SttEngineUnavailableError, match = "did not start in time"): + GgmlSttSidecar._wait_for_server(_fake_alive_process(), port) + finally: + server.shutdown() + + +def test_wait_for_server_accepts_the_whisper_server_page(monkeypatch): + server, port = _serve(b"Whisper.cpp Server") + try: + monkeypatch.setattr(ggml_module, "_SERVER_START_TIMEOUT_SECONDS", 5.0) + GgmlSttSidecar._wait_for_server(_fake_alive_process(), port) + finally: + server.shutdown() + + +def test_probe_requires_the_managed_child_to_be_alive(): + server, port = _serve(b"whisper") + try: + dead = SimpleNamespace(poll = lambda: 0, pid = 999999) + assert GgmlSttSidecar._probe_is_whisper_server(dead, port) is False + assert GgmlSttSidecar._probe_is_whisper_server(_fake_alive_process(), port) is True + finally: + server.shutdown() + + +def test_port_reservation_is_held_until_released(): + reservation, port = GgmlSttSidecar._reserve_free_port() + try: + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + with pytest.raises(OSError): + probe.bind(("127.0.0.1", port)) + finally: + probe.close() + finally: + reservation.close() + # Released right before spawn: the port becomes bindable for the child. + child = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + child.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + child.bind(("127.0.0.1", port)) + finally: + child.close() diff --git a/studio/backend/tests/test_stt_sidecar.py b/studio/backend/tests/test_stt_sidecar.py new file mode 100644 index 0000000000..b138f46331 --- /dev/null +++ b/studio/backend/tests/test_stt_sidecar.py @@ -0,0 +1,1302 @@ +# 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 gc +import io +import json +import sys +import threading +import time +import wave +import weakref +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +import core.inference.stt_sidecar as stt_sidecar_module +from core.inference.stt_sidecar import ( + DEFAULT_STT_MODEL, + STT_MODELS, + SttAudioDecodeError, + SttAudioTooLongError, + SttLanguageError, + SttLoadCancelledError, + SttModelCompatibilityError, + SttModelIdError, + SttModelNotDownloadedError, + SttUnavailableError, + WhisperSttSidecar, + normalize_whisper_language, + resolve_model_id, + resolve_model_repo, + validate_remote_model, +) + +_REAL_DECODE_AUDIO_BOUNDED = stt_sidecar_module._decode_audio_bounded +_REAL_ENSURE_STT_AVAILABLE = stt_sidecar_module.ensure_stt_available +_REAL_SNAPSHOT_IS_COMPLETE = stt_sidecar_module._snapshot_is_complete +_REAL_FIND_COMPLETE_CACHED_SNAPSHOT = stt_sidecar_module._find_complete_cached_snapshot + + +@pytest.fixture(autouse = True) +def stub_audio_decoder(monkeypatch): + """Unit tests below exercise orchestration, not PyAV container parsing.""" + monkeypatch.setattr( + stt_sidecar_module, + "_decode_audio_bounded", + lambda _audio: np.zeros(8000, dtype = np.float32), + ) + monkeypatch.setattr( + "huggingface_hub.snapshot_download", + lambda **_kwargs: "/cached/model", + ) + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + lambda _model: Path("/cached/model"), + ) + # The stubbed snapshot path holds no files; snapshot-integrity tests + # restore the real check. + monkeypatch.setattr(stt_sidecar_module, "_snapshot_is_complete", lambda _snapshot: True) + # transcribe() gates on the runtime up front; treat it as present so these + # orchestration tests run without PyTorch/Transformers/PyAV installed. + # The runtime-specific tests restore the real check. + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", lambda: None) + + +class _CaptureInference: + """Stand-in for the model inference step; records how it was called.""" + + def __init__( + self, + text = "hello", + mutate = None, + ) -> None: + self.text = text + self.mutate = mutate + self.generate_kwargs = None + + def __call__(self, model_id, decoded, generate_kwargs): + self.generate_kwargs = generate_kwargs + if self.mutate is not None: + self.mutate() + return self.text + + +def test_five_curated_whisper_models_are_offered(): + assert STT_MODELS == { + "tiny": "unsloth/whisper-tiny", + "base": "unsloth/whisper-base", + "small": "unsloth/whisper-small", + "large-v3-turbo": "unsloth/whisper-large-v3-turbo", + "large-v3": "unsloth/whisper-large-v3", + } + assert all(repo.startswith(("unsloth/", "unslothai/")) for repo in STT_MODELS.values()) + assert DEFAULT_STT_MODEL in STT_MODELS + + +def test_av_is_required_for_stt_availability(monkeypatch): + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", _REAL_ENSURE_STT_AVAILABLE) + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace()) + monkeypatch.setitem(sys.modules, "transformers", SimpleNamespace()) + monkeypatch.setitem(sys.modules, "av", None) + + assert stt_sidecar_module.is_available() is False + + +def test_transformers_is_required_for_stt_availability(monkeypatch): + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", _REAL_ENSURE_STT_AVAILABLE) + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace()) + monkeypatch.setitem(sys.modules, "av", SimpleNamespace()) + monkeypatch.setitem(sys.modules, "transformers", None) + + assert stt_sidecar_module.is_available() is False + + +@pytest.mark.parametrize("missing", ["transformers", "av"]) +def test_load_rejects_an_incomplete_stt_runtime(monkeypatch, missing): + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", _REAL_ENSURE_STT_AVAILABLE) + for module in ("torch", "transformers", "av"): + monkeypatch.setitem(sys.modules, module, SimpleNamespace()) + monkeypatch.setitem(sys.modules, missing, None) + monkeypatch.setattr( + sidecar, + "_ensure_model_downloaded", + lambda _model: pytest.fail("runtime must be checked before the model cache"), + ) + + with pytest.raises(SttUnavailableError, match = "needs PyTorch, Transformers, and PyAV"): + sidecar.load("small") + + +def test_model_id_accepts_defaults_and_custom_hub_repositories(): + assert resolve_model_id("tiny") == "tiny" + assert resolve_model_id(None) == DEFAULT_STT_MODEL + assert resolve_model_id("large-v3") == "large-v3" + assert resolve_model_id("openai/whisper-medium") == "openai/whisper-medium" + assert resolve_model_repo("tiny") == "unsloth/whisper-tiny" + assert resolve_model_repo("openai/whisper-medium") == "openai/whisper-medium" + + +@pytest.mark.parametrize("model", ["tiny-ish", "owner/model/extra", "../model", "owner/"]) +def test_invalid_custom_model_id_is_rejected(model): + with pytest.raises(SttModelIdError, match = "owner/model"): + resolve_model_id(model) + + +def test_remote_custom_model_validation_requires_whisper_config(monkeypatch): + calls = [] + + class FakeApi: + def __init__(self, token): + calls.append(("token", token)) + + def model_info(self, repo, **kwargs): + calls.append(("model_info", repo, kwargs)) + return SimpleNamespace( + sha = "a" * 40, + config = { + "model_type": "whisper", + "architectures": ["WhisperForConditionalGeneration"], + }, + ) + + monkeypatch.setattr("huggingface_hub.HfApi", FakeApi) + + result = validate_remote_model("owner/custom-whisper", "hf_private") + + assert result == { + "model": "owner/custom-whisper", + "repo": "owner/custom-whisper", + "revision": "a" * 40, + } + assert calls == [ + ("token", "hf_private"), + ( + "model_info", + "owner/custom-whisper", + {"expand": ["config", "sha"], "timeout": 10}, + ), + ] + + +def test_remote_custom_model_validation_rejects_non_whisper(monkeypatch): + class FakeApi: + def __init__(self, token): + assert token is False + + def model_info(self, _repo, **_kwargs): + return SimpleNamespace( + config = { + "model_type": "llama", + "architectures": ["LlamaForCausalLM"], + } + ) + + monkeypatch.setattr("huggingface_hub.HfApi", FakeApi) + + with pytest.raises(SttModelCompatibilityError, match = "not a compatible"): + validate_remote_model("owner/chat-model") + + +def test_remote_custom_model_validation_requires_an_immutable_sha(monkeypatch): + class FakeApi: + def __init__(self, token): + pass + + def model_info(self, _repo, **_kwargs): + return SimpleNamespace(sha = None, config = {"model_type": "whisper"}) + + monkeypatch.setattr("huggingface_hub.HfApi", FakeApi) + + with pytest.raises(SttModelCompatibilityError, match = "immutable revision"): + validate_remote_model("owner/custom-whisper") + + +def test_fast_transcription_uses_greedy_decoding(monkeypatch): + sidecar = WhisperSttSidecar() + infer = _CaptureInference() + monkeypatch.setattr(sidecar, "_transcribe_decoded", infer) + + result = sidecar.transcribe(b"encoded audio", language = "en", fast = True) + + assert result["text"] == "hello" + assert result["duration"] == 0.5 + assert result["model"] == DEFAULT_STT_MODEL + assert infer.generate_kwargs == { + "task": "transcribe", + "condition_on_prev_tokens": False, + "num_beams": 1, + "language": "en", + } + + +def test_accurate_transcription_keeps_beam_search_default(monkeypatch): + sidecar = WhisperSttSidecar() + infer = _CaptureInference() + monkeypatch.setattr(sidecar, "_transcribe_decoded", infer) + + sidecar.transcribe(b"encoded audio") + + assert infer.generate_kwargs == { + "task": "transcribe", + "condition_on_prev_tokens": False, + "num_beams": 5, + } + + +@pytest.mark.parametrize( + ("language", "expected"), + [ + (None, None), + ("auto", None), + ("en-US", "en"), + ("en-GB", "en"), + ("zh-CN", "zh"), + ("ja-JP", "ja"), + ("ko-KR", "ko"), + ("es-ES", "es"), + ("fr-FR", "fr"), + ("de-DE", "de"), + ("it-IT", "it"), + ("pt_BR", "pt"), + ("ru-RU", "ru"), + ("hi-IN", "hi"), + ("ar-SA", "ar"), + ("iw-IL", "he"), + ("nb-NO", "no"), + ], +) +def test_normalize_whisper_language_accepts_bcp47(language, expected): + assert normalize_whisper_language(language) == expected + + +def test_transcription_normalizes_region_qualified_language(monkeypatch): + sidecar = WhisperSttSidecar() + infer = _CaptureInference() + monkeypatch.setattr(sidecar, "_transcribe_decoded", infer) + + sidecar.transcribe(b"encoded audio", language = "fr-FR") + + assert infer.generate_kwargs["language"] == "fr" + + +def test_english_only_model_rejects_non_english_before_decode(monkeypatch, tmp_path): + (tmp_path / "config.json").write_text('{"model_type": "whisper"}') + (tmp_path / "generation_config.json").write_text('{"is_multilingual": false}') + sidecar = WhisperSttSidecar() + + def should_not_decode(_audio): + pytest.fail("English-only language mismatch must be rejected before decode") + + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + lambda _model: tmp_path, + ) + monkeypatch.setattr(stt_sidecar_module, "_decode_audio_bounded", should_not_decode) + + with pytest.raises(SttLanguageError, match = "English-only"): + sidecar.transcribe( + b"encoded audio", + model = "owner/whisper-small.en", + language = "fr-FR", + ) + + +def test_english_only_model_omits_forbidden_generation_controls(monkeypatch): + calls = [] + + class FakeTensor: + def to(self, *_args): + return self + + class FakeProcessor: + def __call__(self, *_args, **_kwargs): + return SimpleNamespace(input_features = FakeTensor()) + + def batch_decode(self, *_args, **_kwargs): + return ["hello"] + + class FakeModel: + dtype = None + device = "cpu" + generation_config = SimpleNamespace(is_multilingual = False) + + def generate(self, _features, **kwargs): + calls.append(kwargs) + return [[1]] + + class NoGrad: + def __enter__(self): + return None + + def __exit__(self, *_args): + return False + + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(no_grad = NoGrad)) + sidecar = WhisperSttSidecar() + monkeypatch.setattr(sidecar, "load", lambda _model: (FakeModel(), FakeProcessor())) + + text = sidecar._transcribe_decoded( + "owner/whisper-small.en", + np.zeros(160, dtype = np.float32), + { + "task": "transcribe", + "language": "en", + "condition_on_prev_tokens": False, + "num_beams": 1, + }, + ) + + assert text == "hello" + assert calls == [{"condition_on_prev_tokens": False, "num_beams": 1}] + + +def test_unknown_language_is_rejected_before_decode_or_model_load(monkeypatch): + sidecar = WhisperSttSidecar() + + def should_not_run(*_args, **_kwargs): + pytest.fail("unknown language must be rejected before expensive work") + + monkeypatch.setattr(stt_sidecar_module, "_known_whisper_languages", lambda: frozenset({"en"})) + monkeypatch.setattr(stt_sidecar_module, "_decode_audio_bounded", should_not_run) + monkeypatch.setattr(sidecar, "_transcribe_decoded", should_not_run) + + with pytest.raises(SttLanguageError, match = "is not supported"): + sidecar.transcribe(b"encoded audio", language = "xx-YY") + + +def test_unknown_language_is_not_reported_as_bad_audio(monkeypatch): + sidecar = WhisperSttSidecar() + infer = _CaptureInference() + monkeypatch.setattr(sidecar, "_transcribe_decoded", infer) + monkeypatch.setattr(stt_sidecar_module, "_known_whisper_languages", lambda: frozenset({"en"})) + + with pytest.raises(SttLanguageError, match = "is not supported"): + sidecar.transcribe(b"encoded audio", language = "xx-YY") + + +def test_transcription_result_keeps_requested_model_id_during_switch(monkeypatch): + sidecar = WhisperSttSidecar() + + # Simulate another request changing the mutable resident-model state after + # this request pinned its own model id. + infer = _CaptureInference(mutate = lambda: setattr(sidecar, "_model_id", "large-v3")) + monkeypatch.setattr(sidecar, "_transcribe_decoded", infer) + + result = sidecar.transcribe(b"encoded audio", model = "small") + + assert result["model"] == "small" + + +def test_inference_failure_propagates(monkeypatch): + sidecar = WhisperSttSidecar() + + def boom(*_args, **_kwargs): + raise RuntimeError("inference failed") + + monkeypatch.setattr(sidecar, "_transcribe_decoded", boom) + + with pytest.raises(RuntimeError, match = "inference failed"): + sidecar.transcribe(b"encoded audio") + + +class _FakeModel: + def to(self, *_args, **_kwargs): + return self + + def eval(self): + return self + + +class _FakeTimer: + def __init__( + self, + interval, + function, + args = (), + kwargs = None, + ): + self.interval = interval + self.function = function + self.args = args + self.kwargs = kwargs or {} + self.cancelled = False + self.daemon = False + self.started = False + + def start(self): + self.started = True + + def cancel(self): + self.cancelled = True + + def fire(self): + self.function(*self.args, **self.kwargs) + + +def _install_fake_torch(monkeypatch): + fake_torch = SimpleNamespace( + float16 = "float16", + float32 = "float32", + device = lambda value: value, + cuda = SimpleNamespace(is_available = lambda: False), + backends = SimpleNamespace(mps = SimpleNamespace(is_available = lambda: False)), + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "av", SimpleNamespace()) + return fake_torch + + +def test_load_uses_model_hub_cache_without_implicit_download(monkeypatch): + calls = [] + _install_fake_torch(monkeypatch) + + class FakeWhisperForConditionalGeneration: + @classmethod + def from_pretrained(cls, repo, **kwargs): + calls.append(("model", repo, kwargs)) + return _FakeModel() + + class FakeWhisperProcessor: + @classmethod + def from_pretrained(cls, repo, **kwargs): + calls.append(("processor", repo, kwargs)) + return object() + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + WhisperForConditionalGeneration = FakeWhisperForConditionalGeneration, + WhisperProcessor = FakeWhisperProcessor, + ), + ) + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("cpu", "float32")) + + WhisperSttSidecar(keep_alive_seconds = 0).load("small") + + assert {(kind, repo) for kind, repo, _ in calls} == { + ("processor", "/cached/model"), + ("model", "/cached/model"), + } + # Never fetch weights implicitly; the Model Hub owns downloads. + assert all(kwargs.get("local_files_only") is True for _, _, kwargs in calls) + # The weight load forces safetensors so a pickle checkpoint cannot execute. + model_kwargs = next(kwargs for kind, _, kwargs in calls if kind == "model") + assert model_kwargs.get("use_safetensors") is True + + +def test_model_cache_preflight_uses_shared_offline_resolver(monkeypatch): + seen = [] + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + lambda model: seen.append(model) or Path("/cached/model"), + ) + + WhisperSttSidecar(keep_alive_seconds = 0)._ensure_model_downloaded("small") + + assert seen == ["small"] + + +def test_model_cache_preflight_reports_missing_snapshot(monkeypatch): + monkeypatch.setattr(stt_sidecar_module, "_find_complete_cached_snapshot", lambda _model: None) + + with pytest.raises(SttModelNotDownloadedError, match = "not downloaded"): + WhisperSttSidecar(keep_alive_seconds = 0)._ensure_model_downloaded("large-v3") + + +def test_load_reports_model_hub_cache_miss(monkeypatch): + _install_fake_torch(monkeypatch) + + class LocalEntryNotFoundError(RuntimeError): + pass + + class MissingWhisperProcessor: + @classmethod + def from_pretrained(cls, *_args, **_kwargs): + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + WhisperForConditionalGeneration = object, + WhisperProcessor = MissingWhisperProcessor, + ), + ) + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("cpu", "float32")) + + with pytest.raises(SttModelNotDownloadedError, match = "not downloaded"): + WhisperSttSidecar(keep_alive_seconds = 0).load("large-v3") + + +def test_unavailable_runtime_is_rejected_before_audio_decode(monkeypatch): + sidecar = WhisperSttSidecar() + + def unavailable() -> None: + raise SttUnavailableError("needs PyTorch, Transformers, and PyAV") + + def should_not_decode(_audio): + pytest.fail("runtime must be checked before audio decode") + + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", unavailable) + monkeypatch.setattr(stt_sidecar_module, "_decode_audio_bounded", should_not_decode) + + with pytest.raises(SttUnavailableError, match = "needs PyTorch"): + sidecar.transcribe(b"encoded audio", model = "small") + + +def test_missing_model_is_rejected_before_audio_decode(monkeypatch): + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + + def missing(_model_id): + raise SttModelNotDownloadedError("not downloaded") + + def should_not_decode(_audio): + pytest.fail("missing models must be rejected before audio decode") + + monkeypatch.setattr(sidecar, "_ensure_model_downloaded", missing, raising = False) + monkeypatch.setattr(stt_sidecar_module, "_decode_audio_bounded", should_not_decode) + + with pytest.raises(SttModelNotDownloadedError, match = "not downloaded"): + sidecar.transcribe(b"encoded audio", model = "large-v3") + + +def test_missing_model_switch_keeps_resident_model(monkeypatch): + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + resident = object() + sidecar._engine = resident + sidecar._model_id = "small" + sidecar._device = "cpu" + + def missing(_model_id): + raise SttModelNotDownloadedError("not downloaded") + + monkeypatch.setattr(sidecar, "_ensure_model_downloaded", missing, raising = False) + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", lambda: None) + monkeypatch.setattr( + sidecar, + "_build_model", + lambda *_args: pytest.fail("cache miss must be detected before model replacement"), + ) + _install_fake_torch(monkeypatch) + + with pytest.raises(SttModelNotDownloadedError, match = "not downloaded"): + sidecar.load("large-v3") + + assert sidecar._engine is resident + assert sidecar.loaded_model == "small" + + +def test_incompatible_custom_model_switch_keeps_resident_model(monkeypatch, tmp_path): + (tmp_path / "config.json").write_text( + '{"model_type": "llama", "architectures": ["LlamaForCausalLM"]}' + ) + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + resident = (object(), object()) + sidecar._engine = resident + sidecar._model_id = "small" + sidecar._device = "cpu" + _install_fake_torch(monkeypatch) + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + lambda _model: tmp_path, + ) + + with pytest.raises(SttModelCompatibilityError, match = "not a compatible"): + sidecar.load("owner/chat-model") + + assert sidecar._engine is resident + assert sidecar.loaded_model == "small" + + +def test_loaded_model_stays_warm_until_idle_timer_fires(monkeypatch): + timers = [] + _install_fake_torch(monkeypatch) + + def make_timer(*args, **kwargs): + timer = _FakeTimer(*args, **kwargs) + timers.append(timer) + return timer + + sidecar = WhisperSttSidecar(keep_alive_seconds = 300) + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", lambda: None) + monkeypatch.setattr(stt_sidecar_module.threading, "Timer", make_timer) + monkeypatch.setattr(sidecar, "_build_model", lambda *_args: (object(), object())) + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("cpu", "float32")) + + sidecar.load("small") + + assert sidecar.loaded_model == "small" + assert timers[-1].interval == 300 + assert timers[-1].started + + timers[-1].fire() + + assert sidecar.loaded_model is None + + +def test_reusing_loaded_model_refreshes_idle_timer(monkeypatch): + timers = [] + _install_fake_torch(monkeypatch) + + def make_timer(*args, **kwargs): + timer = _FakeTimer(*args, **kwargs) + timers.append(timer) + return timer + + sidecar = WhisperSttSidecar(keep_alive_seconds = 300) + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", lambda: None) + monkeypatch.setattr(stt_sidecar_module.threading, "Timer", make_timer) + monkeypatch.setattr(sidecar, "_build_model", lambda *_args: (object(), object())) + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("cpu", "float32")) + + sidecar.load("small") + first = timers[-1] + sidecar.load("small") + + assert first.cancelled + assert timers[-1] is not first + + first.fire() + + assert sidecar.loaded_model == "small" + + +def test_unload_waits_for_inflight_transcription(monkeypatch): + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + started = threading.Event() + release = threading.Event() + + def transcribe(*_args): + started.set() + assert release.wait(timeout = 2) + return "hello" + + monkeypatch.setattr(sidecar, "_transcribe_decoded", transcribe) + transcribe_thread = threading.Thread(target = lambda: sidecar.transcribe(b"audio")) + transcribe_thread.start() + assert started.wait(timeout = 2) + + unload_thread = threading.Thread(target = sidecar.unload) + unload_thread.start() + time.sleep(0.02) + assert unload_thread.is_alive() + + release.set() + transcribe_thread.join(timeout = 2) + unload_thread.join(timeout = 2) + + assert not transcribe_thread.is_alive() + assert not unload_thread.is_alive() + + +def test_new_stt_load_uses_cpu_while_training(monkeypatch): + fake_torch = SimpleNamespace( + float16 = "float16", + float32 = "float32", + cuda = SimpleNamespace(is_available = lambda: True), + backends = SimpleNamespace(mps = SimpleNamespace(is_available = lambda: True)), + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setattr(stt_sidecar_module, "_training_active", lambda: True) + + assert stt_sidecar_module._pick_device() == ("cpu", "float32") + + +def test_new_stt_load_prefers_cuda_when_training_is_idle(monkeypatch): + fake_torch = SimpleNamespace( + float16 = "float16", + float32 = "float32", + cuda = SimpleNamespace(is_available = lambda: True), + backends = SimpleNamespace(mps = SimpleNamespace(is_available = lambda: False)), + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setattr(stt_sidecar_module, "_training_active", lambda: False) + + assert stt_sidecar_module._pick_device() == ("cuda", "float16") + + +def test_new_stt_load_prefers_mps_when_cuda_is_unavailable(monkeypatch): + fake_torch = SimpleNamespace( + float16 = "float16", + float32 = "float32", + cuda = SimpleNamespace(is_available = lambda: False), + backends = SimpleNamespace(mps = SimpleNamespace(is_available = lambda: True)), + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setattr(stt_sidecar_module, "_training_active", lambda: False) + + assert stt_sidecar_module._pick_device() == ("mps", "float32") + + +def test_new_stt_load_uses_cpu_without_accelerators(monkeypatch): + fake_torch = SimpleNamespace( + float16 = "float16", + float32 = "float32", + cuda = SimpleNamespace(is_available = lambda: False), + backends = SimpleNamespace(mps = SimpleNamespace(is_available = lambda: False)), + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setattr(stt_sidecar_module, "_training_active", lambda: False) + + assert stt_sidecar_module._pick_device() == ("cpu", "float32") + + +def test_accelerator_load_failure_retries_on_cpu(monkeypatch): + fake_torch = _install_fake_torch(monkeypatch) + calls = [] + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", lambda: None) + + def build(_repo, device, dtype, _cancel_event): + calls.append((device, dtype)) + if device == "cuda": + raise RuntimeError("accelerator allocation failed") + return object(), object() + + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("cuda", "float16")) + monkeypatch.setattr(sidecar, "_build_model", build) + + sidecar.load("small") + + assert calls == [("cuda", "float16"), ("cpu", fake_torch.float32)] + assert sidecar.device == "cpu" + + +def test_pending_load_can_be_cancelled_without_waiting_for_model_lock(monkeypatch): + _install_fake_torch(monkeypatch) + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + build_started = threading.Event() + release_build = threading.Event() + errors = [] + + def build(_repo, _device, _dtype, _cancel_event): + build_started.set() + assert release_build.wait(timeout = 2) + return object(), object() + + def run_load(): + try: + sidecar.load("small") + except Exception as exc: + errors.append(exc) + + monkeypatch.setattr(stt_sidecar_module, "ensure_stt_available", lambda: None) + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("cpu", "float32")) + monkeypatch.setattr(sidecar, "_build_model", build) + + load_thread = threading.Thread(target = run_load) + load_thread.start() + assert build_started.wait(timeout = 2) + + result = [] + cancel_thread = threading.Thread(target = lambda: result.append(sidecar.cancel_pending_load())) + cancel_thread.start() + cancel_thread.join(timeout = 2) + + assert not cancel_thread.is_alive() + assert result == [True] + assert load_thread.is_alive() + + release_build.set() + load_thread.join(timeout = 2) + + assert not load_thread.is_alive() + assert len(errors) == 1 + assert isinstance(errors[0], SttLoadCancelledError) + assert sidecar.loaded_model is None + assert sidecar.is_loading() is False + + +def _wav_bytes(sample_count: int, sample_rate: int = 16000) -> bytes: + output = io.BytesIO() + with wave.open(output, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(sample_rate) + wav.writeframes(np.zeros(sample_count, dtype = np.int16).tobytes()) + return output.getvalue() + + +def test_bounded_decoder_returns_16khz_float_pcm(): + pytest.importorskip("av") + + decoded = _REAL_DECODE_AUDIO_BOUNDED(_wav_bytes(1600)) + + assert decoded.dtype == np.float32 + assert decoded.shape == (1600,) + + +def test_bounded_decoder_rejects_audio_as_soon_as_sample_cap_is_crossed(monkeypatch): + pytest.importorskip("av") + monkeypatch.setattr(stt_sidecar_module, "_MAX_AUDIO_SECONDS", 1) + + with pytest.raises(SttAudioTooLongError, match = "Audio must"): + _REAL_DECODE_AUDIO_BOUNDED(_wav_bytes(16001)) + + +def test_bounded_decoder_resamples_stereo_48khz_to_mono_16khz(): + pytest.importorskip("av") + output = io.BytesIO() + frames = np.zeros((4800, 2), dtype = np.int16) + with wave.open(output, "wb") as wav: + wav.setnchannels(2) + wav.setsampwidth(2) + wav.setframerate(48000) + wav.writeframes(frames.tobytes()) + + decoded = _REAL_DECODE_AUDIO_BOUNDED(output.getvalue()) + + assert decoded.dtype == np.float32 + assert 1590 <= len(decoded) <= 1610 + + +@pytest.mark.parametrize("audio", [b"", b"not audio", b"RIFF\x00\x00"]) +def test_bounded_decoder_rejects_malformed_audio(audio): + pytest.importorskip("av") + + with pytest.raises(SttAudioDecodeError, match = "Could not decode"): + _REAL_DECODE_AUDIO_BOUNDED(audio) + + +def test_bounded_decoder_rejects_container_without_audio_stream(monkeypatch): + class FakeFFmpegError(Exception): + pass + + class FakeResampler: + def __init__(self, **_kwargs): + pass + + class FakeFifo: + samples = 0 + + class FakeContainer: + streams = SimpleNamespace(audio = []) + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + fake_av = SimpleNamespace( + audio = SimpleNamespace( + resampler = SimpleNamespace(AudioResampler = FakeResampler), + fifo = SimpleNamespace(AudioFifo = FakeFifo), + ), + open = lambda *_args, **_kwargs: FakeContainer(), + ) + monkeypatch.setitem(sys.modules, "av", fake_av) + monkeypatch.setitem( + sys.modules, + "av.error", + SimpleNamespace( + FFmpegError = FakeFFmpegError, + InvalidDataError = FakeFFmpegError, + ), + ) + + with pytest.raises(SttAudioDecodeError, match = "Could not decode"): + _REAL_DECODE_AUDIO_BOUNDED(b"video-only") + + +def test_unload_releases_model_and_device(): + sidecar = WhisperSttSidecar() + sidecar._engine = object() + sidecar._model_id = "small" + sidecar._device = "cpu" + + sidecar.unload() + + assert sidecar.loaded_model is None + assert sidecar.device is None + + +# --------------------------------------------------------------------------- +# Snapshot download tracking +# --------------------------------------------------------------------------- + + +def _write_complete_snapshot(snapshot: Path, *, model_type: str = "whisper") -> None: + snapshot.mkdir(parents = True, exist_ok = True) + (snapshot / "config.json").write_text(json.dumps({"model_type": model_type})) + (snapshot / "preprocessor_config.json").write_text("{}") + (snapshot / "tokenizer.json").write_text("{}") + (snapshot / "model.safetensors").write_bytes(b"weights") + + +def _sibling(name: str, size: int, key: str): + return SimpleNamespace(rfilename = name, size = size, blob_id = key, lfs = None) + + +def test_sha_snapshot_without_main_ref_survives_restart_and_cache_relocation(monkeypatch, tmp_path): + repo = "openai/whisper-tiny.en" + revision = "c" * 40 + studio_home = tmp_path / "studio" + first_cache = tmp_path / "first-hub" + second_cache = tmp_path / "second-hub" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr(stt_sidecar_module, "_snapshot_is_complete", _REAL_SNAPSHOT_IS_COMPLETE) + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + _REAL_FIND_COMPLETE_CACHED_SNAPSHOT, + ) + + first = first_cache / "models--openai--whisper-tiny.en" / "snapshots" / revision + _write_complete_snapshot(first) + monkeypatch.setenv("HF_HUB_CACHE", str(first_cache)) + stt_sidecar_module._write_revision_record(repo, revision) + assert stt_sidecar_module._find_complete_cached_snapshot(repo) == first.resolve() + + second = second_cache / "models--openai--whisper-tiny.en" / "snapshots" / revision + _write_complete_snapshot(second) + monkeypatch.setenv("HF_HUB_CACHE", str(second_cache)) + assert stt_sidecar_module._find_complete_cached_snapshot(repo) == second.resolve() + + +def test_corrupt_or_escaping_revision_record_is_ignored(monkeypatch, tmp_path): + repo = "openai/whisper-tiny.en" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "hub")) + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + _REAL_FIND_COMPLETE_CACHED_SNAPSHOT, + ) + record = stt_sidecar_module._revision_record_path(repo) + record.parent.mkdir(parents = True) + record.write_text(json.dumps({"version": 1, "repo": repo, "revision": "../../outside"})) + assert stt_sidecar_module._find_complete_cached_snapshot(repo) is None + + outside = tmp_path / "outside" + _write_complete_snapshot(outside) + snapshots = tmp_path / "hub" / "models--openai--whisper-tiny.en" / "snapshots" + snapshots.mkdir(parents = True) + (snapshots / ("d" * 40)).symlink_to(outside, target_is_directory = True) + assert stt_sidecar_module._find_complete_cached_snapshot(repo) is None + + +def test_adapter_only_snapshot_is_not_complete(tmp_path): + (tmp_path / "config.json").write_text('{"model_type": "whisper"}') + (tmp_path / "preprocessor_config.json").write_text("{}") + (tmp_path / "tokenizer.json").write_text("{}") + (tmp_path / "adapter_model.safetensors").write_bytes(b"adapter") + + assert _REAL_SNAPSHOT_IS_COMPLETE(tmp_path) is False + + +def test_snapshot_selection_prefers_safetensors_and_excludes_unrelated_files(): + info = SimpleNamespace( + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("preprocessor_config.json", 20, "preprocessor"), + _sibling("tokenizer.json", 30, "tokenizer"), + _sibling("model.safetensors", 100, "safe"), + _sibling("pytorch_model.bin", 110, "torch"), + _sibling("README.md", 1000, "readme"), + ] + ) + + selected = stt_sidecar_module._select_snapshot_files( + info, lambda _name: pytest.fail("unsharded selection must not load an index") + ) + + assert {item.path for item in selected} == { + "config.json", + "preprocessor_config.json", + "tokenizer.json", + "model.safetensors", + } + assert sum(item.size for item in selected) == 160 + + +def test_snapshot_selection_includes_every_indexed_shard(): + info = SimpleNamespace( + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("model.safetensors.index.json", 5, "index"), + _sibling("model-00001-of-00002.safetensors", 50, "shard1"), + _sibling("model-00002-of-00002.safetensors", 60, "shard2"), + _sibling("pytorch_model.bin", 120, "torch"), + ] + ) + + selected = stt_sidecar_module._select_snapshot_files( + info, + lambda name: { + "weight_map": { + "a": "model-00001-of-00002.safetensors", + "b": "model-00002-of-00002.safetensors", + } + }, + ) + + assert {item.path for item in selected} == { + "config.json", + "model.safetensors.index.json", + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + } + + +def test_snapshot_selection_rejects_pickle_only_weights(): + # A custom repo shipping only pytorch_model.bin (pickle) must fail closed: + # selecting it would download a checkpoint that runs code at load time. + info = SimpleNamespace( + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("preprocessor_config.json", 20, "preprocessor"), + _sibling("tokenizer.json", 30, "tokenizer"), + _sibling("pytorch_model.bin", 110, "torch"), + ] + ) + + with pytest.raises(SttModelCompatibilityError, match = "safetensors"): + stt_sidecar_module._select_snapshot_files( + info, lambda _name: pytest.fail("pickle weights must not be selected") + ) + + +def test_snapshot_selection_rejects_safe_index_pointing_at_pickle_shards(): + # A safetensors index can name .bin shards; Transformers dispatches shard + # loading by extension, so those shards would still pickle-load. The index + # is attacker-controlled, so a non-safetensors shard must fail closed. + info = SimpleNamespace( + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("model.safetensors.index.json", 5, "index"), + _sibling("pytorch_model-00001-of-00001.bin", 90, "shard"), + ] + ) + + with pytest.raises(SttModelCompatibilityError, match = "non-safetensors shards"): + stt_sidecar_module._select_snapshot_files( + info, + lambda _name: {"weight_map": {"a": "pytorch_model-00001-of-00001.bin"}}, + ) + + +def test_progress_counts_only_selected_blobs_and_caps_incomplete_files(monkeypatch, tmp_path): + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "hub")) + blobs = tmp_path / "hub" / "models--owner--whisper" / "blobs" + blobs.mkdir(parents = True) + (blobs / "one").write_bytes(b"x" * 10) + (blobs / "two.incomplete").write_bytes(b"x" * 30) + (blobs / "unrelated").write_bytes(b"x" * 1000) + state = stt_sidecar_module._SnapshotDownloadState() + state._repo = "owner/whisper" + state._selected_files = ( + stt_sidecar_module._SelectedHubFile("config.json", 10, "one"), + stt_sidecar_module._SelectedHubFile("model.safetensors", 20, "two"), + ) + state._total_bytes = 30 + state._complete = True + + status = state.status() + + assert status["bytes_total"] == 30 + assert status["bytes_done"] == 30 + + +def test_download_metadata_and_snapshot_use_the_same_revision(monkeypatch, tmp_path): + revision = "e" * 40 + calls = [] + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("preprocessor_config.json", 20, "preprocessor"), + _sibling("tokenizer.json", 30, "tokenizer"), + _sibling("model.safetensors", 100, "safe"), + _sibling("pytorch_model.bin", 110, "torch"), + ] + + class FakeApi: + def __init__(self, token): + pass + + def model_info(self, repo, **kwargs): + calls.append(("info", repo, kwargs)) + return SimpleNamespace(sha = revision, siblings = siblings) + + def fake_snapshot_download(**kwargs): + calls.append(("snapshot", kwargs)) + return str(tmp_path) + + monkeypatch.setattr("huggingface_hub.HfApi", FakeApi) + monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download) + monkeypatch.setattr( + "huggingface_hub.hf_hub_download", + lambda **_kwargs: pytest.fail("unsharded selection must not load an index"), + ) + monkeypatch.setattr(stt_sidecar_module, "_snapshot_is_complete", lambda _path: True) + monkeypatch.setattr(stt_sidecar_module, "_write_revision_record", lambda *_args: None) + state = stt_sidecar_module._SnapshotDownloadState() + + state._run("owner/whisper", None, revision) + + assert calls[0] == ( + "info", + "owner/whisper", + {"revision": revision, "files_metadata": True, "timeout": 30}, + ) + assert calls[1][0] == "snapshot" + assert calls[1][1]["revision"] == revision + assert "model.safetensors" in calls[1][1]["allow_patterns"] + assert "pytorch_model.bin" not in calls[1][1]["allow_patterns"] + + +def test_download_status_is_idle_before_any_download(): + state = stt_sidecar_module._SnapshotDownloadState() + + status = state.status() + + assert status == { + "downloading": False, + "model": None, + "error": None, + "bytes_total": None, + "bytes_done": None, + } + + +def test_download_rejects_a_second_model_while_one_is_in_flight(monkeypatch): + state = stt_sidecar_module._SnapshotDownloadState() + release = threading.Event() + monkeypatch.setattr( + state, + "_run", + lambda repo, token, revision: release.wait(timeout = 5), + ) + + state.start("small") + try: + # Re-requesting the in-flight model is a no-op, not an error. + state.start("small") + with pytest.raises(SttModelIdError, match = "still"): + state.start("tiny") + assert state.status()["downloading"] is True + assert state.status()["model"] == "small" + finally: + release.set() + + +def test_download_failure_is_reported_in_status(monkeypatch): + state = stt_sidecar_module._SnapshotDownloadState() + # Mask huggingface_hub so the import inside _run fails fast. + monkeypatch.setitem(sys.modules, "huggingface_hub", None) + + state.start("small") + state._thread.join(timeout = 5) + + status = state.status() + assert status["downloading"] is False + assert "Download failed" in (status["error"] or "") + + +def test_is_model_downloaded_is_false_for_a_cache_miss(monkeypatch): + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + _REAL_FIND_COMPLETE_CACHED_SNAPSHOT, + ) + monkeypatch.setenv("HF_HUB_CACHE", "/nonexistent/stt-test-cache") + + assert stt_sidecar_module.is_model_downloaded("small") is False + + +def test_sharded_snapshot_with_missing_shard_is_not_downloaded(monkeypatch, tmp_path): + import json + + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "hub")) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + _REAL_FIND_COMPLETE_CACHED_SNAPSHOT, + ) + monkeypatch.setattr(stt_sidecar_module, "_snapshot_is_complete", _REAL_SNAPSHOT_IS_COMPLETE) + snap = tmp_path / "hub" / "models--unsloth--whisper-small" / "snapshots" / ("a" * 40) + snap.mkdir(parents = True) + (snap / "config.json").write_bytes(b"{}") + (snap / "preprocessor_config.json").write_bytes(b"{}") + (snap / "tokenizer.json").write_bytes(b"{}") + index = { + "weight_map": { + "a": "model-00001-of-00002.safetensors", + "b": "model-00002-of-00002.safetensors", + } + } + (snap / "model.safetensors.index.json").write_text(json.dumps(index)) + (snap / "model-00001-of-00002.safetensors").write_bytes(b"w" * 8) + + assert stt_sidecar_module.is_model_downloaded("small") is False + + # Completing the second shard flips the verdict. + (snap / "model-00002-of-00002.safetensors").write_bytes(b"w" * 8) + assert stt_sidecar_module.is_model_downloaded("small") is True + + +@pytest.mark.parametrize("model_id", ["small", "openai/whisper-medium"]) +def test_preflight_rejects_partial_snapshot(monkeypatch, tmp_path, model_id): + # A resolvable snapshot with metadata but no weights must fail preflight, + # not survive until load() after the audio has already been decoded. + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "hub")) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + monkeypatch.setattr( + stt_sidecar_module, + "_find_complete_cached_snapshot", + _REAL_FIND_COMPLETE_CACHED_SNAPSHOT, + ) + monkeypatch.setattr(stt_sidecar_module, "_snapshot_is_complete", _REAL_SNAPSHOT_IS_COMPLETE) + repo = STT_MODELS.get(model_id, model_id) + snapshot = tmp_path / "hub" / f"models--{repo.replace('/', '--')}" / "snapshots" / ("b" * 40) + snapshot.mkdir(parents = True) + (snapshot / "config.json").write_text('{"model_type": "whisper"}') + + with pytest.raises(SttModelNotDownloadedError, match = "not downloaded"): + WhisperSttSidecar(keep_alive_seconds = 0)._ensure_model_downloaded(model_id) + + # Completing the snapshot clears the preflight. + (snapshot / "preprocessor_config.json").write_text("{}") + (snapshot / "tokenizer.json").write_text("{}") + (snapshot / "model.safetensors").write_bytes(b"w" * 8) + WhisperSttSidecar(keep_alive_seconds = 0)._ensure_model_downloaded(model_id) + + +def test_cpu_retry_releases_failed_accelerator_load(monkeypatch): + _install_fake_torch(monkeypatch) + monkeypatch.setattr(stt_sidecar_module, "_pick_device", lambda: ("mps", "float16")) + + class Marker: + pass + + seen = {} + + def fake_build(self, repo, device, dtype, cancel_event): + if device != "cpu": + # The frame local stands in for a partly loaded accelerator model + # kept alive only through the raised traceback. + marker = Marker() + seen["ref"] = weakref.ref(marker) + raise RuntimeError("accelerator load failed") + gc.collect() + seen["alive_during_retry"] = seen["ref"]() is not None + return (_FakeModel(), object()) + + monkeypatch.setattr(WhisperSttSidecar, "_build_model", fake_build) + sidecar = WhisperSttSidecar(keep_alive_seconds = 0) + sidecar.load("small") + + # The failed attempt must be collectable before the CPU model loads, or + # its accelerator memory stays stranded for the whole retry. + assert seen["alive_during_retry"] is False + assert sidecar.device == "cpu" diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 5d74bb7d28..64228cec3c 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -97,6 +97,7 @@ def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path): repo, fn, token = None, + cache_dir = None, ): assert repo == "someone/my-remote-lora" assert fn == "adapter_config.json" @@ -128,6 +129,7 @@ def test_lora_identifier_retries_transient_then_resolves(tmp_path: Path): repo, fn, token = None, + cache_dir = None, ): calls["n"] += 1 if calls["n"] == 1: diff --git a/studio/backend/tests/test_training_config_popover_source.py b/studio/backend/tests/test_training_config_popover_source.py index 4263b012eb..452a3a1ea8 100644 --- a/studio/backend/tests/test_training_config_popover_source.py +++ b/studio/backend/tests/test_training_config_popover_source.py @@ -105,5 +105,6 @@ def test_shared_mapper_matches_backend_config_keys(): "lora_dropout", "use_rslora", "use_loftq", + "use_dora", ): assert key in src, f"run-config mapper lost backend key {key}" diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py index d75b205f35..e7e47478b5 100644 --- a/studio/backend/tests/test_training_pump_resilience.py +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -310,6 +310,80 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): assert b._pump_running is False +def test_interrupted_cancel_clears_in_memory_output_dir(monkeypatch): + # Stop-without-save interrupted before its complete event: /status must not + # keep serving the cleared run's output_dir. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert b._output_dir is None + assert finalized.get("status") == "stopped" + assert finalized.get("output_dir") is None + assert finalized.get("clear_output_dir") is True + + +def test_worker_exit_reuses_terminal_stop_save_error(monkeypatch): + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = False + b._output_dir = "/out/x" + b.current_job_id = "job-x" + b._terminal_finalize_payload = { + "status": "error", + "error_message": "checkpoint failed", + "output_dir": "/out/x", + "clear_output_dir": False, + "resume_blocked": True, + "expected_job_id": "job-x", + } + + b._pump_loop() + + assert b._output_dir == "/out/x" + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + assert finalized.get("resume_blocked") is True + + +def test_dead_worker_crash_preserves_output_dir(monkeypatch): + # A crash (no stop requested) after output_dir was emitted must keep the dir + # in the error finalize: checkpoints under it may still exist. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + + def test_start_training_clears_stale_pump_running_flag(): # A prior pump that died abnormally leaves _pump_running True. The next # start_training must clear it during reset so the start-time watchdog can't @@ -492,3 +566,34 @@ def test_db_run_created_before_pump_consumes_events(monkeypatch): # The pump observed an already-created run; it would be False if the pump # were started before the eager create. assert seen["db_created"] is True + + +def test_startup_flag_reports_training_active_before_proc(): + # Between freeing VRAM and _proc going live, a concurrent STT load must see + # training as active so it does not grab the just-freed GPU. + b = TrainingBackend() + b._spawn_in_progress = True + assert b.is_training_active() is True + + +def test_before_spawn_runs_inside_active_window(monkeypatch): + # The VRAM-freeing hook must run while training already counts as active, or + # an STT load racing it would place Whisper back on the freed GPU. + b = TrainingBackend() + _stub_spawn(monkeypatch) + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_pump_loop", lambda: setattr(b, "_pump_running", False)) + + active_during_free = {} + + def before_spawn(): + active_during_free["value"] = b.is_training_active() + + assert b.start_training("job_active_window", model_name = "m", before_spawn = before_spawn) is True + if b._pump_thread is not None: + b._pump_thread.join(timeout = 2.0) + + assert active_during_free["value"] is True + # The transient flag clears, but the live proc keeps training active. + assert b._spawn_in_progress is False + assert b.is_training_active() is True diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py index 91fdac9961..51425b0428 100644 --- a/studio/backend/tests/test_training_resume.py +++ b/studio/backend/tests/test_training_resume.py @@ -7,6 +7,9 @@ import importlib.util import json from pathlib import Path +import pytest +import torch + _BACKEND = Path(__file__).resolve().parents[1] @@ -25,6 +28,30 @@ def _load_resume_module(): resume = _load_resume_module() +def test_resume_request_accepts_sanitized_null_target_modules(): + from models.training import TrainingStartRequest + request = TrainingStartRequest( + model_name = "unsloth/Qwen3-0.6B", + training_type = "Full Finetuning", + format_type = "alpaca", + target_modules = None, + ) + + assert request.target_modules == [] + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + torch.save({"weight": torch.ones(1)}, checkpoint / "adapter_model.bin") + torch.save({"state": {0: torch.ones(1)}}, checkpoint / "optimizer.pt") + torch.save({"last_epoch": step}, checkpoint / "scheduler.pt") + return checkpoint + + def _stopped_run(**overrides): run = { "status": "stopped", @@ -44,6 +71,36 @@ def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch): assert resume.can_resume_run(_stopped_run()) is True +def test_can_resume_run_allows_errored_run_with_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error")) is True + + +def test_can_resume_run_rejects_errored_run_without_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: False) + + assert resume.can_resume_run(_stopped_run(status = "error")) is False + + +def test_can_resume_run_allows_errored_run_at_final_step(monkeypatch): + # A save-time crash records final_step == total_steps; resuming re-runs the + # final-save path from the checkpoint. + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(status = "error", final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is True + + +def test_can_resume_run_rejects_stopped_run_at_final_step(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is False + + def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) @@ -91,3 +148,444 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path) result = studio_db.list_runs() assert result["runs"][0]["config_json"] == config_json + + +def test_crashed_run_with_persisted_output_dir_is_resumable(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-crash", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-crash", str(out)) + conn = studio_db.get_connection() + conn.execute("UPDATE training_runs SET status = 'error' WHERE id = 'run-crash'") + conn.commit() + conn.close() + + run = studio_db.get_run("run-crash") + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is True + + +def test_checkpoint_discovery_skips_malformed_newest(monkeypatch, tmp_path): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_x" + valid = _write_checkpoint(out, 5) + (_write_checkpoint(out, 8) / "scheduler.pt").unlink() + malformed = out / "checkpoint-10" + malformed.mkdir() + (malformed / "trainer_state.json").write_text(json.dumps({"global_step": 10}), encoding = "utf-8") + (malformed / "adapter_model.bin").write_bytes(b"not a torch archive") + (malformed / "optimizer.pt").write_bytes(b"not a torch archive") + + assert resume.get_resume_checkpoint_path(str(out)) == str(valid) + + +def test_completed_run_keeps_output_dir_and_rejects_stale_cancel(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "completed", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.mark_run_cancel_requested("r") is False + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.get_run("r")["resume_blocked"] == 0 + + +def test_finish_run_clears_output_dir_for_stop_without_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + conn = studio_db.get_connection() + conn.execute( + "UPDATE training_runs SET status = 'running', output_dir = '/out/x', resume_blocked = 0 WHERE id = 'r'" + ) + conn.commit() + conn.close() + studio_db.mark_run_cancel_requested("r") + studio_db.cleanup_orphaned_runs() + assert studio_db.get_run("r")["status"] == "stopped" + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_clears_output_dir_on_cancel_error_finalize(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = "worker failed during cancel", + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_preserves_output_dir_for_interrupted_stop_and_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + + +def test_resumed_errored_run_is_not_offered_again(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + with pytest.raises(RuntimeError, match = "no longer available"): + studio_db.create_run( + id = "run-duplicate", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:01Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + assert studio_db.get_run("run-duplicate") is None + studio_db.finish_run( + id = "run-new", + status = "error", + ended_at = "2026-01-02T00:05:00Z", + final_step = 15, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed again", + ) + + old_run = studio_db.get_run("run-old") + new_run = studio_db.get_run("run-new") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert new_run["resumed_later"] == 0 + assert resume.can_resume_run(new_run) is True + assert studio_db.get_resumable_run_by_output_dir(str(out))["id"] == "run-new" + + +def test_running_continuation_blocks_older_resume(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + + old_run = studio_db.get_run("run-old") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert studio_db.get_resumable_run_by_output_dir(str(out)) is None + + +def test_stop_save_checkpoint_failure_keeps_error_status(monkeypatch, tmp_path): + # A stop-and-save whose checkpoint write failed must finalize as an error so + # history explains the missing resume state (keep_error_status flag). + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-failed-save", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-failed-save" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + } + ) + + run = studio_db.get_run("run-failed-save") + assert run["status"] == "error" + assert "resumable checkpoint" in run["error_message"] + + +def test_can_resume_run_rejects_resume_blocked_run(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error", resume_blocked = 1)) is False + + +def test_stop_save_checkpoint_failure_with_stale_checkpoint_is_not_resumable(monkeypatch, tmp_path): + # A failed stop-and-save must not offer Resume from an older periodic checkpoint; + # that would roll back past the recorded final step. + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-stale-ckpt", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-stale-ckpt", str(out)) + backend = TrainingBackend() + backend.current_job_id = "run-stale-ckpt" + backend._db_run_created = True + backend._should_stop = True + backend._output_dir = str(out) + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + "resume_blocked": True, + } + ) + + run = studio_db.get_run("run-stale-ckpt") + assert run["status"] == "error" + assert run["resume_blocked"] == 1 + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is False + + +def test_user_stop_error_without_checkpoint_ack_is_blocked(monkeypatch, tmp_path): + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-user-stop", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-user-stop" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event({"type": "error", "error": "interrupted"}) + + run = studio_db.get_run("run-user-stop") + assert run["status"] == "error" and run["resume_blocked"] == 1 + + +def test_terminal_fallback_keeps_resumable_when_current_checkpoint_landed(monkeypatch, tmp_path): + # Worker died before its terminal event, but a valid current-step checkpoint + # is on disk: the fallback must keep the run resumable, not block it. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_ok" + _write_checkpoint(out, 7) + + backend = TrainingBackend() + backend.current_job_id = "run-ok" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "stopped" + assert kwargs["resume_blocked"] is False + + +def test_terminal_fallback_blocks_when_no_current_checkpoint(monkeypatch, tmp_path): + # Same path, but only a stale (older-step) checkpoint exists: must block. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_stale" + _write_checkpoint(out, 5) + + backend = TrainingBackend() + backend.current_job_id = "run-stale" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "error" + assert kwargs["resume_blocked"] is True diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 0cd702bce2..cbe2082e82 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -353,7 +353,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): # stopped so the UI leaves "Stopping..." and a new run can start. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) # wedged: still reports alive b._should_stop = True @@ -365,7 +365,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): assert b._proc is None, "the wedged handle must be dropped so is_training_active clears" assert b._progress.is_training is False - assert b._progress.status_message == "Training stopped." + assert "valid current-step checkpoint" in b._progress.status_message assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" assert b.is_training_active() is False @@ -375,7 +375,7 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): # must record it even if the watchdog wins the finalize race against the pump. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) b._should_stop = True @@ -390,6 +390,28 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): assert finstop[0][1] == "/tmp/outputs/run-123" +def test_finalize_after_escalation_clears_output_dir_on_cancel(monkeypatch): + # Stop-without-saving promises no resume: a cancel that escalates through the + # watchdog clears the persisted output_dir, not a checkpoint path. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append((a, k))) + + b._proc = _FakeProc(alive = True) + b._should_stop = True + b._cancel_requested = True + b.current_job_id = "job_c" + b._db_run_created = True + b._output_dir = "/tmp/outputs/run-123" + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + assert finstop and finstop[0][0][0] == "job_c" + assert finstop[0][0][1] is None, "a cancelled run must not record a checkpoint path" + assert finstop[0][1].get("clear_output_dir") is True + assert b._output_dir is None, "/status must stop exposing the cancelled run's dir" + + def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch): # No worker -> nothing to escalate; the watchdog must not spawn. b = TrainingBackend() @@ -409,7 +431,7 @@ def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch): # The escalation finalize must then leave the NEW run untouched, not drop its handle. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) old_proc = _FakeProc(alive = False) # force-terminated worker we were watching new_proc = _FakeProc(alive = True) # a new run already took over @@ -430,7 +452,7 @@ def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch): # finalizes the captured run by id. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) proc = _FakeProc(alive = False) b._proc = proc @@ -451,7 +473,7 @@ def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypat # catch this even though the proc-only guard would not. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet b._proc = old_proc # still the old handle (== target), so proc guard would pass @@ -509,6 +531,7 @@ def _install_fake_db(monkeypatch): recs["insert_ids"].append(job_id), ) fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id")) + fake_db.mark_run_cancel_requested = lambda _run_id: True fake_storage.studio_db = fake_db monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) @@ -518,9 +541,48 @@ def _install_fake_db(monkeypatch): return recs +def test_stop_without_save_creates_missing_row_before_signal(monkeypatch): + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id, b._db_config = "job_missing", {"model_name": "m"} + b._stop_queue = queue.Queue() + assert b.stop_training(save = False) is True + assert [run["id"] for run in recs["created"]] == ["job_missing"] + assert b._stop_queue.get_nowait() == {"type": "stop", "save": False} + + b._cancel_requested = b._should_stop = False + sys.modules["storage.studio_db"].mark_run_cancel_requested = lambda _run_id: False + assert b.stop_training(save = False) is False + assert not b._cancel_requested and b._stop_queue.empty() + + new_queue = queue.Queue() + b.current_job_id, b._db_run_created = "job_old", True + b._cancel_requested = b._should_stop = False + + def _supersede(_run_id): + b.current_job_id = "job_new" + b._stop_queue = new_queue + return True + + sys.modules["storage.studio_db"].mark_run_cancel_requested = _supersede + assert b.stop_training(save = False) is False + assert not b._cancel_requested and new_queue.empty() + + def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): # The watchdog and pump can both finalize; only one call may reach finish_run. recs = _install_fake_db(monkeypatch) + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + attempts = 0 + + def flaky_finish(**kw): + nonlocal attempts + attempts += 1 + if attempts < 3: + raise RuntimeError("database is locked") + recs["finished"].append(kw) + + sys.modules["storage.studio_db"].finish_run = flaky_finish b = TrainingBackend() b.current_job_id = "job_x" b._db_run_created = True @@ -539,6 +601,7 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): t.join(timeout = 5) assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert attempts == 3 assert b._run_finalized is True @@ -646,7 +709,13 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) - b._ensure_db_run_created() + b._run_intent_lock.acquire() + creator = threading.Thread(target = b._ensure_db_run_created) + creator.start() + time.sleep(0.02) + assert b._db_create_in_progress is False + b._run_intent_lock.release() + creator.join(timeout = 5) assert observed["flag_during_create"] is False, "flag must not be published before insert" assert observed["in_progress_during_create"] is True @@ -718,6 +787,7 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b = TrainingBackend() b.current_job_id = "job_old" b._db_run_created = True + b._should_stop = True b._proc = _FakeProc(alive = False) b._progress.is_training = True b._progress.step = 42 @@ -726,7 +796,8 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old") assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id" - assert recs["finished"][0]["status"] == "stopped" + assert recs["finished"][0]["status"] == "error" + assert recs["finished"][0]["resume_blocked"] is True assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" assert b._metric_buffer == [], "the captured batch must be drained" @@ -737,7 +808,7 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): # so the pump's create-then-finalize records the run. Parent state still clears. b = TrainingBackend() called: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: called.append(a)) b._proc = _FakeProc(alive = False) b.current_job_id = "job_q" @@ -785,7 +856,7 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): new_proc = _FakeProc(alive = True) b._proc = old_proc - def hijack(*a): + def hijack(*a, **k): b._proc = new_proc # a new run takes over during the finalize monkeypatch.setattr(b, "_finish_stopped_run", hijack) diff --git a/studio/backend/tests/test_training_vram_coexistence.py b/studio/backend/tests/test_training_vram_coexistence.py index 2bedc46d1f..217caaa4fb 100644 --- a/studio/backend/tests/test_training_vram_coexistence.py +++ b/studio/backend/tests/test_training_vram_coexistence.py @@ -82,6 +82,63 @@ def _patch_backends(inf, llama): return patch.dict(sys.modules, {"core.inference": core_inf, "routes.inference": routes_inf}) +def _fake_stt_sidecar( + *, + model = None, + device = None, + loading = False, +): + sidecar = SimpleNamespace( + loaded_model = model, + device = device, + is_loading = lambda: loading, + ) + sidecar.cancel_pending_load = MagicMock(return_value = loading) + sidecar.wait_for_load_to_settle = MagicMock() + sidecar.unload = MagicMock() + return sidecar + + +def _fake_ggml_sidecar( + *, + model = None, + device = None, + loading = False, +): + ggml = SimpleNamespace( + loaded_model = model, + device = device, + is_loading = lambda: loading, + ) + ggml.cancel_pending_load = MagicMock(return_value = loading) + ggml.wait_for_load_to_settle = MagicMock() + ggml.unload = MagicMock() + return ggml + + +def _patch_stt(sidecar): + stt_module = types.ModuleType("core.inference.stt_sidecar") + stt_module.get_stt_sidecar = lambda: sidecar + # A fresh import of the GGUF sidecar pulls names from the fake module + # above and fails; fake it too so test ordering cannot break that import. + ggml_module = types.ModuleType("core.inference.stt_ggml_sidecar") + empty_ggml = _fake_ggml_sidecar() + ggml_module.get_ggml_stt_sidecar = lambda: empty_ggml + return patch.dict( + sys.modules, + { + "core.inference.stt_sidecar": stt_module, + "core.inference.stt_ggml_sidecar": ggml_module, + }, + ) + + +def _patch_ggml_stt(sidecar): + ggml_module = types.ModuleType("core.inference.stt_ggml_sidecar") + ggml_module.get_ggml_stt_sidecar = lambda: sidecar + return patch.dict(sys.modules, {"core.inference.stt_ggml_sidecar": ggml_module}) + + # โ”€โ”€ summarize_resident_chat โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ @@ -169,6 +226,49 @@ class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase): self.assertTrue(out["any"]) # GGUF still detected +class TestSummarizeResidentStt(_GpuCacheResetMixin, unittest.TestCase): + def test_reports_resident_model(self): + sidecar = _fake_stt_sidecar(model = "small", device = "cuda") + with _patch_stt(sidecar): + out = tv.summarize_resident_stt() + self.assertEqual(out["model"], "small") + self.assertEqual(out["device"], "cuda") + self.assertTrue(out["any"]) + self.assertFalse(out["loading"]) + + def test_reports_inflight_load(self): + sidecar = _fake_stt_sidecar(loading = True) + with _patch_stt(sidecar): + out = tv.summarize_resident_stt() + self.assertTrue(out["any"]) + self.assertTrue(out["loading"]) + + def test_reports_empty_sidecar(self): + with _patch_stt(_fake_stt_sidecar()): + out = tv.summarize_resident_stt() + self.assertFalse(out["any"]) + + def test_reports_resident_gguf_when_transformers_idle(self): + ggml = _fake_ggml_sidecar(model = "small", device = "whisper.cpp") + with _patch_stt(_fake_stt_sidecar()), _patch_ggml_stt(ggml): + out = tv.summarize_resident_stt() + self.assertEqual(out["model"], "small") + self.assertEqual(out["device"], "whisper.cpp") + self.assertTrue(out["any"]) + + def test_resident_transformers_does_not_mask_loading_gguf(self): + # A Transformers model resident on CPU holds no VRAM, but a GGUF + # whisper-server still binding its accelerator backend does; the CPU + # model must not hide that in-flight startup from training admission. + sidecar = _fake_stt_sidecar(model = "small", device = "cpu") + ggml = _fake_ggml_sidecar(loading = True) + with _patch_stt(sidecar), _patch_ggml_stt(ggml): + out = tv.summarize_resident_stt() + self.assertEqual(out["model"], "small") + self.assertTrue(out["loading"]) + self.assertTrue(out["any"]) + + # โ”€โ”€ can_keep_during_training (auto mode) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ @@ -226,12 +326,21 @@ class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase): keep, _, _ = self._run((None, meta)) self.assertFalse(keep) - def test_unload_on_non_cuda(self): + def test_unload_on_non_accelerator(self): keep, info, auto_mock = self._run(([0], {}), device = DeviceType.CPU) self.assertFalse(keep) - self.assertEqual(info["mode"], "non_cuda") + self.assertEqual(info["mode"], "non_accelerator") auto_mock.assert_not_called() + def test_xpu_gets_sized_like_cuda(self): + # XPU is a first-class training backend: the keep-guard must size it, + # not blanket-unload it as a non-accelerator. + meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0} + keep, info, auto_mock = self._run(([0], meta), device = DeviceType.XPU) + self.assertTrue(keep) + self.assertNotEqual(info.get("mode"), "non_accelerator") + auto_mock.assert_called_once() + def test_full_finetuning_forces_16bit_in_estimate(self): meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0} _keep, _info, auto_mock = self._run( @@ -438,5 +547,151 @@ class TestFreeChatModels(_GpuCacheResetMixin, unittest.TestCase): self.assertEqual(freed, ["gguf:gemma.gguf"]) +class TestFreeSttModel(_GpuCacheResetMixin, unittest.TestCase): + def test_unloads_resident_model(self): + sidecar = _fake_stt_sidecar(model = "small", device = "cuda") + with _patch_stt(sidecar): + freed = tv.free_stt_model_for_training(reason = "test") + sidecar.unload.assert_called_once() + self.assertEqual(freed, ["stt:small"]) + + def test_cancels_inflight_load_and_waits_to_settle(self): + sidecar = _fake_stt_sidecar(loading = True) + with _patch_stt(sidecar): + freed = tv.free_stt_model_for_training(reason = "test") + sidecar.cancel_pending_load.assert_called_once() + # The cancelled loader may still hold VRAM; we wait for it to release. + sidecar.wait_for_load_to_settle.assert_called_once() + # No model surfaced after the wait, so nothing to unload. + sidecar.unload.assert_not_called() + self.assertEqual(freed, ["stt:loading"]) + + def test_cancels_inflight_load_then_unloads_settled_model(self): + # A load that finished before observing the cancel leaves a resident + # model behind; it must be unloaded so training reclaims the memory. + sidecar = _fake_stt_sidecar(model = "small", loading = True) + with _patch_stt(sidecar): + freed = tv.free_stt_model_for_training(reason = "test") + sidecar.cancel_pending_load.assert_called_once() + sidecar.wait_for_load_to_settle.assert_called_once() + sidecar.unload.assert_called_once() + self.assertEqual(freed, ["stt:loading"]) + + def test_cancelled_load_still_unloads_gguf_sidecar(self): + # Cancelling a Transformers load must not skip the GGUF sidecar; both + # engines can hold memory at once (engine switch or direct load calls). + sidecar = _fake_stt_sidecar(loading = True) + ggml = _fake_ggml_sidecar(model = "small") + with _patch_stt(sidecar), _patch_ggml_stt(ggml): + freed = tv.free_stt_model_for_training(reason = "test") + sidecar.cancel_pending_load.assert_called_once() + ggml.unload.assert_called_once() + self.assertEqual(freed, ["stt:loading", "stt:small"]) + + def test_leaves_empty_sidecar_alone(self): + sidecar = _fake_stt_sidecar() + with _patch_stt(sidecar): + freed = tv.free_stt_model_for_training(reason = "test") + sidecar.unload.assert_not_called() + self.assertEqual(freed, []) + + def test_cancels_inflight_gguf_load_and_waits_to_settle(self): + # A GGUF whisper-server still in startup has no loaded_model yet, so the + # coordinator must cancel and wait for it, not skip it, before training + # claims the accelerator memory it is binding. + sidecar = _fake_stt_sidecar() # Transformers idle + ggml = _fake_ggml_sidecar(loading = True) + with _patch_stt(sidecar), _patch_ggml_stt(ggml): + freed = tv.free_stt_model_for_training(reason = "test") + ggml.cancel_pending_load.assert_called_once() + ggml.wait_for_load_to_settle.assert_called_once() + ggml.unload.assert_not_called() # nothing surfaced after the wait + self.assertEqual(freed, ["stt:gguf-loading"]) + + +class TestCoordinateModels(_GpuCacheResetMixin, unittest.TestCase): + def _run(self, chat, stt, keep_results): + keep = MagicMock(side_effect = keep_results) + with ( + patch.object(tv, "summarize_resident_chat", return_value = chat), + patch.object(tv, "summarize_resident_stt", return_value = stt), + patch.object( + tv, + "free_stt_model_for_training", + return_value = ["stt:small"], + ) as free_stt, + patch.object( + tv, + "free_chat_models_for_training", + return_value = ["hf:chat"], + ) as free_chat, + ): + freed = tv.coordinate_models_for_training(keep) + return freed, keep, free_stt, free_chat + + def test_keeps_everything_when_training_fits(self): + chat = {"any": True, "loading": False} + stt = {"any": True, "loading": False} + freed, keep, free_stt, free_chat = self._run( + chat, + stt, + [(True, {"usable_gb": 40, "required_gb": 10})], + ) + self.assertEqual(freed, []) + keep.assert_called_once() + free_stt.assert_not_called() + free_chat.assert_not_called() + + def test_frees_stt_before_chat(self): + chat = {"any": True, "loading": False} + stt = {"any": True, "loading": False} + freed, keep, free_stt, free_chat = self._run( + chat, + stt, + [ + (False, {"usable_gb": 8, "required_gb": 10}), + (True, {"usable_gb": 12, "required_gb": 10}), + ], + ) + self.assertEqual(freed, ["stt:small"]) + self.assertEqual(keep.call_count, 2) + free_stt.assert_called_once() + free_chat.assert_not_called() + + def test_frees_chat_when_stt_is_not_enough(self): + chat = {"any": True, "loading": False} + stt = {"any": True, "loading": False} + freed, keep, free_stt, free_chat = self._run( + chat, + stt, + [ + (False, {"usable_gb": 8, "required_gb": 10}), + (False, {"usable_gb": 9, "required_gb": 10}), + ], + ) + self.assertEqual(freed, ["stt:small", "hf:chat"]) + self.assertEqual(keep.call_count, 2) + free_stt.assert_called_once() + free_chat.assert_called_once() + + def test_frees_loading_models_without_probe(self): + chat = {"any": True, "loading": True} + stt = {"any": True, "loading": True} + freed, keep, free_stt, free_chat = self._run(chat, stt, []) + self.assertEqual(freed, ["stt:small", "hf:chat"]) + keep.assert_not_called() + free_stt.assert_called_once() + free_chat.assert_called_once() + + def test_cancels_loading_stt_without_probe(self): + chat = {"any": False, "loading": False} + stt = {"any": True, "loading": True} + freed, keep, free_stt, free_chat = self._run(chat, stt, []) + self.assertEqual(freed, ["stt:small"]) + keep.assert_not_called() + free_stt.assert_called_once() + free_chat.assert_not_called() + + if __name__ == "__main__": unittest.main() diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index a6e6803a5c..acb2ec449b 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -160,6 +160,25 @@ class TestResolveBaseModel: class TestRemoteLoraBase: """_remote_lora_base reads a remote adapter's base from its Hub adapter_config.json.""" + @pytest.fixture(autouse = True) + def _selected_cache_follows_env(self, monkeypatch): + # The cache helpers now read the selected cache (get_hf_cache_paths), + # which snapshots env at import; make it follow the HF_HUB_CACHE these + # tests set so they keep driving the lookup via env. + monkeypatch.setattr( + "utils.transformers_version.get_hf_cache_paths", + lambda: _types.SimpleNamespace( + hub_cache = Path( + os.environ.get("HF_HUB_CACHE") + or os.environ.get("HUGGINGFACE_HUB_CACHE") + or os.path.join( + os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"), + "hub", + ) + ) + ), + ) + @staticmethod def _resp(cfg: dict): class _Resp: @@ -645,6 +664,24 @@ def _hf_response(cfg: dict): class TestConfigJsonHfCacheFallback: """HF hub cache is consulted only offline or after a failed fetch (never stale online).""" + @pytest.fixture(autouse = True) + def _selected_cache_follows_env(self, monkeypatch): + # As above: route the selected-cache lookup through the HF_HUB_CACHE env + # these tests set, since get_hf_cache_paths snapshots env at import. + monkeypatch.setattr( + "utils.transformers_version.get_hf_cache_paths", + lambda: _types.SimpleNamespace( + hub_cache = Path( + os.environ.get("HF_HUB_CACHE") + or os.environ.get("HUGGINGFACE_HUB_CACHE") + or os.path.join( + os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"), + "hub", + ) + ) + ), + ) + def setup_method(self): _config_json_cache.clear() diff --git a/studio/backend/tests/test_whisper_cpp_freshness.py b/studio/backend/tests/test_whisper_cpp_freshness.py new file mode 100644 index 0000000000..69f0c87cee --- /dev/null +++ b/studio/backend/tests/test_whisper_cpp_freshness.py @@ -0,0 +1,156 @@ +# 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 whisper.cpp prebuilt freshness check. + +Pins the whisper-specific version policy: the release-tag parser, the +is_behind decision matrix (with its downgrade guard), and one end-to-end +wiring smoke through the shared freshness flow. The shared marker-walk and +fail-open mechanics are covered by test_llama_cpp_freshness.py. +""" + +from __future__ import annotations + +import json +import sys +import types as _types +from datetime import datetime, timedelta, timezone +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + + +class _NoopLogger: + """structlog-style logger: every method swallows positional + kwargs.""" + + def __getattr__(self, _name): + return lambda *a, **k: None + + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda *a, **k: _NoopLogger() +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: _NoopLogger() +sys.modules.setdefault("structlog", _structlog_stub) + +import pytest + +from utils import whisper_cpp_freshness as fr + + +# Helpers. + + +def _write_marker(install_dir: Path, **overrides) -> Path: + payload = { + "requested_tag": "latest", + "release_tag": "v1.9.1-unsloth.1", + "upstream_tag": "v1.9.1", + "published_repo": "unslothai/whisper.cpp", + "asset": "whisper-v1.9.1-unsloth.1-linux-x64-cpu.tar.gz", + "asset_sha256": None, + "source": "published", + "installed_at_utc": (datetime.now(tz = timezone.utc) - timedelta(days = 1)) + .isoformat() + .replace("+00:00", "Z"), + } + payload.update(overrides) + install_dir.mkdir(parents = True, exist_ok = True) + marker = install_dir / "UNSLOTH_WHISPER_PREBUILT_INFO.json" + marker.write_text(json.dumps(payload)) + return marker + + +def _fake_binary(install_dir: Path) -> Path: + """Stub whisper-server under the canonical cmake install layout.""" + bin_dir = install_dir / "build" / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + bin_path = bin_dir / "whisper-server" + bin_path.write_text("stub\n") + return bin_path + + +@pytest.fixture(autouse = True) +def _reset(monkeypatch, tmp_path): + # Isolate disk cache per-test; never touch the real cache. + monkeypatch.setattr(fr, "_cache_dir", lambda: tmp_path / ".freshness") + fr.reset_caches() + yield + fr.reset_caches() + + +# parse_release_version. + + +def test_parse_release_version(): + assert fr.parse_release_version("v1.9.1-unsloth.2") == (1, 9, 1, 2) + assert fr.parse_release_version("1.10.0") == (1, 10, 0, 0) # no v, no serial + assert fr.parse_release_version(" v2.0.0-unsloth.10 ") == (2, 0, 0, 10) + assert fr.parse_release_version("v1.9") == (1, 9, 0, 0) # padded + assert fr.parse_release_version("nightly") is None + assert fr.parse_release_version(None) is None + assert fr.parse_release_version("") is None + + +# is_behind decision matrix + downgrade guard. + + +def test_is_behind_serial_bump(): + assert fr.is_behind("v1.9.1-unsloth.1", "v1.9.1-unsloth.2") is True + + +def test_is_behind_downgrade_guard(): + # A lower serial or version is never "behind". + assert fr.is_behind("v1.9.1-unsloth.2", "v1.9.1-unsloth.1") is False + assert fr.is_behind("v1.10.0-unsloth.1", "v1.9.1-unsloth.9") is False + + +def test_is_behind_upstream_bump(): + assert fr.is_behind("v1.9.1-unsloth.1", "v1.10.0-unsloth.1") is True + + +def test_is_behind_identical_is_false(): + assert fr.is_behind("v1.9.1-unsloth.1", "v1.9.1-unsloth.1") is False + + +def test_is_behind_unparseable_differs_is_behind(): + assert fr.is_behind("v1.9.1-unsloth.1", "nightly") is True + + +def test_is_behind_missing_side_fails_open(): + assert fr.is_behind(None, "v1.9.1-unsloth.2") is False + assert fr.is_behind("v1.9.1-unsloth.1", None) is False + + +# check_prebuilt_freshness end-to-end. + + +def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(monkeypatch, tmp_path): + _write_marker( + tmp_path, + release_tag = "v1.9.1-unsloth.1", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 10)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(tmp_path) + monkeypatch.setattr(fr, "latest_published_release", lambda *a, **k: "v1.9.1-unsloth.3") + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["has_marker"] is True + assert info["behind"] is True + assert info["stale"] is True + assert info["installed_tag"] == "v1.9.1-unsloth.1" + assert info["latest_tag"] == "v1.9.1-unsloth.3" + + +def test_marker_reader_prefers_install_root_over_packaging_marker(tmp_path): + root_marker = _write_marker(tmp_path, release_tag = "v1.9.1-unsloth.2") + binary = _fake_binary(tmp_path) + (binary.parent / root_marker.name).write_text( + json.dumps({"backend": "slim", "release_tag": "archive-metadata"}) + ) + assert fr.read_install_marker(str(binary))["release_tag"] == "v1.9.1-unsloth.2" diff --git a/studio/backend/tests/test_windows_external_drive_paths.py b/studio/backend/tests/test_windows_external_drive_paths.py index 9686d45c9f..5687612916 100644 --- a/studio/backend/tests/test_windows_external_drive_paths.py +++ b/studio/backend/tests/test_windows_external_drive_paths.py @@ -57,6 +57,18 @@ def test_windows_drive_roots_empty_off_windows(monkeypatch): assert external_media.windows_drive_roots() == [] +def test_macos_volume_roots_lists_readable_mounts(monkeypatch, tmp_path): + volumes = tmp_path / "Volumes" + external = volumes / "External SSD" + unreadable = volumes / "Unavailable" + external.mkdir(parents = True) + unreadable.mkdir() + monkeypatch.setattr(external_media.platform, "system", lambda: "Darwin") + monkeypatch.setattr(external_media.os, "access", lambda path, _mode: Path(path) == external) + + assert external_media.macos_volume_roots(volumes) == [external] + + def test_windows_drive_roots_lists_readable_drives(monkeypatch): _stub_windows(monkeypatch, {"C", "D", "E"}) @@ -204,8 +216,10 @@ def test_browse_allowlist_includes_windows_drive_roots(monkeypatch, tmp_path): ) fake_external_media = SimpleNamespace( linux_run_media_mount_roots = lambda: [], + macos_volume_roots = lambda: [], windows_drive_roots = lambda: [drive_root], ) + fake_paths.external_media = fake_external_media fake_studio_db = SimpleNamespace( list_scan_folders = lambda: [], contains_sensitive_path_component = lambda _p: False, @@ -270,8 +284,10 @@ def test_build_browse_allowlist_reuses_passed_roots(monkeypatch, tmp_path): ) fake_external_media = SimpleNamespace( linux_run_media_mount_roots = _media_roots, + macos_volume_roots = lambda: [], windows_drive_roots = _drive_roots, ) + fake_paths.external_media = fake_external_media fake_studio_db = SimpleNamespace(list_scan_folders = lambda: []) monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 95c9a00534..9035068c01 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -406,10 +406,13 @@ def convert_to_vlm_format( elif _image_lookup is not None and image_data in _image_lookup: # Bare filename โ†’ resolve via HF repo lookup from huggingface_hub import hf_hub_download + from utils.hf_cache_settings import active_hf_hub_cache + local_path = hf_hub_download( dataset_name, _image_lookup[image_data], repo_type = "dataset", + cache_dir = active_hf_hub_cache(), ) image_data = Image.open(local_path).convert("RGB") else: @@ -774,10 +777,13 @@ def convert_sharegpt_with_images_to_vlm_format( return Image.open(BytesIO(f.read())).convert("RGB") elif _image_lookup is not None and image_data in _image_lookup: from huggingface_hub import hf_hub_download + from utils.hf_cache_settings import active_hf_hub_cache + local_path = hf_hub_download( dataset_name, _image_lookup[image_data], repo_type = "dataset", + cache_dir = active_hf_hub_cache(), ) return Image.open(local_path).convert("RGB") else: diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index f7b35e2869..c594e883a8 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -58,6 +58,7 @@ def precache_helper_gguf(): try: from huggingface_hub import HfApi, hf_hub_download from huggingface_hub.utils import disable_progress_bars, enable_progress_bars + from utils.hf_cache_settings import active_hf_hub_cache disable_progress_bars() logging.getLogger("huggingface_hub").setLevel(logging.WARNING) @@ -76,7 +77,11 @@ def precache_helper_gguf(): + (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else "") ) for target in matching: - hf_hub_download(repo_id = repo, filename = target) + hf_hub_download( + repo_id = repo, + filename = target, + cache_dir = active_hf_hub_cache(), + ) logger.info(f"Helper GGUF cached: {len(matching)} file(s)") else: logger.warning(f"No GGUF matching variant '{variant}' in {repo}") diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 62b537fbac..138238533f 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -50,6 +50,11 @@ def export_capability() -> dict: return _hardware.export_capability() +def get_torch_device_str() -> str: + """Return the torch device string ("cuda", "xpu", "cpu") for the detected hardware.""" + return _hardware.get_torch_device_str() + + __all__ = [ "DeviceType", "DEVICE", @@ -75,6 +80,7 @@ __all__ = [ "estimate_required_model_memory_gb", "auto_select_gpu_ids", "prepare_gpu_selection", + "get_torch_device_str", "safe_num_proc", "safe_thread_num_proc", "dataset_map_num_proc", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 9fef53e65e..38ebc0b6d4 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -175,18 +175,64 @@ def detect_hardware() -> DeviceType: Call once at FastAPI lifespan startup; idempotent. Detection order: - 1. CUDA (NVIDIA GPU, requires torch) - 2. MLX (Apple Silicon via MLX framework) - 3. CPU (fallback) + 1. XPU-preferred hint: only on an unambiguous "prefer XPU" signal + (CUDA hidden via ``CUDA_VISIBLE_DEVICES="" / "-1"``, + ``UNSLOTH_FORCE_XPU=1``, or CUDA unavailable) AND a non-empty + ``ZE_AFFINITY_MASK`` AND ``torch.xpu`` reports a device. A stray + inherited mask is not enough: CUDA still wins on hybrid hosts. + 2. CUDA (NVIDIA GPU, requires torch) + 3. XPU (Intel GPU, requires torch with XPU support) + 4. MLX (Apple Silicon via MLX framework) + 5. CPU (fallback) """ global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, IS_ROCM CHAT_ONLY = True # reset -- only CUDA/ROCm/XPU/MLX sets it to False CHAT_ONLY_REASON = None IS_ROCM = False - # --- CUDA / ROCm: try PyTorch --- + # --- CUDA / ROCm / XPU: try PyTorch --- if _has_torch(): import torch + + # --- Explicit-XPU hint --- + # Prefer XPU on UNSLOTH_FORCE_XPU=1, or ZE_AFFINITY_MASK set + CUDA + # hidden/unavailable. A bare mask alone is NOT enough (can leak from + # unrelated Intel tooling); torch.xpu must report a device. + ze_mask = os.environ.get("ZE_AFFINITY_MASK") + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + cuda_hidden = cvd is not None and cvd.strip() in ("", "-1") + force_xpu = os.environ.get("UNSLOTH_FORCE_XPU") == "1" + try: + cuda_unavailable = not torch.cuda.is_available() + except Exception: + cuda_unavailable = True + + prefer_xpu = force_xpu or (bool(ze_mask) and (cuda_hidden or cuda_unavailable)) + if prefer_xpu: + try: + xpu_ok = hasattr(torch, "xpu") and torch.xpu.is_available() + except Exception: + xpu_ok = False + if xpu_ok: + # Forced XPU on a hybrid host: unsloth's device_type picks + # CUDA before XPU and ignores this Studio-only env var, so + # hide CUDA or spawned workers would silently train on CUDA. + if force_xpu and not cuda_hidden and not cuda_unavailable: + os.environ["CUDA_VISIBLE_DEVICES"] = "" + DEVICE = DeviceType.XPU + CHAT_ONLY = False + CHAT_ONLY_REASON = None + device_name = torch.xpu.get_device_name(0) + if force_xpu and not ze_mask: + reason = "UNSLOTH_FORCE_XPU=1" + elif force_xpu: + reason = "UNSLOTH_FORCE_XPU=1 + ZE_AFFINITY_MASK" + else: + reason = "ZE_AFFINITY_MASK hint honoured" + print(f"Hardware detected: XPU -- {device_name} ({reason})") + return DEVICE + + # --- CUDA: NVIDIA GPU --- if torch.cuda.is_available(): DEVICE = DeviceType.CUDA CHAT_ONLY = False @@ -327,9 +373,18 @@ def clear_gpu_cache(): torch.cuda.empty_cache() torch.cuda.ipc_collect() elif device == DeviceType.XPU: - import torch - torch.xpu.synchronize() - torch.xpu.empty_cache() + # Guard synchronize/empty_cache: older torch-xpu builds may lack + # them, and an unguarded AttributeError would propagate to callers. + # torch.xpu has no ipc_collect(), so do not call it here. + try: + import torch + if hasattr(torch, "xpu"): + if hasattr(torch.xpu, "synchronize"): + torch.xpu.synchronize() + if hasattr(torch.xpu, "empty_cache"): + torch.xpu.empty_cache() + except Exception as e: + logger.debug("Failed to clear XPU cache: %s", e) elif device == DeviceType.MLX: # MLX manages memory automatically; gc.collect() above is enough. pass @@ -500,14 +555,27 @@ def get_package_versions() -> Dict[str, Optional[str]]: except PackageNotFoundError: versions[name] = None - # GPU runtime version bundled with torch + # GPU runtime versions bundled with torch (CUDA, ROCm/HIP, Intel XPU) try: import torch + versions["cuda"] = getattr(torch.version, "cuda", None) versions["rocm"] = getattr(torch.version, "hip", None) + # Isolated probe: a broken Intel runtime raising in is_available() + # must not blank the already-read cuda/rocm versions. + try: + if hasattr(torch, "xpu") and torch.xpu.is_available(): + # torch.version.xpu may be None on modern builds; fall back to + # "available" so the UI distinguishes present-but-unknown from + # "package not found". + xpu_ver = getattr(torch.version, "xpu", None) + versions["xpu"] = xpu_ver if xpu_ver is not None else "available" + except Exception: + versions["xpu"] = None except Exception: versions["cuda"] = None versions["rocm"] = None + versions["xpu"] = None return versions @@ -547,6 +615,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] if mod is None: return [] + device = get_device() # free==total is a Windows-ROCm-only quirk. _win_rocm = sys.platform == "win32" and IS_ROCM devices = [] @@ -558,11 +627,30 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] used_bytes: Optional[int] # Prefer mem_get_info (system-wide) so auto-select sees other consumers. if hasattr(mod, "mem_get_info"): - free_bytes, total_bytes = mod.mem_get_info(ordinal) - used_bytes = total_bytes - free_bytes - # free==total is the broken-API sentinel, not an idle GPU. - if _win_rocm and free_bytes == total_bytes: + try: + free_bytes, total_bytes = mod.mem_get_info(ordinal) + used_bytes = total_bytes - free_bytes + except Exception as e: + if device != DeviceType.XPU: + raise + # Arc B580 and Lunar Lake can report properties while + # rejecting free-memory queries. Preserve the usable + # device and its total memory with unknown utilization. + logger.debug( + "XPU free-memory query failed for ordinal %d: %s", + ordinal, + e, + ) used_bytes = None + else: + # free==total is the broken-API sentinel, not an idle GPU. + if _win_rocm and free_bytes == total_bytes: + used_bytes = None + elif device == DeviceType.XPU: + # XPU without mem_get_info: memory_allocated() is process-local + # and misleading for placement, so return None for the + # selector's no-telemetry fallback. + used_bytes = None else: used_bytes = mod.memory_allocated(ordinal) devices.append( @@ -571,7 +659,9 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] "visible_ordinal": ordinal, "name": props.name, "total_gb": round(total_bytes / (1024**3), 2), - "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None, + "used_gb": ( + round(used_bytes / (1024**3), 2) if used_bytes is not None else None + ), } ) except Exception as e: @@ -582,6 +672,43 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] # ========== Live GPU Utilization ========== +def _xpu_hierarchy_is_composite() -> bool: + """Return True iff Level Zero is running in COMPOSITE device hierarchy. + + COMPOSITE: numeric ``ZE_AFFINITY_MASK`` entries address root GPU IDs + (tiles use ``N.M``). FLAT (the oneAPI default; also assumed when + ``ZE_FLAT_DEVICE_HIERARCHY`` is unset): entries address tile/device + handles, so mapping them back to root GPU IDs is unsafe. Only COMPOSITE + gives stable root-ID semantics. + """ + hierarchy = (os.environ.get("ZE_FLAT_DEVICE_HIERARCHY") or "FLAT").strip().upper() + return hierarchy == "COMPOSITE" + + +def _parse_ze_mask_roots(mask: str) -> list[int]: + """Parse a ``ZE_AFFINITY_MASK`` value into an ordered list of root device IDs. + + One root ID per mask token, preserving order and duplicates so logical + ordinals map 1-to-1 to physical root IDs (e.g. ``"0.0,0.1"`` -> ``[0, 0]``, + ``"2.0,0.1,0.2"`` -> ``[2, 0, 0]``); empty list if no parseable digits. + Only meaningful in COMPOSITE hierarchy -- callers needing a stable + root-ID mapping must gate on ``_xpu_hierarchy_is_composite()``. + """ + roots: list[int] = [] + if not mask: + return roots + for token in mask.split(","): + token = token.strip() + if not token: + continue + root = token.split(".", 1)[0] + # isdecimal() (not isdigit()) rejects Unicode superscripts like + # "ยฒ"/"ยณ", which pass isdigit() but crash int() with ValueError. + if root.isdecimal(): + roots.append(int(root)) + return roots + + def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: """Query the appropriate SMI backend (amd-smi or nvidia-smi). @@ -734,6 +861,141 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: return None, None +# 0x1002. NVIDIA's open kernel module also registers KFD nodes (vendor_id 0x10DE); +# a non-AMD node is not a HIP device and must never take an ordinal. +_AMD_PCI_VENDOR_ID = 4098 + + +def _rocm_kfd_gpu_pci_ids() -> list[str]: + """PCI addresses of the GPUs ROCm enumerates, in HIP device order. + + Reads /sys/class/kfd/kfd/topology/nodes//properties, the topology ROCm + itself enumerates from: AMD GPU nodes (simd_count > 0 excludes CPUs, + vendor_id == AMD excludes NVIDIA) in node-id order are HIP's device order, so + position N is ROCm physical device N. Unlike DRM sysfs, an amdgpu adapter HIP + cannot enumerate has no node here, so it never consumes an ordinal. + + Returns [] (disabling the overlay) when KFD is absent, and FAILS CLOSED the + same way on any unreadable node or an AMD node with no location_id: dropping + one would shift every later ordinal and let a similar-capacity GPU pass the + total-size guard while showing another card's usage. + + location_id is the kernel's (bus << 8) | devfn; domain is separate. + """ + nodes: list[tuple[int, str]] = [] + try: + node_dirs = glob.glob("/sys/class/kfd/kfd/topology/nodes/*") + except Exception: + return [] + for node_dir in node_dirs: + m = re.fullmatch(r".*/(\d+)", node_dir) + if m is None: + continue + props: dict[str, int] = {} + try: + with open(os.path.join(node_dir, "properties")) as f: + for line in f: + parts = line.split() + if len(parts) == 2: + try: + props[parts[0]] = int(parts[1]) + except ValueError: + continue + except OSError: + return [] # unreadable node could be a GPU: fail closed, don't shift + if props.get("simd_count", 0) <= 0: + continue # CPU node, not a GPU + if props.get("vendor_id") != _AMD_PCI_VENDOR_ID: + continue # non-AMD GPU node (NVIDIA open driver): not a HIP device + location_id = props.get("location_id") + if location_id is None: + return [] # an AMD GPU we cannot place: fail closed for the whole map + domain = props.get("domain", 0) + bus = (location_id >> 8) & 0xFF + devfn = location_id & 0xFF + bdf = f"{domain:04x}:{bus:02x}:{(devfn >> 3) & 0x1F:02x}.{devfn & 0x7}" + nodes.append((int(m.group(1)), bdf)) + nodes.sort(key = lambda n: n[0]) + return [bdf for _node_id, bdf in nodes] + + +def _rocm_linux_amdgpu_cards() -> list[tuple[str, int, str]]: + """The amdgpu-bound DRM cards in PCI order: ``(pci_bdf, card_no, device_dir)``. + + Membership is by the BOUND DRIVER, not the VRAM sysfs files: an AMD device + with incomplete sysfs support (some APUs expose no mem_info_vram_*) still + consumes a ROCm ordinal, and dropping it would shift every later card down. + PCI order is HIP's default enumeration order, so list position is the ROCm + ordinal; card_no is a stable tiebreak when the BDF cannot be resolved. + + NOTE this is a superset of the ROCm-visible set (a HIP-unsupported amdgpu + adapter appears too), so callers must check the counts agree before assuming + a 1:1 mapping onto torch devices. + """ + if platform.system() != "Linux": + return [] + amd_cards: list[tuple[str, int, str]] = [] + try: + for card_path in glob.glob("/sys/class/drm/card*"): + # Match card exactly so connector nodes (card0-DP-1) are skipped. + m = re.fullmatch(r".*/card(\d+)", card_path) + if m is None: + continue + dev_dir = os.path.join(card_path, "device") + try: + driver = os.path.basename(os.path.realpath(os.path.join(dev_dir, "driver"))) + except OSError: + continue + if driver != "amdgpu": + continue # foreign adapter: not a ROCm device, takes no ordinal + try: + bdf = os.path.basename(os.path.realpath(dev_dir)) + except OSError: + bdf = "" + amd_cards.append((bdf, int(m.group(1)), dev_dir)) + except Exception: + return [] + amd_cards.sort(key = lambda c: (c[0], c[1])) + return amd_cards + + +def _rocm_linux_sysfs_vram_by_pci_gb() -> dict[str, tuple[float, float]]: + """System-wide AMD VRAM via Linux DRM sysfs, keyed by the card's PCI address. + + Reads each card's mem_info_vram_{used,total} (kernel-updated across all + processes) so every GPU gets its own figure, unlike _rocm_linux_sysfs_vram_gb + which sums the host. Keyed by PCI address, not an ordinal, so the caller can + join it to _rocm_kfd_gpu_pci_ids() by identity: DRM card numbers include + foreign adapters and this set includes cards HIP does not enumerate, so any + ordinal from this list alone can be shifted relative to ROCm's. A card with + missing/unreadable/zero-total figures simply has no entry. Empty off Linux. + """ + if platform.system() != "Linux": + return {} + + try: + by_pci: dict[str, tuple[float, float]] = {} + for bdf, _card_no, dev_dir in _rocm_linux_amdgpu_cards(): + if not bdf: + continue + try: + with open(os.path.join(dev_dir, "mem_info_vram_used")) as f: + used_bytes = int(f.read().strip()) + with open(os.path.join(dev_dir, "mem_info_vram_total")) as f: + total_bytes = int(f.read().strip()) + except (OSError, ValueError): + continue + if total_bytes <= 0: + continue + by_pci[bdf.lower()] = ( + round(used_bytes / (1024**3), 2), + round(total_bytes / (1024**3), 2), + ) + return by_pci + except Exception: + return {} + + # โ”€โ”€ Windows AMD/ROCm per-adapter VRAM (issue #7072) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ # amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the # per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so @@ -1222,6 +1484,75 @@ def _reconcile_primary_rocm_unified_memory( _apply_unified_memory_correction(utilization, torch_devices[0]) +def _rocm_visibility_mask_active() -> bool: + """True when any ROCm/CUDA visibility variable filters the device set.""" + for var in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", + ): + value = os.environ.get(var) + if value and value.strip(): + return True + return False + + +def _overlay_system_wide_vram(devices: list[Dict[str, Any]]) -> None: + """Replace process-local torch VRAM with system-wide Linux ROCm figures. + + The torch fallback is process-local, so a model served by the separate + llama-server process reads as ~0 used even with the GPU full (#7072). DRM + sysfs gives per-card figures the kernel updates across all processes. Sources + are matched by the device's PHYSICAL index (never list position), and only + when NO visibility mask is active and the device count equals the host GPU + count; under any mask the index is not a verifiable host ordinal, so torch's + figures are kept. Best-effort, in place: a device with no matching card, or a + unified-memory APU whose sysfs total is below torch's GTT-backed total, keeps + torch's (mirrors _apply_unified_memory_correction). + + Windows is intentionally not overlaid: its per-adapter perf counters cannot be + mapped to ROCm ordinals and miss WDDM shared memory, so the multi-GPU view + keeps torch there rather than risk misattributing another adapter's usage. + """ + if not devices or platform.system() != "Linux": + return + # Match by PCI identity, never list position: index N in KFD topology is ROCm + # physical device N and carries its PCI address, which DRM sysfs keys on too. + # The two gates below verify ``index`` really is a host-physical ordinal + # (torch exposes no PCI id to check directly): + # * No visibility mask -- any mask makes ``index`` container/ROCR-relative + # rather than a host ordinal. + # * Device count == host GPU count -- rules out a device-cgroup container + # that sets no env var yet compacts torch's indices from zero. + pci_by_ordinal = _rocm_kfd_gpu_pci_ids() + if not pci_by_ordinal: + return + if _rocm_visibility_mask_active() or len(devices) != len(pci_by_ordinal): + return + vram_by_pci = _rocm_linux_sysfs_vram_by_pci_gb() + for dev in devices: + index = dev.get("index") + if not isinstance(index, int) or not (0 <= index < len(pci_by_ordinal)): + continue + entry = vram_by_pci.get(pci_by_ordinal[index].lower()) + if entry is None: + continue + used, total = entry + dev_total = dev.get("vram_total_gb") or 0.0 + # Overlay only a device that maps 1:1 to the whole card: torch total must + # match sysfs total within ~10%. A mismatch either way means a different + # memory scope -- a unified-memory APU (sysfs sees only the dedicated + # slice, torch the GTT pool) or a partitioned MI300 (sysfs reports the + # whole card, dwarfing a partition) -- and overlaying would misstate free + # VRAM (a partition would look like it has the whole card free). + if dev_total <= 0 or abs(total - dev_total) > 0.1 * dev_total: + continue + dev["vram_used_gb"] = used + dev["vram_total_gb"] = total + dev["vram_utilization_pct"] = round((used / total) * 100, 1) if total > 0 else None + + def get_visible_gpu_utilization() -> Dict[str, Any]: device = get_device() @@ -1300,6 +1631,13 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: for td in torch_devices: total = td["total_gb"] used = td["used_gb"] + # used=None is a deliberate "telemetry unavailable" signal + # from _torch_get_per_device_info (e.g. XPU without + # mem_get_info); propagate None instead of dividing by it. On + # CUDA/ROCm used is always an int, so this stays byte-identical. + vram_pct = ( + round((used / total) * 100, 1) if used is not None and total > 0 else None + ) devices.append( { "index": td["index"], @@ -1309,14 +1647,18 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: "temperature_c": None, "vram_used_gb": used, "vram_total_gb": total, - "vram_utilization_pct": round((used / total) * 100, 1) - if total > 0 and used is not None - else None, + "vram_utilization_pct": vram_pct, "power_draw_w": None, "power_limit_w": None, "power_utilization_pct": None, } ) + if IS_ROCM and index_kind == "physical": + # Swap process-local torch VRAM for system-wide sysfs so a model + # held by the separate llama-server process shows up (#7072). + # Physical-index only: a relative index (UUID/MIG mask) is not a + # host GPU id. The overlay verifies the rest itself. + _overlay_system_wide_vram(devices) return { "available": True, "backend": _backend_label(device), @@ -1373,6 +1715,82 @@ _visible_gpu_count: Optional[int] = None def _get_parent_visible_gpu_spec() -> Dict[str, Any]: + # On Intel XPU, visibility is controlled by ZE_AFFINITY_MASK (Level Zero), + # not CUDA_VISIBLE_DEVICES. + if get_device() == DeviceType.XPU: + xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK") + composite = _xpu_hierarchy_is_composite() + + if xpu_mask_raw is None: + # COMPOSITE: root GPU IDs are stable physical IDs. + if composite: + return { + "raw": None, + "numeric_ids": list(range(get_physical_gpu_count())), + "supports_explicit_gpu_ids": True, + } + # FLAT (oneAPI default): ordinals are tile/device handles, not + # physical GPU IDs. numeric_ids=None so telemetry uses relative + # ordinals; explicit selection needs ZE_FLAT_DEVICE_HIERARCHY=COMPOSITE. + return { + "raw": None, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + xpu_mask = xpu_mask_raw.strip() + if xpu_mask == "": + return { + "raw": xpu_mask, + "numeric_ids": [], + "supports_explicit_gpu_ids": True, + } + + # Subdevice syntax ("N.M") expands one root into multiple + # logical devices -- not addressable by explicit root-ID selection. + has_subdevice = any("." in token.strip() for token in xpu_mask.split(",") if token.strip()) + if has_subdevice: + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + # FLAT numeric entries are tile handles, not physical GPU IDs. Keep + # numeric_ids unresolved so every telemetry and picker consumer uses + # relative torch ordinals and cannot advertise them as pinnable roots. + if not composite: + tokens = [token.strip() for token in xpu_mask.split(",") if token.strip()] + if tokens and all(token.isdecimal() for token in tokens): + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + # COMPOSITE + pure numeric (subdevice handled above). _parse_ze_mask_roots + # maps to root GPU IDs, dropping non-decimal tokens so "*"/"GPU-uuid" -> []. + roots_with_dupes = _parse_ze_mask_roots(xpu_mask) + if not roots_with_dupes: + # Unparseable mask (e.g. "*", "GPU-uuid") -- cannot map to + # physical root IDs. + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + return { + "raw": xpu_mask, + "numeric_ids": roots_with_dupes, + "supports_explicit_gpu_ids": True, + } + # ROCm uses HIP/ROCR_VISIBLE_DEVICES on top of CUDA_VISIBLE_DEVICES; check # them first. Explicit None checks (not `or`) so "" reads as "no visible GPUs". cuda_visible = None @@ -1429,24 +1847,44 @@ def get_parent_visible_gpu_ids() -> list[int]: return list(parent_visible_ids) if parent_visible_ids is not None else [] -def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: +def resolve_requested_gpu_ids( + gpu_ids: Optional[list[int]], *, is_vulkan: bool = False +) -> list[int]: parent_visible_spec = _get_parent_visible_gpu_spec() parent_visible_ids = get_parent_visible_gpu_ids() physical_gpu_count = get_physical_gpu_count() if gpu_ids is None: - return parent_visible_ids + return [] if is_vulkan else parent_visible_ids requested_ids = list(gpu_ids) if len(requested_ids) == 0: - return parent_visible_ids + return [] if is_vulkan else parent_visible_ids + + if is_vulkan: + # A Vulkan build selects by ggml Vulkan ordinal (--device VulkanN), a separate + # index space from CUDA/ROCm ids that may be empty under CPU-only torch. The + # CUDA parent-visible / physical-count checks below do not apply; only reject + # malformed ordinals (issue #7239). + if len(set(requested_ids)) != len(requested_ids): + raise ValueError(f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed.") + negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0] + if negative_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: GPU IDs must be non-negative. " + f"Rejected IDs: {negative_ids}." + ) + return requested_ids if not parent_visible_spec["supports_explicit_gpu_ids"]: + env_var_name = ( + "ZE_AFFINITY_MASK" if get_device() == DeviceType.XPU else "CUDA_VISIBLE_DEVICES" + ) raise ValueError( f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are " - f"unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries " - f"({parent_visible_spec['raw']!r}). Omit gpu_ids to use the " - "parent-visible devices." + f"unsupported when {env_var_name} uses non-numeric or subdevice " + f"entries ({parent_visible_spec['raw']!r}). Omit gpu_ids to use " + "the parent-visible devices." ) if len(set(requested_ids)) != len(requested_ids): @@ -1885,8 +2323,11 @@ def auto_select_gpu_ids( ) -> tuple[Optional[list[int]], Dict[str, Any]]: metadata: Dict[str, Any] = {"selection_mode": "auto"} - if get_device() != DeviceType.CUDA: - metadata["selection_mode"] = "non_cuda" + # Auto-selection needs per-device free-VRAM telemetry, available on CUDA + # (nvidia-smi) and XPU (torch.xpu) but not MLX/CPU, which fall + # through to inheriting parent visibility. + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + metadata["selection_mode"] = "non_accelerator" return None, metadata required_gb, estimate_metadata = estimate_required_model_memory_gb( @@ -1983,12 +2424,13 @@ def auto_select_gpu_ids( metadata["selection_mode"] = "auto" metadata["selected_gpu_ids"] = selected logger.debug( - "Selected GPUs automatically", - model_name = model_name, - selected_gpu_ids = selected, - usable_gb = metadata["usable_gb"], - required_gb = metadata.get("required_gb"), - multi_gpu_overhead = multi_gpu_overhead, + "Selected GPUs automatically: model=%s selected=%s usable_gb=%s " + "required_gb=%s multi_gpu_overhead=%s", + model_name, + selected, + metadata["usable_gb"], + metadata.get("required_gb"), + multi_gpu_overhead, ) return selected, metadata @@ -2004,12 +2446,13 @@ def auto_select_gpu_ids( metadata["usable_gb"] = round(fallback_usable, 3) metadata["selected_gpu_ids"] = fallback_all logger.warning( - "Falling back to all visible GPUs -- model may not fit", - model_name = model_name, - selected_gpu_ids = fallback_all, - usable_gb = metadata["usable_gb"], - required_gb = metadata.get("required_gb"), - multi_gpu_overhead = multi_gpu_overhead, + "Falling back to all visible GPUs; model may not fit: model=%s " + "selected=%s usable_gb=%s required_gb=%s multi_gpu_overhead=%s", + model_name, + fallback_all, + metadata["usable_gb"], + metadata.get("required_gb"), + multi_gpu_overhead, ) return fallback_all, metadata @@ -2043,10 +2486,10 @@ def prepare_gpu_selection( to a Hugging Face ``device_map`` string) and to ``apply_gpu_ids()`` in the worker subprocess (narrows ``CUDA_VISIBLE_DEVICES`` before torch/CUDA init). """ - if gpu_ids and get_device() != DeviceType.CUDA: + if gpu_ids and get_device() not in (DeviceType.CUDA, DeviceType.XPU): raise ValueError( - f"gpu_ids {list(gpu_ids)} is only supported on CUDA devices, " - f"but the current backend is '{get_device().value}'." + f"gpu_ids {list(gpu_ids)} is only supported on CUDA and Intel XPU " + f"devices, but the current backend is '{get_device().value}'." ) if gpu_ids: @@ -2119,11 +2562,14 @@ def get_physical_gpu_count() -> int: def _backend_visible_devices_env() -> Optional[str]: """Return the raw visibility env string that applies to this backend. - On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over - CUDA_VISIBLE_DEVICES; this mirrors ``_get_parent_visible_gpu_spec`` so + On XPU the control is ``ZE_AFFINITY_MASK`` (not ``CUDA_VISIBLE_DEVICES``); + on ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over + CUDA_VISIBLE_DEVICES. Mirrors ``_get_parent_visible_gpu_spec`` so ``backend_cuda_visible_devices`` reports the value actually narrowing the - visible device set. + visible device set on the current backend. """ + if get_device() == DeviceType.XPU: + return os.environ.get("ZE_AFFINITY_MASK") if IS_ROCM: return _get_parent_visible_gpu_spec().get("raw") return os.environ.get("CUDA_VISIBLE_DEVICES") @@ -2238,6 +2684,43 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count + # Prefer torch.xpu.device_count() on Intel XPU: the Level Zero runtime + # correctly interprets ZE_AFFINITY_MASK semantics (e.g. subdevice syntax + # "0.0,0.1" collapses onto one root GPU). Supersedes the torch fallback below. + if get_device() == DeviceType.XPU: + xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK") + xpu_mask_set = xpu_mask_raw is not None + xpu_visible = (xpu_mask_raw or "").strip() + if xpu_mask_set and xpu_visible == "": + _visible_gpu_count = 0 + return _visible_gpu_count + + try: + import torch + _visible_gpu_count = torch.xpu.device_count() + except Exception as e: + logger.debug( + "torch.xpu.device_count() failed, falling back to mask parsing: %s", + e, + ) + if xpu_visible: + # Fallback: count unique root device IDs from the mask. + # "device.subdevice" notation means "0.0,0.1" is 1 root, not 2. + # Without torch the hierarchy mode is unknown, so root-device + # counting is the conservative choice. + if xpu_visible == "*": + # Documented wildcard: all physical XPUs visible. + _visible_gpu_count = get_physical_gpu_count() + else: + roots = _parse_ze_mask_roots(xpu_visible) + # Non-parseable masks (",,,", "GPU-abc") yield an empty + # roots list, treated as 0 visible devices, not "all + # visible" -- no evidence the whole fleet was intended. + _visible_gpu_count = len(set(roots)) + else: + _visible_gpu_count = get_physical_gpu_count() + return _visible_gpu_count + # _get_parent_visible_gpu_spec() already handles HIP_VISIBLE_DEVICES / # ROCR_VISIBLE_DEVICES on ROCm. visible_spec = _get_parent_visible_gpu_spec() @@ -2251,20 +2734,18 @@ def get_visible_gpu_count() -> int: _visible_gpu_count = len([x for x in raw.split(",") if x.strip()]) return _visible_gpu_count - # No visibility env var set -- try torch, else physical count + # No visibility env var set -- try torch, else physical count. XPU is + # handled by the early return above, so only torch.cuda is needed here. try: import torch - if get_device() == DeviceType.XPU and hasattr(torch, "xpu"): - _visible_gpu_count = torch.xpu.device_count() - else: - _visible_gpu_count = torch.cuda.device_count() + _visible_gpu_count = torch.cuda.device_count() except Exception: _visible_gpu_count = get_physical_gpu_count() return _visible_gpu_count -def apply_gpu_ids(gpu_ids) -> None: +def apply_gpu_ids(gpu_ids, backend: Optional[str] = None) -> None: if gpu_ids is None: return @@ -2280,6 +2761,62 @@ def apply_gpu_ids(gpu_ids) -> None: else: value = str(gpu_ids) + # Intel XPU honors ZE_AFFINITY_MASK, not CUDA_VISIBLE_DEVICES; route XPU + # pinning through it so worker subprocesses are restricted to the intended GPU. + # Decide WITHOUT get_device(): workers call this before detect_hardware(), + # and a lazy detect would probe torch.cuda against the unmasked parent env, + # latching device enumeration before the mask below is written. Pre-detect, + # use env + torch BUILD attributes only (no runtime init, like the ROCm + # mirror below). + _is_xpu = DEVICE == DeviceType.XPU + if backend is not None: + # The spawning parent's detected backend (config["device_backend"]): + # exact and probe-free, so the mask target always matches what + # detect_hardware() decided in the parent, including its XPU + # availability check and CUDA fallback. + _is_xpu = backend == DeviceType.XPU.value + elif DEVICE is None: + # No parent backend passed (direct caller). version.xpu can be None + # on a working XPU build, so also accept torch.xpu._is_compiled() + # (a pure symbol-presence check, no runtime init). UNSLOTH_FORCE_XPU + # counts only on an XPU-capable build: detect_hardware() falls back + # to CUDA when XPU is missing, and the mask target must follow. + try: + import torch as _torch + + _ver = _torch.version + _is_comp = getattr(getattr(_torch, "xpu", None), "_is_compiled", None) + _xpu_build = (callable(_is_comp) and bool(_is_comp())) or ( + getattr(_ver, "xpu", None) is not None + ) + if os.environ.get("UNSLOTH_FORCE_XPU") == "1": + _is_xpu = _xpu_build + else: + # Mirror detect_hardware: hidden CUDA prefers XPU on an + # XPU-capable build (with or without a ZE mask -- detection + # falls through to XPU either way), where writing these ids + # to CUDA_VISIBLE_DEVICES would re-expose the deliberately + # hidden CUDA. + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + _cuda_hidden = _cvd is not None and _cvd.strip() in ("", "-1") + _is_xpu = _xpu_build and ( + _cuda_hidden + or (getattr(_ver, "cuda", None) is None and getattr(_ver, "hip", None) is None) + ) + except Exception as e: + logger.debug( + "apply_gpu_ids: torch XPU probe skipped (%s: %s)", + type(e).__name__, + e, + ) + if _is_xpu: + os.environ["ZE_AFFINITY_MASK"] = value + # Leave inherited CUDA_VISIBLE_DEVICES alone -- clearing it could let + # the worker flip back to CUDA on hybrid hosts. + _visible_gpu_count = None + logger.info("Applied gpu_ids: ZE_AFFINITY_MASK='%s'", value) + return + os.environ["CUDA_VISIBLE_DEVICES"] = value # Keep ROCm visibility env vars in sync. Workers may call apply_gpu_ids() # before detect_hardware() (IS_ROCM still False), so also mirror when the @@ -2324,26 +2861,41 @@ def get_device_map(gpu_ids: Optional[list[int]] = None) -> str: Returns ``"balanced"`` (shard evenly across GPUs) when: - ``gpu_ids`` explicitly lists >1 GPU, **or** - - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and - >1 GPU is visible (fallback: numeric IDs unresolvable, so assume - multi-GPU is intended). + - ``CUDA_VISIBLE_DEVICES``/``ZE_AFFINITY_MASK`` uses non-numeric + identifiers (UUID/MIG/wildcard) and >1 GPU is visible (fallback: + numeric IDs unresolvable, so assume multi-GPU is intended). - Returns ``"sequential"`` (single device) otherwise, including non-CUDA - backends (CPU, MLX). + Returns ``"sequential"`` (single device) otherwise, including CPU/MLX + backends. Use ``prepare_gpu_selection()`` upstream to determine ``gpu_ids`` -- it handles auto-selecting the minimum GPUs needed for a model. """ device = get_device() - if device == DeviceType.CUDA: + if device in (DeviceType.CUDA, DeviceType.XPU): multi_gpu = gpu_ids is not None and len(gpu_ids) > 1 if not multi_gpu: - # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU - # means multi-GPU sharding is intended. parent_visible_spec = _get_parent_visible_gpu_spec() - if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: - multi_gpu = True + if device == DeviceType.CUDA: + # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU + # means multi-GPU sharding is intended. + if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: + multi_gpu = True + elif device == DeviceType.XPU and gpu_ids is None: + # Shard across visible XPU ordinals via HF (no mask rewrite), + # only when no gpu_ids were passed -- an explicit gpu_ids=[0] + # means "use exactly device 0" and must stay sequential. + supports_physical = parent_visible_spec["supports_explicit_gpu_ids"] + has_multiple_numeric = ( + parent_visible_spec["numeric_ids"] is not None + and len(parent_visible_spec["numeric_ids"]) > 1 + ) + has_multiple_unresolved = ( + parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1 + ) + if has_multiple_unresolved or (not supports_physical and has_multiple_numeric): + multi_gpu = True if multi_gpu: return "balanced" @@ -2378,6 +2930,19 @@ def raise_if_offloaded( ) +def get_torch_device_str() -> str: + """ + Return the torch device string for the detected hardware. + E.g. "cuda", "xpu", or "cpu". + """ + device = get_device() + if device == DeviceType.CUDA: + return "cuda" + elif device == DeviceType.XPU: + return "xpu" + return "cpu" + + def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. @@ -2445,7 +3010,32 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: Returns ``None`` on spawn platforms (Windows, macOS) because ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``); only ``num_proc=None`` guarantees in-process execution. + + Also returns ``None`` on XPU once its runtime is initialized in this + process: ``os.fork()`` corrupts the Level-Zero context, making Triton + kernels fail with "Pointer argument doesn't reference XPU device memory". + Pre-init XPU hosts can still parallelize CPU-side preprocessing. """ if sys.platform in ("win32", "darwin"): return None + + if get_device() == DeviceType.XPU: + try: + import torch + except Exception: + # No torch means no active XPU runtime, so CPU-side dataset + # parallelism is still safe. + return safe_num_proc(desired) + + xpu = getattr(torch, "xpu", None) + is_initialized = getattr(xpu, "is_initialized", None) + if callable(is_initialized): + try: + if is_initialized(): + return None + except Exception as e: + # Treat a failing probe as "runtime not touched yet" so + # pre-init CPU preprocessing can still parallelize. + logger.debug("torch.xpu.is_initialized() probe failed: %s", e) + return safe_num_proc(desired) diff --git a/studio/backend/utils/hf_cache_settings.py b/studio/backend/utils/hf_cache_settings.py new file mode 100644 index 0000000000..07d901a3d2 --- /dev/null +++ b/studio/backend/utils/hf_cache_settings.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Live, persisted Hugging Face cache routing for Unsloth Studio. + +Hugging Face reads cache environment variables at import time. Studio therefore +owns an explicit cache snapshot for each operation instead of trying to refresh +``huggingface_hub.constants`` in the long-running API process. +""" + +from __future__ import annotations + +import os +import shutil +import tempfile +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator, Literal, Mapping, Optional + + +CACHE_HOME_SETTING_KEY = "hugging_face_cache_home" +CACHE_HISTORY_SETTING_KEY = "hugging_face_cache_history" +MAX_CACHE_HISTORY = 16 + +CacheSource = Literal["default", "studio", "environment"] + +_CACHE_ENV_KEYS = ( + "HF_HOME", + "HF_HUB_CACHE", + "HUGGINGFACE_HUB_CACHE", + "HF_XET_CACHE", +) +# Imported by storage_roots._setup_cache_env before Studio seeds defaults. +_EXPLICIT_CACHE_ENV = { + key: value.strip() + for key in _CACHE_ENV_KEYS + if (value := os.environ.get(key)) is not None and value.strip() +} +_settings_lock = threading.RLock() +_spawn_env_lock = threading.RLock() + + +@dataclass(frozen = True) +class HuggingFaceCachePaths: + cache_home: Path + hub_cache: Path + xet_cache: Path + source: CacheSource + environment_variable: Optional[str] = None + + @property + def editable(self) -> bool: + return self.source != "environment" + + @property + def is_custom(self) -> bool: + return self.source == "studio" + + def child_env(self, base: Optional[Mapping[str, str]] = None) -> dict[str, str]: + env = dict(os.environ if base is None else base) + # Do not rewrite HF_HOME. It also owns HF's token path, and credentials + # must not be moved onto a removable cache volume. + env["HF_HUB_CACHE"] = str(self.hub_cache) + env["HF_XET_CACHE"] = str(self.xet_cache) + env.pop("HUGGINGFACE_HUB_CACHE", None) + return env + + +def _default_cache_home() -> Path: + xdg = (os.environ.get("XDG_CACHE_HOME") or "").strip() + return (Path(xdg).expanduser() if xdg else Path.home() / ".cache") / "huggingface" + + +def _canonical(path: Path | str) -> Path: + return Path(path).expanduser().resolve(strict = False) + + +def _environment_paths() -> Optional[HuggingFaceCachePaths]: + explicit_home = _EXPLICIT_CACHE_ENV.get("HF_HOME") + explicit_hub = _EXPLICIT_CACHE_ENV.get("HF_HUB_CACHE") or _EXPLICIT_CACHE_ENV.get( + "HUGGINGFACE_HUB_CACHE" + ) + if not explicit_home and not explicit_hub: + return None + explicit_xet = _EXPLICIT_CACHE_ENV.get("HF_XET_CACHE") + default_home = _default_cache_home() + hf_home = _canonical(explicit_home) if explicit_home else default_home + hub = _canonical(explicit_hub) if explicit_hub else hf_home / "hub" + xet = _canonical(explicit_xet) if explicit_xet else hf_home / "xet" + controlling = next( + key + for key in ("HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE", "HF_HOME") + if key in _EXPLICIT_CACHE_ENV + ) + # Settings describes model downloads, so an explicit hub path is the + # displayed/opened location even when HF_HOME points somewhere else for + # credentials or XET data. + display_home = ( + (hub.parent if explicit_hub and hub.name.lower() == "hub" else hub) + if explicit_hub + else hf_home + ) + return HuggingFaceCachePaths(display_home, hub, xet, "environment", controlling) + + +def _stored_cache_home() -> Optional[Path]: + try: + from storage.studio_db import get_app_setting + value = get_app_setting(CACHE_HOME_SETTING_KEY, None) + except Exception: + return None + if not isinstance(value, str) or not value.strip(): + return None + try: + return _canonical(value.strip()) + except (OSError, RuntimeError, ValueError): + return None + + +def get_hf_cache_paths() -> HuggingFaceCachePaths: + env_paths = _environment_paths() + if env_paths is not None: + return env_paths + stored = _stored_cache_home() + if stored is not None: + xet = _EXPLICIT_CACHE_ENV.get("HF_XET_CACHE") + return HuggingFaceCachePaths( + stored, + stored / "hub", + _canonical(xet) if xet else stored / "xet", + "studio", + ) + home = _default_cache_home() + xet = _EXPLICIT_CACHE_ENV.get("HF_XET_CACHE") + return HuggingFaceCachePaths( + home, + home / "hub", + _canonical(xet) if xet else home / "xet", + "default", + ) + + +def active_hf_hub_cache() -> str: + """Return the current hub cache as a string for library call kwargs.""" + + return str(get_hf_cache_paths().hub_cache) + + +@contextmanager +def child_environment_for_spawn(environment: Mapping[str, str]) -> Iterator[None]: + """Apply captured env before spawn imports the child entrypoint. + + Applying variables only inside the multiprocessing target can be too late + for libraries that snapshot environment variables at import. The lock keeps + this short parent-process override atomic through ``Process.start()``. + """ + + with _spawn_env_lock: + missing = object() + saved_environment: dict[str, str | object] = {} + for key, value in environment.items(): + saved_environment[key] = os.environ.get(key, missing) + os.environ[key] = value + try: + yield + finally: + for key, previous in saved_environment.items(): + if previous is missing: + os.environ.pop(key, None) + else: + os.environ[key] = str(previous) + + +def initialize_hf_cache_environment() -> HuggingFaceCachePaths: + """Seed import-time HF variables once during backend startup.""" + + paths = get_hf_cache_paths() + # Preserve an explicit HF_HOME, otherwise keep credentials at the platform + # default while routing cache bytes through the selected home. + if not os.environ.get("HF_HOME", "").strip(): + os.environ["HF_HOME"] = str(_default_cache_home()) + os.environ["HF_HUB_CACHE"] = str(paths.hub_cache) + os.environ["HF_XET_CACHE"] = str(paths.xet_cache) + if "HUGGINGFACE_HUB_CACHE" not in _EXPLICIT_CACHE_ENV: + os.environ.pop("HUGGINGFACE_HUB_CACHE", None) + for directory in (paths.hub_cache, paths.xet_cache): + try: + directory.mkdir(parents = True, exist_ok = True) + except OSError: + pass + return paths + + +def _validate_cache_home(raw_path: str) -> Path: + value = raw_path.strip() + if not value: + raise ValueError("Choose a cache folder.") + candidate = Path(value).expanduser() + if not candidate.is_absolute(): + raise ValueError("The Hugging Face cache folder must be an absolute path.") + try: + resolved = candidate.resolve(strict = False) + except (OSError, RuntimeError, ValueError) as exc: + raise ValueError("The Hugging Face cache folder is invalid.") from exc + + if resolved.parent == resolved: + raise ValueError("Choose a folder inside the filesystem or drive root.") + try: + from hub.storage.scan_folders import ( + contains_sensitive_path_component, + is_denied_system_path, + ) + except ImportError: + contains_sensitive_path_component = is_denied_system_path = None + if is_denied_system_path is not None and is_denied_system_path(str(resolved)): + raise ValueError("System folders cannot be used for model downloads.") + if contains_sensitive_path_component is not None and contains_sensitive_path_component( + str(resolved) + ): + raise ValueError("Credential or config folders cannot be used for model downloads.") + + parent = resolved.parent + if not parent.exists() or not parent.is_dir(): + raise ValueError("The parent folder does not exist.") + try: + resolved.mkdir(exist_ok = True) + if not resolved.is_dir(): + raise ValueError("The selected cache location is not a folder.") + for child in (resolved / "hub", resolved / "xet"): + child.mkdir(exist_ok = True) + with tempfile.NamedTemporaryFile(prefix = ".unsloth-write-test-", dir = child): + pass + except PermissionError as exc: + raise ValueError("Studio does not have permission to write to this folder.") from exc + except OSError as exc: + raise ValueError(f"Studio cannot use this cache folder: {exc}") from exc + return resolved + + +def _stored_history() -> list[Path]: + try: + from storage.studio_db import get_app_setting + raw = get_app_setting(CACHE_HISTORY_SETTING_KEY, []) + except Exception: + raw = [] + if not isinstance(raw, list): + return [] + out: list[Path] = [] + seen: set[str] = set() + for value in raw: + if not isinstance(value, str) or not value.strip(): + continue + try: + path = _canonical(value) + except (OSError, RuntimeError, ValueError): + continue + key = os.path.normcase(str(path)) + if key in seen: + continue + seen.add(key) + out.append(path) + return out[:MAX_CACHE_HISTORY] + + +def set_hf_cache_home(cache_home: Optional[str]) -> HuggingFaceCachePaths: + if _environment_paths() is not None: + raise RuntimeError("The Hugging Face cache location is managed by an environment variable.") + with _settings_lock: + previous = _stored_cache_home() + next_home = _validate_cache_home(cache_home) if cache_home is not None else None + history = _stored_history() + if previous is not None and previous != next_home: + history.insert(0, previous) + deduped: list[str] = [] + seen: set[str] = set() + for path in history: + key = os.path.normcase(str(path)) + if key in seen or path == next_home: + continue + seen.add(key) + deduped.append(str(path)) + if len(deduped) >= MAX_CACHE_HISTORY: + break + from storage.studio_db import upsert_app_settings + + upsert_app_settings( + { + CACHE_HOME_SETTING_KEY: str(next_home) if next_home is not None else None, + CACHE_HISTORY_SETTING_KEY: deduped, + } + ) + # Inventory scans are cached independently from settings. Invalidate after + # persistence so the next request sees both the new active root and history. + from hub.utils.inventory_scan import invalidate_hf_cache_scans + + invalidate_hf_cache_scans() + return get_hf_cache_paths() + + +def known_hf_cache_homes() -> list[Path]: + paths = get_hf_cache_paths() + stored = _stored_cache_home() + candidates: list[Path] = [] + if paths.source != "environment": + candidates.append(paths.cache_home) + elif explicit_home := _EXPLICIT_CACHE_ENV.get("HF_HOME"): + candidates.append(_canonical(explicit_home)) + if stored is not None: + candidates.append(stored) + candidates.extend([*_stored_history(), _default_cache_home()]) + out: list[Path] = [] + seen: set[str] = set() + for candidate in candidates: + try: + canonical = _canonical(candidate) + except (OSError, RuntimeError, ValueError): + continue + key = os.path.normcase(str(canonical)) + if key in seen: + continue + seen.add(key) + out.append(canonical) + return out + + +def known_hf_hub_caches() -> list[Path]: + active = get_hf_cache_paths() + out = [active.hub_cache] + seen = {os.path.normcase(str(_canonical(active.hub_cache)))} + for home in known_hf_cache_homes(): + hub = _canonical(home / "hub") + key = os.path.normcase(str(hub)) + if key not in seen: + seen.add(key) + out.append(hub) + return out + + +def cache_status(paths: Optional[HuggingFaceCachePaths] = None) -> dict: + paths = paths or get_hf_cache_paths() + available = paths.cache_home.is_dir() + writable = available and os.access(paths.cache_home, os.W_OK | os.X_OK) + free_bytes: Optional[int] = None + if available: + try: + free_bytes = int(shutil.disk_usage(paths.cache_home).free) + except OSError: + pass + return { + "cache_home": str(paths.cache_home), + "hub_cache": str(paths.hub_cache), + "xet_cache": str(paths.xet_cache), + "source": paths.source, + "editable": paths.editable, + "is_custom": paths.is_custom, + "available": available, + "writable": writable, + "free_bytes": free_bytes, + "environment_variable": paths.environment_variable, + } diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 2628b99a2d..49872f371e 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -21,6 +21,8 @@ never triggers the heavy load. from __future__ import annotations import threading +from functools import partial +from pathlib import Path from typing import Any, Callable, Optional # Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as @@ -262,13 +264,23 @@ __all__ = [ ] -def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: +def _studio_prepare_for_http( + repo_type: str, + repo_id: str, + *, + cache_dir: Optional[str] = None, +) -> None: """Unsloth's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, not fatal to the retry.""" try: from hub.utils.download_registry import prepare_cache_for_transport - prepare_cache_for_transport(repo_type, repo_id, "http") + prepare_cache_for_transport( + repo_type, + repo_id, + "http", + root = Path(cache_dir) if cache_dir else None, + ) except Exception as exc: try: from loggers import get_logger @@ -293,9 +305,13 @@ def hf_hub_download_with_xet_fallback( grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, force_download: bool = False, + cache_dir: Optional[str] = None, ) -> str: """Single-file download via the shared fallback with Unsloth's marker-aware HTTP-retry prep. ``force_download`` re-fetches a newer blob over a cached one (Unsloth's model-update path).""" + if cache_dir is None: + from utils.hf_cache_settings import get_hf_cache_paths + cache_dir = str(get_hf_cache_paths().hub_cache) return _shared_hf_hub_download_with_xet_fallback( repo_id, filename, @@ -308,11 +324,18 @@ def hf_hub_download_with_xet_fallback( grace_period = grace_period, on_status = on_status, force_download = force_download, - prepare_for_http_fn = _studio_prepare_for_http, + cache_dir = cache_dir, + prepare_for_http_fn = partial(_studio_prepare_for_http, cache_dir = cache_dir), ) def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: """Whole-repo download via the shared fallback with Unsloth's marker-aware HTTP-retry prep.""" - kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) + if kwargs.get("cache_dir") is None: + from utils.hf_cache_settings import get_hf_cache_paths + kwargs["cache_dir"] = str(get_hf_cache_paths().hub_cache) + kwargs.setdefault( + "prepare_for_http_fn", + partial(_studio_prepare_for_http, cache_dir = kwargs["cache_dir"]), + ) return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/backend/utils/hidden_models.py b/studio/backend/utils/hidden_models.py index 20d0bb966e..e7c3181d71 100644 --- a/studio/backend/utils/hidden_models.py +++ b/studio/backend/utils/hidden_models.py @@ -9,6 +9,7 @@ which eagerly loads the model-config/checkpoint stack, and without importing from __future__ import annotations +import json import re from pathlib import Path from typing import Optional @@ -31,6 +32,61 @@ _DEFAULT_EMBEDDING_REPO_IDS = { # fallback for Studio's static default embedder only; configured custom repos # remain exact-match-only. _DEFAULT_EMBEDDING_PATH_BASENAMES = {"bge-small-en-v1.5"} +# Curated Whisper dictation checkpoints (STT, never chat), hidden from the chat +# inventory and pickers: Transformers safetensors repos (unsloth/whisper-*) and +# their GGUF companions (unslothai/whisper-*-GGUF). Custom checkpoints are caught +# by config below, but the GGUF companions carry a raw .bin (no config.json), so +# they must be listed here by id or they leak into chat pickers. +_HIDDEN_STT_REPO_IDS = frozenset( + { + "unsloth/whisper-tiny", + "unsloth/whisper-base", + "unsloth/whisper-small", + "unsloth/whisper-large-v3-turbo", + "unsloth/whisper-large-v3", + "unslothai/whisper-tiny-GGUF", + "unslothai/whisper-base-GGUF", + "unslothai/whisper-small-GGUF", + "unslothai/whisper-large-v3-turbo-GGUF", + "unslothai/whisper-large-v3-GGUF", + } +) + + +def _config_is_whisper(path: Path) -> bool: + """True if a config.json declares a Whisper model.""" + try: + with open(path, "r", encoding = "utf-8") as file: + config = json.load(file) + except Exception: + return False + if not isinstance(config, dict): + return False + model_type = config.get("model_type") + if isinstance(model_type, str) and model_type.strip().lower() == "whisper": + return True + architectures = config.get("architectures") + return isinstance(architectures, list) and any( + isinstance(name, str) and name == "WhisperForConditionalGeneration" + for name in architectures + ) + + +def _path_is_whisper_model(value: str) -> bool: + """Inspect an existing local model path's config; never hides name-only matches.""" + if _HF_REPO_ID_RE.fullmatch(value.strip()): + return False + path = Path(value).expanduser() + try: + if path.is_file(): + path = path.parent + candidates = [path / "config.json"] + snapshots = path / "snapshots" + if snapshots.is_dir(): + candidates.extend(child / "config.json" for child in snapshots.iterdir()) + except OSError: + return False + return any(_config_is_whisper(candidate) for candidate in candidates) def _safe_resolve(path: Path) -> Optional[str]: @@ -79,11 +135,11 @@ def _path_basename_is_default_embedder(value: str) -> bool: def is_hidden_model(*values: str | None) -> bool: """True if any id/path is the RAG embedding model (the effective embedder - or its GGUF companion repo) 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. + or its GGUF companion repo), the llama.cpp install validation probe + (ggml-org/models / stories260K), or a curated/custom Whisper dictation + model, 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. Hub repo ids are matched EXACTLY (case-insensitive full "owner/name"), so a custom embedder with a generic basename like "org/model" cannot substring @@ -97,6 +153,7 @@ def is_hidden_model(*values: str | None) -> bool: hidden_repo_ids = { _PROBE_REPO_ID.lower(), *(repo_id.lower() for repo_id in _DEFAULT_EMBEDDING_REPO_IDS), + *(repo_id.lower() for repo_id in _HIDDEN_STT_REPO_IDS), } exact_paths: list[str] = [] for model in { @@ -135,6 +192,9 @@ def is_hidden_model(*values: str | None) -> bool: return True if _path_contains_repo_id(v, hidden_repo_ids): return True + # Custom Whisper checkpoints keep no curated repo id, so match by config. + if _path_is_whisper_model(v): + return True if exact_paths: resolved = _safe_resolve(Path(v).expanduser()) if resolved and resolved.lower() in exact_paths: diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 05eb08067c..a264e06c85 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -5,7 +5,10 @@ from pathlib import Path from typing import Dict, Any, Optional +from functools import lru_cache import json +import math +import os import yaml import structlog from loggers import get_logger @@ -160,3 +163,137 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]: } return inference_config + + +# โ”€โ”€ Effective sampling resolution for `unsloth run` / `unsloth start` โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# +# Per-model recommended sampling is applied to a request only for the fields the +# client omitted; an operator can pin a field from the CLI via UNSLOTH_SAMPLING_* +# (a hard override that wins even over an explicit client value). Precedence per +# field: operator pin -> client explicit -> per-model recommendation -> the static +# schema default (mirroring ChatCompletionRequest, so behavior is unchanged when +# nothing is recommended or pinned). + +# field -> (env var, static default, min, max, is_int) +_SAMPLING_FIELDS = { + "temperature": ("UNSLOTH_SAMPLING_TEMPERATURE", 0.6, 0.0, 2.0, False), + "top_p": ("UNSLOTH_SAMPLING_TOP_P", 0.95, 0.0, 1.0, False), + "top_k": ("UNSLOTH_SAMPLING_TOP_K", 20, -1, 100, True), + "min_p": ("UNSLOTH_SAMPLING_MIN_P", 0.01, 0.0, 1.0, False), + "repetition_penalty": ("UNSLOTH_SAMPLING_REPETITION_PENALTY", 1.0, 1.0, 2.0, False), + "presence_penalty": ("UNSLOTH_SAMPLING_PRESENCE_PENALTY", 0.0, 0.0, 2.0, False), +} + +# Public, ordered tuple of the sampling fields callers resolve. +SAMPLING_FIELD_NAMES = tuple(_SAMPLING_FIELDS) + +# Fields the Studio Chat UI adopts as *per-model recommendations* from the backend +# `.inference` block. Its frontend `mergeBackendRecommendedInference` +# (presets/preset-policy.ts) seeds exactly these five and never reads repetition_penalty, +# so the server auto-recommends the same five for request parity. repetition_penalty stays a +# manual-only knob (client-sent or an UNSLOTH_SAMPLING_REPETITION_PENALTY operator pin), +# matching the UI where it is never auto-filled per model. +_UI_RECOMMENDED_FIELDS = ("temperature", "top_p", "top_k", "min_p", "presence_penalty") + + +def _clean_sampling_value(field: str, val: Any): + """Coerce ``val`` to the field's numeric type when it is a finite, in-range number, else None. + + Rejects bool, non-numeric, NaN/inf, and out-of-range values so neither a bad operator env + var nor a malformed model recommendation can reach llama-server. NaN matters because + ``nan < lo`` and ``nan > hi`` are both False, so a plain range check would let it through. + Coerce before the finiteness check: ``math.isfinite`` and ``float()`` raise ``OverflowError`` + on an int too big for a C double (an oversized UNSLOTH_SAMPLING_TOP_K would otherwise 500 the + request), while an in-range int is range-checked exactly and ``int()`` rejects a NaN/inf that + reached an int field. + """ + if isinstance(val, bool) or not isinstance(val, (int, float)): + return None + _env, _default, lo, hi, is_int = _SAMPLING_FIELDS[field] + try: + val = int(val) if is_int else float(val) + except (ValueError, OverflowError): + # int(nan)/int(inf) and float(oversized_int) raise; treat them as unusable. + return None + # After coercion an int is always finite; only a float can still be NaN/inf. + if isinstance(val, float) and not math.isfinite(val): + return None + if val < lo or val > hi: + return None + return val + + +def _operator_sampling_override(field: str): + """Operator-pinned value for a sampling field from UNSLOTH_SAMPLING_*, or None. + + An unparseable, non-finite, or out-of-range value is ignored so a bad env var can never + reach llama-server; the field then falls back to the client / recommended value. + """ + _env, _default, _lo, _hi, is_int = _SAMPLING_FIELDS[field] + raw = os.environ.get(_env) + if raw is None or raw.strip() == "": + return None + try: + val = int(raw) if is_int else float(raw) + except (TypeError, ValueError): + return None + return _clean_sampling_value(field, val) + + +@lru_cache(maxsize = 128) +def _recommended_sampling(model_id: str) -> Dict[str, Any]: + """Per-model recommended sampling, resolved through the SAME path the Studio Chat UI uses. + + The Chat UI seeds its sampling from the ``.inference`` block of the load/status responses, + which is exactly :func:`load_inference_config` (model-specific YAML -> family defaults + (inference_defaults.json) -> default.yaml). Sourcing recommendations here keeps the values + the server applies to a request identical to what the UI shows for the same model. Only the + fields the UI actually adopts (:data:`_UI_RECOMMENDED_FIELDS`) are recommended; each value + is validated (finite + in range) before use. Cached by model id. + """ + if not model_id: + return {} + try: + cfg = load_inference_config(model_id) or {} + except Exception as e: + logger.debug(f"Could not load recommended sampling for '{model_id}': {e}") + return {} + recommended: Dict[str, Any] = {} + for field in _UI_RECOMMENDED_FIELDS: + cleaned = _clean_sampling_value(field, cfg.get(field)) + if cleaned is not None: + recommended[field] = cleaned + return recommended + + +def resolve_effective_sampling( + model_id: Optional[str], + explicit: Dict[str, Any], + *, + fill_defaults: bool = True, +) -> Dict[str, Any]: + """Resolve the effective sampling params for a request. + + ``explicit`` maps each field in :data:`SAMPLING_FIELD_NAMES` to the client-sent + value, or ``None`` when the client omitted it. Precedence (highest first): an + operator ``UNSLOTH_SAMPLING_*`` pin, then the client's explicit value, then the + per-model recommendation, then the static schema default. + + When ``fill_defaults`` is False a field with no operator pin, client value, or + per-model recommendation is omitted from the result instead of set to the static + schema default, so a raw proxy body (``/v1/completions``) keeps llama-server's own + default for that field rather than being forced onto this schema's value. + """ + recommended = _recommended_sampling(model_id or "") + effective: Dict[str, Any] = {} + for field, (_env, default, _lo, _hi, _int) in _SAMPLING_FIELDS.items(): + override = _operator_sampling_override(field) + if override is not None: + effective[field] = override + elif explicit.get(field) is not None: + effective[field] = explicit[field] + elif field in recommended: + effective[field] = recommended[field] + elif fill_defaults: + effective[field] = default + return effective diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index 7d077bfa3b..a184fdb3e9 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -7,28 +7,28 @@ Reads UNSLOTH_PREBUILT_INFO.json (written by install_llama_prebuilt.py) and compares the installed release tag against the latest on GitHub. Surfaced via main.py:lifespan() and /api/inference/status. Fails open on any missing data so we never show a misleading banner. + +The mechanics (marker walk-up, GitHub fetch, memo + disk cache, report +skeleton) live in utils.prebuilt.freshness_flow; this module keeps the +llama version policy and the per-module caches its tests patch. """ from __future__ import annotations -import json -import os import re -import time -from datetime import datetime, timezone +from datetime import datetime from pathlib import Path from typing import Optional import structlog +from utils.prebuilt import freshness_flow as _flow + logger = structlog.get_logger(__name__) # 3 days matches Unsloth's typical llama.cpp release cadence. STALENESS_THRESHOLD_DAYS = 3 -# 24h TTL keeps the GitHub call off the hot path and within rate limits. -_RELEASE_CACHE_TTL_SECONDS = 24 * 60 * 60 - _INSTALL_MARKER_NAME = "UNSLOTH_PREBUILT_INFO.json" _marker_cache: dict[str, Optional[dict]] = {} @@ -49,203 +49,60 @@ def _cache_dir() -> Path: def read_install_marker(binary_path: Optional[str]) -> Optional[dict]: """Walk up from binary_path to find UNSLOTH_PREBUILT_INFO.json. None = no marker (source build / custom path) or invalid JSON.""" - if not binary_path: - return None - cached = _marker_cache.get(binary_path) - if cached is not None or binary_path in _marker_cache: - return cached - p = Path(binary_path) - marker: Optional[dict] = None - # Cover all _find_llama_server_binary layouts (binary is 1-4 dirs deep): - for parent in p.parents[:5]: - candidate = parent / _INSTALL_MARKER_NAME - if candidate.is_file(): - try: - marker = json.loads(candidate.read_text(encoding = "utf-8")) - except (OSError, json.JSONDecodeError) as exc: - logger.debug( - "failed to parse install marker", - path = str(candidate), - error = str(exc), - ) - marker = None - break - _marker_cache[binary_path] = marker - return marker - - -def _cache_path_for(repo: str) -> Path: - safe = repo.replace("/", "__") - return _cache_dir() / f"{safe}.json" + return _flow.read_install_marker( + binary_path, + marker_name = _INSTALL_MARKER_NAME, + cache = _marker_cache, + log_message = "failed to parse install marker", + ) def _load_disk_cache(repo: str) -> Optional[tuple[float, Optional[str]]]: - path = _cache_path_for(repo) - try: - payload = json.loads(path.read_text(encoding = "utf-8")) - except (OSError, json.JSONDecodeError): - return None - ts = payload.get("fetched_at") - tag = payload.get("latest_tag") - if not isinstance(ts, (int, float)): - return None - return float(ts), tag if isinstance(tag, str) else None + return _flow.load_disk_cache(repo, _cache_dir()) def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None: - path = _cache_path_for(repo) - try: - path.parent.mkdir(parents = True, exist_ok = True) - tmp = path.with_suffix(".tmp") - tmp.write_text( - json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag}), - encoding = "utf-8", - ) - tmp.replace(path) - except OSError as exc: - logger.debug("freshness cache write failed", repo = repo, error = str(exc)) + _flow.save_disk_cache( + repo, latest_tag, _cache_dir(), log_message = "freshness cache write failed" + ) def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]: - """Newest published release tag for `repo`, by publish time. - - Resolves "latest" the way install_llama_prebuilt.py does (newest - non-draft/non-prerelease by ``published_at``), NOT via GitHub's - ``/releases/latest`` pointer. That pointer sorts by commit date and can lag - behind the build the installer actually installs, so detection and apply - disagreed -- the cause of the downgrade/sticky banner. None on any failure - (offline, rate-limited, etc).""" - import urllib.error - import urllib.request - - url = f"https://api.github.com/repos/{repo}/releases?per_page=30" - headers = { - "Accept": "application/vnd.github+json", - "User-Agent": "unsloth-studio-freshness-check", - } - token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") - if token: - headers["Authorization"] = f"Bearer {token}" - req = urllib.request.Request(url, headers = headers) - try: - with urllib.request.urlopen(req, timeout = timeout) as resp: - data = json.loads(resp.read().decode("utf-8")) - except ( - urllib.error.URLError, - urllib.error.HTTPError, - OSError, - json.JSONDecodeError, - ) as exc: - logger.debug("freshness fetch failed", repo = repo, error = str(exc)) - return None - if not isinstance(data, list): - return None - published = [ - r - for r in data - if isinstance(r, dict) - and not r.get("draft") - and not r.get("prerelease") - and isinstance(r.get("tag_name"), str) - and r.get("tag_name") - ] - if not published: - return None - newest = max(published, key = lambda r: r.get("published_at") or "") - return newest["tag_name"] + """Newest published release tag for `repo`, by publish time (see + freshness_flow for why this is not GitHub's /releases/latest pointer).""" + return _flow.fetch_latest_release_tag(repo, timeout, log_message = "freshness fetch failed") def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]: """Latest release tag for `repo`. Memo + disk-cached (24h TTL). None when offline and never previously cached.""" - if not repo: - return None - now = time.time() - if not force_refresh: - memo = _release_memo.get(repo) - if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS: - return memo[1] - disk = _load_disk_cache(repo) - if disk and now - disk[0] < _RELEASE_CACHE_TTL_SECONDS: - _release_memo[repo] = disk - return disk[1] - latest = _fetch_latest_release_tag(repo) - if latest is None: - # Keep last-good disk value rather than poisoning with None. - disk = _load_disk_cache(repo) - if disk: - _release_memo[repo] = disk - return disk[1] - return None - _release_memo[repo] = (now, latest) - _save_disk_cache(repo, latest) - return latest + return _flow.latest_published_release( + repo, + force_refresh = force_refresh, + memo = _release_memo, + cache_dir = lambda: _cache_dir(), + fetch = lambda r: _fetch_latest_release_tag(r), + save = lambda r, tag: _save_disk_cache(r, tag), + ) def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]: """Asset name -> size (bytes) for the newest published release of `repo`, selected exactly like _fetch_latest_release_tag. None on any failure.""" - import urllib.error - import urllib.request - - url = f"https://api.github.com/repos/{repo}/releases?per_page=30" - headers = { - "Accept": "application/vnd.github+json", - "User-Agent": "unsloth-studio-freshness-check", - } - token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") - if token: - headers["Authorization"] = f"Bearer {token}" - req = urllib.request.Request(url, headers = headers) - try: - with urllib.request.urlopen(req, timeout = timeout) as resp: - data = json.loads(resp.read().decode("utf-8")) - except ( - urllib.error.URLError, - urllib.error.HTTPError, - OSError, - json.JSONDecodeError, - ) as exc: - logger.debug("freshness asset fetch failed", repo = repo, error = str(exc)) - return None - if not isinstance(data, list): - return None - published = [ - r - for r in data - if isinstance(r, dict) - and not r.get("draft") - and not r.get("prerelease") - and isinstance(r.get("tag_name"), str) - and r.get("tag_name") - ] - if not published: - return None - newest = max(published, key = lambda r: r.get("published_at") or "") - assets: dict[str, int] = {} - for a in newest.get("assets") or []: - name, size = a.get("name"), a.get("size") - if isinstance(name, str) and isinstance(size, int): - assets[name] = size - return assets + return _flow.fetch_latest_release_assets( + repo, timeout, log_message = "freshness asset fetch failed" + ) def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]: """Newest-release asset sizes for `repo`, memoized (24h TTL). None when offline and never fetched. In-memory only -- a restart simply re-fetches.""" - if not repo: - return None - now = time.time() - if not force_refresh: - memo = _assets_memo.get(repo) - if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS: - return memo[1] - assets = _fetch_latest_release_assets(repo) - if assets is None: - memo = _assets_memo.get(repo) - return memo[1] if memo else None - _assets_memo[repo] = (now, assets) - return assets + return _flow.latest_release_assets( + repo, + force_refresh = force_refresh, + memo = _assets_memo, + fetch = lambda r: _fetch_latest_release_assets(r), + ) def update_download_size_bytes( @@ -290,16 +147,7 @@ def update_download_size_bytes( def _parse_installed_at(value: object) -> Optional[datetime]: - if not isinstance(value, str) or not value: - return None - s = value.replace("Z", "+00:00") if value.endswith("Z") else value - try: - dt = datetime.fromisoformat(s) - except ValueError: - return None - if dt.tzinfo is None: - dt = dt.replace(tzinfo = timezone.utc) - return dt + return _flow.parse_installed_at(value) def parse_base_build(tag: object) -> Optional[int]: @@ -350,64 +198,27 @@ def check_prebuilt_freshness( behind = installed genuinely older than latest (see is_behind). stale = behind AND age >= threshold. Fails open on missing data (behind/stale stay False).""" - out: dict = { - "has_marker": False, - "stale": False, - "behind": False, - "installed_tag": None, - "latest_tag": None, - "installed_at_utc": None, - "age_days": None, - "published_repo": None, - "threshold_days": int(threshold_days), - } - marker = read_install_marker(binary_path) - if not marker: - return out - out["has_marker"] = True - # Display prefers the normalized base ("tag"); comparison below prefers the - # full "release_tag" -- deliberately opposite fallbacks. - out["installed_tag"] = marker.get("tag") or marker.get("release_tag") - out["installed_at_utc"] = marker.get("installed_at_utc") - out["published_repo"] = marker.get("published_repo") - # The marker records both a normalized base tag ("tag", e.g. b9596) and the - # full release tag ("release_tag", e.g. b9596-mix-). Compare against the - # FULL identity, since GitHub /releases/latest returns the full tag_name -- - # comparing the normalized base against the full latest is what produced the - # permanent "downgrade" banner on every mix release. - installed_full = marker.get("release_tag") or marker.get("tag") - repo = out["published_repo"] - if not repo or not installed_full: - return out - latest = latest_published_release(repo) - out["latest_tag"] = latest - out["behind"] = is_behind(installed_full, latest) - if not out["behind"]: - return out - - installed_at = _parse_installed_at(out["installed_at_utc"]) - if installed_at is None: - return out - now = now or datetime.now(tz = timezone.utc) - age_seconds = (now - installed_at).total_seconds() - out["age_days"] = max(0, int(age_seconds // 86400)) - if age_seconds >= threshold_days * 86400: - out["stale"] = True - return out + # full release tag ("release_tag", e.g. b9596-mix-). Display prefers the + # normalized base; comparison uses the FULL identity, since GitHub + # /releases/latest returns the full tag_name -- comparing the normalized base + # against the full latest is what produced the permanent "downgrade" banner + # on every mix release. Deliberately opposite fallbacks. + return _flow.check_freshness( + binary_path, + threshold_days = threshold_days, + now = now, + read_marker = lambda p: read_install_marker(p), + latest_release = lambda repo: latest_published_release(repo), + behind = lambda installed, latest: is_behind(installed, latest), + display_tag = lambda marker: marker.get("tag") or marker.get("release_tag"), + compare_tag = lambda marker: marker.get("release_tag") or marker.get("tag"), + ) def format_stale_warning(info: dict) -> str: """Human-readable one-liner for stale prebuilt info.""" - age = info.get("age_days") - installed = info.get("installed_tag") or "unknown" - latest = info.get("latest_tag") or "unknown" - age_str = f"{age} day{'s' if age != 1 else ''}" if age is not None else "some time" - return ( - f"llama.cpp prebuilt is {age_str} behind: installed " - f"{installed}, latest {latest}. Run `unsloth studio update` " - f"to refresh." - ) + return _flow.format_stale_warning(info, component = "llama.cpp") def reset_caches(*, drop_disk: bool = False) -> None: @@ -420,13 +231,8 @@ def reset_caches(*, drop_disk: bool = False) -> None: (see its last-good fallback) and the banner could linger. Dropping the disk cache makes latest read as None in that offline case, so the banner fails open (off) instead of pointing at the just-replaced build.""" - _marker_cache.clear() - _release_memo.clear() - _assets_memo.clear() - if drop_disk: - import shutil - - # _cache_dir() is a dedicated freshness-only subdir; it is re-created on - # the next _save_disk_cache. ignore_errors so a missing/locked dir is a - # no-op rather than breaking an otherwise successful install. - shutil.rmtree(_cache_dir(), ignore_errors = True) + _flow.reset_caches( + (_marker_cache, _release_memo, _assets_memo), + drop_disk = drop_disk, + cache_dir = lambda: _cache_dir(), + ) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 67733bde35..174e6ef4dc 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -17,17 +17,22 @@ Design notes: thread; callers poll get_update_status() for the job state. - Everything fails open: a missing marker / offline GitHub / source build just reports update_available=False and never blocks the app. +- The mechanics (managed-root resolution, local-link detection, the resolve + probe, the streamed installer run) live in utils.prebuilt.update_flow; this + module keeps the llama policy and the job dict its callers poll. +- This is the single main update item: whisper.cpp piggybacks on it. Status + folds in a whisper sub-status (update_available becomes the union) and apply + chains a whisper phase after the llama phase when whisper is behind (see + update_flow.run_chained_update and whisper_cpp_update.chained_phase_plan). """ from __future__ import annotations -import json import os import re import subprocess import sys import threading -import time from pathlib import Path from typing import Optional @@ -43,7 +48,7 @@ from utils.llama_cpp_freshness import ( reset_caches, update_download_size_bytes, ) -from utils.process_lifetime import child_popen_kwargs +from utils.prebuilt import update_flow as _flow logger = structlog.get_logger(__name__) @@ -51,33 +56,18 @@ DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" _INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate # Background job state. Single in-flight update at a time, guarded by _job_lock. -_JOB_IDLE = "idle" -_JOB_RUNNING = "running" -_JOB_SUCCESS = "success" -_JOB_ERROR = "error" +_JOB_IDLE = _flow.JOB_IDLE +_JOB_RUNNING = _flow.JOB_RUNNING +_JOB_SUCCESS = _flow.JOB_SUCCESS +_JOB_ERROR = _flow.JOB_ERROR _job_lock = threading.Lock() -_job: dict = { - "state": _JOB_IDLE, - "message": "", - "from_tag": None, - "to_tag": None, - "reload_required": None, - "error": None, - "progress": None, - "started_at": None, - "finished_at": None, -} +_job: dict = _flow.new_job() -# Matches the installer's download progress lines, e.g. -# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s". -_PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(") -# The download dominates the update; extract/validate fill the last slice. -_DOWNLOAD_PROGRESS_CEILING = 0.95 - - -def _utcnow() -> str: - return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) +_utcnow = _flow.utcnow +_is_under = _flow.is_under +_is_external_link = _flow.is_external_link +_rocm_install_args = _flow.rocm_install_args def _find_binary() -> Optional[str]: @@ -94,37 +84,19 @@ def _find_binary() -> Optional[str]: def _install_dir_for(binary_path: Optional[str]) -> Optional[Path]: """The directory holding UNSLOTH_PREBUILT_INFO.json -- i.e. the install root - install_llama_prebuilt.py wrote and the one we re-install into. Walks up from - the binary the same way read_install_marker() does.""" - if not binary_path: - return None - p = Path(binary_path) - for parent in p.parents[:5]: - if (parent / _INSTALL_MARKER_NAME).is_file(): - return parent - return None + install_llama_prebuilt.py wrote and the one we re-install into.""" + return _flow.install_dir_for(binary_path, marker_name = _INSTALL_MARKER_NAME) def _installer_script() -> Optional[Path]: - """Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then - searches up from this file for both ``/install_llama_prebuilt.py`` and - ``/studio/install_llama_prebuilt.py`` so it works in the dev tree and - in an installed Unsloth layout.""" - env = os.environ.get("UNSLOTH_LLAMA_INSTALLER") - if env and Path(env).is_file(): - return Path(env) - here = Path(__file__).resolve() - for up in here.parents: - for cand in (up / "install_llama_prebuilt.py", up / "studio" / "install_llama_prebuilt.py"): - if cand.is_file(): - return cand - return None + """Locate install_llama_prebuilt.py (UNSLOTH_LLAMA_INSTALLER wins).""" + return _flow.find_installer_script( + env_var = "UNSLOTH_LLAMA_INSTALLER", script_name = "install_llama_prebuilt.py" + ) # Markerless (source-build) installs have no UNSLOTH_PREBUILT_INFO.json, so we -# ask the installer whether an official prebuilt now exists for this host. Memo -# is 24h; only successful answers are cached so a network blip retries. -_RESOLVE_TTL_SECONDS = 24 * 60 * 60 +# ask the installer whether an official prebuilt now exists for this host. _resolve_memo: dict = {} @@ -132,39 +104,12 @@ def _resolve_prebuilt_for_host(*, force_refresh: bool = False) -> Optional[dict] """Run install_llama_prebuilt.py --resolve-prebuilt (no download) and return {prebuilt_available, repo, release_tag, llama_tag, asset, install_kind} or None. Fail-open: any error -> None so a source build never blocks the app.""" - now = time.time() - if not force_refresh and _resolve_memo: - if now - _resolve_memo.get("at", 0.0) < _RESOLVE_TTL_SECONDS: - return _resolve_memo.get("value") - script = _installer_script() - if script is None: - return None - value: Optional[dict] = None - try: - proc = subprocess.run( - [ - sys.executable, - str(script), - "--resolve-prebuilt", - "latest", - "--output-format", - "json", - ], - capture_output = True, - text = True, - timeout = 60, - ) - out = (proc.stdout or "").strip() - if proc.returncode == 0 and out: - parsed = json.loads(out.splitlines()[-1]) - if isinstance(parsed, dict): - value = parsed - except Exception as exc: # pragma: no cover - subprocess/json defensive - logger.debug("llama update: resolve-prebuilt failed", error = str(exc)) - value = None - if value is not None: # cache real answers; let failures retry next poll - _resolve_memo.update(at = now, value = value) - return value + return _flow.resolve_prebuilt_for_host( + force_refresh = force_refresh, + memo = _resolve_memo, + installer_script = lambda: _installer_script(), + log_message = "llama update: resolve-prebuilt failed", + ) def _installed_build_number(binary: Optional[str]) -> Optional[int]: @@ -218,38 +163,16 @@ def get_installed_llama_version() -> Optional[str]: return f"b{n}" if n is not None else None -def _is_under(path: Path, root: Path) -> bool: - try: - p, r = path.resolve(), root.resolve() - except (OSError, ValueError): - p, r = path, root - return p == r or r in p.parents - - def _llama_install_root(binary: Optional[str]) -> Optional[Path]: """The Unsloth-managed llama.cpp root the active binary lives under, or None - when the binary is unmanaged. Installing anywhere the active binary is not - would not replace what _find_llama_server_binary runs (which prefers a pinned - LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we - refuse rather than silently install into an inactive or foreign tree.""" - marked = _install_dir_for(binary) - if marked is not None: - return marked - if not binary: - return None - # LLAMA_SERVER_PATH is an explicit user pin that always wins in discovery; - # never auto-replace its tree (even a user's own llama.cpp checkout). - if os.environ.get("LLAMA_SERVER_PATH"): - return None - p = Path(binary) - env = os.environ.get("UNSLOTH_LLAMA_CPP_PATH") - if env and _is_under(p, Path(env)): - return Path(env) - for parent in p.parents: - if parent.name == "llama.cpp": - return parent - # PATH / system / custom install: not a managed tree, so do not offer. - return None + when the binary is unmanaged (see update_flow.managed_install_root).""" + return _flow.managed_install_root( + binary, + marker_root = _install_dir_for(binary), + server_path_var = "LLAMA_SERVER_PATH", + cpp_path_var = "UNSLOTH_LLAMA_CPP_PATH", + dir_name = "llama.cpp", + ) def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: @@ -324,69 +247,83 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: } -def _is_external_link(path: Optional[Path]) -> bool: - """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink - or a Windows directory junction / reparse point. Such a link resolves into - the user's own llama.cpp checkout, so Unsloth must never auto-update it.""" - if path is None: - return False - try: - if os.path.islink(path): - return True - except OSError: - return False - if os.name == "nt": - try: - import stat - attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] - return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) - except (OSError, AttributeError): - return False - return False - - def _active_install_is_local_link(binary: Optional[str]) -> bool: """True when the active llama-server resolves through a --with-llama-cpp-dir - local link at the canonical llama.cpp directory. An update would write - through that link into the user's own checkout (or fail), so the install is - treated as externally managed: no update is offered or applied. Checks only - up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root - above it can't trip a false positive.""" - if not binary: - return False - for parent in Path(binary).parents: - if _is_external_link(parent): - return True - if parent.name == "llama.cpp": - break - return False + local link at the canonical llama.cpp directory (see + update_flow.active_install_is_local_link).""" + return _flow.active_install_is_local_link(binary, dir_name = "llama.cpp") def _local_link_status() -> dict: """Status payload for a local-link install: unmanaged, no update offered.""" - with _job_lock: - job = dict(_job) - return { - "supported": False, - "update_available": False, - "stale": False, - "installed_tag": None, - "latest_tag": None, - "published_repo": None, - "installed_at_utc": None, - "age_days": None, - "source_build": False, - "local_link": True, - "update_size_bytes": None, - "job": job, + return _flow.local_link_status(_job, _job_lock) + + +def _whisper_chain_status( + *, force_refresh: bool = False, paired_llama_will_update: bool = False +) -> Optional[dict]: + """Whisper's piggyback plan for the combined update item (see + whisper_cpp_update.chained_phase_plan). None disables the piggyback -- + fail-open so whisper can never break the llama status or apply.""" + try: + from utils import whisper_cpp_update + return whisper_cpp_update.chained_phase_plan( + force_refresh = force_refresh, + paired_llama_will_update = paired_llama_will_update, + ) + except Exception as exc: # pragma: no cover - defensive + logger.debug("llama update: whisper piggyback probe failed", error = str(exc)) + return None + + +def _merge_whisper_status(status: dict, *, force_refresh: bool = False) -> dict: + """Fold the whisper sub-status into the llama status payload: the llama + update item is the single UI surface, so update_available becomes the union + (llama behind OR whisper behind) while llama_update_available keeps the + llama-only flag. All pre-existing top-level fields are preserved.""" + status["llama_update_available"] = bool(status.get("update_available")) + plan = _whisper_chain_status( + force_refresh = force_refresh, + paired_llama_will_update = status["llama_update_available"], + ) + if plan is None: + status["whisper"] = None + status["update_component"] = "llama" if status["llama_update_available"] else None + return status + sub = plan.get("status") or {} + status["whisper"] = { + "update_available": bool(plan.get("update_available")), + "installed_tag": sub.get("installed_tag"), + "latest_tag": sub.get("latest_tag"), + "update_size_bytes": sub.get("update_size_bytes"), + "skip_reason": plan.get("skip_reason"), } + whisper_update_available = bool(plan.get("update_available")) + if whisper_update_available: + status["update_available"] = True + status["update_component"] = ( + "llama" + if status["llama_update_available"] + else "whisper" + if whisper_update_available + else None + ) + return status def get_update_status(*, force_refresh: bool = False) -> dict: - """Report whether a newer prebuilt exists plus the current job state. + """Report whether an update is available plus the current job state. - force_refresh bypasses the 24h release cache for an explicit "check now". + This is the single main update item: llama.cpp drives it and the whisper + piggyback is folded in (see _merge_whisper_status). force_refresh bypasses + the 24h release cache for an explicit "check now". """ + status = _llama_only_status(force_refresh = force_refresh) + return _merge_whisper_status(status, force_refresh = force_refresh) + + +def _llama_only_status(*, force_refresh: bool = False) -> dict: + """The llama.cpp half of get_update_status (no whisper sub-status).""" binary = _find_binary() # A --with-llama-cpp-dir local link is the user's own tree; never offer to # replace it. Bail before any network/freshness work. @@ -456,33 +393,19 @@ def get_update_status(*, force_refresh: bool = False) -> dict: } -def _rocm_install_args(asset: Optional[str]) -> list[str]: - """Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh. - The installer probe can miss the gfx arch on amd-smi-only hosts; per-gfx - ROCm bundles carry the family in the name (rocm-gfx110X), version-tagged - bundles only rocm/hip.""" - if not asset: - return [] - low = asset.lower() - if "rocm" not in low and "hip" not in low: - return [] - gfx = re.search(r"-gfx[0-9a-z]+", low) - if gfx: - # _normalize_forwarded_gfx accepts the family form (gfx110x -> gfx110X). - return ["--rocm-gfx", gfx.group(0).lstrip("-")] - return ["--has-rocm"] - - -def _run_update( +def _run_llama_phase( install_dir: Path, repo: str, asset: Optional[str], script: Path, - pin_release_tag: Optional[str] = None, + pin_release_tag: Optional[str], + set_progress, force_cpu: bool = False, -) -> None: - """Worker: put the backend into a maintenance state, run the installer for - the latest prebuilt, then refresh caches so the next load uses the new build. +) -> dict: + """The llama phase of a chained update: put the backend into a maintenance + state, run the installer for the latest prebuilt, then refresh caches so the + next load uses the new build. Returns {to_tag, reload_required, message}; + raises on failure. pin_release_tag pins the installer to that exact published release instead of letting it re-resolve "latest" itself (see start_update for why).""" @@ -530,7 +453,6 @@ def _run_update( if force_cpu: cmd.append("--force-cpu") logger.info("llama update: installing", cmd = " ".join(cmd)) - # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm # box would otherwise re-route and silently replace the Vulkan build. @@ -538,44 +460,12 @@ def _run_update( # _rocm_install_args). if asset and "vulkan" in asset.lower(): env["UNSLOTH_FORCE_VULKAN"] = "1" - proc = subprocess.Popen( + _flow.stream_installer( cmd, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - text = True, - env = env, - **child_popen_kwargs(), + env, + set_progress = set_progress, + timeout_seconds = _INSTALL_TIMEOUT_SECONDS, ) - timed_out = threading.Event() - - def _kill_on_timeout() -> None: - timed_out.set() - proc.kill() - - watchdog = threading.Timer(_INSTALL_TIMEOUT_SECONDS, _kill_on_timeout) - watchdog.daemon = True - watchdog.start() - tail_lines: list[str] = [] - try: - assert proc.stdout is not None - for line in proc.stdout: - tail_lines.append(line) - if len(tail_lines) > 80: - del tail_lines[0] - m = _PROGRESS_LINE_RE.search(line) - if m is None: - continue - fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING - with _job_lock: - _job["progress"] = max(_job.get("progress") or 0.0, fraction) - returncode = proc.wait() - finally: - watchdog.cancel() - if timed_out.is_set(): - raise RuntimeError(f"installer timed out after {_INSTALL_TIMEOUT_SECONDS}s") - if returncode != 0: - tail = "".join(tail_lines).strip()[-1500:] - raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}") # Drop stale caches so the banner re-checks the swapped marker. # If GitHub is offline, latest stays unknown and the banner fails open. @@ -597,29 +487,18 @@ def _run_update( ): raise RuntimeError(f"pinned release {pin_release_tag} but installer produced {new_tag}") - with _job_lock: - _job.update( - state = _JOB_SUCCESS, - message = ( - f"Updated llama.cpp to {new_tag}." - + (" Reload your model to use it." if model_was_active else "") - ), - to_tag = new_tag, - reload_required = model_was_active, - error = None, - progress = 1.0, - finished_at = _utcnow(), - ) logger.info("llama update: success", to_tag = new_tag) + return { + "to_tag": new_tag, + "reload_required": model_was_active, + "message": ( + f"Updated llama.cpp to {new_tag}." + + (" Reload your model to use it." if model_was_active else "") + ), + } except Exception as exc: logger.warning("llama update: failed", error = str(exc)) - with _job_lock: - _job.update( - state = _JOB_ERROR, - message = "llama.cpp update failed.", - error = str(exc), - finished_at = _utcnow(), - ) + raise finally: # Always clear maintenance state. if backend is not None: @@ -629,50 +508,58 @@ def _run_update( pass -def start_update() -> dict: - """Kick off a background update. Idempotent: a second call while one is - running returns the in-flight job rather than starting another.""" +# Combined-job progress split when both phases run (download sizes: the llama +# bundle dwarfs the whisper one); normalized to 0..1 when a phase is skipped. +_LLAMA_PHASE_WEIGHT = 0.7 +_WHISPER_PHASE_WEIGHT = 0.3 + + +def _plan_llama_phase() -> dict: + """Decide how the llama phase of a combined update runs. Returns {"spec"} + when llama should install, else {"skip_reason", "refusal"}: skip_reason + marks the phase skipped inside a chained job, refusal is the started=False + response when the whisper phase has nothing to run either.""" binary = _find_binary() # Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt # here would write through the link into the user's own checkout (or fail) # and silently drop the link the flag created. if _active_install_is_local_link(binary): return { - "started": False, - "reason": "local_link", - "message": ( - "llama.cpp is a local directory linked with --with-llama-cpp-dir; " - "Unsloth won't replace it. Update your own llama.cpp checkout instead." - ), - "job": get_update_status()["job"], + "skip_reason": "local_link", + "refusal": { + "started": False, + "reason": "local_link", + "message": ( + "llama.cpp is a local directory linked with --with-llama-cpp-dir; " + "Unsloth won't replace it. Update your own llama.cpp checkout instead." + ), + }, } marker = read_install_marker(binary) script = _installer_script() if script is None: return { - "started": False, - "reason": "installer_missing", - "message": "install_llama_prebuilt.py could not be located.", - "job": get_update_status()["job"], + "skip_reason": "installer_missing", + "refusal": { + "started": False, + "reason": "installer_missing", + "message": "install_llama_prebuilt.py could not be located.", + }, } - # A job already in flight wins over any freshness re-check below (and skips - # its network call). The final lock block re-checks to close the TOCTOU. - with _job_lock: - if _job["state"] == _JOB_RUNNING: - return {"started": False, "reason": "already_running", "job": dict(_job)} - if marker: # Mirror the detection guard: a direct POST or a stale banner must not # start an install when the latest is not actually newer (force a fresh # check so a stale 24h cache can't wrongly block a real update either). - status = get_update_status(force_refresh = True) + status = _llama_only_status(force_refresh = True) if not status.get("update_available"): return { - "started": False, - "reason": "up_to_date", - "message": "The installed llama.cpp build is already at the latest prebuilt.", - "job": status["job"], + "skip_reason": "up_to_date", + "refusal": { + "started": False, + "reason": "up_to_date", + "message": "The installed llama.cpp build is already at the latest prebuilt.", + }, } install_dir = _install_dir_for(binary) repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO @@ -693,20 +580,27 @@ def start_update() -> dict: src = _source_build_status(binary, force_refresh = True) if binary else None if src is None: return { - "started": False, - "reason": "no_prebuilt_available", - "message": ( - "No official llama.cpp prebuilt is available for this host, " - "so the source build cannot be swapped automatically." - ), - "job": get_update_status()["job"], + "skip_reason": "no_prebuilt_available", + "refusal": { + "started": False, + "reason": "no_prebuilt_available", + "message": ( + "No official llama.cpp prebuilt is available for this host, " + "so the source build cannot be swapped automatically." + ), + }, } if not src.get("update_available"): return { - "started": False, - "reason": "up_to_date", - "message": "The installed llama.cpp build is already at or newer than the latest prebuilt.", - "job": get_update_status()["job"], + "skip_reason": "up_to_date", + "refusal": { + "started": False, + "reason": "up_to_date", + "message": ( + "The installed llama.cpp build is already at or newer than the " + "latest prebuilt." + ), + }, } res = _resolve_prebuilt_for_host() install_dir = _llama_install_root(binary) @@ -721,31 +615,116 @@ def start_update() -> dict: if install_dir is None: return { - "started": False, - "reason": "no_install_dir", - "message": "Could not determine the llama.cpp install directory.", - "job": get_update_status()["job"], + "skip_reason": "no_install_dir", + "refusal": { + "started": False, + "reason": "no_install_dir", + "message": "Could not determine the llama.cpp install directory.", + }, } + return { + "spec": { + "install_dir": install_dir, + "repo": repo, + "asset": asset, + "script": script, + "pin_release_tag": pin_release_tag, + "from_tag": from_tag, + "force_cpu": force_cpu, + } + } + + +def start_update() -> dict: + """Kick off a background update job. The job chains the llama phase (the + existing flow) with a whisper phase that runs only when whisper is actually + behind; either phase no-ops cleanly when its component is current or + unmanaged. Idempotent: a second call while one is running returns the + in-flight job rather than starting another.""" + # A job already in flight wins over any freshness re-check below (and skips + # its network calls). The final lock block re-checks to close the TOCTOU. + with _job_lock: + if _job["state"] == _JOB_RUNNING: + return {"started": False, "reason": "already_running", "job": dict(_job)} + + llama_plan = _plan_llama_phase() + llama_spec = llama_plan.get("spec") + whisper_plan = _whisper_chain_status( + force_refresh = True, + paired_llama_will_update = llama_spec is not None, + ) + whisper_spec = (whisper_plan or {}).get("phase") + if llama_spec is None and whisper_spec is None: + # Nothing to run in either phase: answer with the llama refusal so the + # existing reasons (local_link / up_to_date / ...) keep their meaning. + refusal = dict(llama_plan["refusal"]) + with _job_lock: + refusal["job"] = dict(_job) + return refusal + + whisper_run = None + if whisper_spec is not None: + from utils import whisper_cpp_update as _whisper + whisper_run = lambda set_progress: _whisper.run_chained_phase(whisper_spec, set_progress) + + phases = [ + { + "name": "llama", + "weight": _LLAMA_PHASE_WEIGHT, + "failure_message": "llama.cpp update failed.", + "skip_reason": llama_plan.get("skip_reason"), + "run": ( + ( + lambda set_progress: _run_llama_phase( + llama_spec["install_dir"], + llama_spec["repo"], + llama_spec["asset"], + llama_spec["script"], + llama_spec["pin_release_tag"], + set_progress, + force_cpu = llama_spec.get("force_cpu", False), + ) + ) + if llama_spec + else None + ), + }, + { + "name": "whisper", + "weight": _WHISPER_PHASE_WEIGHT, + "failure_message": "whisper.cpp update failed.", + # The sidecar reload is whisper-internal; it must not trip the + # job-level reload flag the chat frontend resyncs on. + "affects_job_reload": False, + "skip_reason": (whisper_plan or {}).get("skip_reason") or "unavailable", + "run": whisper_run, + }, + ] + running = " + ".join( + name for name, spec in (("llama.cpp", llama_spec), ("whisper.cpp", whisper_spec)) if spec + ) with _job_lock: if _job["state"] == _JOB_RUNNING: return {"started": False, "reason": "already_running", "job": dict(_job)} _job.update( state = _JOB_RUNNING, - message = "Downloading and installing the latest llama.cpp prebuilt...", - from_tag = from_tag, + message = f"Downloading and installing the latest {running} prebuilt...", + from_tag = (llama_spec or {}).get("from_tag"), to_tag = None, reload_required = None, error = None, progress = 0.0, started_at = _utcnow(), finished_at = None, + phases = None, ) job_snapshot = dict(_job) thread = threading.Thread( - target = _run_update, - args = (install_dir, repo, asset, script, pin_release_tag, force_cpu), + target = _flow.run_chained_update, + args = (phases,), + kwargs = {"job": _job, "job_lock": _job_lock}, name = "llama-cpp-update", daemon = True, ) @@ -755,15 +734,4 @@ def start_update() -> dict: def _reset_job_for_tests() -> None: """Test-only: return the job tracker to idle.""" - with _job_lock: - _job.update( - state = _JOB_IDLE, - message = "", - from_tag = None, - to_tag = None, - reload_required = None, - error = None, - progress = None, - started_at = None, - finished_at = None, - ) + _flow.reset_job(_job, _job_lock) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index 50b3cd3513..749f2c9234 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,10 +50,14 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} +_STRING_CACHE: Dict[Tuple[_CacheKey, str], Optional[str]] = {} + # GGUF header dims for the staged/deferred-load UI: context_length, layer_count # (block_count), and moe_layer_count (block_count minus leading dense layers; 0 # if not MoE). One cached pass fills all three so the staged sheet can size every -# slider before the model loads. None = unreadable / not a GGUF. +# slider before the model loads. None = unreadable / not a GGUF. The native +# training context length (``{arch}.context_length``) the UI shows before a model +# loads is read from here via read_gguf_context_length. _DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {} @@ -408,6 +412,83 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]: return result +def _parse_gguf_string(path: str, wanted_key: str) -> Optional[str]: + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + return sbytes.decode("utf-8", "replace") + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"_parse_gguf_string: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"_parse_gguf_string: parse failure on {path}: {e}") + return None + return None + + +def _read_gguf_string(path: str, wanted_key: str) -> Optional[str]: + fkey = _cache_key(path) + if fkey is None: + return None + ckey = (fkey, wanted_key) + with _CACHE_LOCK: + if ckey in _STRING_CACHE: + return _STRING_CACHE[ckey] + result = _parse_gguf_string(path, wanted_key) + with _CACHE_LOCK: + while len(_STRING_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _STRING_CACHE.pop(next(iter(_STRING_CACHE))) + except StopIteration: + break + _STRING_CACHE[ckey] = result + return result + + +def read_gguf_chat_template(path: str) -> Optional[str]: + template = _read_gguf_string(path, "tokenizer.chat_template") + if isinstance(template, str) and template.strip(): + return template + return None + + def read_mmproj_audio_capability(path: str) -> Optional[bool]: """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index dadf103cea..4897f05ce4 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -37,6 +37,7 @@ import yaml from utils.native_path_leases import child_env_without_native_path_secret +from utils.hf_cache_settings import active_hf_hub_cache, get_hf_cache_paths from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) @@ -493,6 +494,7 @@ def load_model_config( trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + cache_dir = active_hf_hub_cache(), ) if not use_auth: @@ -503,6 +505,7 @@ def load_model_config( trust_remote_code = trust_remote_code, token = None, local_files_only = local_files_only, + cache_dir = active_hf_hub_cache(), ) # Default auth (cached tokens) @@ -510,6 +513,7 @@ def load_model_config( model_name, trust_remote_code = trust_remote_code, local_files_only = local_files_only, + cache_dir = active_hf_hub_cache(), ) @@ -624,6 +628,7 @@ def _raw_config_has_vision_config( filename = "config.json", token = hf_token, local_files_only = local_files_only, + cache_dir = active_hf_hub_cache(), ) ) config = json.loads(config_path.read_text()) @@ -770,7 +775,7 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) capture_output = True, text = True, timeout = 60, - env = child_env_without_native_path_secret(), + env = get_hf_cache_paths().child_env(child_env_without_native_path_secret()), **_windows_hidden_subprocess_kwargs(), ) @@ -1249,6 +1254,77 @@ def _iter_gguf_files(directory: Path, recursive: bool = False): yield f +_GGUF_SPLIT_FILE_RE = re.compile( + r"^(?P.+)-(?P\d{5})-of-(?P\d{5})\.gguf$", + re.IGNORECASE, +) + + +def _colocated_first_split_shard(path: Path) -> tuple[Optional[Path], bool]: + """Return shard 1 and whether every shard is beside *path*.""" + match = _GGUF_SPLIT_FILE_RE.match(path.name) + if match is None: + return None, False + + prefix = match.group("prefix").casefold() + total_text = match.group("total") + total = int(total_text) + if total < 1: + return None, False + + first: Optional[Path] = None + indices: set[int] = set() + try: + siblings = path.parent.iterdir() + for sibling in siblings: + sibling_match = _GGUF_SPLIT_FILE_RE.match(sibling.name) + if ( + sibling_match is None + or sibling_match.group("prefix").casefold() != prefix + or sibling_match.group("total") != total_text + ): + continue + try: + if not sibling.is_file(): + continue + except OSError: + continue + index = int(sibling_match.group("index")) + if not 1 <= index <= total: + continue + indices.add(index) + if index == 1: + first = sibling + except OSError: + return None, False + + return first, first is not None and len(indices) == total + + +def _local_gguf_load_path(path: Path) -> Path: + """Choose a loadable local path while preserving complete symlink sets.""" + if _GGUF_SPLIT_FILE_RE.match(path.name) is None: + return path.absolute() + + first, complete = _colocated_first_split_shard(path) + if complete and first is not None: + return first.absolute() + + try: + is_symlink = path.is_symlink() + except OSError: + is_symlink = False + if is_symlink: + try: + target = path.resolve() + except OSError: + return (first or path).absolute() + target_first, _ = _colocated_first_split_shard(target) + return (target_first or target).absolute() + + return (first or path).absolute() + + def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]: """Find the mmproj GGUF for a model. @@ -1434,7 +1510,7 @@ def detect_gguf_model(path: str) -> Optional[str]: except OSError: is_dir = False # stat() unavailable in the lock window if not is_dir: - return str(p.absolute()) # absolute() keeps symlink names readable + return str(_local_gguf_load_path(p)) # Directory named "*.gguf": fall through to the dir scan below. # Case 2: directory containing .gguf files (skip mmproj / MTP drafter) @@ -1452,7 +1528,7 @@ def detect_gguf_model(path: str) -> Optional[str]: gguf_files.append(f) gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True) if gguf_files: - return str(gguf_files[0].resolve()) + return str(_local_gguf_load_path(gguf_files[0])) return None @@ -1643,19 +1719,20 @@ def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str return str(gguf_dir) -def _iter_hf_cache_snapshots(repo_id: str): +def _iter_hf_cache_snapshots(repo_id: str, cache_dir: Optional[str | Path] = None): """Yield HF cache snapshot dirs for *repo_id*, newest first. Empty if HF_HUB_CACHE is missing, the repo isn't cached, or has no snapshots. Repo name match is case-insensitive to handle casing drift between download time and lookup. """ - try: - from huggingface_hub import constants as hf_constants - except Exception: - return - - cache_dir = Path(hf_constants.HF_HUB_CACHE) + if cache_dir is None: + try: + from utils.hf_cache_settings import get_hf_cache_paths + cache_dir = get_hf_cache_paths().hub_cache + except Exception: + return + cache_dir = Path(cache_dir) target = f"models--{repo_id.replace('/', '--')}".lower() repo_dirs: list[Path] = [] try: @@ -1879,7 +1956,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: For sharded GGUFs (multiple files sharing a quant label), returns the first shard (sorted by name), which is what ``llama-server -m`` expects. - Returns the resolved absolute path, or ``None`` if no match. + Returns the absolute path, or ``None`` if no match. """ p = _resolve_gguf_dir(Path(directory)) if p is None: @@ -1900,7 +1977,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: matches.append(f) matches.sort() if matches: - return str(matches[0].resolve()) + return str(_local_gguf_load_path(matches[0])) return None @@ -1997,6 +2074,7 @@ def download_gguf_file( repo_id = repo_id, filename = filename, token = hf_token, + cache_dir = active_hf_hub_cache(), ) return local_path @@ -2005,6 +2083,24 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} +# Bound the Hub lookup so a DNS-dead session fails fast to the cache instead of hanging on retries. +_HUB_MODEL_INFO_TIMEOUT = 15.0 + + +def _embedding_marker_in_hf_cache(model_name: str) -> bool: + """True when model_name's cached snapshot carries a modules.json (the ST marker). + Cache-only, no network; used offline and as a fallback when the Hub lookup times out.""" + from utils.utils import hf_cache_snapshot_dir + + snapshot = hf_cache_snapshot_dir(model_name) + if snapshot is None: + return False + try: + return (snapshot / "modules.json").is_file() + except OSError: + return False + + def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """Detect embedding/sentence-transformer models via HF metadata. @@ -2019,6 +2115,15 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: Returns: True if embedding model, else False (default for local paths or errors). """ + from utils.utils import hf_env_offline + + # Offline (remote repo): reclassify from the local cache on every call, before/without the + # memo. An online lookup can memoize True from tags with no weights cached, so trusting it once + # the session goes offline would accept a repo _get() cannot load; a cached negative can also be + # invalidated by later cache materialization. The cache probe is local-only, so it's cheap. + if not is_local_path(model_name) and hf_env_offline(): + return _embedding_marker_in_hf_cache(model_name) + cache_key = (model_name, hf_token) if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] @@ -2033,7 +2138,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: try: from huggingface_hub import model_info as hf_model_info - info = hf_model_info(model_name, token = hf_token) + info = hf_model_info(model_name, token = hf_token, timeout = _HUB_MODEL_INFO_TIMEOUT) tags = set(info.tags or []) pipeline_tag = info.pipeline_tag or "" @@ -2054,9 +2159,11 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb except Exception as e: + # Timeout or transient network error: fall back to the local cache marker, don't hard-fail. logger.warning(f"Could not determine if {model_name} is embedding model: {e}") - _embedding_detection_cache[cache_key] = False - return False + is_emb = _embedding_marker_in_hf_cache(model_name) + _embedding_detection_cache[cache_key] = is_emb + return is_emb def _has_model_weight_files(model_dir: Path) -> bool: @@ -2416,7 +2523,10 @@ def get_base_model_from_lora_identifier( for _attempt in range(2): # one retry: a transient blip must not skip the base try: cfg_path = hf_hub_download( - identifier, "adapter_config.json", token = hf_token if hf_token else None + identifier, + "adapter_config.json", + token = hf_token if hf_token else None, + cache_dir = active_hf_hub_cache(), ) except (EntryNotFoundError, RepositoryNotFoundError): # No adapter_config.json -> not a resolvable LoRA; caller scans the identifier. @@ -2796,7 +2906,12 @@ class ModelConfig: try: from huggingface_hub import hf_hub_download - config_path = hf_hub_download(identifier, "adapter_config.json", token = hf_token) + config_path = hf_hub_download( + identifier, + "adapter_config.json", + token = hf_token, + cache_dir = active_hf_hub_cache(), + ) with open(config_path, "r") as f: adapter_config = json.load(f) base_model = adapter_config.get("base_model_name_or_path") diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index 08671cfe39..3ed7faa7c2 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -15,6 +15,7 @@ import base64 import binascii import hashlib import hmac +import importlib import json import os import stat as _stat_module @@ -35,7 +36,7 @@ _USED_NONCES: dict[str, int] = {} _REDACTION_LOCK = threading.Lock() _NATIVE_PATH_REDACTIONS: list[str] = [] _NATIVE_PATH_LABELS: dict[str, str] = {} -_NATIVE_PATH_ENV_LOCK = threading.Lock() +_NATIVE_PATH_ENV_LOCK = threading.RLock() _SECRET_INIT_LOCK = threading.Lock() _CACHED_LEASE_SECRET: bytes | None = None _SCRUB_REFCOUNT = 0 @@ -80,7 +81,9 @@ def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) - return cleaned -def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: +def run_without_native_path_secret( + target: Callable[..., Any] | str, *args: Any, **kwargs: Any +) -> Any: """Run a multiprocessing child target without the native path lease secret.""" # Runs in the spawned child: bind it to the parent's death (Linux), since @@ -96,6 +99,11 @@ def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwa os.environ.pop(LEASE_SECRET_ENV, None) _CACHED_LEASE_SECRET = None _SCRUB_SAVED_SECRET = None + if isinstance(target, str): + function_name, environment, *args = args + for key, value in environment.items(): + os.environ[key] = value + target = getattr(importlib.import_module(target), function_name) return target(*args, **kwargs) @@ -107,10 +115,9 @@ def native_path_secret_removed_for_child_start() -> Iterator[None]: _SCRUB_SAVED_SECRET = os.environ.pop(LEASE_SECRET_ENV, None) _CACHED_LEASE_SECRET = None _SCRUB_REFCOUNT += 1 - try: - yield - finally: - with _NATIVE_PATH_ENV_LOCK: + try: + yield + finally: _SCRUB_REFCOUNT -= 1 if _SCRUB_REFCOUNT == 0 and _SCRUB_SAVED_SECRET is not None: os.environ[LEASE_SECRET_ENV] = _SCRUB_SAVED_SECRET diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py index 0ea0477cc7..1a0d2d2746 100644 --- a/studio/backend/utils/paths/external_media.py +++ b/studio/backend/utils/paths/external_media.py @@ -131,6 +131,29 @@ def linux_run_media_mount_roots( return roots +def macos_volume_roots(base: Path | str = "/Volumes") -> list[Path]: + """Readable mounted volumes for the macOS folder browser.""" + + if platform.system() != "Darwin": + return [] + base_path = Path(base) + try: + entries = list(base_path.iterdir()) + except OSError: + return [] + roots: list[Path] = [] + for entry in entries: + if is_sensitive_path_component(entry.name): + continue + try: + resolved = entry.resolve() + if resolved.is_dir() and os.access(resolved, os.R_OK | os.X_OK): + roots.append(resolved) + except (OSError, RuntimeError, ValueError): + continue + return roots + + def _active_windows_drive_bitmask() -> int: """Active-logical-drive bitmask from ``GetLogicalDrives`` (bit 0 = ``A:``), or ``0`` when unavailable. diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py index e8dabc8954..65541661f1 100644 --- a/studio/backend/utils/paths/path_utils.py +++ b/studio/backend/utils/paths/path_utils.py @@ -122,15 +122,8 @@ def is_model_cached(model_name: str) -> bool: def _hf_hub_cache_dir() -> Path: """Return HF cache root honoring HF_HUB_CACHE when available.""" - try: - from huggingface_hub.constants import HF_HUB_CACHE - return Path(HF_HUB_CACHE) - except Exception as exc: - logger.debug( - "Could not read huggingface_hub HF_HUB_CACHE, using default hub path: %s", - exc, - ) - return Path.home() / ".cache" / "huggingface" / "hub" + from utils.hf_cache_settings import get_hf_cache_paths + return get_hf_cache_paths().hub_cache def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 35b8c57e9b..cea3cc61e3 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -277,27 +277,15 @@ def well_known_model_dirs() -> list[Path]: def _setup_cache_env() -> None: """Set cache env vars for HuggingFace, uv, and vLLM. - Respects the standard HF cache chain (explicit HF_HOME / HF_HUB_CACHE, - then XDG_CACHE_HOME, then ~/.cache/huggingface) and only sets vars the - user hasn't, so explicit overrides are honored. A user-set HF_HOME also - seeds HF_HUB_CACHE / HF_XET_CACHE (HF defaults them to $HF_HOME/hub and - $HF_HOME/xet); without this, models download to and load from the standard - cache even when HF_HOME points elsewhere, and both the Xet and HTTP-fallback - download paths inherit the same wrong root. + Explicit Hugging Face environment variables take precedence over Studio's + stored location. Studio seeds import-time variables once, while each later + worker receives its own captured cache location. """ root = cache_root() - xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() - # HUGGINGFACE_HUB_CACHE is HF's legacy alias for HF_HUB_CACHE; honor it. - if "HF_HUB_CACHE" not in os.environ and os.environ.get("HUGGINGFACE_HUB_CACHE"): - os.environ["HF_HUB_CACHE"] = os.environ["HUGGINGFACE_HUB_CACHE"] - # Seed the hub/xet caches from HF_HOME when set, else the platform default. - # Strip so a blank/whitespace HF_HOME falls back instead of making " /hub". - hf_home = (os.environ.get("HF_HOME") or "").strip() - hf_base = Path(hf_home).expanduser() if hf_home else xdg_cache / "huggingface" + from utils.hf_cache_settings import initialize_hf_cache_environment + + initialize_hf_cache_environment() defaults: dict[str, str] = { - "HF_HOME": str(hf_base), - "HF_HUB_CACHE": str(hf_base / "hub"), - "HF_XET_CACHE": str(hf_base / "xet"), "UV_CACHE_DIR": str(root / "uv"), "VLLM_CACHE_ROOT": str(root / "vllm"), } diff --git a/studio/backend/utils/prebuilt/__init__.py b/studio/backend/utils/prebuilt/__init__.py new file mode 100644 index 0000000000..c41cd1150d --- /dev/null +++ b/studio/backend/utils/prebuilt/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend-importable prebuilt helpers. + +The installers reuse install_llama_prebuilt.py directly; this package holds the +backend-side shapes the studio/ scripts cannot provide (the backend runs with +studio/backend as its sys.path root): runtime_libs (wheel CUDA dirs), child_env +(secret scrubbing + WSL ROCm dirs), freshness_flow and update_flow (the shared +mechanics behind the *_cpp_freshness / *_cpp_update twins). +""" diff --git a/studio/backend/utils/prebuilt/child_env.py b/studio/backend/utils/prebuilt/child_env.py new file mode 100644 index 0000000000..b6b7a40df7 --- /dev/null +++ b/studio/backend/utils/prebuilt/child_env.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Child-process environment hygiene for the managed ggml servers. + +Secret-env scrubbing and the WSL2 ROCm library-dir probe, shared by the STT +sidecar (and any future launcher of a downloaded binary). Kept in sync with +install_llama_prebuilt.py's scrub_env / _wsl_system_rocm_lib_dirs; the backend +cannot import the studio/ installer scripts, so this copy stays importable with +only the backend root on sys.path. +""" + +from __future__ import annotations + +import os +import re +from typing import Mapping + +SECRET_ENV_EXACT = frozenset( + { + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "GH_TOKEN", + "GITHUB_TOKEN", + "WANDB_API_KEY", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "GOOGLE_APPLICATION_CREDENTIALS", + "AZURE_CLIENT_SECRET", + "KUBECONFIG", + "SSH_AUTH_SOCK", + } +) +# Case-insensitive substring markers for names we do not enumerate (no bare "KEY"). +SECRET_ENV_MARKERS = ( + "TOKEN", + "SECRET", + "PASSWORD", + "PASSWD", + "PASSPHRASE", + "CREDENTIAL", + "PRIVATE_KEY", + "API_KEY", +) +# Proxy / index URLs embed creds in their value; the offline server never needs them. +SECRET_ENV_URL_NAMES = frozenset( + { + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "FTP_PROXY", + "RSYNC_PROXY", + "PIP_INDEX_URL", + "PIP_EXTRA_INDEX_URL", + "UV_INDEX_URL", + "UV_DEFAULT_INDEX", + "UV_EXTRA_INDEX_URL", + } +) +# Also drop values with URL userinfo creds (scheme://user:secret@host). +URL_USERINFO_RE = re.compile(r"://[^/@\s]+@") + + +def is_secret_env_name(name: str) -> bool: + upper = name.upper() + return ( + upper in SECRET_ENV_EXACT + or upper in SECRET_ENV_URL_NAMES + or any(marker in upper for marker in SECRET_ENV_MARKERS) + ) + + +def scrub_env(env: Mapping[str, str]) -> dict[str, str]: + """Copy of ``env`` without secret-bearing names or URL-userinfo values.""" + return { + k: v + for k, v in env.items() + if not is_secret_env_name(k) and not URL_USERINFO_RE.search(v or "") + } + + +# Filesystem pointers a downloaded binary could follow to on-disk credential +# stores (token caches under $HF_HOME, ~/.netrc, XDG config). Dropped, not +# repointed; the offline inference server needs none. Mirrors the cred-location +# list of the tools bypass env (core/inference/tools.py). +CRED_LOCATION_ENV_NAMES = frozenset( + { + "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", + "BASH_ENV", + "GIT_CONFIG_GLOBAL", + "GIT_CONFIG_SYSTEM", + "GIT_ASKPASS", + "SSH_ASKPASS", + "HOMEDRIVE", + "HOMEPATH", + } +) +# Home dirs are repointed (not dropped): loaders and SDKs expect them present, +# but they must not resolve to the user's real profile with its token caches. +HOME_ENV_NAMES = ("HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA") + + +def isolate_home(env: dict[str, str], scratch_dir: str) -> dict[str, str]: + """Repoint home/profile vars at ``scratch_dir`` and drop credential-store + pointers so a compromised downloaded server cannot read token caches or cred + files through the environment. Mutates and returns ``env``.""" + os.makedirs(scratch_dir, exist_ok = True) + for name in HOME_ENV_NAMES: + if name in env: + env[name] = scratch_dir + for name in CRED_LOCATION_ENV_NAMES: + env.pop(name, None) + return env + + +def wsl_system_rocm_lib_dirs() -> list[str]: + """System ROCm lib dir(s) to load before a bundle's HIP on WSL2. Strict no-op + off WSL (needs /dev/dxg, a "microsoft" /proc/version, and a librocdxg).""" + try: + if not os.path.exists("/dev/dxg"): + return [] + with open("/proc/version", encoding = "utf-8", errors = "replace") as fh: + if "microsoft" not in fh.read().lower(): + return [] + except OSError: + return [] + dirs: list[str] = [] + for d in ("/opt/rocm/lib", "/opt/rocm/lib64"): + if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists( + os.path.join(d, "librocdxg.so.1") + ): + dirs.append(d) + return dirs diff --git a/studio/backend/utils/prebuilt/freshness_flow.py b/studio/backend/utils/prebuilt/freshness_flow.py new file mode 100644 index 0000000000..b90ebf776c --- /dev/null +++ b/studio/backend/utils/prebuilt/freshness_flow.py @@ -0,0 +1,325 @@ +# 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 mechanics of the llama.cpp / whisper.cpp prebuilt freshness checks. + +The component modules (utils.llama_cpp_freshness / utils.whisper_cpp_freshness) +keep their public names, per-module caches, and version-comparison policy; +everything mechanical (marker walk-up, GitHub release fetch, memo + disk cache, +the freshness report skeleton) lives here, parameterized by call-time callables +so the modules' monkeypatch seams keep working. +""" + +from __future__ import annotations + +import json +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Optional + +import structlog + +logger = structlog.get_logger(__name__) + +# 24h TTL keeps the GitHub call off the hot path and within rate limits. +RELEASE_CACHE_TTL_SECONDS = 24 * 60 * 60 + + +def read_install_marker( + binary_path: Optional[str], + *, + marker_name: str, + cache: dict[str, Optional[dict]], + log_message: str, +) -> Optional[dict]: + """Walk up from binary_path to find the install marker JSON. + None = no marker (source build / custom path) or invalid JSON.""" + if not binary_path: + return None + cached = cache.get(binary_path) + if cached is not None or binary_path in cache: + return cached + p = Path(binary_path) + marker: Optional[dict] = None + # Cover all managed binary layouts (binary is 1-4 dirs deep). + for parent in p.parents[:5]: + candidate = parent / marker_name + if candidate.is_file(): + try: + marker = json.loads(candidate.read_text(encoding = "utf-8")) + except (OSError, json.JSONDecodeError) as exc: + logger.debug(log_message, path = str(candidate), error = str(exc)) + marker = None + break + cache[binary_path] = marker + return marker + + +def cache_path_for(repo: str, cache_dir: Path) -> Path: + safe = repo.replace("/", "__") + return cache_dir / f"{safe}.json" + + +def load_disk_cache(repo: str, cache_dir: Path) -> Optional[tuple[float, Optional[str]]]: + path = cache_path_for(repo, cache_dir) + try: + payload = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, json.JSONDecodeError): + return None + ts = payload.get("fetched_at") + tag = payload.get("latest_tag") + if not isinstance(ts, (int, float)): + return None + return float(ts), tag if isinstance(tag, str) else None + + +def save_disk_cache( + repo: str, latest_tag: Optional[str], cache_dir: Path, *, log_message: str +) -> None: + path = cache_path_for(repo, cache_dir) + try: + path.parent.mkdir(parents = True, exist_ok = True) + tmp = path.with_suffix(".tmp") + tmp.write_text( + json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag}), + encoding = "utf-8", + ) + tmp.replace(path) + except OSError as exc: + logger.debug(log_message, repo = repo, error = str(exc)) + + +def _fetch_newest_published_release( + repo: str, timeout: float, *, log_message: str +) -> Optional[dict]: + """Newest published (non-draft/non-prerelease) release object for `repo`, by + ``published_at``. + + Resolves "latest" the way the installers do, NOT via GitHub's + ``/releases/latest`` pointer, which sorts by commit date and can lag the + build the installer installs (detection and apply then disagree -- the + downgrade/sticky-banner bug). None on any failure (offline, rate-limited).""" + import os + import urllib.error + import urllib.request + + url = f"https://api.github.com/repos/{repo}/releases?per_page=30" + headers = { + "Accept": "application/vnd.github+json", + "User-Agent": "unsloth-studio-freshness-check", + } + token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, headers = headers) + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + data = json.loads(resp.read().decode("utf-8")) + except ( + urllib.error.URLError, + urllib.error.HTTPError, + OSError, + json.JSONDecodeError, + ) as exc: + logger.debug(log_message, repo = repo, error = str(exc)) + return None + if not isinstance(data, list): + return None + published = [ + r + for r in data + if isinstance(r, dict) + and not r.get("draft") + and not r.get("prerelease") + and isinstance(r.get("tag_name"), str) + and r.get("tag_name") + ] + if not published: + return None + return max(published, key = lambda r: r.get("published_at") or "") + + +def fetch_latest_release_tag( + repo: str, + timeout: float = 5.0, + *, + log_message: str, +) -> Optional[str]: + """Newest published release tag for `repo`, by publish time. None on failure.""" + newest = _fetch_newest_published_release(repo, timeout, log_message = log_message) + return newest["tag_name"] if newest else None + + +def fetch_latest_release_assets( + repo: str, + timeout: float = 5.0, + *, + log_message: str, +) -> Optional[dict[str, int]]: + """Asset name -> size (bytes) for the newest published release of `repo`, + selected exactly like fetch_latest_release_tag. None on any failure.""" + newest = _fetch_newest_published_release(repo, timeout, log_message = log_message) + if newest is None: + return None + assets: dict[str, int] = {} + for a in newest.get("assets") or []: + name, size = a.get("name"), a.get("size") + if isinstance(name, str) and isinstance(size, int): + assets[name] = size + return assets + + +def latest_published_release( + repo: str, + *, + force_refresh: bool, + memo: dict[str, tuple[float, Optional[str]]], + cache_dir: Callable[[], Path], + fetch: Callable[[str], Optional[str]], + save: Callable[[str, Optional[str]], None], +) -> Optional[str]: + """Latest release tag for `repo`. Memo + disk-cached (24h TTL). + None when offline and never previously cached.""" + if not repo: + return None + now = time.time() + if not force_refresh: + cached = memo.get(repo) + if cached and now - cached[0] < RELEASE_CACHE_TTL_SECONDS: + return cached[1] + disk = load_disk_cache(repo, cache_dir()) + if disk and now - disk[0] < RELEASE_CACHE_TTL_SECONDS: + memo[repo] = disk + return disk[1] + latest = fetch(repo) + if latest is None: + # Keep the last-good disk value rather than poison it with None. + disk = load_disk_cache(repo, cache_dir()) + if disk: + memo[repo] = disk + return disk[1] + return None + memo[repo] = (now, latest) + save(repo, latest) + return latest + + +def latest_release_assets( + repo: str, + *, + force_refresh: bool, + memo: dict[str, tuple[float, dict[str, int]]], + fetch: Callable[[str], Optional[dict[str, int]]], +) -> Optional[dict[str, int]]: + """Newest-release asset sizes for `repo`, memoized (24h TTL). None when + offline and never fetched. In-memory only -- a restart re-fetches.""" + if not repo: + return None + now = time.time() + if not force_refresh: + cached = memo.get(repo) + if cached and now - cached[0] < RELEASE_CACHE_TTL_SECONDS: + return cached[1] + assets = fetch(repo) + if assets is None: + cached = memo.get(repo) + return cached[1] if cached else None + memo[repo] = (now, assets) + return assets + + +def parse_installed_at(value: object) -> Optional[datetime]: + if not isinstance(value, str) or not value: + return None + s = value.replace("Z", "+00:00") if value.endswith("Z") else value + try: + dt = datetime.fromisoformat(s) + except ValueError: + return None + if dt.tzinfo is None: + dt = dt.replace(tzinfo = timezone.utc) + return dt + + +def check_freshness( + binary_path: Optional[str], + *, + threshold_days: int, + now: Optional[datetime], + read_marker: Callable[[Optional[str]], Optional[dict]], + latest_release: Callable[[str], Optional[str]], + behind: Callable[[Optional[str], Optional[str]], bool], + display_tag: Callable[[dict], Any], + compare_tag: Callable[[dict], Any], +) -> dict: + """Freshness report skeleton shared by both components; the component's + marker-tag choice and is_behind policy come in as callables. Fails open on + missing data (behind/stale stay False).""" + out: dict = { + "has_marker": False, + "stale": False, + "behind": False, + "installed_tag": None, + "latest_tag": None, + "installed_at_utc": None, + "age_days": None, + "published_repo": None, + "threshold_days": int(threshold_days), + } + marker = read_marker(binary_path) + if not marker: + return out + out["has_marker"] = True + out["installed_tag"] = display_tag(marker) + out["installed_at_utc"] = marker.get("installed_at_utc") + out["published_repo"] = marker.get("published_repo") + + installed_full = compare_tag(marker) + repo = out["published_repo"] + if not repo or not installed_full: + return out + latest = latest_release(repo) + out["latest_tag"] = latest + out["behind"] = behind(installed_full, latest) + if not out["behind"]: + return out + + installed_at = parse_installed_at(out["installed_at_utc"]) + if installed_at is None: + return out + now = now or datetime.now(tz = timezone.utc) + age_seconds = (now - installed_at).total_seconds() + out["age_days"] = max(0, int(age_seconds // 86400)) + if age_seconds >= threshold_days * 86400: + out["stale"] = True + return out + + +def format_stale_warning(info: dict, *, component: str) -> str: + """Human-readable one-liner for stale prebuilt info.""" + age = info.get("age_days") + installed = info.get("installed_tag") or "unknown" + latest = info.get("latest_tag") or "unknown" + age_str = f"{age} day{'s' if age != 1 else ''}" if age is not None else "some time" + return ( + f"{component} prebuilt is {age_str} behind: installed " + f"{installed}, latest {latest}. Run `unsloth studio update` " + f"to refresh." + ) + + +def reset_caches( + caches: tuple[dict, ...], *, drop_disk: bool, cache_dir: Callable[[], Path] +) -> None: + """Drop the in-memory freshness caches; with drop_disk also the on-disk 24h + release cache (see the component modules for why).""" + for cache in caches: + cache.clear() + if drop_disk: + import shutil + + # cache_dir() is a dedicated freshness-only subdir, re-created on the next + # save_disk_cache. ignore_errors so a missing/locked dir is a no-op rather + # than breaking an otherwise successful install. + shutil.rmtree(cache_dir(), ignore_errors = True) diff --git a/studio/backend/utils/prebuilt/runtime_libs.py b/studio/backend/utils/prebuilt/runtime_libs.py new file mode 100644 index 0000000000..6e51fb8246 --- /dev/null +++ b/studio/backend/utils/prebuilt/runtime_libs.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""CUDA runtime dirs shipped inside Python wheels, for the STT sidecar's child env. + +Kept in sync with install_llama_prebuilt.py's python_runtime_dirs; the backend +cannot import the studio/ installer scripts, so this small copy stays importable +with only the backend root on sys.path. +""" + +from __future__ import annotations + +import site +import sys +from pathlib import Path +from typing import Iterable + + +def dedupe_existing_dirs(paths: Iterable[str | Path]) -> list[str]: + unique: list[str] = [] + seen: set[str] = set() + for raw in paths: + if not raw: + continue + try: + path = Path(raw).expanduser() + if not path.is_dir(): + continue + resolved = str(path.resolve()) + except (OSError, ValueError): + continue + if resolved in seen: + continue + seen.add(resolved) + unique.append(resolved) + return unique + + +def python_runtime_dirs() -> list[str]: + """CUDA runtime dirs shipped inside Python wheels (torch + nvidia-* wheels).""" + candidates: list[Path] = [] + search_roots = [Path(entry) for entry in sys.path if entry] + try: + search_roots.extend(Path(path) for path in site.getsitepackages()) + except Exception: + pass + try: + user_site = site.getusersitepackages() + if user_site: + search_roots.append(Path(user_site)) + except Exception: + pass + + for root in search_roots: + if not root.is_dir(): + continue + candidates.extend(root.glob("nvidia/*/lib")) # Linux convention + candidates.extend(root.glob("nvidia/*/bin")) # legacy modular Windows wheels + candidates.extend(root.glob("nvidia/*/bin/x86_64")) # CUDA 13 Windows wheel layout + candidates.extend(root.glob("nvidia/*/bin/x64")) + candidates.extend(root.glob("nvidia/*/Library/bin")) # conda-style repacks + candidates.extend(root.glob("nvidia/*/Library/bin/x86_64")) + candidates.extend(root.glob("nvidia/*/Library/bin/x64")) + candidates.extend(root.glob("torch/lib")) + return dedupe_existing_dirs(candidates) diff --git a/studio/backend/utils/prebuilt/update_flow.py b/studio/backend/utils/prebuilt/update_flow.py new file mode 100644 index 0000000000..74af0c18f9 --- /dev/null +++ b/studio/backend/utils/prebuilt/update_flow.py @@ -0,0 +1,447 @@ +# 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 mechanics of the llama.cpp / whisper.cpp in-app prebuilt updates. + +The component modules (utils.llama_cpp_update / utils.whisper_cpp_update) keep +their public names, job dicts, and update policy (version comparison, pinning, +pre/post install steps); everything mechanical (managed-root resolution, +local-link detection, the resolve probe, the streamed installer run) lives here, +parameterized so the modules' monkeypatch seams keep working. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Callable, Optional + +import structlog + +from utils.process_lifetime import child_popen_kwargs + +logger = structlog.get_logger(__name__) + +# Markerless (source-build) resolve answers are memoized for 24h; only +# successful answers are cached so a network blip retries. +RESOLVE_TTL_SECONDS = 24 * 60 * 60 + +# Matches the installer's download progress lines, e.g. +# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s". +PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(") +# The download dominates the update; extract/validate fill the last slice. +DOWNLOAD_PROGRESS_CEILING = 0.95 + + +class InstallerExit(RuntimeError): + """Installer subprocess exited nonzero; carries the exit code so phase + runners can special-case contractual codes (whisper's 2 = unavailable).""" + + def __init__(self, returncode: int, message: str) -> None: + super().__init__(message) + self.returncode = returncode + + +JOB_IDLE = "idle" +JOB_RUNNING = "running" +JOB_SUCCESS = "success" +JOB_ERROR = "error" + +# Per-phase states inside a chained job's "phases" breakdown. +PHASE_PENDING = "pending" +PHASE_RUNNING = "running" +PHASE_SUCCESS = "success" +PHASE_ERROR = "error" +PHASE_SKIPPED = "skipped" + +_IDLE_JOB_FIELDS = dict( + state = JOB_IDLE, + message = "", + from_tag = None, + to_tag = None, + reload_required = None, + error = None, + progress = None, + started_at = None, + finished_at = None, + phases = None, +) + + +def new_job() -> dict: + """A fresh idle job-state dict (one per component module).""" + return dict(_IDLE_JOB_FIELDS) + + +def reset_job(job: dict, job_lock: threading.Lock) -> None: + """Return a job tracker to idle (test seam).""" + with job_lock: + job.update(_IDLE_JOB_FIELDS) + + +def utcnow() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + + +def is_under(path: Path, root: Path) -> bool: + try: + p, r = path.resolve(), root.resolve() + except (OSError, ValueError): + p, r = path, root + return p == r or r in p.parents + + +def install_dir_for(binary_path: Optional[str], *, marker_name: str) -> Optional[Path]: + """The directory holding the install marker: the install root the installer + wrote and the one we re-install into. Walks up from the binary like the + freshness marker reader does.""" + if not binary_path: + return None + p = Path(binary_path) + for parent in p.parents[:5]: + if (parent / marker_name).is_file(): + return parent + return None + + +def find_installer_script(*, env_var: str, script_name: str) -> Optional[Path]: + """Locate the installer script. Honours the env override, then searches up + from this file for both ``/