diff --git a/install.ps1 b/install.ps1 index 17834d1443..07ea7613f3 100644 --- a/install.ps1 +++ b/install.ps1 @@ -3966,13 +3966,32 @@ exit 1 (Get-CanonicalDir -Path (Join-Path $env:USERPROFILE ".unsloth\studio"))) } - # Explicit staging root, shared default cache, or the custom Unsloth home's tree. + # The master root storage_roots.unsloth_home() reads. llama.cpp, node and whisper.cpp sit + # BESIDE studio\ under it, so deriving them from $StudioHome would put them one level below + # where every runtime resolver looks. + function Get-MasterRootOverride { + if ([string]::IsNullOrWhiteSpace($env:UNSLOTH_HOME)) { return $null } + $value = $env:UNSLOTH_HOME.Trim() + if ($value -eq "~") { + $value = $env:USERPROFILE + } elseif ($value -like "~/*" -or $value -like "~\*") { + $value = (Join-Path $env:USERPROFILE $value.Substring(1).TrimStart('/', '\')) + } + return (Get-CanonicalDir -Path $value) + } + + # Explicit staging root, the master root, the shared default cache, or the custom Unsloth + # home's tree. function Get-ManagedLlamaCppDir { param([AllowNull()][string]$StagingRoot = $null) if ($StagingRoot) { return (Join-Path $StagingRoot "llama.cpp") } + $masterRoot = Get-MasterRootOverride + if ($masterRoot) { + return (Join-Path $masterRoot "llama.cpp") + } if (-not (Test-StudioHomeIsCustom)) { return (Join-Path $env:USERPROFILE ".unsloth\llama.cpp") } @@ -3988,8 +4007,9 @@ exit 1 $dir = Get-ManagedLlamaCppDir -StagingRoot $StagingRoot if ((Get-LlamaCppInstallReadState -Path $dir) -ne "Denied") { return $null } Write-StudioLine "" - # A denied custom home cannot be claimed as an Unsloth-managed cache. - $homeIsCustom = Test-StudioHomeIsCustom + # A denied custom home cannot be claimed as an Unsloth-managed cache. Computed rather + # than read off $RuntimeRootIsCustom: this runs beside the line that defines it. + $homeIsCustom = (Test-StudioHomeIsCustom) -or [bool](Get-MasterRootOverride) # Preserve user-supplied wording when either override names this tree, or # names a build inside it: moving or deleting this folder takes that build # with it, and the later --with-llama-cpp-dir check then aborts on a path diff --git a/scripts/build_whisper_cpp.sh b/scripts/build_whisper_cpp.sh index 19fe6d9605..775e939b84 100755 --- a/scripts/build_whisper_cpp.sh +++ b/scripts/build_whisper_cpp.sh @@ -2,6 +2,7 @@ # Build whisper.cpp's whisper-server for Unsloth's GGUF dictation engine. # Installs into the managed Unsloth 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 (master root) # /whisper.cpp/build/bin/whisper-server (custom home) # ~/.unsloth/whisper.cpp/build/bin/whisper-server (default) # Usage: @@ -16,7 +17,25 @@ 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:-${UNSLOTH_HOME:-}}}" +# Stripped and tilde-expanded before it can outrank anything, as setup.sh and the Python +# resolvers do: ${VAR:-} only treats the EMPTY string as unset, so a whitespace-only value +# would win here and name a relative " /whisper.cpp" the backend never looks in. +_root_value() { + _rv=$(printf '%s' "${1:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + case "$_rv" in + "~") _rv="$HOME" ;; + "~/"*) _rv="$HOME/${_rv#'~/'}" ;; + esac + printf '%s' "$_rv" +} + +# UNSLOTH_HOME first: whisper.cpp is a SIBLING of studio/ under the master root, and the CLI +# exports UNSLOTH_STUDIO_HOME=/studio beside it, so taking that one would install a level +# below where stt_ggml_sidecar._managed_whisper_cpp_dir() looks. +_STUDIO_HOME_ALIAS="${STUDIO_HOME:-}" # read before the name below is reassigned +STUDIO_HOME="$(_root_value "${UNSLOTH_HOME:-}")" +[ -n "$STUDIO_HOME" ] || STUDIO_HOME="$(_root_value "${UNSLOTH_STUDIO_HOME:-}")" +[ -n "$STUDIO_HOME" ] || STUDIO_HOME="$(_root_value "$_STUDIO_HOME_ALIAS")" CUSTOM_STUDIO_HOME=false if [ -n "$STUDIO_HOME" ]; then CUSTOM_STUDIO_HOME=true diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index fd9c0e41fe..a811cbed35 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -124,6 +124,48 @@ Environment: } } + # rmdir, as uninstall.sh uses for the same two paths: empty or not at all, never recursed + # into. Not "list it, then _RemovePath": -ErrorAction SilentlyContinue makes an enumeration + # failure look like an empty directory, so a master root whose ACL denies listing was deleted + # recursively with the user's files in it, and a separate check and delete let a file + # arriving in between be taken by the -Recurse. Directory.Delete($path, $false) is one call. + function _RemoveDirIfEmpty { + param([string]$Path) + if ([string]::IsNullOrWhiteSpace($Path)) { return } + try { + if (-not (Test-Path -LiteralPath $Path -PathType Container)) { return } + [System.IO.Directory]::Delete($Path, $false) + _Substep "removed: $Path" "Green" + } catch { + # Non-empty, unreadable or held open. All three mean "not ours to take". + _Substep "keeping non-empty directory: $Path" "Yellow" + } + } + + # The exact shape prebuilt_core.py leaves: a component lock name, ".stale.", and the pid it + # moved aside. The leading dot alone also matched ".backup.install.lock.stale.copy", which in + # a user-chosen master root is theirs. uninstall.sh applies the same shape. + $script:StaleLockPattern = '^\.(llama\.cpp|node|whisper\.cpp|sd\.cpp)\.install\.lock\.stale\.[0-9]+$' + + # An install lock, and only an install lock. prebuilt_core.install_lock creates these with + # O_CREAT | O_EXCL, so a lock is always a regular file, while _RemovePath deletes recursively + # and in a user-chosen master root would take a whole tree carrying one of these fixed names + # with none of the owner-marker proof its neighbours require. Same rule as uninstall.sh. + function _RemoveLockFile { + param([string]$Path) + if ([string]::IsNullOrWhiteSpace($Path)) { return } + if (-not (Test-Path -LiteralPath $Path)) { return } + # A reparse point at a lock name is not a lock either, and following one would delete + # whatever it points at. + $item = Get-Item -LiteralPath $Path -Force -ErrorAction SilentlyContinue + if ($null -ne $item -and -not $item.PSIsContainer -and + -not ($item.Attributes -band [IO.FileAttributes]::ReparsePoint)) { + _RemovePath $Path + } else { + _Substep "keeping non-file at an install-lock path: $Path" "Yellow" + } + } + # LOCALAPPDATA / APPDATA are dropped in service and CI contexts, so fall back to the known # folder; $null only if that fails too, which callers treat as incomplete cleanup. function _AppDataRoot { @@ -567,6 +609,122 @@ Environment: # Discover non-default Unsloth roots from env vars + studio.conf files. Mirrors install.ps1's # precedence: UNSLOTH_STUDIO_HOME wins and STUDIO_HOME is ignored when both are set, or # uninstalling install A would also delete install B from a stale STUDIO_HOME. + # The master root UNSLOTH_HOME names, or $null. studio\ is its child and llama.cpp, node and + # whisper.cpp are its other children, so removing the Studio root alone strands them. Trimmed + # and tilde-expanded like storage_roots.unsloth_home() and studio\setup.ps1. + function _MasterRoot { + $noteStudio = $null + $raw = $env:UNSLOTH_HOME + if ([string]::IsNullOrWhiteSpace($raw)) { + # The note setup.ps1 leaves in the Studio tree, when this run has no UNSLOTH_HOME of + # its own: a one-command `$env:UNSLOTH_HOME = 'D:\portable'` leaves nothing in a + # later environment, so without it an uninstall stranded the runtimes. The deny list + # and the marker gates still apply, so a stale note licenses nothing. + # + # Only the Studio root this run is about, in studio_root()'s precedence and with no + # falling through: with two installs on one box, a legacy note otherwise won here + # while the removal worked on the named tree, so the run deleted one Studio and the + # OTHER install's runtime children. Mirrors _master_root in uninstall.sh. + $noteRoots = @() + foreach ($override in @($env:UNSLOTH_STUDIO_HOME, $env:STUDIO_HOME)) { + if (-not [string]::IsNullOrWhiteSpace($override)) { + $noteRoots = @(_ExpandTilde $override.Trim()) + break + } + } + if (-not $noteRoots) { + # No override names one, so the root install.ps1 recorded in the default-mode + # studio.conf, then the legacy tree. A master-root install puts Studio at + # \studio, and that conf is how this script finds it with a bare + # environment. Read directly rather than through _CustomStudioRoots, which calls + # back into this function. + if ($env:LOCALAPPDATA) { + $confRoot = _RootFromConf (Join-Path $env:LOCALAPPDATA "Unsloth Studio\studio.conf") + if ($confRoot) { $noteRoots += $confRoot } + } + if ($env:USERPROFILE) { + $noteRoots += (Join-Path $env:USERPROFILE ".unsloth\studio") + } + } + foreach ($noteRoot in $noteRoots) { + $notePath = Join-Path $noteRoot "share\.unsloth-master-root" + if (-not (Test-Path -LiteralPath $notePath -PathType Leaf)) { continue } + try { + # One line, and REFUSED when there is a second: storage_roots.py and the CLI + # strip the WHOLE file, so a two-line note is not a directory to them and + # they decline it, while accepting line 1 here let a note no runtime reader + # honours authorise removing the master root's runtime children. Two lines + # are read so the second can be seen; a trailing newline is not one. + # -Encoding UTF8: 5.1 decodes BOM-less input with the ANSI code page, so a + # master root holding non-ASCII came back mangled and its runtime siblings + # were never found. + $noteLines = @(Get-Content -LiteralPath $notePath -TotalCount 2 -Encoding UTF8 -ErrorAction Stop) + } catch { continue } + if ($noteLines.Count -gt 1 -and -not [string]::IsNullOrWhiteSpace($noteLines[1])) { continue } + $line = if ($noteLines.Count -ge 1) { $noteLines[0] } else { $null } + if (-not [string]::IsNullOrWhiteSpace($line)) { + $raw = $line + $noteStudio = $noteRoot + break + } + } + } + if ([string]::IsNullOrWhiteSpace($raw)) { return $null } + $expanded = _ExpandTilde $raw.Trim() + $norm = $null + # The provider path first, exactly as setup.ps1's Get-CanonicalDir resolves it. + # [IO.Path]::GetFullPath anchors a RELATIVE root at [Environment]::CurrentDirectory, + # which PowerShell does not keep in step with its own location, so after a Set-Location + # a relative UNSLOTH_HOME named one install to setup and another here. Nothing downstream + # notices: the marker spares trees that are not Unsloth's, not another Unsloth install. + try { + $norm = $ExecutionContext.SessionState.Path.GetUnresolvedProviderPathFromPSPath($expanded) + } catch { $norm = $null } + try { $norm = [System.IO.Path]::GetFullPath($(if ($norm) { $norm } else { $expanded })).TrimEnd('\','/') } + catch { return $null } + if (-not $norm) { return $null } + # A note must describe the tree it was found in: copy a Studio tree from master root A to + # B and uninstall B, and the copied note still names A, whose runtimes carry the same + # owner markers -- so the gates below would delete the ORIGINAL install's. Containment, + # not an exact \studio match, since the flat layout is supported. An explicit + # UNSLOTH_HOME skips the check: that is the user speaking, not a file on disk. Mirrors + # uninstall.sh and storage_roots._recorded_master_root(). + if ($noteStudio) { + $here = $null + try { + $here = $ExecutionContext.SessionState.Path.GetUnresolvedProviderPathFromPSPath($noteStudio) + } catch { $here = $null } + try { $here = [System.IO.Path]::GetFullPath($(if ($here) { $here } else { $noteStudio })).TrimEnd('\','/') } + catch { $here = $null } + if (-not $here) { return $null } + # A note found in the ordinary %USERPROFILE%\.unsloth\studio install is declined + # outright, whatever it records, which is the rule storage_roots' + # _is_legacy_studio_tree and the CLI already apply. Containment alone is not enough: + # the profile directory CONTAINS the legacy tree, so a note naming it passed, + # _CustomStudioRoots then derived "\studio", and uninstalling the legacy + # install accepted a SECOND marked install there as a root to remove. The deny list + # below refuses the profile as a MASTER root, not the Studio candidate derived from + # it. Mirrors uninstall.sh. + if ($env:USERPROFILE) { + $legacyStudio = $null + try { + $legacyStudio = [System.IO.Path]::GetFullPath( + (Join-Path $env:USERPROFILE ".unsloth\studio")).TrimEnd('\','/') + } catch { $legacyStudio = $null } + if ($legacyStudio -and ($here -ieq $legacyStudio)) { return $null } + } + if (-not ($here -ieq $norm -or $here.StartsWith($norm + [System.IO.Path]::DirectorySeparatorChar, + [System.StringComparison]::OrdinalIgnoreCase))) { + return $null + } + } + # The default root is left to the blocks that own it, which remove it unconditionally. + if ($env:USERPROFILE -and ($norm -ieq (Join-Path $env:USERPROFILE ".unsloth").TrimEnd('\','/'))) { + return $null + } + return $norm + } + function _CustomStudioRoots { $seen = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase) $defaultRoot = $null @@ -586,10 +744,18 @@ Environment: } $envRoot = $null - if ($env:UNSLOTH_STUDIO_HOME) { - $envRoot = $env:UNSLOTH_STUDIO_HOME - } elseif ($env:STUDIO_HOME) { - $envRoot = $env:STUDIO_HOME + # IsNullOrWhiteSpace, not truthiness: storage_roots.studio_root() trims these, so a + # whitespace-only override is unset to every resolver, while a bare truthy test called it + # present and suppressed the master-root branch below, leaving \studio. + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { + $envRoot = $env:UNSLOTH_STUDIO_HOME.Trim() + } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { + $envRoot = $env:STUDIO_HOME.Trim() + } else { + # Last, as in storage_roots.studio_root(): UNSLOTH_HOME names the tree, and the two + # above name this exact directory, so either of them wins outright. + $master = _MasterRoot + if ($master) { $envRoot = (Join-Path $master "studio") } } if ($envRoot) { $expandedEnv = _ExpandTilde $envRoot @@ -918,7 +1084,23 @@ Environment: } # Also stop anything holding a handle on the exact paths we delete (llama-server, # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. - $stopRoots = @($ownedRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode, $defaultWhisperCpp) + @($defaultSdCppToStop | Where-Object { $_ }) + @($customSdCppToStop) + # Resolved here, before the stop pass: a loaded llama-server.exe under the master root keeps + # its own tree locked, and the removal below would exhaust its retries and leave it installed. + $masterRootToStop = _MasterRoot + $masterChildrenToStop = @() + if ($masterRootToStop -and -not (_IsUnsafeRoot $masterRootToStop)) { + foreach ($childName in @("llama.cpp", "node", "whisper.cpp", "stable-diffusion.cpp")) { + $childPath = Join-Path $masterRootToStop $childName + # Only the trees this uninstall is going to delete, so an unmarked neighbour's + # process is never killed. + if (Test-Path -LiteralPath (Join-Path $childPath ".unsloth-studio-owned") -PathType Leaf) { + $masterChildrenToStop += $childPath + } + } + } + # $ownedRoots, not $knownRoots: see the note where the two lists are built. The master + # children are already marker-gated above, so they carry their own proof of ownership. + $stopRoots = @($ownedRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode, $defaultWhisperCpp) + @($defaultSdCppToStop | Where-Object { $_ }) + @($customSdCppToStop) + @($masterChildrenToStop) # The reparse expansion turns one path into generic subdirectories of wherever it points # (node, bin, unsloth_studio), so it is gated for the same reason. _StopProcessesLockingRoots -Roots ($stopRoots + @(_ManagedPathsUnderReparseTargets $ownedRoots)) @@ -937,6 +1119,18 @@ Environment: _Substep "refusing to remove non-Unsloth path: $r" "Yellow" continue } + # The flat layout, both variables naming one directory: _RemoveRootRecordingDb takes a + # marked Studio root WHOLE, so without this the user-chosen master root went with the + # install, which is the hole in the rule every other child below is marker-gated for. + # Kept rather than pruned, since data left behind is recoverable. Mirrors uninstall.sh. + $flatMaster = $masterRootToStop + if ($flatMaster -and ($r.TrimEnd('\', '/') -ieq $flatMaster.TrimEnd('\', '/'))) { + _Substep "keeping $r`: UNSLOTH_HOME and the Studio root name the same directory," "Yellow" + _Substep "so removing it would take whatever else you keep there. Delete it by hand" "Yellow" + _Substep "once you have checked what is in it." "Yellow" + $script:RemoveFailed = $true + continue + } _RemoveRootRecordingDb $r # Native diffusion now installs UNDER the custom root, so the removal above already took # it. Older builds put it BESIDE the root at \stable-diffusion.cpp, which removing @@ -971,6 +1165,45 @@ Environment: # Unsloth's %TEMP% before the sweep ever looked at its owner.pid. $preservedTemp = @(_RemoveStudioPrivateTempTrees -Paths $privateTempDirs -PrimaryPath $primaryPrivateTemp) if ($defaultDataDir) { _RemoveDataDirKeepingWslIcon $defaultDataDir -Preserve $preservedTemp } + # The master root's own children, marker-gated rather than removed outright like the + # ~/.unsloth ones below: is a directory the user chose. The locks and .staging are + # ours by name (prebuilt_core.py) and carry no marker. Mirrors scripts/uninstall.sh. + # $masterRootToStop, not a fresh _MasterRoot: that resolver can read its answer from a note + # inside a Studio tree the loop above has already removed, after which a second call returns + # nothing and the runtime siblings are stranded. + $masterRoot = $masterRootToStop + if ($masterRoot -and (_IsUnsafeRoot $masterRoot)) { + _Substep "refusing to remove unsafe path: $masterRoot" "Yellow" + $masterRoot = $null + } + if ($masterRoot) { + foreach ($childName in @("llama.cpp", "node", "whisper.cpp", "stable-diffusion.cpp")) { + $childPath = Join-Path $masterRoot $childName + if ((Test-Path -LiteralPath $childPath) -and + -not (Test-Path -LiteralPath (Join-Path $childPath ".unsloth-studio-owned") -PathType Leaf)) { + _Substep "keeping $childName without Unsloth owner marker: $childPath" "Yellow" + } else { + _RemovePath $childPath + } + } + foreach ($lockName in @(".llama.cpp.install.lock", ".node.install.lock", + ".whisper.cpp.install.lock", ".sd.cpp.install.lock")) { + _RemoveLockFile (Join-Path $masterRoot $lockName) + } + # The prebuilt installers SHARE \.staging and prune it only when empty, so + # anything left in it here is not ours: remove the directory only, never its contents. + _RemoveDirIfEmpty (Join-Path $masterRoot ".staging") + if (Test-Path -LiteralPath $masterRoot) { + # $script:StaleLockPattern, not a glob: the shape has to be the installer's own, + # name and numeric pid both, or a user's file in their own root is taken. + foreach ($stale in @(Get-ChildItem -LiteralPath $masterRoot -Force -ErrorAction SilentlyContinue | + Where-Object { $_.Name -match $script:StaleLockPattern })) { + _RemoveLockFile $stale.FullName + } + } + # Only when nothing of the user's is left. + _RemoveDirIfEmpty $masterRoot + } # Shared llama.cpp build + cache, siblings of studio under ~/.unsloth in default mode. if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp } # "stable-diffusion.cpp" is exactly what a git clone of leejet/stable-diffusion.cpp produces @@ -992,15 +1225,16 @@ Environment: # (prebuilt_core.py), and a stray lock keeps ~/.unsloth from being pruned below. if ($defaultUnslothHome) { foreach ($lockName in @(".llama.cpp.install.lock", ".node.install.lock", ".whisper.cpp.install.lock")) { - _RemovePath (Join-Path $defaultUnslothHome $lockName) + _RemoveLockFile (Join-Path $defaultUnslothHome $lockName) } # Taking over an abandoned lock renames it to .stale. before unlinking; a crash # between the two strands the rename, so sweep any leftovers. -Force sees dotted names. if (Test-Path -LiteralPath $defaultUnslothHome) { - # -like, not -Filter: the Win32 filter is unreliable for names with several dots. + # -match, not -Filter: the Win32 filter is unreliable for names with several dots, + # and only the installer's exact shape is ours to remove. foreach ($stale in @(Get-ChildItem -LiteralPath $defaultUnslothHome -Force -ErrorAction SilentlyContinue | - Where-Object { $_.Name -like "*.install.lock.stale.*" })) { - _RemovePath $stale.FullName + Where-Object { $_.Name -match $script:StaleLockPattern })) { + _RemoveLockFile $stale.FullName } } } @@ -1105,6 +1339,27 @@ Environment: } catch { _Substep "could not update user PATH: $($_.Exception.Message)" "Yellow" } + # Clear the persisted Inductor cache path when it still names a tree this run deleted: + # setup.ps1 writes TORCHINDUCTOR_CACHE_DIR to the USER environment, so every later PyTorch + # process on this account inherits it and compiles into the removed directory. Only a value + # inside a root this run owned; a user's own directory and the shared C:\tc fallback stay. + try { + $persistedCache = [Environment]::GetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', 'User') + if (-not [string]::IsNullOrWhiteSpace($persistedCache)) { + $expandedCache = [Environment]::ExpandEnvironmentVariables($persistedCache).TrimEnd('\', '/') + foreach ($r in $ownedRoots) { + if (-not $r) { continue } + $rNorm = $r.TrimEnd('\', '/') + if ($expandedCache -ieq $rNorm -or $expandedCache -ilike "$rNorm\*") { + [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', [NullString]::Value, 'User') + _Substep "cleared TORCHINDUCTOR_CACHE_DIR: $persistedCache" "Green" + break + } + } + } + } catch { + _Substep "could not clear TORCHINDUCTOR_CACHE_DIR: $($_.Exception.Message)" "Yellow" + } # Remove HKCU\Software\Unsloth (PathBackup lives here; install.ps1 owns it). try { Remove-Item -LiteralPath 'HKCU:\Software\Unsloth' -Recurse -Force -ErrorAction SilentlyContinue diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index 93e5cb0fec..bf391ee9ec 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -341,6 +341,22 @@ _remove_path() { fi } +# An install lock, and only an install lock. prebuilt_core.install_lock creates these with +# O_CREAT | O_EXCL, so a lock is always a regular file: a directory or a link at that name is +# the user's, and in a master root they chose _remove_path would rm -rf it. -L before -f, which +# is true for a link to a file. uninstall.ps1's _RemoveLockFile excludes reparse points for the +# same reason, and the two halves of one rule have to agree. +_remove_lock_file() { + _rlf="$1" + if [ -L "$_rlf" ]; then + echo " keeping link at an install-lock path: $_rlf" >&2 + elif [ -f "$_rlf" ]; then + _remove_path "$_rlf" + elif [ -e "$_rlf" ]; then + echo " keeping non-file at an install-lock path: $_rlf" >&2 + fi +} + # $1 override, $2 default. A relative override is invalid per XDG and dropped by dirs (which # Tauri resolves through), so honouring one would spare the real data and rm -rf under our cwd. _xdg_dir() { @@ -462,6 +478,112 @@ _custom_studio_data_dirs() { # 3. Env-mode studio.conf at $/share/studio.conf (discovered via 1) # install.sh writes UNSLOTH_EXE='/unsloth_studio/bin/unsloth', so the install root is three # dirnames up. Each discovered non-default root is printed on its own line, de-duplicated. +# The master root UNSLOTH_HOME names, or empty. studio/ is its child and llama.cpp, node and +# whisper.cpp are its other children, so removing the Studio root alone strands them. Stripped +# and tilde-expanded like storage_roots.unsloth_home() and studio/setup.sh, or a padded value +# would name a directory neither install nor uninstall agrees on. +_master_root() { + _mr_from_note= + _mr=$(printf '%s' "${UNSLOTH_HOME:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + # The note setup.sh leaves in the Studio tree, when this run has no UNSLOTH_HOME of its own: + # `UNSLOTH_HOME=/mnt/portable unsloth studio update` leaves nothing in the environment, so + # without it an uninstall stranded the runtimes. The deny list and the marker gate still + # apply, so a stale note licenses nothing the environment could not. + # + # Only the Studio root this run is about, in studio_root()'s precedence and with no falling + # through: with two installs on one box, a legacy note otherwise won here while the removal + # worked on the named tree, so the run deleted one Studio and the OTHER install's runtimes. + if [ -z "$_mr" ]; then + # Trimmed as setup.sh trimmed them when it chose where to write the note: appending to + # a padded UNSLOTH_STUDIO_HOME looked somewhere that does not exist, so the note was + # never found while the loop below still removed the tree. + _mr_ush=$(printf '%s' "${UNSLOTH_STUDIO_HOME:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + _mr_sh=$(printf '%s' "${STUDIO_HOME:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + if [ -n "$_mr_ush" ]; then + _mr_roots="$_mr_ush" + elif [ -n "$_mr_sh" ]; then + _mr_roots="$_mr_sh" + else + # No override, so the root install.sh recorded in the default-mode studio.conf, then + # the legacy tree: with a bare environment that conf is the only way to find a + # /studio at all. Read here, not through _custom_studio_roots, which calls + # back into this function. UNSLOTH_EXE='/unsloth_studio/bin/unsloth', so the + # root is three dirnames up; apostrophes arrive escaped as '\''. + _mr_roots="" + _mr_dconf="$HOME/.local/share/unsloth/studio.conf" + if [ -f "$_mr_dconf" ]; then + _mr_exe=$(sed -n "s/^UNSLOTH_EXE='\(.*\)'\$/\1/p" "$_mr_dconf" | head -n1) + _mr_exe=$(printf '%s' "$_mr_exe" | sed "s/'\\\\''/'/g") + [ -n "$_mr_exe" ] && _mr_roots=$(dirname "$(dirname "$(dirname "$_mr_exe")")") + fi + _mr_roots="$_mr_roots +$HOME/.unsloth/studio" + fi + _mr_saved_ifs=$IFS + IFS=' +' + for _mr_studio in $_mr_roots; do + IFS=$_mr_saved_ifs + [ -n "$_mr_studio" ] || continue + _mr_conf="${_mr_studio}/share/.unsloth-master-root" + [ -f "$_mr_conf" ] || continue + # One line, and REFUSED when there is a second. storage_roots.py and the CLI both + # strip the WHOLE file, so a two-line note is not a directory to them and they + # decline it, while taking line 1 here authorised removing /llama.cpp, node, + # whisper.cpp and stable-diffusion.cpp: a note no runtime reader honours must not + # license a delete. A trailing newline is not a second line. + if [ -n "$(sed -n '2,$p' "$_mr_conf" 2>/dev/null | tr -d '[:space:]')" ]; then + continue + fi + _mr=$(head -n 1 "$_mr_conf" 2>/dev/null | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') \ + || _mr="" + if [ -n "$_mr" ]; then + _mr_from_note=1 + break + fi + done + IFS=$_mr_saved_ifs + fi + [ -n "$_mr" ] || return 0 + # shellcheck disable=SC2088 + case "$_mr" in + "~") _mr="$HOME" ;; + "~/"*) _mr="$HOME/${_mr#'~/'}" ;; + esac + # shellcheck disable=SC1007 + # || _mr_canon="": a bare assignment takes the substitution's exit status under set -e, and + # a root already gone or on a disconnected drive is the ordinary case here, which killed the + # whole run. bash hides this by clearing errexit inside command substitution; dash and POSIX + # `sh`, which the advertised `| sh` one-liner uses, do not. + _mr_canon=$(CDPATH= cd -P -- "$_mr" 2>/dev/null && pwd -P) || _mr_canon="" + [ -n "$_mr_canon" ] && _mr="$_mr_canon" + # A note must describe the tree it was found in: copy a Studio tree from master root A to B + # and uninstall B, and the copied note still names A, whose runtimes carry the same owner + # markers -- so the gates below would delete the ORIGINAL install's. Containment, not an + # exact /studio match, since the flat layout is supported. An explicit UNSLOTH_HOME + # skips all of this: that is the user speaking, not a file on disk. + if [ "${_mr_from_note:-}" = 1 ]; then + _mr_here=$(CDPATH= cd -P -- "$_mr_studio" 2>/dev/null && pwd -P) || _mr_here="" + # A note found in the ordinary ~/.unsloth/studio install is declined outright, whatever + # it records, which is the rule storage_roots._is_legacy_studio_tree and the CLI already + # apply. Containment alone is not enough here: $HOME contains the legacy tree, so a note + # naming it passed, _custom_studio_roots then emitted "$HOME/studio", and uninstalling + # the legacy install accepted a SECOND marked install there as a root to remove, taking + # its database and outputs with it. The deny list does not catch that -- it refuses + # $HOME as a master root, not the separate Studio candidate derived from it. + _mr_legacy_studio=$(CDPATH= cd -P -- "$HOME/.unsloth/studio" 2>/dev/null && pwd -P) \ + || _mr_legacy_studio="$HOME/.unsloth/studio" + [ "$_mr_here" != "$_mr_legacy_studio" ] || return 0 + case "$_mr_here" in + "$_mr") : ;; + "$_mr"/*) : ;; + *) return 0 ;; + esac + fi + case "$_mr" in "$HOME/.unsloth"|/|"") return 0 ;; esac + printf '%s\n' "$_mr" +} + _custom_studio_roots() { # $1 = "lexical": skip the canonicalization (see the legacy sd.cpp sibling below). Reset on # every call, so a plain call is never affected by a preceding lexical one. @@ -500,12 +622,22 @@ _custom_studio_roots() { } # Mirror install.sh's precedence: UNSLOTH_STUDIO_HOME wins, STUDIO_HOME is ignored when both # are set, or uninstalling install A could also delete install B from a leftover STUDIO_HOME. - if [ -n "${UNSLOTH_STUDIO_HOME:-}" ]; then - _emit "$UNSLOTH_STUDIO_HOME" - _from_conf "$UNSLOTH_STUDIO_HOME/share/studio.conf" - elif [ -n "${STUDIO_HOME:-}" ]; then - _emit "$STUDIO_HOME" - _from_conf "$STUDIO_HOME/share/studio.conf" + # Trimmed, as storage_roots.studio_root() trims them: a whitespace-only override is unset to + # every resolver, but a bare -n test called it present and suppressed the master-root branch + # below, leaving /studio installed. + _ush=$(printf '%s' "${UNSLOTH_STUDIO_HOME:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + _sh=$(printf '%s' "${STUDIO_HOME:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + if [ -n "$_ush" ]; then + _emit "$_ush" + _from_conf "$_ush/share/studio.conf" + elif [ -n "$_sh" ]; then + _emit "$_sh" + _from_conf "$_sh/share/studio.conf" + elif [ -n "$(_master_root)" ]; then + # Last, as in storage_roots.studio_root(): UNSLOTH_HOME names the tree, and the two + # above name this exact directory, so either of them wins outright. + _emit "$(_master_root)/studio" + _from_conf "$(_master_root)/studio/share/studio.conf" fi # Default-mode conf. _from_conf "$HOME/.local/share/unsloth/studio.conf" @@ -593,6 +725,9 @@ _unsloth_uninstall_main() { _pkill_studio echo "Removing data and install directories..." + # Resolved ONCE, before anything is deleted: _master_root can read its answer from a note + # inside a Studio tree the loop below removes, after which the runtime siblings are stranded. + _MASTER_ROOT_SAVED="$(_master_root)" _custom_studio_roots | while IFS= read -r _custom_root; do [ -n "$_custom_root" ] || continue if _is_unsafe_root "$_custom_root"; then @@ -607,6 +742,23 @@ _unsloth_uninstall_main() { echo " refusing to remove non-Unsloth path: $_custom_root" >&2 continue fi + # The flat layout, both variables naming one directory: _remove_root_recording_db takes + # a marked Studio root WHOLE, so without this the user-chosen master root went with the + # install, which is the hole in the rule that every other child below is marker-gated + # for. Kept rather than pruned, since data left behind is recoverable and printed. + # Canonicalised on both sides, or a symlinked path compares unequal to itself. + _crf_canon=$(CDPATH= cd -P -- "$_custom_root" 2>/dev/null && pwd -P) || _crf_canon="" + [ -n "$_crf_canon" ] || _crf_canon="$_custom_root" + _crf_master="$_MASTER_ROOT_SAVED" + if [ -n "$_crf_master" ] && [ "$_crf_canon" = "$_crf_master" ]; then + echo " keeping $_custom_root: UNSLOTH_HOME and the Studio root name the same" >&2 + echo " directory, so removing it would take whatever else you keep there." >&2 + echo " Delete it by hand once you have checked what is in it." >&2 + _set_marker "$_REMOVE_FAILED_FLAG" + unset _crf_canon _crf_master + continue + fi + unset _crf_canon _crf_master _remove_root_recording_db "$_custom_root" # Native diffusion now installs UNDER the custom root, so the removal above already took # it. Older builds put it BESIDE the root at /stable-diffusion.cpp, which removing @@ -645,6 +797,52 @@ _unsloth_uninstall_main() { _remove_path "$_lex_sd_cpp" fi done + # The master root's own children, marker-gated and deny-listed rather than removed outright + # like the ~/.unsloth ones below: is a directory the user chose. The locks and + # .staging are ours by name (prebuilt_core.py) and carry no marker. + _mr_root="$_MASTER_ROOT_SAVED" + if [ -n "$_mr_root" ]; then + if _is_unsafe_root "$_mr_root"; then + echo " refusing to remove unsafe path: $_mr_root" >&2 + else + for _mr_child in llama.cpp node whisper.cpp stable-diffusion.cpp; do + _mr_path="$_mr_root/$_mr_child" + if _is_unsafe_root "$_mr_path"; then + echo " refusing to remove unsafe path: $_mr_path" >&2 + # -L as well as -e: -e is false for a dangling symlink, so one named llama.cpp + # with its target volume unmounted fell through to _remove_path, which treats + # -L as present and unlinks it. In a root the user chose, that link is theirs. + elif { [ -e "$_mr_path" ] || [ -L "$_mr_path" ]; } \ + && [ ! -f "$_mr_path/.unsloth-studio-owned" ]; then + echo " keeping $_mr_child without Unsloth owner marker: $_mr_path" >&2 + else + _remove_path "$_mr_path" + fi + done + for _mr_lock in .llama.cpp.install.lock .node.install.lock \ + .whisper.cpp.install.lock .sd.cpp.install.lock; do + _remove_lock_file "$_mr_root/$_mr_lock" + done + # The prebuilt installers SHARE /.staging and prune it only when empty, so + # anything left in it here is not ours. rmdir, not _remove_path: in a user-chosen + # root a recursive delete would take files an install was content to leave. + rmdir "$_mr_root/.staging" 2>/dev/null || true + # The exact shape prebuilt_core.py leaves: a component lock name, ".stale.", and the + # pid. A bare .*.install.lock.stale.* also matched .backup.install.lock.stale.copy. + for _mr_lock in .llama.cpp.install.lock .node.install.lock \ + .whisper.cpp.install.lock .sd.cpp.install.lock; do + for _mr_stale in "$_mr_root/$_mr_lock".stale.*; do + case "${_mr_stale##*.stale.}" in + ''|*[!0-9]*) continue ;; + esac + { [ -e "$_mr_stale" ] || [ -L "$_mr_stale" ]; } && _remove_lock_file "$_mr_stale" + done + done + # Only when nothing of the user's is left; rmdir refuses a non-empty directory. + rmdir "$_mr_root" 2>/dev/null || true + fi + fi + # end master-root children # Same gate as a custom root: an ungated run takes a hand-made ~/.unsloth/studio, and then # ~/.unsloth via the empty-dir prune below. # -e OR -L: -e follows a link and misses a dangling one, which _remove_path would still @@ -685,14 +883,23 @@ _unsloth_uninstall_main() { _remove_path "$HOME/.unsloth/whisper.cpp" # Prebuilt install locks: every prebuilt serializes on /..install.lock # (prebuilt_core.py), and a stray lock keeps ~/.unsloth from being pruned below. - _remove_path "$HOME/.unsloth/.llama.cpp.install.lock" - _remove_path "$HOME/.unsloth/.node.install.lock" - _remove_path "$HOME/.unsloth/.whisper.cpp.install.lock" + _remove_lock_file "$HOME/.unsloth/.llama.cpp.install.lock" + _remove_lock_file "$HOME/.unsloth/.node.install.lock" + _remove_lock_file "$HOME/.unsloth/.whisper.cpp.install.lock" # Taking over an abandoned lock renames it to .stale. before unlinking # (install_node_prebuilt.py); a crash between the two strands the rename, and a stranded one # blocks the rmdir below. Unmatched globs stay literal, hence the existence test. - for _stale in "$HOME"/.unsloth/.*.install.lock.stale.*; do - [ -e "$_stale" ] && _remove_path "$_stale" + # Same shape restriction as the master-root sweep above: the help text promises that + # anything else kept under ~/.unsloth is left in place, and .backup.install.lock.stale.copy + # is dotted too. + for _lock in .llama.cpp.install.lock .node.install.lock \ + .whisper.cpp.install.lock .sd.cpp.install.lock; do + for _stale in "$HOME/.unsloth/$_lock".stale.*; do + case "${_stale##*.stale.}" in + ''|*[!0-9]*) continue ;; + esac + { [ -e "$_stale" ] || [ -L "$_stale" ]; } && _remove_lock_file "$_stale" + done done # ROCm-on-WSL helper artifacts (librocdxg clone, smoke-test venv); removing them frees the rmdir. _remove_path "$HOME/.unsloth/librocdxg" diff --git a/studio/backend/core/inference/diffusion_compile_cache.py b/studio/backend/core/inference/diffusion_compile_cache.py index a5a416882f..b10bcae615 100644 --- a/studio/backend/core/inference/diffusion_compile_cache.py +++ b/studio/backend/core/inference/diffusion_compile_cache.py @@ -57,6 +57,7 @@ import dataclasses import hashlib import json import os +import shutil import tempfile import threading import time @@ -75,7 +76,7 @@ _ENV_DIR = "UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR" _ENV_SAVE = "UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE" _ENV_SYNC = "UNSLOTH_DIFFUSION_COMPILE_CACHE_SYNC" -_DEFAULT_ROOT = Path.home() / ".cache" / "unsloth" / "diffusion_compile_cache" +_LEGACY_ROOT = Path.home() / ".cache" / "unsloth" / "diffusion_compile_cache" _MANIFEST_NAME = "manifest.json" # The pre-content-addressing bundle name. Still read (a manifest without a "bundle" key names it, which is every @@ -133,14 +134,61 @@ def _save_enabled(mode: str) -> bool: return (os.environ.get(_ENV_SAVE) or "").strip().lower() not in ("0", "off", "false", "no") +def _portable_mode() -> bool: + try: + from utils.paths.storage_roots import portable_mode + except ImportError: + return False + try: + return bool(portable_mode()) + except Exception: # noqa: BLE001 - never fail a cache lookup over this + return False + + +def _default_root() -> Path: + """Resolved per call, not a module constant: this module is imported before startup sets + UNSLOTH_STUDIO_HOME, which storage_roots reads.""" + try: + from utils.paths.storage_roots import cache_root as studio_cache_root + except ImportError: + return _LEGACY_ROOT + return studio_cache_root() / "diffusion_compile_cache" + + def sync_saves() -> bool: """Whether ``save_async`` must write inline instead of handing off to the worker.""" return (os.environ.get(_ENV_SYNC) or "").strip().lower() in ("1", "on", "true", "yes") def cache_root() -> Path: + """The root every bundle, manifest and inductor artifact is WRITTEN under.""" root = os.environ.get(_ENV_DIR) - return Path(root) if root else _DEFAULT_ROOT + if root: + return Path(root) + return _default_root() + + +def legacy_cache_root() -> Optional[Path]: + """The pre-relocation root, when it is still worth READING old bundles from. + + Read-only on purpose: returning it as the write root would pin an upgraded install to the home + directory forever. Skipped under an explicit dir override (it names one exact directory) and + in portable mode (the host's home is not part of the install). + """ + if os.environ.get(_ENV_DIR) or _portable_mode(): + return None + # Equal when storage_roots is unavailable and the default IS the legacy root. + if _LEGACY_ROOT == _default_root(): + return None + # This root is the HOST's home, not the install, so it is the one that may be on a mount the + # new cache does not depend on. Path.exists raises for EACCES and EIO before 3.14, and an + # optional migration source we cannot inspect is a miss, never a failed generation. + try: + if not _LEGACY_ROOT.exists(): + return None + except OSError: + return None + return _LEGACY_ROOT def _triton_version() -> Optional[str]: @@ -339,17 +387,70 @@ def begin( _collect_superseded(cdir, logger) # Try an exact-match load. A miss/mismatch is normal and non-fatal. - if published is not None and ctx.bundle.exists(): + # Guarded like _load_from_legacy's probe: Path.exists() raises rather than returning False when a parent denies + # traversal, and this branch moves the write root to a directory the process may not own. Unguarded, that + # exception left begin() before the legacy fallback below could run. + try: + pair_present = published is not None and ctx.bundle.exists() + except OSError as exc: + _warn(logger, f"compile-cache: cannot check the bundle at {ctx.dir}: {exc}") + pair_present = False + if pair_present: ctx.hit = _try_load(ctx, logger) - if ctx.hit and mode != "on": - # Loaded artifacts == on-disk artifacts, so nothing to save. A new static-compile shape re-dirties via - # register_shape; mode "on" keeps saving. - ctx.saved = True - else: + if not ctx.hit: + # An install that predates the relocation may still hold this key under the old root. + ctx.hit = _load_from_legacy(ctx, logger) + if not ctx.hit: _info(logger, f"compile-cache: no bundle for key {key} (will compile locally)") + elif mode != "on": + # Loaded artifacts == on-disk artifacts, so nothing to save. A new static-compile shape + # re-dirties via register_shape; mode "on" keeps saving. + ctx.saved = True return ctx +def _load_from_legacy(ctx: CacheContext, logger: Any) -> bool: + """Load the same key's bundle from the pre-relocation root, then migrate it. + + The key already covers every portability dimension, so a legacy bundle under it is the same + artifact this run would have written. The copy stops the read fallback becoming permanent. + Best-effort, and skipped when saving is off, since that mode promises a read-only cache.""" + root = legacy_cache_root() + if root is None: + return False + ldir = root / ctx.key + manifest_path = ldir / _MANIFEST_NAME + # The manifest is the commit point here too, so it decides which legacy bundle is live; a + # pre-content-addressing one names cache.bin. + published = _read_manifest(manifest_path) + if published is None: + return False + bundle = _manifest_bundle(ldir, published) + try: + if not bundle.exists(): + return False + except OSError: + # Same reason as legacy_cache_root: a pair we cannot even stat is a miss. + return False + if not _try_load(ctx, logger, bundle = bundle, manifest_path = manifest_path): + return False + if _save_enabled(ctx.mode): + try: + ctx.dir.mkdir(parents = True, exist_ok = True) + # Under the name the copied manifest names, which is not ctx.bundle unless the legacy + # pair predates content addressing. Published like begin() does, so the context goes + # on naming the live bundle in the write root. + migrated = _manifest_bundle(ctx.dir, published) + _atomic_copy(bundle, migrated) + # Bundle first: a manifest without one reads as a miss, not an unservable hit. + _atomic_copy(manifest_path, ctx.manifest_path) + ctx.bundle = migrated + _info(logger, f"compile-cache: migrated legacy bundle for key {ctx.key}") + except OSError as exc: + _warn(logger, f"compile-cache: could not migrate legacy bundle: {exc}") + return True + + def register_shape(ctx: Optional[CacheContext], shape: Any, *, static: bool) -> None: """Record a generation's (width, height, batch) against the bundle coverage. @@ -372,20 +473,37 @@ def register_shape(ctx: Optional[CacheContext], shape: Any, *, static: bool) -> pass -def _try_load(ctx: CacheContext, logger: Any) -> bool: +def _try_load( + ctx: CacheContext, + logger: Any, + *, + bundle: Optional[Path] = None, + manifest_path: Optional[Path] = None, +) -> bool: + """Validate and load a bundle pair, defaulting to the context's own (write-root) one.""" + bundle = bundle if bundle is not None else ctx.bundle + manifest_path = manifest_path if manifest_path is not None else ctx.manifest_path try: - manifest = json.loads(ctx.manifest_path.read_text(encoding = "utf-8")) + manifest = json.loads(manifest_path.read_text(encoding = "utf-8")) except Exception as exc: # noqa: BLE001 _warn(logger, f"compile-cache: unreadable manifest: {exc}") return False + # A manifest that decoded but is not an object. json.loads happily returns [] or null, and + # the .get() below would then raise AttributeError out of a function whose whole contract is + # that a bad cache entry is a miss. The legacy root makes this reachable: the bundle being + # validated was written by an older build, on a disk this run has never checked. + if not isinstance(manifest, dict): + _warn(logger, "compile-cache: manifest is not an object; ignoring") + return False + # Exact-match guard (defence in depth: torch also validates internally on load). if manifest.get("env") != ctx.env_fp or manifest.get("model") != ctx.model_fp: _warn(logger, "compile-cache: fingerprint mismatch; falling back to local compile") return False try: - data = ctx.bundle.read_bytes() + data = bundle.read_bytes() except Exception as exc: # noqa: BLE001 _warn(logger, f"compile-cache: cannot read bundle: {exc}") return False @@ -443,6 +561,30 @@ def _atomic_write(path: Path, data: bytes) -> None: pass +def _atomic_copy(src: Path, dst: Path) -> None: + """``shutil.copyfile`` into a temp file in *dst*'s directory, then ``os.replace``. + + Same rule and same reason as _atomic_write: a plain copyfile onto the live name is visible + while it is still partial, and the migration below runs on the same interruptible path as a + save. A torn manifest costs only a miss, but a miss here means the cold compile the migration + exists to avoid, and two backends migrating one key would otherwise interleave their writes + into a single destination file. + """ + tmp: Optional[str] = None + try: + fd, tmp = tempfile.mkstemp(dir = str(dst.parent), prefix = f".{dst.name}.", suffix = _TEMP_SUFFIX) + os.close(fd) + shutil.copyfile(src, tmp) + os.replace(tmp, dst) + tmp = None + finally: + if tmp is not None: + try: + os.unlink(tmp) + except OSError: + pass + + def _collect_superseded(cdir: Path, logger: Any) -> list[str]: """Delete bundles no manifest names any more. Returns the names removed. Never raises. diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e7f8a4ff29..5983cdc0a3 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8478,8 +8478,8 @@ class LlamaCppBackend: @staticmethod def _resolved_studio_root_and_is_legacy() -> "tuple[Optional[Path], bool]": - """Resolve the Unsloth install root and classify it as the legacy - ~/.unsloth/studio root vs. a custom (env/venv-inferred) root. + """Resolve the directory llama.cpp hangs off and classify it as the + legacy ~/.unsloth/studio root vs. a custom (env/venv-inferred) root. Returns (resolved_root, is_legacy). On any import/resolution failure the root is treated as legacy and resolved_root is None -- callers must read @@ -8488,8 +8488,16 @@ class LlamaCppBackend: (cleanup) so the two never disagree on which root is legacy. """ try: - from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 + from utils.paths.storage_roots import ( # noqa: WPS433 + studio_root as _sr, + unsloth_home as _uh, + ) + # llama.cpp is a sibling of studio/ under the master root, the path run.py and + # main.py export, so a portable install is never the legacy layout. + master = _uh() + if master is not None: + return master, False resolved = _sr() legacy_studio = Path.home() / ".unsloth" / "studio" try: diff --git a/studio/backend/core/inference/stt_ggml_sidecar.py b/studio/backend/core/inference/stt_ggml_sidecar.py index 960f64115f..7c3be28955 100644 --- a/studio/backend/core/inference/stt_ggml_sidecar.py +++ b/studio/backend/core/inference/stt_ggml_sidecar.py @@ -120,15 +120,20 @@ def resolve_ggml_model_id(model: Optional[str]) -> str: def _managed_whisper_cpp_dir() -> Path: - """`/whisper.cpp` in custom mode, else `~/.unsloth/whisper.cpp`. + """`/whisper.cpp` when set, else `/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 + from utils.paths.storage_roots import studio_root, unsloth_home + # setup.sh installs whisper.cpp at , beside studio/. No env var bridges it. + master = unsloth_home() + if master is not None: + return master / "whisper.cpp" resolved = studio_root() legacy_studio = Path.home() / ".unsloth" / "studio" try: diff --git a/studio/backend/main.py b/studio/backend/main.py index ef471ebf30..794b66949a 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -230,10 +230,18 @@ try: _STUDIO_ROOT_RESOLVED = _studio_root().resolve() except (OSError, ValueError): _STUDIO_ROOT_RESOLVED = _studio_root() -if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: +from utils.paths.storage_roots import unsloth_home as _unsloth_home + +_MASTER_ROOT = _unsloth_home() +# A master root pointed at the legacy path still owns runtimes beside it, so the equality alone +# would skip the export and leave unsloth_zoo on ~/.unsloth/llama.cpp. +if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT or _MASTER_ROOT is not None: if not os.environ.get("UNSLOTH_STUDIO_HOME"): os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED) - _MANAGED_LLAMA_CPP_PATH = _STUDIO_ROOT_RESOLVED / "llama.cpp" + # Same derivation as run.py: the runtimes sit at the master root, beside studio/, and marking + # the wrong directory managed would make the bundled path look like an immutable override. + _MANAGED_ROOT = _MASTER_ROOT or _STUDIO_ROOT_RESOLVED + _MANAGED_LLAMA_CPP_PATH = _MANAGED_ROOT / "llama.cpp" if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_MANAGED_LLAMA_CPP_PATH) # A CLI/desktop launcher may already have exported Unsloth's own install path. diff --git a/studio/backend/run.py b/studio/backend/run.py index 4526304b6a..d3da78db36 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1241,10 +1241,18 @@ try: _STUDIO_ROOT_RESOLVED = _studio_root().resolve() except (OSError, ValueError): _STUDIO_ROOT_RESOLVED = _studio_root() -if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: +from utils.paths.storage_roots import unsloth_home as _unsloth_home + +_MASTER_ROOT = _unsloth_home() +# A master root pointed at the legacy path still owns runtimes beside it, so the equality alone +# would skip the export and leave unsloth_zoo on ~/.unsloth/llama.cpp. +if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT or _MASTER_ROOT is not None: if not os.environ.get("UNSLOTH_STUDIO_HOME"): os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED) - _MANAGED_LLAMA_CPP_PATH = _STUDIO_ROOT_RESOLVED / "llama.cpp" + # The runtimes sit at the master root, beside studio/; deriving from the Studio root would + # pin a path one level too deep for every worker. + _MANAGED_ROOT = _MASTER_ROOT or _STUDIO_ROOT_RESOLVED + _MANAGED_LLAMA_CPP_PATH = _MANAGED_ROOT / "llama.cpp" if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_MANAGED_LLAMA_CPP_PATH) # The CLI and generated launchers can export this path before run.py starts. diff --git a/studio/backend/tests/test_diffusion_compile_cache.py b/studio/backend/tests/test_diffusion_compile_cache.py index 6665490935..988ed675c3 100644 --- a/studio/backend/tests/test_diffusion_compile_cache.py +++ b/studio/backend/tests/test_diffusion_compile_cache.py @@ -14,6 +14,7 @@ from __future__ import annotations import hashlib import json import types +from pathlib import Path import pytest @@ -342,6 +343,227 @@ def test_restore_inductor_dir(monkeypatch, tmp_path, fake_megacache): assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == "/tmp/prior-inductor" # restored +# -------------------------------------------------------------------------- legacy root +def _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) -> tuple: + """Write a bundle under a fake pre-relocation root, then hand back an upgraded install: + ``(legacy_root, studio_home, legacy_bundle)``, with the environment already pointing at a + non-portable install whose new default root is empty. The bundle is content-addressed, so its + name comes from the seeding save rather than from a constant.""" + legacy = tmp_path / "legacy" / "diffusion_compile_cache" + studio_home = tmp_path / "studio" + monkeypatch.setenv(cc._ENV_MODE, "auto") + monkeypatch.delenv(cc._ENV_SAVE, raising = False) + monkeypatch.setenv(cc._ENV_DIR, str(legacy)) + seeded = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert cc.save(seeded) is True + + monkeypatch.delenv(cc._ENV_DIR, raising = False) + monkeypatch.delenv("UNSLOTH_HOME", raising = False) + monkeypatch.delenv("UNSLOTH_PORTABLE", raising = False) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.setattr(cc, "_LEGACY_ROOT", legacy) + return legacy, studio_home, seeded.bundle + + +def test_legacy_bundle_is_read_but_the_new_root_takes_the_writes( + monkeypatch, tmp_path, fake_megacache +): + # An upgraded install stays warm without the home directory becoming the write root: + # begin() points TORCHINDUCTOR_CACHE_DIR at ctx.dir and save() writes every later bundle. + legacy, studio_home, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + import os + + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + + assert ctx.hit is True # warm: the legacy bundle still counts + assert str(ctx.dir).startswith(str(studio_home)) + assert str(os.environ["TORCHINDUCTOR_CACHE_DIR"]).startswith(str(studio_home)) + assert legacy not in ctx.dir.parents + # The read fallback migrates: the pair now lives in the write root, byte-identical. + assert ctx.bundle.read_bytes() == legacy_bundle.read_bytes() + assert ctx.bundle.parent == ctx.dir + assert ctx.manifest_path.exists() + + +def test_an_interrupted_migration_leaves_no_partial_pair(monkeypatch, tmp_path, fake_megacache): + """The migration copies through a temp file and renames, as every other write here does. + + A plain copyfile onto the live name is visible while it is still partial. The manifest is the + commit point, so a torn one is read as a miss and costs the cold compile the migration exists + to avoid; a second backend migrating the same key would interleave its writes into the same + destination. Interruption is injected at the manifest copy, the later of the two. + """ + legacy, studio_home, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + + import shutil as _shutil + + real_copyfile = _shutil.copyfile + calls = {"n": 0} + + def exploding_copyfile(src, dst, *args, **kwargs): + calls["n"] += 1 + if calls["n"] == 2: + # Half a manifest on disk, then death, exactly as interpreter exit kills the thread. + Path(dst).write_bytes(Path(src).read_bytes()[:3]) + raise OSError("interrupted") + return real_copyfile(src, dst, *args, **kwargs) + + monkeypatch.setattr(cc.shutil, "copyfile", exploding_copyfile) + + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + + # The read still served this run from the legacy pair. + assert ctx.hit is True + # Nothing partial is published under a name a later run reads: the temp file the copy died + # inside is not the live manifest. + if ctx.manifest_path.exists(): + assert ctx.manifest_path.read_bytes() == (legacy / ctx.key / cc._MANIFEST_NAME).read_bytes() + leftovers = [q.name for q in ctx.dir.iterdir() if q.name.endswith(cc._TEMP_SUFFIX)] + assert leftovers == [], leftovers + + +def test_migrated_bundle_serves_the_next_run_without_the_legacy_root( + monkeypatch, tmp_path, fake_megacache +): + legacy, _, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + assert cc.begin(transformer = _transformer(), **_BEGIN_KW).hit is True + + import shutil + + shutil.rmtree(legacy) # the old cache gets cleaned up: the warm start must survive + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx.hit is True + + +def test_load_only_mode_reads_legacy_without_writing_to_it(monkeypatch, tmp_path, fake_megacache): + legacy, _, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + monkeypatch.setenv(cc._ENV_SAVE, "0") + + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + + assert ctx.hit is True + assert not ( + ctx.dir / legacy_bundle.name + ).exists() # a read-only cache stays read-only, both roots + assert legacy_bundle.exists() + + +def test_key_absent_from_legacy_never_writes_into_it(monkeypatch, tmp_path, fake_megacache): + # The old fallback swapped the WHOLE root, so even an unheld key wrote into the home dir. + legacy, studio_home, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + before = {p.name for p in legacy.iterdir()} + + other = dict(_BEGIN_KW, family = "qwen-image") + ctx = cc.begin(transformer = _transformer(("QwenImageTransformerBlock",)), **other) + + assert ctx.hit is False + assert cc.save(ctx) is True + assert str(ctx.bundle).startswith(str(studio_home)) + assert {p.name for p in legacy.iterdir()} == before + + +def test_portable_mode_never_falls_back_to_the_home_directory(monkeypatch, tmp_path): + # begin() points TORCHINDUCTOR_CACHE_DIR inside this root, so a fallback here would write + # GBs into the host machine's home directory. + monkeypatch.delenv(cc._ENV_DIR, raising = False) + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + monkeypatch.setenv("UNSLOTH_HOME", str(tmp_path / "portable")) + legacy = tmp_path / "legacy" / "diffusion_compile_cache" + legacy.mkdir(parents = True) + monkeypatch.setattr(cc, "_LEGACY_ROOT", legacy) + + root = cc.cache_root() + + assert root != legacy + assert str(root).startswith(str(tmp_path / "portable")) + assert cc.legacy_cache_root() is None # not even read: not part of the install + + +def test_explicit_dir_override_ignores_the_legacy_root(monkeypatch, tmp_path): + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path / "chosen")) + legacy = tmp_path / "legacy" / "diffusion_compile_cache" + legacy.mkdir(parents = True) + monkeypatch.setattr(cc, "_LEGACY_ROOT", legacy) + + assert cc.cache_root() == tmp_path / "chosen" + assert cc.legacy_cache_root() is None + + +def test_an_unreadable_legacy_root_is_a_miss_not_a_failure(monkeypatch, tmp_path): + # The legacy root is the HOST's home, so it can be on a mount the new cache does not need. + # Path.exists raises for EACCES and EIO before 3.14, and begin() does not catch around here. + monkeypatch.delenv(cc._ENV_DIR, raising = False) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + monkeypatch.delenv("UNSLOTH_HOME", raising = False) + monkeypatch.delenv("UNSLOTH_PORTABLE", raising = False) + legacy = tmp_path / "legacy" / "diffusion_compile_cache" + monkeypatch.setattr(cc, "_LEGACY_ROOT", legacy) + + def _raise(self, *a, **k): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr(Path, "exists", _raise) + + assert cc.legacy_cache_root() is None + + +def test_an_unreadable_legacy_bundle_pair_is_a_miss_not_a_failure(monkeypatch, tmp_path): + legacy = tmp_path / "legacy" / "diffusion_compile_cache" + legacy.mkdir(parents = True) + monkeypatch.setattr(cc, "legacy_cache_root", lambda: legacy) + ctx = cc.CacheContext( + key = "abc", + dir = tmp_path / "new" / "abc", + bundle = tmp_path / "new" / "abc" / "cache.bin", + manifest_path = tmp_path / "new" / "abc" / "manifest.json", + env_fp = "e", + model_fp = "m", + mode = "auto", + ) + + def _raise(self, *a, **k): + raise OSError(5, "Input/output error") + + monkeypatch.setattr(Path, "exists", _raise) + + assert cc._load_from_legacy(ctx, None) is False + + +def test_a_manifest_that_is_not_an_object_is_a_miss(monkeypatch, tmp_path, fake_megacache): + """json.loads returns [] for "[]" and None for "null", and .get() on either raises. + + _try_load's whole contract is that a bad cache entry is a miss, and the legacy root makes + this reachable in a way it was not before: the manifest being validated was written by an + older build, on a disk this run has never checked. + """ + legacy, _, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + for payload in ("[]", "null", '"a string"', "42"): + for key_dir in legacy.iterdir(): + (key_dir / cc._MANIFEST_NAME).write_text(payload, encoding = "utf-8") + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx.hit is False, payload + + +def test_an_unreadable_write_root_falls_back_to_legacy(monkeypatch, tmp_path, fake_megacache): + """Path.exists() raises rather than answering False when a parent denies traversal. + + This branch moves the write root, so it can land on a directory the process does not own on + some machine. Unguarded, that exception left begin() before the legacy fallback could run, + which is the fallback the relocation depends on for a warm start. + """ + legacy, studio_home, legacy_bundle = _seed_legacy_bundle(monkeypatch, tmp_path, fake_megacache) + real_exists = Path.exists + + def _raise_under_studio(self, *a, **k): + if str(self).startswith(str(studio_home)): + raise PermissionError(13, "Permission denied") + return real_exists(self, *a, **k) + + monkeypatch.setattr(Path, "exists", _raise_under_studio) + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx.hit is True + + # ------------------------------------------------------------------ background save worker # save_async hands the write to a single module-level daemon worker so a generation stops paying for it. These drive # the worker directly, gating the fake save_cache_artifacts on Events so ordering is asserted, never slept on. diff --git a/studio/backend/tests/test_hf_cache_settings.py b/studio/backend/tests/test_hf_cache_settings.py index 506b552a4c..2874e0f64a 100644 --- a/studio/backend/tests/test_hf_cache_settings.py +++ b/studio/backend/tests/test_hf_cache_settings.py @@ -3,7 +3,9 @@ from __future__ import annotations +import json import os +import subprocess import sys import threading import time @@ -315,3 +317,174 @@ def test_diffusion_loader_calls_pin_the_cache_dir(): assert ( "cache_dir" in window or "kwargs" in window ), f"{rel}:{index} calls {call} without a pinned cache_dir" + + +# _stored_cache_home skips the database read when nothing uses one, so `unsloth train` does not +# build a studio.db on a machine that never opened Studio. Only a positively observed absence may +# license that skip. Driven in a subprocess: the skip needs storage.studio_db out of sys.modules. +_GUARD_PROBE = """ +import json, os, sys + +sys.path.insert(0, os.environ["FAKE_STORAGE"]) +sys.path.insert(1, os.environ["BACKEND_DIR"]) +from utils import hf_cache_settings + +assert "storage.studio_db" not in sys.modules, "the probe must exercise the skip" +stored = hf_cache_settings._stored_cache_home() +print(json.dumps({"stored": None if stored is None else str(stored)})) +""" + +_FAKE_STUDIO_DB = """ +import os +from pathlib import Path + + +def get_app_setting(key, fallback = None): + Path(os.environ["READ_WITNESS"]).write_text("read", encoding = "utf-8") + if key == "hugging_face_cache_home": + return os.environ["STORED_CACHE_HOME"] + return fallback +""" + + +def _run_guard_probe( + tmp_path, + studio_home: Path, + stored: Path, + *, + probe: str | None = None, + extra_env: dict[str, str] | None = None, +) -> tuple[str | None, bool]: + """Return (_stored_cache_home() answer, whether the database was read).""" + fake = tmp_path / "fake_storage" + (fake / "storage").mkdir(parents = True, exist_ok = True) + (fake / "storage" / "__init__.py").write_text("", encoding = "utf-8") + (fake / "storage" / "studio_db.py").write_text(_FAKE_STUDIO_DB, encoding = "utf-8") + witness = tmp_path / "read_witness" + witness.unlink(missing_ok = True) + + environment = dict(os.environ) + environment.update( + BACKEND_DIR = _BACKEND_DIR, + FAKE_STORAGE = str(fake), + READ_WITNESS = str(witness), + STORED_CACHE_HOME = str(stored), + UNSLOTH_STUDIO_HOME = str(studio_home), + PYTHONPATH = "", + ) + environment.update(extra_env or {}) + for key in ("HF_HOME", "HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE", "HF_XET_CACHE"): + environment.pop(key, None) + result = subprocess.run( + [sys.executable, "-c", probe if probe is not None else _GUARD_PROBE], + capture_output = True, + text = True, + env = environment, + timeout = 120, + ) + assert result.returncode == 0, result.stderr + return json.loads(result.stdout.strip().splitlines()[-1])["stored"], witness.exists() + + +def test_absent_studio_db_skips_the_database_read(tmp_path): + # The skip 912024e84 added must survive the tightening below: no studio.db answers None + # WITHOUT a connection, which is what stops the CLI creating a 250 KB database. + studio_home = tmp_path / "root" / "studio" + studio_home.mkdir(parents = True) + + answer, was_read = _run_guard_probe(tmp_path, studio_home, tmp_path / "chosen") + + assert answer is None + assert not was_read, "an absent database must not be opened" + + +@pytest.mark.parametrize("fixture", ["not_a_directory", "symlink_loop", "unreadable_parent"]) +def test_uninspectable_studio_db_keeps_the_stored_cache_home(tmp_path, fixture): + # Path.exists reports ENOTDIR and ELOOP as absence on every release, and swallows EACCES + # from 3.14. Reading any as "no database" discards the cache home chosen in Settings. + chosen = tmp_path / "chosen" + studio_home = tmp_path / "root" / "studio" + if fixture == "not_a_directory": + studio_home.parent.mkdir(parents = True) + studio_home.write_text("", encoding = "utf-8") + elif fixture == "symlink_loop": + studio_home.parent.mkdir(parents = True) + studio_home.symlink_to(studio_home) + else: + studio_home.mkdir(parents = True) + (studio_home / "studio.db").write_bytes(b"") + os.chmod(studio_home, 0o000) + + try: + answer, was_read = _run_guard_probe(tmp_path, studio_home, chosen) + finally: + if fixture == "unreadable_parent": + os.chmod(studio_home, 0o755) + + # unreadable_parent is the case 3.14 newly breaks; the other two hold on every release. + assert was_read, "a database we could not inspect must still be read" + assert answer == str(chosen) + + +_WINDOWS_SHAPED_GUARD_PROBE = """ +import errno, json, os, sys + +# Windows reports a path whose PARENT is a file as ERROR_PATH_NOT_FOUND, which Python raises as +# FileNotFoundError; POSIX raises NotADirectoryError. Reproduce the error shape, since the Linux +# runners cannot produce it, and do it before hf_cache_settings is imported. +_NOT_A_DIR = os.environ["NOT_A_DIRECTORY"] +_real_stat = os.stat + + +def _stat(path, *args, **kwargs): + text = os.fspath(path) + if text != _NOT_A_DIR and text.startswith(_NOT_A_DIR + os.sep): + raise FileNotFoundError(errno.ENOENT, "The system cannot find the path specified", text) + return _real_stat(path, *args, **kwargs) + + +os.stat = _stat +sys.path.insert(0, os.environ["FAKE_STORAGE"]) +sys.path.insert(1, os.environ["BACKEND_DIR"]) +from utils import hf_cache_settings + +assert "storage.studio_db" not in sys.modules, "the probe must exercise the skip" +stored = hf_cache_settings._stored_cache_home() +print(json.dumps({"stored": None if stored is None else str(stored)})) +""" + + +def test_a_studio_home_that_is_a_file_still_reads_the_database_on_windows(tmp_path): + """The POSIX arm of this is test_uninspectable_studio_db_keeps_the_stored_cache_home + [not_a_directory], which passes there because ENOTDIR is its own exception type. Windows + reports the same situation as FileNotFoundError, so the skip read it as "no database stored" + and discarded the cache home chosen in Settings -- on Windows and nowhere else. Caught by the + cross-platform leg, held here by reproducing the error shape rather than the platform. + """ + studio_home = tmp_path / "root" / "studio" + studio_home.parent.mkdir(parents = True) + studio_home.write_text("", encoding = "utf-8") + chosen = tmp_path / "chosen" + + answer, was_read = _run_guard_probe( + tmp_path, + studio_home, + chosen, + probe = _WINDOWS_SHAPED_GUARD_PROBE, + extra_env = {"NOT_A_DIRECTORY": str(studio_home)}, + ) + + assert was_read, "a studio home we could not inspect was read as no database" + assert answer == str(chosen) + + +def test_a_studio_home_that_does_not_exist_still_skips_the_database_read(tmp_path): + """The other half. The fix above must not widen the skip into the case it exists for: a + machine that has never opened Studio has no studio/ directory at all, and falling through + there builds the 250 KB database this guard was added to avoid.""" + studio_home = tmp_path / "never-created" / "studio" + + answer, was_read = _run_guard_probe(tmp_path, studio_home, tmp_path / "chosen") + + assert answer is None + assert not was_read, "an absent studio home must not open a database" diff --git a/studio/backend/tests/test_setup_cache_env_containment.py b/studio/backend/tests/test_setup_cache_env_containment.py new file mode 100644 index 0000000000..3d0894d65c --- /dev/null +++ b/studio/backend/tests/test_setup_cache_env_containment.py @@ -0,0 +1,1446 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Everything Unsloth writes should sit under one directory (issue #8865).""" + +import contextlib +import errno +import importlib.util +import json +import os +import platform +import re +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_STORAGE_ROOTS_PATH = Path(__file__).resolve().parent.parent / "utils/paths/storage_roots.py" + +_ALWAYS_PINNED = ( + "UV_CACHE_DIR", + "VLLM_CACHE_ROOT", + "UNSLOTH_COMPILE_LOCATION", + "TORCHINDUCTOR_CACHE_DIR", + "TRITON_CACHE_DIR", + "TRITON_DUMP_DIR", + "TORCH_EXTENSIONS_DIR", + "CUDA_CACHE_PATH", + "MPLCONFIGDIR", + "NUMBA_CACHE_DIR", + "DATA_DESIGNER_HOME", + "DATA_DESIGNER_MANAGED_ASSETS_PATH", +) + +# Shared user data / large re-downloads: portable mode only. +_PORTABLE_ONLY = ("HF_DATASETS_CACHE", "HF_ASSETS_CACHE", "TORCH_HOME") + +_HF_ENV = ("HF_HOME", "HF_HUB_CACHE", "HF_XET_CACHE", "HUGGINGFACE_HUB_CACHE") + + +@pytest.fixture(autouse = True) +def _clean_env(monkeypatch, tmp_path): + for key in _ALWAYS_PINNED + _PORTABLE_ONLY + _HF_ENV: + monkeypatch.delenv(key, raising = False) + for key in ("UNSLOTH_HOME", "UNSLOTH_PORTABLE", "STUDIO_HOME"): + monkeypatch.delenv(key, raising = False) + # Not pinned, but both change what Triton resolves under the test home. + for key in ("TRITON_HOME", "TRITON_OVERRIDE_DIR"): + monkeypatch.delenv(key, raising = False) + # _default_cache_home reads this before ~/.cache, and CI runners set it. + monkeypatch.delenv("XDG_CACHE_HOME", raising = False) + # Same for matplotlib's config dir, whose contents decide whether MPLCONFIGDIR is ours. + monkeypatch.delenv("XDG_CONFIG_HOME", raising = False) + # Empty home: a real ~/.data-designer would change what the resolver pins. + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.setenv("USERPROFILE", str(tmp_path / "home")) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + + +@pytest.fixture(autouse = True) +def _restore_hf_cache_settings_module(): + # Left popped, the next import builds a second module object for a later test to read. + import utils + + name = "utils.hf_cache_settings" + saved = sys.modules.get(name) + saved_attr = getattr(utils, "hf_cache_settings", None) + try: + yield + finally: + if saved is not None: + sys.modules[name] = saved + else: + sys.modules.pop(name, None) + if saved_attr is not None: + utils.hf_cache_settings = saved_attr + else: + with contextlib.suppress(AttributeError): + del utils.hf_cache_settings + + +def _load_storage_roots(): + # The resolver snapshots explicit env vars 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) + return module + + +def test_regenerable_caches_are_pinned_under_the_studio_root(tmp_path): + sr = _load_storage_roots() + + sr._setup_cache_env() + + root = str(tmp_path / "studio") + for key in _ALWAYS_PINNED: + value = os.environ.get(key) + assert value, f"{key} was not pinned" + assert value.startswith(root), f"{key} escaped the studio root: {value}" + + +# Pinned to a path a compiler is handed on a command line: cpp_builder.py joins with spaces, +# interpolates unquoted and reparses with shlex.split, so a root containing a space splits +# mid-path and the C++ build fails outright. +_TOOLCHAIN_PINNED = ( + "TORCHINDUCTOR_CACHE_DIR", + "TORCH_EXTENSIONS_DIR", + "TRITON_CACHE_DIR", + "TRITON_DUMP_DIR", + "CUDA_CACHE_PATH", +) + + +def test_a_spaced_root_leaves_the_compiler_caches_to_their_own_defaults(monkeypatch, tmp_path): + """ "C:\\Users\\First Last" is an ordinary Windows account name, so the DEFAULT Studio root + contains a space for a large share of installs. Before this file pinned these, Inductor used + its own whitespace-free temporary directory and the build worked; pinning it into a spaced + root broke torch.compile outright. Unset is the behaviour that shipped, so that is the + fallback: the rest of the caches, which nobody pastes into a command line, still move.""" + spaced = tmp_path / "my home" / "studio" + spaced.mkdir(parents = True) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(spaced)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + for key in _TOOLCHAIN_PINNED: + assert key not in os.environ, f"{key} was pinned to a path with a space" + # Non-vacuity, and the point of the guard being narrow: everything else is still contained. + for key in ("UV_CACHE_DIR", "NUMBA_CACHE_DIR", "MPLCONFIGDIR", "UNSLOTH_COMPILE_LOCATION"): + assert os.environ[key].startswith(str(spaced.parent)), key + + +def test_a_root_without_spaces_still_pins_the_compiler_caches(monkeypatch, tmp_path): + """The other half of the guard: it must fire on whitespace and nothing else.""" + plain = tmp_path / "plain_home" / "studio" + plain.mkdir(parents = True) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(plain)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + for key in _TOOLCHAIN_PINNED: + assert os.environ.get(key, "").startswith(str(plain)), key + + +def test_an_explicit_spaced_compiler_cache_is_left_alone(monkeypatch, tmp_path): + """Only a default we invented is ours to withhold. A caller who set the variable chose it, + and silently dropping it would send their cache somewhere they did not ask for.""" + chosen = tmp_path / "their choice" + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == str(chosen) + + +def test_default_install_leaves_the_shared_hf_cache_alone(monkeypatch, tmp_path): + home = tmp_path / "home" + monkeypatch.setenv("HOME", str(home)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["HF_HUB_CACHE"] == str(home / ".cache" / "huggingface" / "hub") + for key in _PORTABLE_ONLY: + assert key not in os.environ + + +def test_portable_mode_moves_the_hf_and_torch_caches_under_the_root(monkeypatch, tmp_path): + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + master = tmp_path / "portable" + monkeypatch.setenv("UNSLOTH_HOME", str(master)) + sr = _load_storage_roots() + + assert sr.portable_mode() is True + assert sr.studio_root() == master / "studio" + sr._setup_cache_env() + + root = str(master) + for key in _ALWAYS_PINNED + _PORTABLE_ONLY + ("HF_HUB_CACHE", "HF_XET_CACHE"): + value = os.environ.get(key) + assert value, f"{key} was not pinned" + assert value.startswith(root), f"{key} escaped the portable root: {value}" + + +def test_portable_mode_still_leaves_hf_home_alone(monkeypatch, tmp_path): + # HF_HOME owns the token path; credentials stay off a removable volume. + home = tmp_path / "home" + monkeypatch.setenv("HOME", str(home)) + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + monkeypatch.setenv("UNSLOTH_HOME", str(tmp_path / "portable")) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["HF_HOME"] == str(home / ".cache" / "huggingface") + + +def test_unsloth_portable_alone_enables_portable_mode(monkeypatch, tmp_path): + monkeypatch.setenv("UNSLOTH_PORTABLE", "1") + sr = _load_storage_roots() + + assert sr.portable_mode() is True + sr._setup_cache_env() + + assert os.environ["TORCH_HOME"].startswith(str(tmp_path / "studio")) + + +@pytest.mark.parametrize( + "value", + ("0", "false", "False", "FALSE", "off", "OFF", "no", "No", " false ", ""), +) +def test_unsloth_portable_off_values_do_not_enable_portable_mode(monkeypatch, value): + monkeypatch.setenv("UNSLOTH_PORTABLE", value) + sr = _load_storage_roots() + + assert sr.portable_mode() is False + + +# Every spelling the installers refuse: an intent the shell never acted on, and a typo. +_UNRECOGNIZED_PORTABLE = ("enabled", "flase", "2", "bogus", "y", "n", "disabled", "-1") + + +@pytest.mark.parametrize("value", _UNRECOGNIZED_PORTABLE) +def test_unrecognized_unsloth_portable_does_not_enable_portable_mode(monkeypatch, tmp_path, value): + # A rejected value must not turn portable mode on: the caches would move for this launch + # and move back on the next one. + monkeypatch.setenv("UNSLOTH_PORTABLE", value) + sr = _load_storage_roots() + + assert sr.portable_mode() is False + sr._setup_cache_env() + + for key in _PORTABLE_ONLY: + assert key not in os.environ, f"{key} was redirected by UNSLOTH_PORTABLE={value!r}" + + +@pytest.mark.parametrize("value", _UNRECOGNIZED_PORTABLE + ("0", "false", "off", "no")) +def test_a_portable_root_stays_portable_whatever_unsloth_portable_says( + monkeypatch, tmp_path, value +): + # The root is what makes an install portable; no UNSLOTH_PORTABLE value may veto it. + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + master = tmp_path / "portable" + monkeypatch.setenv("UNSLOTH_HOME", str(master)) + monkeypatch.setenv("UNSLOTH_PORTABLE", value) + sr = _load_storage_roots() + + assert sr.portable_mode() is True + assert sr.studio_root() == master / "studio" + sr._setup_cache_env() + + for key in _PORTABLE_ONLY: + assert os.environ[key].startswith(str(master)), f"{key} escaped the portable root" + + +@pytest.mark.parametrize("value", _UNRECOGNIZED_PORTABLE) +def test_an_on_disk_portable_root_outranks_an_unrecognized_value(monkeypatch, value): + # unsloth_home() also resolves from install.sh's root marker, which a venv-activated + # launch has instead of an environment. + monkeypatch.setenv("UNSLOTH_PORTABLE", value) + sr = _load_storage_roots() + monkeypatch.setattr(sr, "unsloth_home", lambda: Path("/opt/unsloth-portable")) + + assert sr.portable_mode() is True + + +class _RecordingLogger: + def __init__(self): + self.warnings = [] + + def warning(self, message, *args): + self.warnings.append(message % args if args else message) + + +def test_an_unrecognized_value_is_reported_once_not_once_per_call(monkeypatch): + # portable_mode() runs on every cache-var lookup, so a per-call warning floods the log. + monkeypatch.setenv("UNSLOTH_PORTABLE", "enabled") + sr = _load_storage_roots() + recorder = _RecordingLogger() + monkeypatch.setattr(sr, "logger", recorder) + + for _ in range(200): + assert sr.portable_mode() is False + + assert len(recorder.warnings) == 1 + warning = recorder.warnings[0] + assert "enabled" in warning + # Naming only the rejection leaves the user guessing at the spelling. + for accepted in ("1", "true", "yes", "on", "0", "false", "off", "no"): + assert accepted in warning + + +@pytest.mark.parametrize("value", ("1", "true", "TRUE", " on ", "0", "false", "off", "no", "")) +def test_accepted_spellings_are_silent(monkeypatch, value): + monkeypatch.setenv("UNSLOTH_PORTABLE", value) + sr = _load_storage_roots() + recorder = _RecordingLogger() + monkeypatch.setattr(sr, "logger", recorder) + + sr.portable_mode() + + assert recorder.warnings == [] + + +def test_conflicting_roots_are_reported_once_per_conflict(monkeypatch, tmp_path): + # studio_root() runs many times per request, so a per-call warning floods the log and + # adds synchronous log I/O for the life of the backend. + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "elsewhere" / "studio")) + monkeypatch.setenv("UNSLOTH_HOME", str(tmp_path / "portable")) + sr = _load_storage_roots() + recorder = _RecordingLogger() + monkeypatch.setattr(sr, "logger", recorder) + + for _ in range(200): + sr.studio_root() + + assert len(recorder.warnings) == 1 + assert str(tmp_path / "elsewhere" / "studio") in recorder.warnings[0] + assert str(tmp_path / "portable") in recorder.warnings[0] + + # A different pair of roots is a different mistake and must still be heard. + monkeypatch.setenv("UNSLOTH_HOME", str(tmp_path / "other-portable")) + for _ in range(200): + sr.studio_root() + + assert len(recorder.warnings) == 2 + assert str(tmp_path / "other-portable") in recorder.warnings[1] + + +@pytest.mark.parametrize("layout", ("nested", "flat")) +def test_a_self_contained_layout_never_warns(monkeypatch, tmp_path, layout): + # Both supported shapes; silencing the repeat must not silence the whole diagnostic. + master = tmp_path / "portable" + monkeypatch.setenv("UNSLOTH_HOME", str(master)) + monkeypatch.setenv( + "UNSLOTH_STUDIO_HOME", + str(master / "studio") if layout == "nested" else str(master), + ) + sr = _load_storage_roots() + recorder = _RecordingLogger() + monkeypatch.setattr(sr, "logger", recorder) + + for _ in range(50): + sr.studio_root() + + assert recorder.warnings == [] + + +def test_explicit_env_beats_the_pinned_default(monkeypatch, tmp_path): + chosen = tmp_path / "elsewhere" / "inductor" + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == str(chosen) + + +def test_blank_inherited_value_counts_as_unset(monkeypatch, tmp_path): + # "" would send the library to the working directory or the system temp dir. + monkeypatch.setenv("MPLCONFIGDIR", " ") + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["MPLCONFIGDIR"] == str(tmp_path / "studio" / "cache" / "matplotlib") + + +def test_studio_home_outranks_unsloth_home(monkeypatch, tmp_path): + explicit = tmp_path / "explicit" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(explicit)) + monkeypatch.setenv("UNSLOTH_HOME", str(tmp_path / "portable")) + sr = _load_storage_roots() + + assert sr.studio_root() == explicit.resolve() + + +def _use_data_designer(home: Path) -> None: + """Write what a Studio Data Designer session leaves behind.""" + (home / "managed-assets").mkdir(parents = True, exist_ok = True) + (home / "model_configs.yaml").write_text("models: []\n", encoding = "utf-8") + (home / "managed-assets" / "seeds.parquet").write_bytes(b"PAR1") + + +def test_a_used_managed_home_does_not_flip_when_the_legacy_dir_is_deleted(tmp_path): + # Deleting and recreating ~/.data-designer used to toggle which home a run read. + managed = tmp_path / "studio" / "data-designer" + _use_data_designer(managed) + legacy = tmp_path / "home" / ".data-designer" + + seen = [] + for exists in (False, True, False, True): + if exists: + legacy.mkdir(parents = True, exist_ok = True) + elif legacy.exists(): + shutil.rmtree(legacy) + for key in ("DATA_DESIGNER_HOME", "DATA_DESIGNER_MANAGED_ASSETS_PATH"): + os.environ.pop(key, None) + sr = _load_storage_roots() + sr._setup_cache_env() + seen.append(os.environ.get("DATA_DESIGNER_HOME")) + + assert seen == [str(managed)] * 4 + + +def _fail_stat_on(monkeypatch, target: Path, error: OSError) -> None: + """Make every stat of *target* raise *error*, and leave every other path alone. + + chmod covers EACCES; a failing mount answering EIO has no on-disk equivalent an unprivileged + user can set up. Both os.stat and os.lstat are patched so the predicates this replaced see + the same failure the fix does. + """ + + def denying(real): + def deny(path, *args, **kwargs): + if isinstance(path, (str, os.PathLike)) and Path(path) == target: + raise error + return real(path, *args, **kwargs) + + return deny + + for name in ("stat", "lstat"): + monkeypatch.setattr(os, name, denying(getattr(os, name))) + + +def _assert_the_resolver_really_ran(tmp_path: Path) -> None: + """A decline only means something beside a pin that still happened. + + "MPLCONFIGDIR is absent" is equally true when no pinning code exists at all, so on its own + every decline case passed against a reverted implementation. Naming a sibling that is pinned + unconditionally makes the whole family falsifiable without a control test per case. + """ + assert os.environ["TORCHINDUCTOR_CACHE_DIR"].startswith(str(tmp_path / "studio")) + assert os.environ["NUMBA_CACHE_DIR"].startswith(str(tmp_path / "studio")) + + +def _mpl_managed(tmp_path: Path) -> Path: + return tmp_path / "studio" / "cache" / "matplotlib" + + +def _write_rc(directory: Path) -> None: + directory.mkdir(parents = True, exist_ok = True) + (directory / "matplotlibrc").write_text("figure.dpi: 222\n", encoding = "utf-8") + + +def _write_style(stylelib: Path) -> None: + stylelib.mkdir(parents = True, exist_ok = True) + (stylelib / "house.mplstyle").write_text("axes.facecolor: black\n", encoding = "utf-8") + + +@contextlib.contextmanager +def _denied(*paths: Path): + """chmod 000 for the duration of the call under test, restored however it ends.""" + for path in paths: + path.chmod(0o000) + try: + yield + finally: + for path in paths: + path.chmod(0o755) + + +_NEEDS_CHMOD = pytest.mark.skipif( + os.name == "nt" or os.geteuid() == 0, + reason = "chmod 000 denies neither root nor Windows", +) +_XDG_ONLY = pytest.mark.skipif( + not sys.platform.startswith(("linux", "freebsd")), + reason = "XDG config base is the Linux/FreeBSD branch", +) + + +def _mpl_user_rc(tmp_path, monkeypatch): + _write_rc(_matplotlib_config_dir(tmp_path / "home")) + + +def _mpl_user_styles(tmp_path, monkeypatch): + _write_style(_matplotlib_config_dir(tmp_path / "home") / "stylelib") + + +def _mpl_empty_config(tmp_path, monkeypatch): + _matplotlib_config_dir(tmp_path / "home").mkdir(parents = True) + + +def _mpl_xdg_rc(tmp_path, monkeypatch): + _write_rc(tmp_path / "xdg" / "matplotlib") + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "xdg")) + + +def _mpl_denied_config(tmp_path, monkeypatch): + config = _matplotlib_config_dir(tmp_path / "home") + _write_rc(config) + return _denied(config.parent) + + +def _mpl_denied_styles(tmp_path, monkeypatch): + styles = _matplotlib_config_dir(tmp_path / "home") / "stylelib" + _write_style(styles) + + @contextlib.contextmanager + def guard(): + with _denied(styles): + # Path.glob suppresses the scandir error, which is what read an unreadable stylelib + # as empty; without this the case could pass while chmod did nothing. + assert list(styles.glob("*.mplstyle")) == [] + yield + + return guard() + + +def _mpl_failing_volume(tmp_path, monkeypatch): + config = _matplotlib_config_dir(tmp_path / "home") + _write_rc(config) + _fail_stat_on(monkeypatch, config / "matplotlibrc", OSError(errno.EIO, "Input/output error")) + + +def _mpl_managed_styles_and_user_rc(tmp_path, monkeypatch): + _write_style(_mpl_managed(tmp_path) / "stylelib") + _write_rc(_matplotlib_config_dir(tmp_path / "home")) + + +def _mpl_empty_managed_and_user_rc(tmp_path, monkeypatch): + (_mpl_managed(tmp_path) / "stylelib").mkdir(parents = True) + _write_rc(_matplotlib_config_dir(tmp_path / "home")) + + +def _mpl_linked_styles(tmp_path, monkeypatch): + config = _matplotlib_config_dir(tmp_path / "home") + config.mkdir(parents = True) + elsewhere = tmp_path / "styles-volume" + elsewhere.mkdir() # empty, which is the whole point + (config / "stylelib").symlink_to(elsewhere, target_is_directory = True) + + +def _mpl_dangling_styles(tmp_path, monkeypatch): + config = _matplotlib_config_dir(tmp_path / "home") + config.mkdir(parents = True) + (config / "stylelib").symlink_to(tmp_path / "never-mounted", target_is_directory = True) + + +def _mpl_plain_empty_styles(tmp_path, monkeypatch): + (_matplotlib_config_dir(tmp_path / "home") / "stylelib").mkdir(parents = True) + + +@pytest.mark.parametrize( + "prepare, pinned", + [ + # MPLCONFIGDIR moves the CONFIG directory as well as the cache, so a pin wins only when + # there is nothing of the user's at matplotlib's own directory to hide. Each "must still + # pin" row is the control for the row above it: matplotlib creates the config dir and + # stylelib on import, so existence alone must never count as configuration. + pytest.param(_mpl_user_rc, False, id = "a user matplotlibrc"), + pytest.param(_mpl_user_styles, False, id = "a user style library"), + pytest.param(_mpl_empty_config, True, id = "an empty config dir"), + pytest.param(_mpl_xdg_rc, False, id = "an rc under XDG_CONFIG_HOME", marks = _XDG_ONLY), + pytest.param( + _mpl_denied_config, False, id = "a config dir we may not read", marks = _NEEDS_CHMOD + ), + pytest.param( + _mpl_denied_styles, False, id = "a stylelib we may not read", marks = _NEEDS_CHMOD + ), + pytest.param(_mpl_failing_volume, False, id = "an rc on a failing volume"), + pytest.param(_mpl_managed_styles_and_user_rc, True, id = "our own styles, already ours"), + pytest.param(_mpl_empty_managed_and_user_rc, False, id = "an empty managed dir"), + pytest.param(_mpl_linked_styles, False, id = "a stylelib linked to an empty volume"), + pytest.param(_mpl_dangling_styles, False, id = "a dangling stylelib link"), + pytest.param(_mpl_plain_empty_styles, True, id = "a plain empty stylelib"), + ], +) +def test_mplconfigdir_is_pinned_only_when_it_would_hide_nothing( + tmp_path, monkeypatch, prepare, pinned +): + guard = prepare(tmp_path, monkeypatch) or contextlib.nullcontext() + sr = _load_storage_roots() + + with guard: + sr._setup_cache_env() + + if pinned: + assert os.environ["MPLCONFIGDIR"] == str(_mpl_managed(tmp_path)) + else: + assert "MPLCONFIGDIR" not in os.environ + _assert_the_resolver_really_ran(tmp_path) + + +def _dd_managed(tmp_path: Path) -> Path: + return tmp_path / "studio" / "data-designer" + + +def _dd_legacy(tmp_path: Path) -> Path: + return tmp_path / "home" / ".data-designer" + + +def _dd_nothing(tmp_path, monkeypatch): + (tmp_path / "home").mkdir(parents = True, exist_ok = True) + + +def _dd_legacy_in_use(tmp_path, monkeypatch): + _use_data_designer(_dd_legacy(tmp_path)) + + +def _dd_managed_in_use(tmp_path, monkeypatch): + _use_data_designer(_dd_managed(tmp_path)) + _dd_legacy(tmp_path).mkdir(parents = True) + + +def _dd_managed_untouched(tmp_path, monkeypatch): + (_dd_managed(tmp_path) / "managed-assets").mkdir(parents = True) + _dd_legacy(tmp_path).mkdir(parents = True) + + +def _dd_managed_untouched_legacy_in_use(tmp_path, monkeypatch): + (_dd_managed(tmp_path) / "managed-assets").mkdir(parents = True) + _use_data_designer(_dd_legacy(tmp_path)) + + +def _dd_managed_unlistable(tmp_path, monkeypatch): + managed = _dd_managed(tmp_path) + _use_data_designer(managed) + _dd_legacy(tmp_path).mkdir(parents = True) + real_iterdir = Path.iterdir + + def deny(self): + if self in (managed, managed / "managed-assets"): + raise PermissionError(13, "Permission denied") + return real_iterdir(self) + + monkeypatch.setattr(Path, "iterdir", deny) + + +def _dd_managed_assets_denied(tmp_path, monkeypatch): + managed = _dd_managed(tmp_path) + _use_data_designer(managed) + _dd_legacy(tmp_path).mkdir(parents = True) + return _denied(managed / "managed-assets") + + +def _dd_redirected_assets(tmp_path, monkeypatch): + managed = _dd_managed(tmp_path) + managed.mkdir(parents = True) + elsewhere = tmp_path / "big-disk" / "assets" + elsewhere.mkdir(parents = True) # empty, which is the whole point + (managed / "managed-assets").symlink_to(elsewhere, target_is_directory = True) + _use_data_designer(_dd_legacy(tmp_path)) + + +def _dd_legacy_denied(tmp_path, monkeypatch): + _use_data_designer(_dd_legacy(tmp_path)) + return _denied(tmp_path / "home") + + +def _dd_legacy_failing_volume(tmp_path, monkeypatch): + legacy = _dd_legacy(tmp_path) + _use_data_designer(legacy) + _fail_stat_on(monkeypatch, legacy, OSError(errno.EIO, "Input/output error")) + + +def _dd_legacy_symlink_loop(tmp_path, monkeypatch): + legacy = _dd_legacy(tmp_path) + legacy.parent.mkdir(parents = True, exist_ok = True) + legacy.symlink_to(legacy) + # ELOOP reads as absence through Path.exists() on every release, which is the bug. + assert Path(legacy).exists() is False + + +@pytest.mark.parametrize( + "prepare, pinned", + [ + # Data Designer's home is not a cache: repointing one that holds yaml configs and + # multi-GB parquet hides them behind a re-seeded default. Our own home wins once it has + # been USED, since the legacy probe re-runs every launch and a standalone run creating + # ~/.data-designer would otherwise take the work written under the Studio root. + pytest.param(_dd_nothing, True, id = "no legacy home at all"), + pytest.param(_dd_legacy_in_use, False, id = "a legacy home in use"), + pytest.param(_dd_managed_in_use, True, id = "our home in use, legacy empty"), + pytest.param(_dd_managed_untouched, False, id = "our home untouched, legacy empty"), + pytest.param( + _dd_managed_untouched_legacy_in_use, False, id = "our home untouched, legacy in use" + ), + pytest.param(_dd_managed_unlistable, True, id = "our home we may not list"), + pytest.param( + _dd_managed_assets_denied, True, id = "our managed-assets denied", marks = _NEEDS_CHMOD + ), + pytest.param(_dd_redirected_assets, True, id = "managed-assets redirected by a link"), + pytest.param( + _dd_legacy_denied, False, id = "a legacy home we may not read", marks = _NEEDS_CHMOD + ), + pytest.param(_dd_legacy_failing_volume, False, id = "a legacy home on a failing volume"), + pytest.param(_dd_legacy_symlink_loop, False, id = "a legacy home that is a symlink loop"), + ], +) +def test_the_data_designer_home_is_pinned_only_when_it_would_hide_nothing( + tmp_path, monkeypatch, prepare, pinned +): + guard = prepare(tmp_path, monkeypatch) or contextlib.nullcontext() + sr = _load_storage_roots() + + with guard: + sr._setup_cache_env() + + managed = _dd_managed(tmp_path) + if pinned: + assert os.environ["DATA_DESIGNER_HOME"] == str(managed) + assert os.environ["DATA_DESIGNER_MANAGED_ASSETS_PATH"] == str(managed / "managed-assets") + else: + assert "DATA_DESIGNER_HOME" not in os.environ + assert "DATA_DESIGNER_MANAGED_ASSETS_PATH" not in os.environ + _assert_the_resolver_really_ran(tmp_path) + + +def test_managed_assets_follow_an_explicit_data_designer_home(monkeypatch, tmp_path): + chosen = tmp_path / "mine" / ".data-designer" + monkeypatch.setenv("DATA_DESIGNER_HOME", str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["DATA_DESIGNER_HOME"] == str(chosen) + assert "DATA_DESIGNER_MANAGED_ASSETS_PATH" not in os.environ + + +def test_an_explicit_triton_home_keeps_its_own_cache_dir(monkeypatch, tmp_path): + chosen = tmp_path / "mine" / "triton-home" + monkeypatch.setenv("TRITON_HOME", str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["TRITON_HOME"] == str(chosen) + assert "TRITON_CACHE_DIR" not in os.environ + + +def test_triton_keeps_reading_the_default_override_dir(tmp_path): + # TRITON_HOME would move ~/.triton/override, which holds user files, along with the cache, + # so a TRITON_KERNEL_OVERRIDE=1 run would compile something else instead. + pytest.importorskip("triton") + override = tmp_path / "home" / ".triton" / "override" / "0123456789abcdef" + override.mkdir(parents = True) + (override / "kernel.ttir").write_text("// hand-tuned\n", encoding = "utf-8") + sr = _load_storage_roots() + + sr._setup_cache_env() + + # A fresh interpreter: torch may have imported Triton and pinned a cache dir already. + probe = subprocess.run( + [ + sys.executable, + "-c", + "import json; from triton import knobs;" + "print(json.dumps({'cache': knobs.cache.dir," + "'dump': knobs.cache.dump_dir," + "'override': knobs.cache.override_dir}))", + ], + env = dict(os.environ), + capture_output = True, + text = True, + check = True, + ) + result = json.loads(probe.stdout.strip().splitlines()[-1]) + + assert Path(result["override"]) == override.parent + assert (Path(result["override"]) / override.name / "kernel.ttir").is_file() + assert result["cache"] == str(tmp_path / "studio" / "cache" / "triton") + assert result["dump"] == str(tmp_path / "studio" / "cache" / "triton-dump") + + +def test_the_macos_matplotlib_config_dir_matches_matplotlibs_own(monkeypatch, tmp_path): + # macOS is matplotlib's "other platforms" branch: ~/.matplotlib, not Application Support. + pytest.importorskip("matplotlib") + monkeypatch.setattr(sys, "platform", "darwin") + sr = _load_storage_roots() + + ours = sr._matplotlib_config_dir() + + # sys.platform is read inside _get_config_or_cache_dir, so a fresh interpreter can be + # asked what it would do on a Mac. + probe = subprocess.run( + [ + sys.executable, + "-c", + "import sys, matplotlib; sys.platform = 'darwin';" + "print(matplotlib._get_config_or_cache_dir(matplotlib._get_xdg_config_dir))", + ], + env = {**os.environ, "HOME": str(tmp_path / "home"), "MPLCONFIGDIR": ""}, + capture_output = True, + text = True, + check = True, + ) + theirs = Path(probe.stdout.strip().splitlines()[-1]) + + assert ours is not None + assert ours.resolve() == theirs.resolve() + + +def _matplotlib_config_dir(home: Path) -> Path: + """Where matplotlib on THIS platform would look for a user matplotlibrc. + + _get_config_or_cache_dir: XDG on Linux and FreeBSD, a non-empty ~/.matplotlib ahead of + %LOCALAPPDATA% on Windows, ~/.matplotlib everywhere else. The Linux branch alone was hardcoded + here, so on a mac runner these tests wrote the rc into a directory matplotlib never reads: the + guard correctly found nothing to strand, pinned MPLCONFIGDIR, and nine tests failed for a + reason that was in them rather than in the code they cover. The two tests above hold this + derivation against real matplotlib on each spoofed platform, so it cannot quietly drift into + agreeing with storage_roots.py and nothing else. + + Windows resolves to ~/.matplotlib rather than %LOCALAPPDATA%: the callers write a file into + whatever this returns, LOCALAPPDATA on a runner is outside tmp_path, and a non-empty + ~/.matplotlib is the branch matplotlib itself prefers. + """ + if sys.platform.startswith(("linux", "freebsd")): + base = (os.environ.get("XDG_CONFIG_HOME") or "").strip() + return (Path(base) if base else home / ".config") / "matplotlib" + return home / ".matplotlib" + + +def test_matplotlib_reads_the_config_the_pin_would_have_hidden(tmp_path): + pytest.importorskip("matplotlib") + config = _matplotlib_config_dir(tmp_path / "home") + (config / "stylelib").mkdir(parents = True) + (config / "matplotlibrc").write_text("figure.dpi: 222\n", encoding = "utf-8") + (config / "stylelib" / "house.mplstyle").write_text("axes.facecolor: black\n", encoding = "utf-8") + sr = _load_storage_roots() + + sr._setup_cache_env() + + # A fresh interpreter: matplotlib caches the config dir on first read. + probe = subprocess.run( + [ + sys.executable, + "-c", + "import json, matplotlib, matplotlib.style;" + "print(json.dumps({'rc': matplotlib.matplotlib_fname()," + "'dpi': matplotlib.rcParams['figure.dpi']," + "'style': 'house' in matplotlib.style.available}))", + ], + env = dict(os.environ), + capture_output = True, + text = True, + check = True, + ) + result = json.loads(probe.stdout.strip().splitlines()[-1]) + + assert result["rc"] == str(config / "matplotlibrc") + assert result["dpi"] == 222.0 + assert result["style"] is True + + +@pytest.mark.skipif( + not sys.platform.startswith(("linux", "freebsd")), + reason = "XDG config base is the Linux/FreeBSD branch", +) +def test_an_xdg_config_dir_is_read_without_a_resolvable_home(monkeypatch, tmp_path): + # _get_xdg_config_dir reads XDG_CONFIG_HOME before it needs a home, so bailing out on + # Path.home() pinned over a real matplotlibrc. + config = tmp_path / "xdg" / "matplotlib" + config.mkdir(parents = True) + (config / "matplotlibrc").write_text("figure.dpi: 222\n", encoding = "utf-8") + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "xdg")) + sr = _load_storage_roots() + + def no_home(): + raise RuntimeError("Could not determine home directory.") + + monkeypatch.setattr(Path, "home", staticmethod(no_home)) + + assert sr._matplotlib_defaults(tmp_path / "studio" / "cache") == {} + + +def test_only_a_positive_absence_counts_as_nothing_at_a_path(tmp_path): + # The one rule every "pin only when there is nothing to strand" probe shares. + sr = _load_storage_roots() + present = tmp_path / "present" + present.write_text("x", encoding = "utf-8") + loop = tmp_path / "loop" + loop.symlink_to(loop) + dangling = tmp_path / "dangling" + dangling.symlink_to(tmp_path / "not-mounted-yet" / "data") + + assert sr._nothing_at(tmp_path / "absent") is True + assert sr._nothing_at(present) is False + # Both of these answer False through Path.exists(), which is the bug. + assert Path(loop).exists() is False and sr._nothing_at(loop) is False + assert Path(dangling).exists() is False and sr._nothing_at(dangling) is False + + +def test_nothing_at_reports_an_uninspectable_directory_as_holding_something(monkeypatch, tmp_path): + sr = _load_storage_roots() + styles = tmp_path / "stylelib" + styles.mkdir() + real_scandir = os.scandir + + def deny( + path = ".", + *args, + **kwargs, + ): + if isinstance(path, (str, os.PathLike)) and Path(path) == styles: + raise OSError(errno.EIO, "Input/output error") + return real_scandir(path, *args, **kwargs) + + monkeypatch.setattr(os, "scandir", deny) + + assert sr._nothing_at(styles, ending = ".mplstyle") is False + assert sr._nothing_at(tmp_path / "absent", ending = ".mplstyle") is True + + +def _fake_torch_on_path( + tmp_path, + label, + version, + cuda = None, + hip = None, + debug = False, +): + """A torch that has a version.py and explodes if anything imports it. version.py is written + the way recent torch generates it, annotations and all, so the parser meets the real shape.""" + pkg = tmp_path / label / "torch" + pkg.mkdir(parents = True) + (pkg / "__init__.py").write_text("raise AssertionError('torch was imported')\n") + (pkg / "version.py").write_text( + "from typing import Optional\n\n" + f"__version__ = {version!r}\n" + f"debug = {debug!r}\n" + f"cuda: Optional[str] = {cuda!r}\n" + "git_version = '5811a8d7da873dd699ff6687092c225caffcf1bb'\n" + f"hip: Optional[str] = {hip!r}\n" + "xpu: Optional[str] = None\n" + ) + return str(tmp_path / label) + + +@contextlib.contextmanager +def _only_torch(entry): + saved_path = list(sys.path) + saved_module = sys.modules.pop("torch", None) + sys.path.insert(0, entry) + importlib.invalidate_caches() + try: + yield + finally: + sys.path[:] = saved_path + if saved_module is not None: + sys.modules["torch"] = saved_module + importlib.invalidate_caches() + + +def test_torch_extension_cache_keeps_an_abi_folder(tmp_path): + # torch appends py_ to its DEFAULT root only, so a pin carries it. + sr = _load_storage_roots() + + sr._setup_cache_env() + + pinned = Path(os.environ["TORCH_EXTENSIONS_DIR"]) + assert str(pinned).startswith(str(tmp_path / "studio")) + assert pinned.parent.name == "torch-extensions" + assert pinned.name != "torch-extensions", "extension cache is shared across runtimes" + assert pinned.name.startswith(f"py{sys.version_info.major}{sys.version_info.minor}") + + +def _interpreter_prefix() -> str: + """The interpreter-and-ABI part of the tag, spelled out here rather than imported. + + Reusing storage_roots' own helper would make every assertion below agree with whatever + it happens to produce, so the exact-string tests would stop being able to catch a change + in the tag at all. + """ + abi = getattr(sys, "abiflags", "") + host = re.sub(r"[^A-Za-z0-9.]+", "-", f"{sys.platform}-{platform.machine() or 'unknown'}") + return f"py{sys.version_info.major}{sys.version_info.minor}{abi}_{host}" + + +def test_torch_extension_cache_separates_incompatible_builds(tmp_path): + sr = _load_storage_roots() + + tags = [] + for label, version, cuda in (("a", "2.9.1+cu128", "12.8"), ("b", "2.9.1+cu126", "12.6")): + with _only_torch(_fake_torch_on_path(tmp_path, label, version, cuda = cuda)): + tags.append(sr._torch_runtime_tag()) + + assert tags[0] != tags[1], f"two torch builds shared one extension dir: {tags[0]}" + assert "cu128" in tags[0] and "cu126" in tags[1] + # Path-safe: no '+' or other separators survive into the directory name. + assert all(part.replace(".", "").replace("-", "").replace("_", "").isalnum() for part in tags) + + +def test_torch_extension_tag_survives_a_missing_torch(tmp_path): + # First launch, before the venv has torch: still isolated by interpreter. + sr = _load_storage_roots() + empty = tmp_path / "empty" + empty.mkdir() + + with _only_torch(str(empty)): + saved = list(sys.path) + try: + sys.path[:] = [str(empty)] + tag = sr._torch_runtime_tag() + finally: + sys.path[:] = saved + + assert tag == _interpreter_prefix() + + +def test_torch_extension_cache_separates_builds_sharing_one_version_string(tmp_path): + # conda-forge's CPU and CUDA packages of one release share a __version__, differing only + # in a conda build string. + sr = _load_storage_roots() + + tags = [] + for label, cuda in (("cpu", None), ("cu126", "12.6"), ("cu128", "12.8")): + with _only_torch(_fake_torch_on_path(tmp_path, label, "2.9.1", cuda = cuda)): + tags.append(sr._torch_runtime_tag()) + + assert len(set(tags)) == 3, f"builds with different CUDA ABIs shared one dir: {tags}" + prefix = _interpreter_prefix() + assert tags[0] == f"{prefix}_cpu_2.9.1" + assert tags[1] == f"{prefix}_cu126_2.9.1" + assert tags[2] == f"{prefix}_cu128_2.9.1" + + +def test_torch_extension_cache_separates_a_rocm_build_from_a_cpu_build(tmp_path): + # torch names a ROCm build 'cpu', since version.cuda is unset; hip is read first. + sr = _load_storage_roots() + + tags = [] + for label, hip in (("cpu", None), ("rocm", "6.4.43484-123eb5128")): + with _only_torch(_fake_torch_on_path(tmp_path, label, "2.9.1", hip = hip)): + tags.append(sr._torch_runtime_tag()) + + assert tags[0] != tags[1], f"a ROCm build shared the CPU extension dir: {tags[0]}" + assert "rocm6.4.43484-123eb5128" in tags[1] + + +def test_torch_extension_cache_separates_two_host_architectures(tmp_path, monkeypatch): + # An arm64 python and a Rosetta x86_64 python on ONE Mac agree on every other field, so + # torch's py_ folder gives them one directory and the .so fails to load. + sr = _load_storage_roots() + + tags = [] + for machine in ("arm64", "x86_64"): + monkeypatch.setattr(platform, "machine", lambda m = machine: m) + with _only_torch(_fake_torch_on_path(tmp_path, machine, "2.9.1")): + tags.append(sr._torch_runtime_tag()) + + assert tags[0] != tags[1], f"two host architectures shared one extension dir: {tags[0]}" + assert "arm64" in tags[0] and "x86-64" in tags[1] + + +def test_torch_extension_cache_survives_an_unnameable_architecture(tmp_path, monkeypatch): + # platform.machine() returns "" when the platform cannot answer; no empty segment. + sr = _load_storage_roots() + + monkeypatch.setattr(platform, "machine", lambda: "") + with _only_torch(_fake_torch_on_path(tmp_path, "blank", "2.9.1")): + tag = sr._torch_runtime_tag() + + assert "__" not in tag and not tag.endswith("_") + assert "unknown" in tag + + +def test_torch_extension_cache_separates_a_debug_build(tmp_path): + # A debug build keeps the soname of a release one but not its ABI. + sr = _load_storage_roots() + + tags = [] + for label, debug in (("release", False), ("debug", True)): + entry = _fake_torch_on_path(tmp_path, label, "2.9.1+cu128", cuda = "12.8", debug = debug) + with _only_torch(entry): + tags.append(sr._torch_runtime_tag()) + + assert tags[0] != tags[1], f"a debug build shared the release extension dir: {tags[0]}" + assert tags[1].endswith("_debug") + + +def test_torch_runtime_tag_never_imports_torch(tmp_path): + # Runs before torch exists in a fresh venv, and importing it on the startup path would cost + # seconds. The fake package raises on import, so a skip is a failure. + sr = _load_storage_roots() + entry = _fake_torch_on_path(tmp_path, "guard", "2.9.1+cu128", cuda = "12.8") + + with _only_torch(entry): + tag = sr._torch_runtime_tag() + assert "torch" not in sys.modules, "torch was imported to build the cache tag" + + assert "cu128" in tag + + +def test_a_managed_matplotlibrc_is_not_displaced_by_a_later_legacy_one(tmp_path): + """The legacy probe re-runs every launch, so on its own it hands a matplotlibrc written + HERE to a ~/.config/matplotlib created later by some other tool. + + Measured across four launches before the fix: dpi 177 from the managed config, then 222 once + a legacy rc appeared, then 177 again when it was removed, with the managed rc present + throughout. A style that changes on a later launch and changes back is worse than either + choice made once. _data_designer_defaults already applies this rule to its own home. + """ + managed = tmp_path / "studio" / "cache" / "matplotlib" + managed.mkdir(parents = True) + (managed / "matplotlibrc").write_text("figure.dpi: 177\n", encoding = "utf-8") + config = _matplotlib_config_dir(tmp_path / "home") + config.mkdir(parents = True) + (config / "matplotlibrc").write_text("figure.dpi: 222\n", encoding = "utf-8") + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["MPLCONFIGDIR"] == str(managed) + + +def test_a_blank_toolchain_override_is_dropped_on_a_spaced_root(monkeypatch, tmp_path): + """ "blank counts as unset" has to hold for a root we refuse to pin, too. + + Inductor distinguishes an absent TORCHINDUCTOR_CACHE_DIR from a present one, so a leftover + " " becomes a relative compiler path and is then split by the very unquoted command + construction the whitespace refusal exists to avoid. + """ + spaced = tmp_path / "My Studio" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(spaced)) + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", " ") + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert "TORCHINDUCTOR_CACHE_DIR" not in os.environ + + +def test_an_unusable_managed_inductor_path_is_not_published(monkeypatch, tmp_path): + """torch treats the value as authoritative, so a path it cannot use fails every compile. + + Unset, it would have used its own temporary cache instead. The placement error is swallowed + on purpose so startup survives, which is what let an unusable path stay published. + """ + cache = tmp_path / "studio" / "cache" + cache.mkdir(parents = True) + # A regular file where the directory should go: mkdir raises FileExistsError, which the + # best-effort handler treats as "already there". + (cache / "torchinductor").write_text("not a directory", encoding = "utf-8") + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert "TORCHINDUCTOR_CACHE_DIR" not in os.environ + # The other pins are unaffected: only the one that could not be made is withheld. + assert os.environ["CUDA_CACHE_PATH"] == str(cache / "cuda") + + +@pytest.mark.skipif( + os.name == "nt" or os.geteuid() == 0, reason = "chmod 555 denies neither root nor Windows" +) +def test_a_read_only_managed_cache_is_not_published(tmp_path): + """A directory that exists but cannot be written to is the same defect as a missing one. + + mkdir passes exist_ok = False, so an existing directory takes the FileExistsError path and + proves nothing about writability. torch only falls back to its own temporary cache when the + variable is UNSET (torch/_inductor/runtime/cache_dir_utils.py), so a read-only pin fails + every compile. + """ + cache = tmp_path / "studio" / "cache" + for name in ("torchinductor", "triton", "cuda"): + (cache / name).mkdir(parents = True) + (cache / name).chmod(0o555) + sr = _load_storage_roots() + try: + sr._setup_cache_env() + + for key in ("TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR", "CUDA_CACHE_PATH"): + assert key not in os.environ, key + # Withheld, not blanked: an empty value is a relative path to the library reading it. + assert os.environ["NUMBA_CACHE_DIR"] == str(cache / "numba") + finally: + for name in ("torchinductor", "triton", "cuda"): + (cache / name).chmod(0o755) + + +def test_the_write_probe_leaves_nothing_behind(tmp_path): + """The probe creates a file in the directory it is testing; it must not accumulate.""" + sr = _load_storage_roots() + + sr._setup_cache_env() + + inductor = Path(os.environ["TORCHINDUCTOR_CACHE_DIR"]) + assert inductor.is_dir() + assert list(inductor.iterdir()) == [] + + +def test_a_usable_managed_inductor_path_is_still_published(tmp_path): + """The rule above must not withhold the ordinary case.""" + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == str( + tmp_path / "studio" / "cache" / "torchinductor" + ) + + +@pytest.mark.skipif(os.name == "nt" or os.geteuid() == 0, reason = "chmod 000 denies neither") +def test_a_note_that_could_not_be_read_is_not_cached_as_a_missing_one(monkeypatch, tmp_path): + """A miss is a fact about the install; a failure to LOOK is a fact about one instant. + + _recorded_master_root() memoises per studio path, and the backend is a long-lived process. + Caching an EACCES on share/ -- a mount that came back, a permission fixed a second later -- + pinned that process to the legacy runtime paths for its whole lifetime with the trees it + should have found sitting right beside it, and nothing in production calls + forget_recorded_master_root() to get back out. + """ + master = tmp_path / "root" + studio = master / "studio" + share = studio / "share" + share.mkdir(parents = True) + (share / ".unsloth-master-root").write_text(str(master) + "\n", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio)) + sr = _load_storage_roots() + + share.chmod(0o000) + try: + assert sr.unsloth_home() is None + finally: + share.chmod(0o755) + + assert sr.unsloth_home() == master, "a transient read failure was cached as a missing note" + + +def test_a_genuinely_absent_note_is_still_cached(monkeypatch, tmp_path): + """The other half: the fix must not turn the memo off. unsloth_home() runs on every + cache-var lookup, so re-stat'ing a note that is simply not there is a syscall per call.""" + studio = tmp_path / "root" / "studio" + (studio / "share").mkdir(parents = True) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio)) + sr = _load_storage_roots() + + assert sr.unsloth_home() is None + # Written after the miss was cached: a cached answer is the point, and + # forget_recorded_master_root() is the way back out for an installer that writes one later. + (studio / "share" / ".unsloth-master-root").write_text( + str(tmp_path / "root") + "\n", + encoding = "utf-8", + ) + assert sr.unsloth_home() is None + sr.forget_recorded_master_root() + assert sr.unsloth_home() == tmp_path / "root" + + +def test_a_write_probe_that_cannot_close_its_handle_is_a_no_not_a_crash(monkeypatch, tmp_path): + """_usable_dir answers yes or no. Every other failure in it returns False; an OSError from + os.close escaping instead would take the whole _setup_cache_env() -- and the backend start + that called it -- down with it.""" + sr = _load_storage_roots() + + real_close = os.close + + def _failing_close(fd): + real_close(fd) + raise OSError(errno.EIO, "close failed") + + monkeypatch.setattr(sr.os, "close", _failing_close) + assert sr._usable_dir(str(tmp_path)) is False + + monkeypatch.setattr(sr.os, "close", real_close) + assert sr._usable_dir(str(tmp_path)) is True + + +def _note_ancestor_of_a_legacy_tree(tmp_path, monkeypatch): + """A note naming $HOME, which passes containment because ~/.unsloth/studio is inside it. + + process.rs scrubs UNSLOTH_HOME and UNSLOTH_PORTABLE from every managed spawn so Tauri uses + the legacy root whatever the environment says. A note is a FILE, which no env_remove reaches, + so this is how both variables came back and moved the packaged app's caches. + """ + home = tmp_path / "home" + studio = home / ".unsloth" / "studio" + _write_note(studio, home) + monkeypatch.setenv("HOME", str(home)) + monkeypatch.setenv("USERPROFILE", str(home)) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio)) + for scrubbed in ("UNSLOTH_HOME", "UNSLOTH_PORTABLE", "STUDIO_HOME"): + monkeypatch.delenv(scrubbed, raising = False) + return None + + +def _note_naming_the_legacy_default(tmp_path, monkeypatch): + """The reader has to decline it, not just the writer: setup only sweeps an old note when it + runs again, and the user who hit this set UNSLOTH_HOME once, for one command.""" + home = tmp_path / "home" + studio = home / ".unsloth" / "studio" + _write_note(studio, home / ".unsloth") + monkeypatch.setenv("HOME", str(home)) + monkeypatch.setenv("USERPROFILE", str(home)) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio)) + return None + + +def _note_naming_a_real_master_root(tmp_path, monkeypatch): + """The other side of the rule: a recorded root that is not the legacy default and that + contains this Studio tree is exactly what the note is for.""" + master = tmp_path / "portable" + _write_note(master / "studio", master) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(master / "studio")) + return master + + +def _an_explicit_legacy_unsloth_home(tmp_path, monkeypatch): + """Only the RECORDED root is declined. unsloth_home() returns an explicit UNSLOTH_HOME before + it reads any note, and narrowing that would change what the variable means.""" + home = tmp_path / "home" + (home / ".unsloth" / "studio" / "share").mkdir(parents = True) + monkeypatch.setenv("HOME", str(home)) + monkeypatch.setenv("USERPROFILE", str(home)) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(home / ".unsloth" / "studio")) + monkeypatch.setenv("UNSLOTH_HOME", str(home / ".unsloth")) + return home / ".unsloth" + + +def _write_note(studio: Path, root: Path) -> None: + (studio / "share").mkdir(parents = True, exist_ok = True) + (studio / "share" / ".unsloth-master-root").write_text(str(root) + "\n", encoding = "utf-8") + + +@pytest.mark.parametrize( + "prepare", + [ + pytest.param(_note_ancestor_of_a_legacy_tree, id = "a note naming an ancestor"), + pytest.param(_note_naming_the_legacy_default, id = "a note naming the legacy default"), + pytest.param(_note_naming_a_real_master_root, id = "a note naming a real master root"), + pytest.param(_an_explicit_legacy_unsloth_home, id = "an explicit legacy UNSLOTH_HOME"), + ], +) +def test_a_master_root_is_honoured_only_when_a_reader_should_honour_it( + tmp_path, monkeypatch, prepare +): + """`prepare` returns the root that must be honoured, or None when it must be declined. + + The hub cache is asserted either way, since declining is only meaningful if the shared cache + really stays shared, and honouring is only meaningful if the caches really move. + """ + expected = prepare(tmp_path, monkeypatch) + sr = _load_storage_roots() + + assert sr.unsloth_home() == expected + assert sr.portable_mode() is (expected is not None) + + sr._setup_cache_env() + + hub = Path(os.environ["HF_HUB_CACHE"]) + if expected is None: + assert hub == Path(os.environ["HOME"]) / ".cache" / "huggingface" / "hub" + assert "TORCH_HOME" not in os.environ + else: + assert expected in hub.parents + assert Path(os.environ["TORCH_HOME"]).is_relative_to(expected) + + +# 25 of 113 cases stayed green with the production hunks reverted: each asserts a variable is +# ABSENT, which is equally true when no pinning code exists, so they were assertions rather than +# negative controls. Pairing a decline with the pins that must STILL happen makes them falsifiable. +_DECLINE_SIBLINGS = ( + "TORCHINDUCTOR_CACHE_DIR", + "NUMBA_CACHE_DIR", + "CUDA_CACHE_PATH", + "UV_CACHE_DIR", + "UNSLOTH_COMPILE_LOCATION", +) + + +@pytest.mark.parametrize("declined", ["MPLCONFIGDIR", "DATA_DESIGNER_HOME"]) +def test_declining_one_pin_never_means_the_resolver_did_nothing(tmp_path, declined): + """A decline is only meaningful next to a pin. If the resolver were absent, or returned early, + the sibling assertions below would fail and the decline would stop proving anything.""" + home = tmp_path / "home" + if declined == "MPLCONFIGDIR": + config = _matplotlib_config_dir(home) + config.mkdir(parents = True) + (config / "matplotlibrc").write_text("figure.dpi: 123\n", encoding = "utf-8") + else: + (home / ".data-designer" / "recipes").mkdir(parents = True) + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert declined not in os.environ, f"{declined} displaced user files" + root = str(tmp_path / "studio") + for sibling in _DECLINE_SIBLINGS: + value = os.environ.get(sibling) + assert value, f"{sibling} unpinned: the decline of {declined} proves nothing" + assert value.startswith(root), f"{sibling} escaped the studio root: {value}" + + +def test_an_uninspectable_probe_declines_and_still_pins_everything_else(tmp_path): + """Same rule for the os.lstat/os.scandir failure paths, which are the subtlest of the 25: + an unreadable matplotlib config must leave MPLCONFIGDIR alone AND leave the rest pinned.""" + config = _matplotlib_config_dir(tmp_path / "home") + config.parent.mkdir(parents = True, exist_ok = True) + config.write_text("", encoding = "utf-8") # a file where a directory belongs: ENOTDIR + sr = _load_storage_roots() + + sr._setup_cache_env() + + assert "MPLCONFIGDIR" not in os.environ + for sibling in _DECLINE_SIBLINGS: + assert os.environ.get(sibling, "").startswith(str(tmp_path / "studio")), sibling + + +def test_a_file_where_the_config_dir_belongs_declines_the_pin_on_windows_too(monkeypatch, tmp_path): + """_nothing_at reads a FileNotFoundError as absence, and on Windows that error also means + "a parent component is a file" -- the case POSIX reports as NotADirectoryError and declines. + So with a file where ~/.matplotlib belongs, POSIX declined the pin and Windows took it, + hiding whatever the user had underneath. Caught by the windows-latest leg. + + Reproduced by the error shape rather than the platform, since the Linux runners cannot raise + it; the POSIX arm of the same fixture is the test directly above. + """ + sr = _load_storage_roots() + config = _matplotlib_config_dir(tmp_path / "home") + config.parent.mkdir(parents = True, exist_ok = True) + config.write_text("", encoding = "utf-8") + + real_lstat = sr.os.lstat + + def windows_shaped_lstat(p, *a, **k): + text = os.fspath(p) + if text != str(config) and text.startswith(str(config) + os.sep): + raise FileNotFoundError(2, "The system cannot find the path specified", text) + return real_lstat(p, *a, **k) + + monkeypatch.setattr(sr.os, "lstat", windows_shaped_lstat) + + assert ( + sr._nothing_at(config / "matplotlibrc") is False + ), "a file where the config dir belongs read as 'nothing there'" + # A genuinely empty, genuinely present directory must still read as empty, or the fix above + # would decline every pin and the guard would stop guarding anything. + empty = tmp_path / "really-empty" + empty.mkdir() + assert sr._nothing_at(empty / "matplotlibrc") is True + # And a path with nothing along it at all is still absence, which is the default install. + assert sr._nothing_at(tmp_path / "never" / "created" / "matplotlibrc") is True 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 c0bf61d914..b3ea8e90d8 100644 --- a/studio/backend/tests/test_setup_cache_env_hf_home.py +++ b/studio/backend/tests/test_setup_cache_env_hf_home.py @@ -37,6 +37,9 @@ _STORAGE_ROOTS_PATH = Path(__file__).resolve().parent.parent / "utils/paths/stor def _isolate_studio_home(monkeypatch, tmp_path): # Keep _setup_cache_env's UV/VLLM mkdirs out of the real ~/.unsloth/studio. monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + # _setup_cache_env writes os.environ directly, so a portable test would leak these forward. + for key in ("UNSLOTH_HOME", "UNSLOTH_PORTABLE", "TORCH_HOME"): + monkeypatch.delenv(key, raising = False) @pytest.fixture(autouse = True) @@ -73,10 +76,27 @@ def _load_storage_roots(): def _clear_hf_env(monkeypatch): - for key in ("HF_HOME", "HF_HUB_CACHE", "HF_XET_CACHE", "HUGGINGFACE_HUB_CACHE"): + for key in ( + "HF_HOME", + "HF_HUB_CACHE", + "HF_XET_CACHE", + "HUGGINGFACE_HUB_CACHE", + "HF_DATASETS_CACHE", + "HF_ASSETS_CACHE", + "HF_MODULES_CACHE", + ): monkeypatch.delenv(key, raising = False) +def _portable_install(monkeypatch, tmp_path): + """Turn the fixture's studio home into a portable one and return its root.""" + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + monkeypatch.delenv("UNSLOTH_PORTABLE", raising = False) + master = tmp_path / "portable" + monkeypatch.setenv("UNSLOTH_HOME", str(master)) + return master + + def test_custom_hf_home_seeds_hub_and_xet(monkeypatch, tmp_path): _clear_hf_env(monkeypatch) custom = tmp_path / "shared" / "huggingface" @@ -127,6 +147,179 @@ def test_whitespace_hf_home_falls_back_to_default(monkeypatch, tmp_path): assert os.environ["HF_HUB_CACHE"] == str(tmp_path / "xdg" / "huggingface" / "hub") +def test_explicit_hf_home_keeps_the_datasets_and_assets_caches(monkeypatch, tmp_path): + # A user who names one HF_HOME gets one Hugging Face cache: assets and datasets derive + # from it, so pinning either under the portable root splits it across two volumes. + _clear_hf_env(monkeypatch) + master = _portable_install(monkeypatch, tmp_path) + chosen = tmp_path / "bigdisk" / "huggingface" + monkeypatch.setenv("HF_HOME", str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_HUB_CACHE"] == str(chosen / "hub") + assert os.environ["HF_XET_CACHE"] == str(chosen / "xet") + assert "HF_DATASETS_CACHE" not in os.environ + assert "HF_ASSETS_CACHE" not in os.environ + assert "HF_MODULES_CACHE" not in os.environ + # Containment is only given up for the root the user named. + assert os.environ["TORCH_HOME"].startswith(str(master)) + + +def test_a_dedicated_cache_var_still_outranks_an_explicit_hf_home(monkeypatch, tmp_path): + _clear_hf_env(monkeypatch) + _portable_install(monkeypatch, tmp_path) + monkeypatch.setenv("HF_HOME", str(tmp_path / "bigdisk" / "huggingface")) + mine = tmp_path / "mine" / "datasets" + monkeypatch.setenv("HF_DATASETS_CACHE", str(mine)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_DATASETS_CACHE"] == str(mine) + + +def test_portable_mode_without_an_explicit_hf_home_still_contains_them(monkeypatch, tmp_path): + # The other side: with no HF_HOME of the user's own, both derive from the host copy + # Unsloth leaves behind and would still write outside the volume. + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + master = _portable_install(monkeypatch, tmp_path) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_DATASETS_CACHE"].startswith(str(master)) + assert os.environ["HF_ASSETS_CACHE"].startswith(str(master)) + # transformers derives this from HF_HOME, which stays on the host, and appends it to + # sys.path, so a trust_remote_code load would leave generated modules outside the volume. + assert os.environ["HF_MODULES_CACHE"].startswith(str(master)) + + +def test_a_normal_install_leaves_the_modules_cache_where_transformers_looks(monkeypatch, tmp_path): + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.delenv("UNSLOTH_HOME", raising = False) + monkeypatch.delenv("UNSLOTH_PORTABLE", raising = False) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert "HF_MODULES_CACHE" not in os.environ + + +def test_an_explicit_modules_cache_outranks_the_portable_default(monkeypatch, tmp_path): + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + _portable_install(monkeypatch, tmp_path) + mine = tmp_path / "mine" / "modules" + monkeypatch.setenv("HF_MODULES_CACHE", str(mine)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_MODULES_CACHE"] == str(mine) + + +@pytest.mark.parametrize("hub_variable", ["HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE"]) +def test_a_hub_only_override_still_contains_the_xet_cache(monkeypatch, tmp_path, hub_variable): + # huggingface_hub derives HF_XET_CACHE from HF_HOME, never from HF_HUB_CACHE, so a hub-only + # override would keep a portable install writing Xet data outside the volume. + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "home" / ".cache")) + master = _portable_install(monkeypatch, tmp_path) + chosen = tmp_path / "bigdisk" / "hub" + monkeypatch.setenv(hub_variable, str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_HUB_CACHE"] == str(chosen) + assert os.environ["HF_XET_CACHE"].startswith(str(master)) + + +def test_an_explicit_xet_cache_outranks_the_portable_default(monkeypatch, tmp_path): + # Containment must not collapse into "always redirect": a named Xet cache is what was asked. + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + _portable_install(monkeypatch, tmp_path) + mine = tmp_path / "bigdisk" / "xet" + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "bigdisk" / "hub")) + monkeypatch.setenv("HF_XET_CACHE", str(mine)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_XET_CACHE"] == str(mine) + + +def test_a_normal_install_leaves_the_xet_cache_in_the_host_home(monkeypatch, tmp_path): + # Portable mode's promise alone: a normal install keeps the platform default so chunks + # shared with plain huggingface_hub still hit. + _clear_hf_env(monkeypatch) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "bigdisk" / "hub")) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import os + + assert os.environ["HF_XET_CACHE"] == str(tmp_path / "xdg" / "huggingface" / "xet") + + +def test_the_libraries_really_derive_these_caches_from_hf_home(monkeypatch, tmp_path): + # Leaving them unset is only correct if huggingface_hub and datasets derive them from + # HF_HOME. Asked of a fresh interpreter: both snapshot their constants at import. + pytest.importorskip("datasets") + _clear_hf_env(monkeypatch) + _portable_install(monkeypatch, tmp_path) + chosen = tmp_path / "bigdisk" / "huggingface" + monkeypatch.setenv("HF_HOME", str(chosen)) + sr = _load_storage_roots() + + sr._setup_cache_env() + + import json + import os + import subprocess + + probe = subprocess.run( + [ + sys.executable, + "-c", + "import json, huggingface_hub.constants as c, datasets.config as d;" + "print(json.dumps({'assets': c.HF_ASSETS_CACHE, 'hub': c.HF_HUB_CACHE," + "'datasets': str(d.HF_DATASETS_CACHE)}))", + ], + env = dict(os.environ), + capture_output = True, + text = True, + check = True, + ) + result = json.loads(probe.stdout.strip().splitlines()[-1]) + + assert Path(result["assets"]) == chosen / "assets" + assert Path(result["datasets"]) == chosen / "datasets" + assert Path(result["hub"]) == chosen / "hub" + + def test_unwritable_hf_home_does_not_crash(monkeypatch, tmp_path): # HF_HOME under a regular file -> mkdir fails; startup must not crash and the # env var is still set (HF surfaces a clear error later, at download time). diff --git a/studio/backend/utils/hf_cache_settings.py b/studio/backend/utils/hf_cache_settings.py index 62effd8e31..73a72e0084 100644 --- a/studio/backend/utils/hf_cache_settings.py +++ b/studio/backend/utils/hf_cache_settings.py @@ -11,6 +11,8 @@ owns an explicit cache snapshot for each operation instead of trying to refresh from __future__ import annotations import os +import stat +import sys import shutil import tempfile import threading @@ -91,7 +93,10 @@ def _environment_paths() -> Optional[HuggingFaceCachePaths]: 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" + # huggingface_hub derives HF_XET_CACHE from HF_HOME, never from HF_HUB_CACHE, so a hub-only + # override would leave the chunk and shard caches in the host home. + xet_home = hf_home if explicit_home else (_portable_cache_home() or hf_home) + xet = _canonical(explicit_xet) if explicit_xet else xet_home / "xet" controlling = next( key for key in ("HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE", "HF_HOME") @@ -108,7 +113,54 @@ def _environment_paths() -> Optional[HuggingFaceCachePaths]: return HuggingFaceCachePaths(display_home, hub, xet, "environment", controlling) +def _absence_is_real(path: Path) -> bool: + """Whether a FileNotFoundError for *path* proves the file is not there. + + On POSIX it always does: a parent component that is a file raises NotADirectoryError, and a + parent that cannot be traversed raises PermissionError, both of which are their own types. + Windows collapses every one of those into ERROR_PATH_NOT_FOUND, which Python surfaces as + FileNotFoundError, so on Windows the exception means "not there, OR somewhere above it is not + a directory, OR I could not look" -- and the caller must only skip its read for the first. + + Decided by the nearest ancestor that can actually be stat'ed: a directory means the tree is + real and the file genuinely is not in it; anything else means a component is a file or a + reparse point and the path could never have existed; an ancestor that cannot be inspected at + all is not proof either way, so it reads as "not proven". Running out of ancestors means + nothing along the path exists, which is the machine that has never opened Studio -- the case + the skip is FOR, so the absence is real there. + """ + current = os.path.dirname(os.fspath(path)) + while current: + try: + info = os.stat(current) + except FileNotFoundError: + parent = os.path.dirname(current) + if parent == current: + return True + current = parent + continue + except OSError: + return False + return stat.S_ISDIR(info.st_mode) + return True + + def _stored_cache_home() -> Optional[Path]: + # get_app_setting CREATES and migrates studio.db, so an unconditional read built one on + # machines that had never opened Studio. os.stat, not Path.exists: only a positively observed + # absence may skip the read, and from 3.14 the predicates report EACCES and EIO as False. + # stat, not lstat: sqlite follows symlinks. + try: + if "storage.studio_db" not in sys.modules: + from utils.paths.storage_roots import studio_db_path + database = studio_db_path() + try: + os.stat(database) + except FileNotFoundError: + if _absence_is_real(database): + return None + except Exception: # noqa: BLE001 - fall through to the read on any doubt + pass try: from storage.studio_db import get_app_setting from utils.account_context import OWNER, run_as @@ -147,6 +199,22 @@ def configured_cache_key() -> str: return "default" +def _portable_cache_home() -> Optional[Path]: + """The HF cache home a portable install uses, else None: a normal install keeps the platform + default so models shared with other tools are not re-downloaded. Lazy import, since + storage_roots reaches into this module during startup.""" + try: + from utils.paths.storage_roots import cache_root, portable_mode + except ImportError: + return None + if not portable_mode(): + return None + try: + return _canonical(cache_root() / "huggingface") + except (OSError, RuntimeError, ValueError): + return None + + def get_hf_cache_paths() -> HuggingFaceCachePaths: env_paths = _environment_paths() if env_paths is not None: @@ -160,7 +228,8 @@ def get_hf_cache_paths() -> HuggingFaceCachePaths: _canonical(xet) if xet else stored / "xet", "studio", ) - home = _default_cache_home() + # Ranks below env vars and Settings: portable mode is a default, not an override. + home = _portable_cache_home() or _default_cache_home() xet = _EXPLICIT_CACHE_ENV.get("HF_XET_CACHE") return HuggingFaceCachePaths( home, diff --git a/studio/backend/utils/node_runtime.py b/studio/backend/utils/node_runtime.py index 6283125ce4..5cb136edf2 100644 --- a/studio/backend/utils/node_runtime.py +++ b/studio/backend/utils/node_runtime.py @@ -38,12 +38,16 @@ def _version_meets_floor(version: str) -> bool: def managed_node_dir() -> Path: - """Isolated Node install dir. Mirrors ``_find_llama_server_binary``: shares a parent with llama.cpp, ```` in custom mode, else legacy ``~/.unsloth``.""" + """Isolated Node install dir. Mirrors ``_find_llama_server_binary``: shares a parent with llama.cpp, ``UNSLOTH_HOME`` when set, else ```` in custom mode, else legacy ``~/.unsloth``.""" legacy_node = Path.home() / ".unsloth" / "node" try: # Lazy import (mirrors _find_llama_server_binary) so this module stays importable even if utils.paths cannot be loaded. - from utils.paths.storage_roots import studio_root + from utils.paths.storage_roots import studio_root, unsloth_home + # setup.sh installs Node at /node, beside studio/. + master = unsloth_home() + if master is not None: + return master / "node" resolved = studio_root() legacy_studio = Path.home() / ".unsloth" / "studio" try: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 7edcb09ba1..05ab40da79 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -3,10 +3,13 @@ from __future__ import annotations +import importlib.util import json import ntpath import os +import platform import re +import stat as stat_module import sys import threading from pathlib import Path, PurePosixPath, PureWindowsPath @@ -53,21 +56,196 @@ def _infer_studio_home_from_venv() -> Path | None: return candidate -def studio_root() -> Path: - """Unsloth install root. +def _resolved(value: str) -> Path: + try: + return Path(value).expanduser().resolve() + except (OSError, ValueError): + return Path(value).expanduser() - Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix - inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins if - both are set (specific signal beats generic alias). + +MASTER_ROOT_NOTE = ".unsloth-master-root" + +_recorded_master_roots: dict[str, Path | None] = {} +_recorded_master_lock = threading.Lock() + + +def _studio_root_without_master() -> Path: + """The studio root as studio_root() would give it with the master root out of the picture. + + studio_root() asks unsloth_home(), so the note reader cannot ask studio_root() back. """ override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() if not override: override = (os.environ.get("STUDIO_HOME") or "").strip() if override: + return _resolved(override) + inferred = _infer_studio_home_from_venv() + if inferred is not None: + return inferred + return Path.home() / ".unsloth" / "studio" + + +def _is_legacy_studio_tree(studio: Path) -> bool: + """Whether the tree a note was read from is the ordinary `~/.unsloth/studio` install. + + A legacy-rooted install has no master root, so a note there says something no reader needs; + both uninstallers already refuse the value it produces. Two things it broke: a one-command + `UNSLOTH_HOME=$HOME/.unsloth` left a note that kept `portable_mode()` true for good, moving + the one cache this module promises not to move; and a note naming any ANCESTOR of the tree + (`$HOME` passes containment) reintroduced through the filesystem what `process.rs` scrubs + from every managed spawn, which no `env_remove` can reach. Hence keyed on the TREE, not on + the recorded value. An explicit UNSLOTH_HOME never gets here: unsloth_home() returns the + override first. + """ + try: + return studio.resolve() == (Path.home() / ".unsloth" / "studio").resolve() + except (OSError, RuntimeError, ValueError): + return False + + +def _recorded_master_root() -> Path | None: + """The master root setup recorded inside the studio tree, or None. + + UNSLOTH_HOME is settable for a single command (`UNSLOTH_HOME=/mnt/portable unsloth studio + update`) and nothing persists it, so without this every later launch resolved node, llama.cpp + and whisper.cpp somewhere other than the siblings of studio/ setup created. setup already + writes the root to share/.unsloth-master-root for the uninstallers. + + Honoured only while it still describes reality: the root must exist, the Studio directory the + note was read from must lie INSIDE it, and it must not be the legacy default. Containment + rather than an exact /studio match, since the flat layout (both variables naming one + directory) is supported; the uninstallers apply the same rule. + """ + studio = _studio_root_without_master() + key = str(studio) + with _recorded_master_lock: + if key in _recorded_master_roots: + return _recorded_master_roots[key] + found: Path | None = None + # A miss is cached; a failure to LOOK is not. EACCES on share/ or EIO off a mount that came + # back is a fact about one instant, and cached it pinned the backend to the legacy runtime + # paths for its whole lifetime, with forget_recorded_master_root() the only way out and + # nothing calling it in production. + definitive = True + try: + recorded = (studio / "share" / MASTER_ROOT_NOTE).read_text(encoding = "utf-8").strip() + except (FileNotFoundError, NotADirectoryError, ValueError, UnicodeDecodeError): + recorded = "" + except OSError: + recorded = "" + definitive = False + if recorded: + master = _resolved(recorded) try: - return Path(override).expanduser().resolve() + here = studio.resolve() + if ( + master.is_dir() + and (here == master or master in here.parents) + and not _is_legacy_studio_tree(studio) + ): + found = master except (OSError, ValueError): - return Path(override).expanduser() + found = None + if definitive: + with _recorded_master_lock: + _recorded_master_roots[key] = found + return found + + +def forget_recorded_master_root() -> None: + """Drop the cached note reading. For tests, and for an installer that writes the note into + a tree this process has already looked at.""" + with _recorded_master_lock: + _recorded_master_roots.clear() + + +def unsloth_home() -> Path | None: + """The master root, or None. STUDIO_HOME is its studio/ child; llama.cpp, node and + whisper.cpp are SIBLINGS of studio/, the layout setup.sh already gives UNSLOTH_HOME.""" + override = (os.environ.get("UNSLOTH_HOME") or "").strip() + if override: + return _resolved(override) + return _recorded_master_root() + + +_PORTABLE_ON_VALUES = ("1", "true", "yes", "on") +_PORTABLE_OFF_VALUES = ("0", "false", "off", "no") + +# portable_mode() runs on every cache-var lookup, so warn once per process, not once per call. +_warned_unrecognized_portable = False + + +def _warn_unrecognized_portable(raw: str) -> None: + global _warned_unrecognized_portable + if _warned_unrecognized_portable: + return + _warned_unrecognized_portable = True + # Not "to leave it off": an off value declines to turn portable mode on by itself, it does + # not veto the master root that carries it. + logger.warning( + "Ignoring UNSLOTH_PORTABLE=%r: expected one of %s to turn portable mode on, or one of " + "%s to leave that choice to UNSLOTH_HOME, which turns it on when it names a master root.", + raw, + "/".join(_PORTABLE_ON_VALUES), + "/".join(_PORTABLE_OFF_VALUES), + ) + + +def portable_mode() -> bool: + """Whether this install keeps everything under one directory. Implied by UNSLOTH_HOME, and + settable on its own so an existing install can opt in.""" + # Case-folded: UNSLOTH_PORTABLE=FALSE read as "on" moves the caches out from under a user + # who asked for the opposite. + raw = (os.environ.get("UNSLOTH_PORTABLE") or "").strip() + value = raw.lower() + if value in _PORTABLE_ON_VALUES: + return True + if value and value not in _PORTABLE_OFF_VALUES: + _warn_unrecognized_portable(raw) + # Unrecognized means no opinion, as an off value does: neither vetoes a real master root. + return unsloth_home() is not None + + +# Warned once per distinct pair rather than per call: studio_root() runs many times a request, +# and both variables can be rebound inside one process. +_warned_root_conflicts: set[tuple[str, str]] = set() + + +def _warn_root_conflict(resolved: Path, master: Path) -> None: + key = (str(resolved), str(master)) + if key in _warned_root_conflicts: + return + _warned_root_conflicts.add(key) + # Not fatal: failing here would break a resolver called at import time. + logger.warning( + "UNSLOTH_STUDIO_HOME (%s) is outside UNSLOTH_HOME (%s); this " + "install is not self-contained.", + resolved, + master, + ) + + +def studio_root() -> Path: + """Unsloth install root. + + Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then UNSLOTH_HOME's studio/ child, then + sys.prefix inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME outranks both: it + names this exact directory, while the others only name the tree it sits in. + """ + override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() + if not override: + override = (os.environ.get("STUDIO_HOME") or "").strip() + if override: + resolved = _resolved(override) + master = unsloth_home() + # Path.parents excludes the path itself, so a flat layout (both naming one root) would + # otherwise warn on every call. + if master is not None and master != resolved and master not in resolved.parents: + _warn_root_conflict(resolved, master) + return resolved + master = unsloth_home() + if master is not None: + return master / "studio" inferred = _infer_studio_home_from_venv() if inferred is not None: return inferred @@ -510,6 +688,367 @@ def well_known_model_dirs() -> list[Path]: return _existing_dirs(candidates, resolve = True) +def _user_set_hf_home() -> bool: + """Whether HF_HOME was set by the user rather than seeded by Unsloth. + + initialize_hf_cache_environment fills a blank HF_HOME first, so the import-time snapshot is + the only record of who chose it. + """ + try: + from utils.hf_cache_settings import _EXPLICIT_CACHE_ENV + except ImportError: + return False + return bool(_EXPLICIT_CACHE_ENV.get("HF_HOME")) + + +def _portable_cache_defaults(root: Path) -> dict[str, str]: + """Cache vars that only move under the root in portable mode, since they hold shared user data + or large re-downloads. The hub and xet caches are handled inside hf_cache_settings, which must + agree with the Settings UI. HF_HOME never moves: credentials should not follow a cache onto a + removable volume. + """ + if not portable_mode(): + return {} + if _user_set_hf_home(): + # Assets, datasets and modules derive from an explicit HF_HOME: pinning them splits one + # deliberately chosen cache across two volumes. + return {"TORCH_HOME": str(root / "torch")} + return { + "HF_DATASETS_CACHE": str(root / "huggingface" / "datasets"), + "HF_ASSETS_CACHE": str(root / "huggingface" / "assets"), + # transformers.utils.hub reads this at import and appends it to sys.path, so a + # trust_remote_code load leaves generated modules on the host without it. + "HF_MODULES_CACHE": str(root / "huggingface" / "modules"), + "TORCH_HOME": str(root / "torch"), + } + + +def _triton_cache_defaults(root: Path) -> dict[str, str]: + """Triton's regenerable directories, named one at a time. + + Not TRITON_HOME (triton-lang/triton#4265): it would take ~/.triton/override with it, and a + TRITON_KERNEL_OVERRIDE=1 run would silently fall back to the compiler's own output. The + dedicated variables outrank the derivation (triton/knobs.py cache_knobs). + """ + if (os.environ.get("TRITON_HOME") or "").strip(): + # Whoever moved the whole tree meant the cache with it; TRITON_CACHE_DIR would outrank it. + return {} + return { + "TRITON_CACHE_DIR": str(root / "triton"), + # A sibling, not a child: dumps are asked for by hand and should outlive a cache wipe. + "TRITON_DUMP_DIR": str(root / "triton-dump"), + } + + +def _nothing_at(path: Path, *, ending: str = "") -> bool: + """Whether *path* positively holds nothing, the only state that licenses a pin. + + Not Path.exists/is_dir/glob: they report ENOTDIR, ELOOP, EBADF (and from 3.14 EACCES and EIO) + as absence, so a directory we merely cannot inspect would read as empty and the redirect would + hide the user's files. os.lstat and os.scandir raise on all of it. lstat rather than stat: a + dangling symlink is still something the user put there. *ending* asks instead whether the + DIRECTORY holds no entry with that suffix; pass it lowercase, as Windows glob matched + case-insensitively. + """ + try: + if not ending: + os.lstat(path) + return False + # lstat, not scandir, at the directory itself: scandir follows a link, so an empty or + # unmounted target would read as an empty directory and hide what is placed there. + if stat_module.S_ISLNK(os.lstat(path).st_mode) or _is_reparse_point(path): + return False + with os.scandir(path) as entries: + return not any(entry.name.lower().endswith(ending) for entry in entries) + except FileNotFoundError: + # Not absence on its own: Windows collapses NotADirectoryError and PermissionError into + # ERROR_PATH_NOT_FOUND, so with a FILE where ~/.matplotlib belongs POSIX declined the pin + # and Windows took it. + return _nothing_above(path) + except (OSError, ValueError): + # Could not look. Declining a pin costs a shared cache directory; taking one wrongly + # hides the configuration or datasets underneath it. + return False + + +def _nothing_above(path: Path) -> bool: + """Whether a FileNotFoundError for *path* really means nothing is there. + + Decided by the nearest ancestor that can be stat'ed: a directory means real absence, anything + else means a component is a file or a reparse point, and an ancestor that cannot be inspected + is proof of neither. Same rule as hf_cache_settings._absence_is_real. + """ + current = os.path.dirname(os.fspath(path)) + while current: + try: + info = os.stat(current) + except FileNotFoundError: + parent = os.path.dirname(current) + if parent == current: + return True + current = parent + continue + except OSError: + return False + return stat_module.S_ISDIR(info.st_mode) + return True + + +def _matplotlib_config_dir() -> Path | None: + """Where matplotlib reads matplotlibrc and stylelib/ from when MPLCONFIGDIR is unset, or None + when this machine has no such directory. Mirrors _get_config_or_cache_dir: XDG config base on + Linux/FreeBSD, %LOCALAPPDATA% on Windows but keeping a pre-existing ~/.matplotlib there. + + None means matplotlib falls back to a temporary directory, so a pin can strand nothing. The + Windows branch is matplotlib 3.11's and the pinned 3.10.9 uses ~/.matplotlib; that + disagreement only ever costs us the pin, it never hides a file matplotlib reads. + """ + # XDG_CONFIG_HOME ahead of Path.home(), as _get_xdg_config_dir does: an install that sets it + # has a config dir even with no resolvable home. + if sys.platform.startswith(("linux", "freebsd")): + base = (os.environ.get("XDG_CONFIG_HOME") or "").strip() + if base: + return Path(base) / "matplotlib" + try: + home = Path.home() + except (OSError, RuntimeError): + return None + if sys.platform.startswith(("linux", "freebsd")): + return home / ".config" / "matplotlib" + if sys.platform == "win32": + legacy = home / ".matplotlib" + # Not is_dir(): an unreadable ~/.matplotlib raises before 3.14 and reads as absent from + # 3.14, which sends the probe to a %LOCALAPPDATA% that is empty on exactly the machines + # where the old directory holds the rc. + if not _nothing_at(legacy): + return legacy + local_app_data = (os.environ.get("LOCALAPPDATA") or "").strip() + return Path(local_app_data) / "matplotlib" if local_app_data else legacy + return home / ".matplotlib" + + +def _matplotlib_defaults(root: Path) -> dict[str, str]: + """MPLCONFIGDIR, unless matplotlib's own directory holds user configuration. + + The one variable moves the CONFIG directory as well as the cache, so pinning it would drop a + user matplotlibrc and every custom style, silently changing the loss plots + core/training/training.py draws. matplotlib creates the directory on import, so contents + decide, not existence. + """ + managed = root / "matplotlib" + pinned = {"MPLCONFIGDIR": str(managed)} + # Our own configuration first: the legacy probe re-runs every launch, so on its own a + # ~/.config/matplotlib created later by another tool takes over a matplotlibrc written HERE, + # and the plot style flips back and forth across launches. + if not ( + _nothing_at(managed / "matplotlibrc") + and _nothing_at(managed / "stylelib", ending = ".mplstyle") + ): + return pinned + config_dir = _matplotlib_config_dir() + if config_dir is not None and not ( + _nothing_at(config_dir / "matplotlibrc") + and _nothing_at(config_dir / "stylelib", ending = ".mplstyle") + ): + return {} + return pinned + + +def _is_reparse_point(path: Path) -> bool: + """A Windows junction, which is_symlink() answers False for. + + os.path.isjunction arrived in 3.12 and the Studio venv can be 3.11, so its absence means + "no junctions to worry about" rather than an error. + """ + isjunction = getattr(os.path, "isjunction", None) + if isjunction is None: + return False + try: + return bool(isjunction(path)) + except (OSError, ValueError): + return False + + +def _data_designer_in_use(home: Path) -> bool: + """Whether the managed Data Designer home holds work worth keeping. + + _setup_cache_env creates this directory and its managed-assets child on the first launch, so + existence alone would pin a home nobody has written to. A home we cannot list counts as in + use: dropping the pin would hide the recipes here behind a re-seeded ~/.data-designer. + """ + try: + entries = list(home.iterdir()) + except FileNotFoundError: + return False + except (OSError, ValueError): + return True + for entry in entries: + try: + if entry.name != "managed-assets": + return True + # A link here is the user redirecting their assets, which is state; is_dir() follows + # it, so an empty target read as the untouched layout and took the redirect with it. + if entry.is_symlink() or _is_reparse_point(entry): + return True + if not entry.is_dir(): + return True + if any(entry.iterdir()): + return True + except FileNotFoundError: + continue + except (OSError, ValueError): + return True + return False + + +def _data_designer_defaults(root: Path) -> dict[str, str]: + """Data Designer's home, unless the user already has one. + + Not a cache: repointing an existing ~/.data-designer hides its yaml configs and multi-GB + parquet behind a re-seeded default. MANAGED_ASSETS_PATH derives from the home, so it is only + ours to set when the home is. data_designer.config.utils.constants reads both at IMPORT time. + """ + if (os.environ.get("DATA_DESIGNER_HOME") or "").strip(): + return {} + home = root.parent / "data-designer" + pinned = { + "DATA_DESIGNER_HOME": str(home), + "DATA_DESIGNER_MANAGED_ASSETS_PATH": str(home / "managed-assets"), + } + # Our own populated home first: the legacy probe re-runs every launch, so a ~/.data-designer + # created later by a standalone run would otherwise take over the recipes written here. + if _data_designer_in_use(home): + return pinned + try: + legacy = Path.home() / ".data-designer" + except (OSError, RuntimeError): + # No home, so the pin can hide nothing: data_designer's default is off the same call. + return pinned + return pinned if _nothing_at(legacy) else {} + + +def _path_safe(value: str) -> str: + """A directory-name-safe rendering of a build field.""" + return re.sub(r"[^A-Za-z0-9.]+", "-", value) + + +def _torch_version_fields() -> dict[str, str]: + """The build identity torch.version exposes, read without importing torch. + + Runs on a startup path that executes before torch exists in a fresh venv, so it stays a file + read. torch/version.py assigns only literals, but whether it annotates them + (``cuda: Optional[str] = ...``) varies by release, hence a regex rather than ast or exec. + """ + origin = getattr(importlib.util.find_spec("torch"), "origin", None) + if not origin: + return {} + text = (Path(origin).parent / "version.py").read_text(encoding = "utf-8") + found = re.findall( + r"""^(__version__|debug|cuda|hip|xpu)\s*(?::[^=\n]+)?=\s*([^\s#]+)""", + text, + re.MULTILINE, + ) + return {name: value.strip("'\"") for name, value in found} + + +def _torch_accelerator_tag(fields: dict[str, str]) -> str: + """torch's own cu_str, widened to the runtimes it declines to name: cpp_extension picks 'cpu' + whenever version.cuda is unset, filing a ROCm build beside a real CPU one. hip first, as + torch main prioritises ROCm.""" + for field, prefix in (("hip", "rocm"), ("cuda", "cu"), ("xpu", "xpu")): + value = fields.get(field) + if not value or value == "None": + continue + # torch spells the CUDA version without dots: 12.8 -> cu128. + return prefix + _path_safe(value.replace(".", "") if field == "cuda" else value) + return "cpu" + + +def _torch_runtime_tag() -> str: + """Name the extension cache after the runtime that builds into it. + + torch.utils.cpp_extension._get_build_directory appends a ``py_`` folder to + the DEFAULT root only, never to a TORCH_EXTENSIONS_DIR we supply, so pinning a flat path drops + the isolation that keeps a py313/cu128 build from being loaded by a py312/cu126 one. The + accelerator comes from the generated cuda/hip fields, not from a local version segment: + conda-forge's CPU and CUDA packages of one release share a __version__. __version__ stays in + the tag anyway, since segments like +cpu.cxx11.abi mark ABI splits no other field records. + """ + tag = f"py{sys.version_info.major}{sys.version_info.minor}{getattr(sys, 'abiflags', '')}" + # Architecture too, which torch's py_ omits: an arm64 and a Rosetta x86_64 + # python agree on every other field, so ninja reads the other's build as up to date. + tag += "_" + _path_safe(f"{sys.platform}-{platform.machine() or 'unknown'}") + try: + fields = _torch_version_fields() + except (ImportError, OSError, ValueError, AttributeError): + # No torch yet, or a half-built tree; the interpreter tag alone still isolates. + return tag + if not fields: + return tag + tag += "_" + _torch_accelerator_tag(fields) + version = fields.get("__version__") + if version: + tag += "_" + _path_safe(version) + if fields.get("debug") == "True": + # A debug build keeps the soname of a release one but not its ABI. + tag += "_debug" + return tag + + +# Caches whose path is pasted into a compiler command line, unquoted, by somebody else's code: +# torch/_inductor/cpp_builder.py joins the g++ invocation with spaces and reparses it with +# shlex.split, so a root containing a space splits in two and the build fails outright +# ("g++: fatal error: input file .../my is the same as output file"). The quoting is PyTorch's +# to fix; skipping the pin just returns to the temporary directory Inductor used before. Not a +# rare shape: "C:\Users\First Last" is an ordinary Windows account name. +_TOOLCHAIN_PATH_KEYS = frozenset( + { + "TORCHINDUCTOR_CACHE_DIR", + "TORCH_EXTENSIONS_DIR", + "TRITON_CACHE_DIR", + "TRITON_DUMP_DIR", + "TRITON_HOME", + "CUDA_CACHE_PATH", + } +) + + +def _usable_dir(value: str) -> bool: + """Whether a path we generated is actually a directory the toolchain can compile into. + + A real create, not is_dir() alone. torch falls back to its own temporary directory only when + the variable is UNSET, then calls os.makedirs(exist_ok = True), which succeeds on an existing + read-only directory and leaves every later write to fail + (torch/_inductor/runtime/cache_dir_utils.py), so publishing an unwritable path is worse than + publishing none. Same rule install.sh states for the uv cache probe. + """ + try: + if not Path(value).is_dir(): + return False + except (OSError, ValueError): + return False + try: + handle, probe = tempfile.mkstemp(dir = value, prefix = ".unsloth-write-probe.") + except (OSError, ValueError): + return False + # Guarded like the unlink below: an EIO or ENOSPC on close escaping here would take the + # whole backend start with it, for a probe whose job is to answer yes or no. + try: + os.close(handle) + except OSError: + return False + try: + os.unlink(probe) + except OSError: + pass + return True + + +def _toolchain_unsafe(key: str, value: str) -> bool: + """Whether pinning *key* to *value* would hand a compiler a path it cannot parse.""" + return key in _TOOLCHAIN_PATH_KEYS and any(ch.isspace() for ch in value) + + def _setup_cache_env() -> None: """Set cache env vars for HuggingFace, uv, and vLLM. @@ -529,11 +1068,37 @@ def _setup_cache_env() -> None: # cache landed in the user home. Must be set before unsloth_zoo.compiler imports: it reads the value at import # time and puts it on sys.path. "UNSLOTH_COMPILE_LOCATION": str(root.parent / "compiled_cache"), + # Regenerable and process-scoped. Shared user data (HF hub cache, torch.hub checkpoints) + # stays where the other tools look, except in portable mode. + "TORCHINDUCTOR_CACHE_DIR": str(root / "torchinductor"), + # Tagged: torch only inserts its own ABI folder when TORCH_EXTENSIONS_DIR is unset, so a + # flat pin would let two runtimes sharing this root import each other's .so. + "TORCH_EXTENSIONS_DIR": str(root / "torch-extensions" / _torch_runtime_tag()), + # NVIDIA's JIT compile cache; ~/.nv/ComputeCache otherwise. + "CUDA_CACHE_PATH": str(root / "cuda"), + "NUMBA_CACHE_DIR": str(root / "numba"), } + defaults.update(_matplotlib_defaults(root)) + defaults.update(_triton_cache_defaults(root)) + defaults.update(_data_designer_defaults(root)) + defaults.update(_portable_cache_defaults(root)) for key, value in defaults.items(): # Blank counts as unset: an inherited KEY= would otherwise pin the cache to "", which puts an empty entry on # sys.path and sends the compiler to the system temp directory instead. if not (os.environ.get(key) or "").strip(): + # An explicit value is still honoured above: the caller chose it, and only a default + # we invented is ours to withhold. + if _toolchain_unsafe(key, value): + logger.debug( + "leaving %s unset: %s contains whitespace, which the C++ builders " + "paste into a command line unquoted", + key, + value, + ) + # Popped, not blanked: Inductor reads " " as a relative path and hands it to + # the very command line being refused here. + os.environ.pop(key, None) + continue os.environ[key] = value # Best-effort: a non-writable custom HF_HOME must not crash startup try: @@ -550,6 +1115,14 @@ def _setup_cache_env() -> None: (Path(value) / CACHE_MARKER).touch(exist_ok = True) except (OSError, ImportError): pass + # A toolchain path we invented and could not make is worse than none: torch treats + # it as authoritative and every compile fails, where unset would have used the + # library's temporary cache. Only these keys -- UNSLOTH_COMPILE_LOCATION unset falls + # back to a CWD-relative name, the bug that pin exists to fix, and the data roots + # have no library default to fall back to. + if key in _TOOLCHAIN_PATH_KEYS and not _usable_dir(value): + logger.debug("leaving %s unset: %s is not a usable directory", key, value) + os.environ.pop(key, None) def setup_cache_env() -> None: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 99044e7deb..73cde3c11c 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -763,13 +763,32 @@ function Test-UnslothCmdShimFile { return ($text -like "*unsloth-studio-managed-launcher*" -and $text -like "*from unsloth_cli import app*") } -# Explicit staging root, shared default cache, or the custom Unsloth home's tree. +# The master root storage_roots.unsloth_home() reads. llama.cpp, node and whisper.cpp sit +# BESIDE studio\ under it, so deriving them from $StudioHome would put them one level below +# where every runtime resolver looks. +function Get-MasterRootOverride { + if ([string]::IsNullOrWhiteSpace($env:UNSLOTH_HOME)) { return $null } + $value = $env:UNSLOTH_HOME.Trim() + if ($value -eq "~") { + $value = $env:USERPROFILE + } elseif ($value -like "~/*" -or $value -like "~\*") { + $value = (Join-Path $env:USERPROFILE $value.Substring(1).TrimStart('/', '\')) + } + return (Get-CanonicalDir -Path $value) +} + +# Explicit staging root, the master root, the shared default cache, or the custom Unsloth +# home's tree. function Get-ManagedLlamaCppDir { param([AllowNull()][string]$StagingRoot = $null) if ($StagingRoot) { return (Join-Path $StagingRoot "llama.cpp") } + $masterRoot = Get-MasterRootOverride + if ($masterRoot) { + return (Join-Path $masterRoot "llama.cpp") + } if (-not (Test-StudioHomeIsCustom)) { return (Join-Path $env:USERPROFILE ".unsloth\llama.cpp") } @@ -785,8 +804,9 @@ function Invoke-ManagedLlamaCppPreflight { $dir = Get-ManagedLlamaCppDir -StagingRoot $StagingRoot if ((Get-LlamaCppInstallReadState -Path $dir) -ne "Denied") { return $null } Write-StudioLine "" - # A denied custom home cannot be claimed as an Unsloth-managed cache. - $homeIsCustom = Test-StudioHomeIsCustom + # A denied custom home cannot be claimed as an Unsloth-managed cache. Computed rather + # than read off $RuntimeRootIsCustom: this runs beside the line that defines it. + $homeIsCustom = (Test-StudioHomeIsCustom) -or [bool](Get-MasterRootOverride) # Preserve user-supplied wording when either override names this tree, or # names a build inside it: moving or deleting this folder takes that build # with it, and the later --with-llama-cpp-dir check then aborts on a path @@ -2792,9 +2812,170 @@ $StudioStaleMarker = ".unsloth-studio-stale" $NoTorchMarker = ".unsloth-no-torch" $LegacyStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" $StudioHomeIsCustom = Test-StudioHomeIsCustom +# Ownership applies to node\, llama.cpp\ and whisper.cpp\ whenever a master root moves them, +# even with $StudioHome left at the legacy path: $StudioHomeIsCustom is false there, and false +# is what licenses the installers to replace and delete without the Unsloth-owned marker. The +# Studio home itself, and the venvs under it, keep the other flag. Mirrors setup.sh. +$RuntimeRootIsCustom = $StudioHomeIsCustom +# Keyed on where the runtimes LAND, not on whether a master root was named: UNSLOTH_HOME set to +# the root an install already uses moves nothing, and calling that custom would demand an owner +# marker from a legacy source-built .unsloth\llama.cpp that predates markers. The comparison +# setup.sh makes; taking any non-empty master root instead was a real divergence. +# +# Staging is excluded for the same reason, not as an exception to it: Get-ManagedLlamaCppDir +# gives $StageRoot precedence over the master root, so during a staged update the master root is +# not where anything lands. ASSIGNED, not merely raised, or a custom Studio home whose runtimes +# land in the legacy root keeps the flag true and demands markers from exactly the pre-marker +# tree this comparison exists to spare. +$_masterRootForOwnership = Get-MasterRootOverride +if ($_masterRootForOwnership -and -not $StageRoot) { + # Canonicalised the same way Get-MasterRootOverride canonicalises its answer, or a + # junctioned profile compares unequal to itself. Case-insensitively, since two Windows + # paths differing only in case are one directory. + $_legacyRuntimeRoot = Get-CanonicalDir -Path (Join-Path $env:USERPROFILE ".unsloth") + $RuntimeRootIsCustom = ($_masterRootForOwnership -ine $_legacyRuntimeRoot) +} $LlamaCppDir = Get-ManagedLlamaCppDir -StagingRoot $StageRoot $UnslothHome = Split-Path -Parent $LlamaCppDir +# Record the master root inside the Studio tree, for the uninstaller: UNSLOTH_HOME can be set +# for a single command, and without the note a later uninstall removed the Studio tree and +# stranded multi-gigabyte runtimes beside it. +# +# Only for a master root (the other branches derive the root from paths the uninstaller already +# knows), and only when a reader will honour it. setup.sh's _master_root_note_is_honoured holds +# the same two rules for the same reasons; keep them together. +function Test-MasterRootNoteIsHonoured { + param([string]$Root, [string]$StudioRoot) + $norm = Get-CanonicalDir -Path $Root + if (-not $norm) { return $false } + $here = Get-CanonicalDir -Path $StudioRoot + if (-not $here) { return $false } + # The legacy default is the one root every reader finds without help, and both uninstallers + # refuse it outright. Recording it turned a one-command UNSLOTH_HOME=%USERPROFILE%\.unsloth + # into a permanent portable install: every later bare launch read it back and HF_HUB_CACHE + # moved off the shared Hugging Face cache this change promises not to move. + $legacy = Get-CanonicalDir -Path (Join-Path $env:USERPROFILE ".unsloth") + if ($legacy -and ($norm -ieq $legacy)) { return $false } + # Keyed on the TREE as well, because that is what the readers key on: storage_roots' + # _is_legacy_studio_tree and the CLI both decline ANY note found in the legacy Studio tree, + # whatever it records. Without this, UNSLOTH_HOME=%USERPROFILE% passed containment (the + # legacy tree is inside the profile) and is not the legacy root, so the note was written, + # warned about nothing, and was then honoured by nobody. + $legacyStudio = Get-CanonicalDir -Path (Join-Path $env:USERPROFILE ".unsloth\studio") + if ($legacyStudio -and ($here -ieq $legacyStudio)) { + Write-StudioLine " note: the managed runtimes were installed under $norm, but Studio is the" -ForegroundColor Yellow + Write-StudioLine " default install at $here, where no reader honours a recorded root." -ForegroundColor Yellow + Write-StudioLine " That root cannot be recorded, so a later launch or uninstall will" -ForegroundColor Yellow + Write-StudioLine " not find them. Set UNSLOTH_STUDIO_HOME to $norm\studio, or re-set" -ForegroundColor Yellow + Write-StudioLine " UNSLOTH_HOME whenever you run Unsloth." -ForegroundColor Yellow + return $false + } + # A note must describe the tree it is written into: every reader requires the Studio + # directory to lie INSIDE the root it names, so a tree copied between master roots cannot aim + # a removal at the original install. install.ps1 does not read UNSLOTH_HOME yet, so a master + # root on an ordinary install leaves Studio at the legacy path with the runtimes elsewhere, + # which is that shape: say so, rather than leave a file that behaves as no note at all. + $sep = [System.IO.Path]::DirectorySeparatorChar + if (($here -ieq $norm) -or $here.StartsWith($norm + $sep, [System.StringComparison]::OrdinalIgnoreCase)) { + return $true + } + Write-StudioLine " note: the managed runtimes were installed under $norm, but Studio lives at" -ForegroundColor Yellow + Write-StudioLine " $here, which is outside it. That root cannot be recorded, so a later" -ForegroundColor Yellow + Write-StudioLine " launch or uninstall will not find them. Set UNSLOTH_STUDIO_HOME to" -ForegroundColor Yellow + Write-StudioLine " $norm\studio, or re-set UNSLOTH_HOME whenever you run Unsloth." -ForegroundColor Yellow + return $false +} +if ((Get-MasterRootOverride) -and -not $StageRoot -and + (Test-MasterRootNoteIsHonoured -Root $UnslothHome -StudioRoot $StudioHome)) { + try { + $noteDir = Join-Path $StudioHome "share" + if (-not (Test-Path -LiteralPath $noteDir -PathType Container)) { + [void][System.IO.Directory]::CreateDirectory($noteDir) + } + # Staged then renamed: a reader that caught a half-written note would name a truncated + # path, and this note licenses deletions. + $notePath = Join-Path $noteDir ".unsloth-master-root" + $noteBytes = [System.Text.Encoding]::UTF8.GetBytes($UnslothHome + [Environment]::NewLine) + # Nothing to do when the note already says this, which is every run after the first. + # Move-Item -Force is NOT an atomic replace: the provider deletes the destination then + # moves the source, so an interruption leaves NO note and a backend that looks next + # caches "no master root". Re-entering that window to rewrite identical bytes is the + # avoidable half. + $noteCurrent = $null + try { + if (Test-Path -LiteralPath $notePath -PathType Leaf) { + $noteCurrent = [System.IO.File]::ReadAllBytes($notePath) + } + } catch { $noteCurrent = $null } + $noteSame = $false + if ($null -ne $noteCurrent -and $noteCurrent.Length -eq $noteBytes.Length) { + $noteSame = $true + for ($i = 0; $i -lt $noteBytes.Length; $i++) { + if ($noteCurrent[$i] -ne $noteBytes[$i]) { $noteSame = $false; break } + } + } + if (-not $noteSame) { + # An unpredictable name opened CreateNew, not "$notePath.$PID" through WriteAllText: + # the predictable pair lets anyone who can write share/ precreate that name as a link + # and have this write land on its target. Move-Item -Force replaces a link at the + # final path rather than following it. + $noteTmp = Join-Path $noteDir (".unsloth-master-root." + [System.IO.Path]::GetRandomFileName()) + try { + $noteStream = [System.IO.File]::Open( + $noteTmp, + [System.IO.FileMode]::CreateNew, + [System.IO.FileAccess]::Write, + [System.IO.FileShare]::None + ) + try { + $noteStream.Write($noteBytes, 0, $noteBytes.Length) + } finally { + $noteStream.Dispose() + } + # Move-Item -Force, not [IO.File]::Move with an overwrite flag: that overload is + # .NET Core only and setup.ps1 still runs under Windows PowerShell 5.1. + Move-Item -LiteralPath $noteTmp -Destination $notePath -Force + $noteTmp = $null + } finally { + # The staging file is unreadable to every reader but still occupies share/, and + # nothing else ever removes it. Cleared whether the move threw or not. + if ($noteTmp) { Remove-Item -LiteralPath $noteTmp -Force -ErrorAction SilentlyContinue } + } + } + } catch { + Write-StudioLine " note: could not record the master root for uninstall: $($_.Exception.Message)" -ForegroundColor DarkGray + } +} + +# Clear a note an earlier run left naming the legacy default root, whatever this run was asked +# to do: the gate above stops new ones, but an install that already has one keeps reading it +# back and no later update would reach the write block to correct it. Narrow on purpose -- both +# uninstallers already refuse this value, so removing it takes nothing a reader could act on. +try { + $staleNote = Join-Path (Join-Path $StudioHome "share") ".unsloth-master-root" + if (Test-Path -LiteralPath $staleNote -PathType Leaf) { + # -Encoding UTF8: the writer emits BOM-less UTF-8 and Windows PowerShell 5.1 decodes + # BOM-less input with the ANSI code page, so a root holding non-ASCII read back + # mangled and never matched the value this sweep exists to clear. + $staleValue = (Get-Content -LiteralPath $staleNote -TotalCount 1 -Encoding UTF8 -ErrorAction Stop) + if ($staleValue) { $staleValue = $staleValue.Trim() } + if ($staleValue) { + if ($staleValue -eq "~") { $staleValue = $env:USERPROFILE } + elseif ($staleValue -like "~/*" -or $staleValue -like "~\*") { + $staleValue = (Join-Path $env:USERPROFILE $staleValue.Substring(1).TrimStart('/', '\')) + } + $staleNorm = Get-CanonicalDir -Path $staleValue + $legacyNorm = Get-CanonicalDir -Path (Join-Path $env:USERPROFILE ".unsloth") + if ($staleNorm -and $legacyNorm -and ($staleNorm -ieq $legacyNorm)) { + Remove-Item -LiteralPath $staleNote -Force -ErrorAction SilentlyContinue + } + } + } +} catch { + # A note this run cannot read is one it has no business deleting. +} + $WithLlamaCppDir = $null $llamaPreflightFailure = Invoke-ManagedLlamaCppPreflight -StagingRoot $StageRoot if ($llamaPreflightFailure) { @@ -4434,8 +4615,12 @@ $NodeSource = $null if (-not $IsPipInstall) { # Put Node beside the Unsloth root. OXC can still need npm when the # frontend build is skipped. + $_masterRoot = Get-MasterRootOverride if ($StageRoot) { $NodeParent = $StageRoot + } elseif ($_masterRoot) { + # Same derivation as Get-ManagedLlamaCppDir: beside studio\, not inside it. + $NodeParent = $_masterRoot } else { if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $NodeOverride = $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $NodeOverride = $env:STUDIO_HOME.Trim() } @@ -4701,11 +4886,34 @@ if ($NeedNodeForSetup) { } elseif ($NodeSource -eq "bundled") { New-Item -ItemType Directory -Force -Path $NodeParent -ErrorAction SilentlyContinue | Out-Null # Minimal ownership guard; never os.replace over a user-owned dir. - if ($NodeOverride -and (Test-Path -LiteralPath $NodeDir -PathType Container)) { + # $RuntimeRootIsCustom as well as $NodeOverride: the master-root branch above leaves + # $NodeOverride null, so \node reached install_node_prebuilt.py's whole-directory + # replacement with none of the ownership evidence a custom UNSLOTH_STUDIO_HOME requires, + # and the master root is the user's directory too. Or, not replacing: a custom Studio + # home keeps the guard it had. + # + # Anything at the path counts as occupied, not just a directory: _swap_into_place renames + # whatever it finds out of the way, so a file or symlink named `node` was displaced by a + # guard that only asked about containers. Ownership evidence lives inside a directory, so + # a non-directory can never carry it and is refused outright. + # Get-Item -Force, not Test-Path: under Windows PowerShell 5.1 Test-Path reports whether + # the TARGET resolves, so a dangling link named `node` reads as absent and + # install_node_prebuilt.py installs through it. + $nodeEntry = if ($NodeOverride -or $RuntimeRootIsCustom) { + Get-Item -LiteralPath $NodeDir -Force -ErrorAction SilentlyContinue + } else { $null } + if ($nodeEntry) { $nodeOwnedMarker = Join-Path $NodeDir ".unsloth-studio-owned" $nodeMeta = Join-Path $NodeDir "UNSLOTH_NODE_PREBUILT_INFO.json" - if (-not (Test-Path -LiteralPath $nodeOwnedMarker) -and -not (Test-Path -LiteralPath $nodeMeta)) { - Write-StudioLine "[ERROR] $NodeDir already exists and is not an Unsloth-owned Node install." -ForegroundColor Red + # From the entry, not a second probe: a dangling link is a container to neither, and + # a reparse point carrying no ownership evidence is not an install of ours. + $nodeIsDir = $nodeEntry.PSIsContainer -and + -not ($nodeEntry.Attributes -band [IO.FileAttributes]::ReparsePoint) + if (-not $nodeIsDir -or ( + -not (Test-Path -LiteralPath $nodeOwnedMarker) -and + -not (Test-Path -LiteralPath $nodeMeta))) { + $what = if ($nodeIsDir) { "an Unsloth-owned Node install" } else { "a directory" } + Write-StudioLine "[ERROR] $NodeDir already exists and is not $what." -ForegroundColor Red Write-StudioLine " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow Exit-SetupFailure "$NodeDir is not an Unsloth-owned Node install" } @@ -4742,7 +4950,9 @@ if ($NeedNodeForSetup) { Write-StudioLine " Install Node >= 20.19 (with npm >= 11) from https://nodejs.org/ and re-run, or check your network." -ForegroundColor Yellow Exit-SetupFailure "Could not install an isolated Node runtime" } - if ($NodeOverride -and (Test-Path -LiteralPath $NodeDir -PathType Container)) { + # Same condition as the guard above, or the tree this run just created under a master + # root stays unmarked and the uninstaller declines to remove it as somebody else's. + if (($NodeOverride -or $RuntimeRootIsCustom) -and (Test-Path -LiteralPath $NodeDir -PathType Container)) { New-Item -ItemType File -Force -Path (Join-Path $NodeDir ".unsloth-studio-owned") -ErrorAction SilentlyContinue | Out-Null } $env:PATH = "$NodeDir;" + $env:PATH @@ -5023,26 +5233,41 @@ function Assert-StudioOwnedOrAbsent { [Parameter(Mandatory = $true)][string]$Label, # whisper.cpp is non-fatal by contract, so it needs the denial handed back # rather than exited on. Only this mode returns a value. - [switch]$NonFatal + [switch]$NonFatal, + # The runtime children pass $RuntimeRootIsCustom; everything under the + # Studio home keeps $StudioHomeIsCustom. + [AllowNull()][object]$IsCustom = $null ) + $isCustomRoot = $StudioHomeIsCustom + if ($null -ne $IsCustom) { $isCustomRoot = [bool]$IsCustom } # Denied is not Absent: a root we cannot read cannot be proven ours, and - # returning here would let the caller replace it. Both stops stay gated on - # $StudioHomeIsCustom, as before; a default-home denial is reported by the - # phase that owns the path. + # returning here would let the caller replace it; a default-home denial is + # reported by the phase that owns the path. $pathState = Get-PathState -Path $Path -PathType Container if ($pathState -ne "Present") { - if ($StudioHomeIsCustom -and $pathState -eq "Denied") { + if ($isCustomRoot -and $pathState -eq "Denied") { if ($NonFatal) { return "Denied" } Exit-PathAccessDenied -Path $Path -Label $Label -OwnershipUnverified } - return + # A Container probe answers false for a regular file and for a link whose target is gone, + # and the caller then deleted the path or let install_*_prebuilt.py replace it; both + # shapes are things the user put in a directory they chose. Get-Item -Force still sees + # the directory entry where Test-Path does not, which is the dangling link under 5.1. + if (-not $isCustomRoot) { return } + if (-not (Get-Item -LiteralPath $Path -Force -ErrorAction SilentlyContinue)) { return } + # Ownership evidence lives inside a directory, so a non-directory can never carry it + # and is refused outright rather than asked for markers. Fatal even under -NonFatal: + # that mode rescues a denial, not somebody else's file. + Write-StudioLine "[ERROR] $Path already exists and is not a directory." -ForegroundColor Red + Write-StudioLine " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow + Exit-SetupFailure "$Label path is not an Unsloth-owned install: $Path" } $markerState = Get-PathState -Path (Join-Path $Path $StudioOwnedMarker) -PathType Leaf - if ($StudioHomeIsCustom -and $markerState -eq "Denied") { + if ($isCustomRoot -and $markerState -eq "Denied") { if ($NonFatal) { return "Denied" } Exit-PathAccessDenied -Path $Path -Label $Label -OwnershipUnverified } - if ($StudioHomeIsCustom -and $markerState -ne "Present") { + if ($isCustomRoot -and $markerState -ne "Present") { $adoptState = Get-StudioAdoptableState -Path $Path if ($adoptState -eq "Denied") { if ($NonFatal) { return "Denied" } @@ -7221,9 +7446,24 @@ if ($script:UnslothVerbose) { # Triton/inductor filenames are long and can hit Windows MAX_PATH (260). With long # paths on, cache under Unsloth home; else use a short drive-root dir for headroom. -if ($StageRoot -or $LongPathsEnabled) { +# +# Persisted to the USER environment below, so every later Studio process inherits it and +# _setup_cache_env's fill-if-unset default never applies on Windows: it has to name the +# directory the resolver would have chosen, or the containment does not happen here at all. +# +# Two things outrank that. Long paths off keeps the short drive-root directory, since Inductor's +# filenames hit MAX_PATH and a contained cache that cannot be written is worse than an +# uncontained one. And a path containing a space is refused as +# storage_roots._TOOLCHAIN_PATH_KEYS refuses one: cpp_builder.py pastes it into a compiler +# command line unquoted, and "C:\Users\First Last" is an ordinary account name. +$TorchCacheDir = $null +if ($StageRoot) { $TorchCacheDir = Join-Path $RuntimeRoot "TORCHINDUCTOR_CACHE_DIR" -} else { +} elseif ($LongPathsEnabled) { + $candidate = Join-Path (Join-Path $StudioHome "cache") "torchinductor" + if ($candidate -notmatch '\s') { $TorchCacheDir = $candidate } +} +if (-not $TorchCacheDir) { $TorchCacheDir = "C:\tc" } if (-not (Test-Path -LiteralPath $TorchCacheDir)) { [System.IO.Directory]::CreateDirectory($TorchCacheDir) | Out-Null } @@ -8465,8 +8705,8 @@ if ($LocalLlamaCppSrc) { throw } } - if ($StudioHomeIsCustom) { - Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + if ($RuntimeRootIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" -IsCustom $RuntimeRootIsCustom } # The destination is about to be deleted and replaced, so a denial here # must stop rather than throw raw: under a default home nothing above @@ -8519,7 +8759,7 @@ if ($LocalLlamaCppLinked) { # Keep this late guard as defense in depth before the prebuilt installer. $llamaDirState = Get-LlamaCppInstallReadState -Path $LlamaCppDir if ($llamaDirState -eq "Denied") { - Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install" -OwnershipUnverified:$StudioHomeIsCustom + Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install" -OwnershipUnverified:$RuntimeRootIsCustom } if ($llamaDirState -eq "Readable") { substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" @@ -8563,8 +8803,8 @@ if ($LocalLlamaCppLinked) { } substep "installing prebuilt llama.cpp bundle (preferred path)..." # install_llama_prebuilt.py's os.replace() would displace a custom-home llama.cpp first. - if ($StudioHomeIsCustom) { - Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + if ($RuntimeRootIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" -IsCustom $RuntimeRootIsCustom } $prebuiltArgs = @( "$PSScriptRoot\install_llama_prebuilt.py", @@ -8643,7 +8883,7 @@ if ($LocalLlamaCppLinked) { } else { step "llama.cpp" "prebuilt installed and validated" } - if ($StudioHomeIsCustom -and (Test-Path -LiteralPath $LlamaCppDir -PathType Container)) { + if ($RuntimeRootIsCustom -and (Test-Path -LiteralPath $LlamaCppDir -PathType Container)) { Mark-StudioOwned -Path $LlamaCppDir } $installedRelease = Get-InstalledLlamaPrebuiltRelease -InstallDir $LlamaCppDir @@ -8730,8 +8970,8 @@ if ($env:WHISPER_SERVER_PATH -or $env:UNSLOTH_WHISPER_CPP_PATH) { substep "whisper.cpp: using a user-configured binary/dir; skipping managed install" } elseif ($env:UNSLOTH_SKIP_WHISPER_INSTALL -eq "1") { substep "whisper.cpp: install skipped (UNSLOTH_SKIP_WHISPER_INSTALL=1)" -} elseif ($StudioHomeIsCustom -and (Test-Path -LiteralPath $WhisperInstaller) -and - (Assert-StudioOwnedOrAbsent -Path $WhisperCppDir -Label "whisper.cpp install" -NonFatal) -eq "Denied") { +} elseif ($RuntimeRootIsCustom -and (Test-Path -LiteralPath $WhisperInstaller) -and + (Assert-StudioOwnedOrAbsent -Path $WhisperCppDir -Label "whisper.cpp install" -NonFatal -IsCustom $RuntimeRootIsCustom) -eq "Denied") { # Never fatal, per the phase header: the guard below would exit the whole run # on an unreadable tree, taking llama.cpp down with it. Only the denial is # caught here; an unowned tree still stops. @@ -8739,8 +8979,8 @@ if ($env:WHISPER_SERVER_PATH -or $env:UNSLOTH_WHISPER_CPP_PATH) { } elseif (Test-Path -LiteralPath $WhisperInstaller) { # The installer's atomic activation replaces the whole directory, so the # custom-home ownership guard must run first (mirrors the llama block). - if ($StudioHomeIsCustom) { - Assert-StudioOwnedOrAbsent -Path $WhisperCppDir -Label "whisper.cpp install" + if ($RuntimeRootIsCustom) { + Assert-StudioOwnedOrAbsent -Path $WhisperCppDir -Label "whisper.cpp install" -IsCustom $RuntimeRootIsCustom } $whisperArgs = @($WhisperInstaller, "--install-dir", $WhisperCppDir) if ($env:UNSLOTH_WHISPER_RELEASE_TAG) { @@ -8780,7 +9020,7 @@ if ($env:WHISPER_SERVER_PATH -or $env:UNSLOTH_WHISPER_CPP_PATH) { } else { step "whisper.cpp" "prebuilt installed" } - if ($StudioHomeIsCustom -and (Test-PathQuiet $WhisperCppDir "Container")) { + if ($RuntimeRootIsCustom -and (Test-PathQuiet $WhisperCppDir "Container")) { Mark-StudioOwned -Path $WhisperCppDir } } elseif ($whisperExit -eq 3) { @@ -8873,7 +9113,7 @@ $NeedRebuild = $false $llamaBinState = if ($LocalLlamaCppLinked) { "Absent" } else { Get-PathState -Path $LlamaServerBin -PathType Leaf } if ($llamaBinState -eq "Denied") { # Nothing proved this tree is ours here, so do not advise deleting it. - Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install" -OwnershipUnverified:$StudioHomeIsCustom + Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install" -OwnershipUnverified:$RuntimeRootIsCustom } if ($llamaBinState -eq "Present") { $CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt" @@ -8883,7 +9123,7 @@ if ($llamaBinState -eq "Present") { $cachedCuda = Select-String -LiteralPath $CmakeCacheFile -Pattern 'GGML_CUDA:BOOL=ON' -Quiet } catch { if (-not (Test-AccessDeniedError $_)) { throw } - Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install" -OwnershipUnverified:$StudioHomeIsCustom + Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install" -OwnershipUnverified:$RuntimeRootIsCustom } if ($HasNvidiaSmi -and -not $cachedCuda) { Write-StudioLine " Existing llama-server is CPU-only but GPU is available -- rebuilding" -ForegroundColor Yellow @@ -9098,8 +9338,8 @@ if ($LocalLlamaCppLinked) { # why: in-place git mutation (remote set-url, checkout -B, clean -fdx) # rewrites $LlamaCppDir; mirror the prebuilt and temp-dir-swap guards # so an unrelated workspace .git tree is never silently overwritten. - if ($StudioHomeIsCustom) { - Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + if ($RuntimeRootIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" -IsCustom $RuntimeRootIsCustom } Write-StudioLine " Syncing llama.cpp to $ResolvedSourceRef..." -ForegroundColor Gray # Always sync the remote URL so switching between default/fork sources works @@ -9172,7 +9412,7 @@ if ($LocalLlamaCppLinked) { } } # In-place git-sync must mark the tree, else a later Assert-StudioOwnedOrAbsent exits. - if ($BuildOk -and $StudioHomeIsCustom) { + if ($BuildOk -and $RuntimeRootIsCustom) { Mark-StudioOwned -Path $LlamaCppDir } } else { @@ -9398,7 +9638,7 @@ if ($LocalLlamaCppLinked) { # Swap temp build dir into final location (only if we built in a temp dir) if ($BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) { - Assert-StudioOwnedOrAbsent -Path $OriginalLlamaCppDir -Label "llama.cpp install" + Assert-StudioOwnedOrAbsent -Path $OriginalLlamaCppDir -Label "llama.cpp install" -IsCustom $RuntimeRootIsCustom if ((Get-PathState -Path $OriginalLlamaCppDir) -ne "Absent") { Remove-Item -LiteralPath $OriginalLlamaCppDir -Recurse -Force -ErrorAction SilentlyContinue # Any unreadable or locked child survives the removal, and Move-Item @@ -9467,7 +9707,7 @@ if ($LocalLlamaCppLinked) { $llamaCppItem = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue $llamaCppIsLink = $llamaCppItem -and ($llamaCppItem.Attributes -band [System.IO.FileAttributes]::ReparsePoint) if (-not $llamaCppIsLink -and ( - -not $StudioHomeIsCustom -or + -not $RuntimeRootIsCustom -or (Test-PathQuiet (Join-Path $LlamaCppDir $StudioOwnedMarker) "Leaf") -or (Test-StudioOwnedAdoptable $LlamaCppDir) )) { diff --git a/studio/setup.sh b/studio/setup.sh index 2747c9565f..72e52b5deb 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1443,6 +1443,63 @@ _STUDIO_HOME_IS_CUSTOM=false if [ "$_studio_home_canon" != "$_LEGACY_STUDIO_HOME" ]; then _STUDIO_HOME_IS_CUSTOM=true fi +# The master root storage_roots.unsloth_home() reads. llama.cpp, node and whisper.cpp sit BESIDE +# studio/ under it, so installing them under $STUDIO_HOME would put them one level below where +# every runtime resolver looks. Captured here because section 7 assigns over UNSLOTH_HOME. +# Stripped before anything else, like _studio_override above: " " counts as unset and " /opt/uns " +# names /opt/uns, which is what the Python resolver and the CLI both see. +_MASTER_ROOT=$(printf '%s' "${UNSLOTH_HOME:-}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') +if [ -n "$_MASTER_ROOT" ]; then + case "$_MASTER_ROOT" in + "~") _MASTER_ROOT="$HOME" ;; + "~/"*) _MASTER_ROOT="$HOME/${_MASTER_ROOT#'~/'}" ;; + esac + # Absolute against the CALLER's directory even when it does not exist yet, as Path.resolve() + # does: node is chosen before the first `cd "$SCRIPT_DIR"` and llama.cpp after it, so a + # relative value names two different directories and matches the backend's neither. + case "$_MASTER_ROOT" in + /*) ;; + *) _MASTER_ROOT="$PWD/$_MASTER_ROOT" ;; + esac + if [ -d "$_MASTER_ROOT" ]; then + # Keep the expanded value when it cannot be canonicalized, as the Python resolver does: + # dropping it here would send the runtimes to a root nothing else agrees on. + _master_root_canon=$(CDPATH= cd -P -- "$_MASTER_ROOT" 2>/dev/null && pwd -P) || _master_root_canon="" + [ -z "$_master_root_canon" ] || _MASTER_ROOT="$_master_root_canon" + unset _master_root_canon + fi +fi + +# Ownership applies to node/, llama.cpp/ and whisper.cpp/ whenever a master root moves them, +# even with STUDIO_HOME left at the legacy path, where _STUDIO_HOME_IS_CUSTOM is false and +# "false" licenses the installers to os.replace() and rm -rf unmarked trees. The Studio home +# and its venvs keep the other flag. +_RUNTIME_ROOT_IS_CUSTOM="$_STUDIO_HOME_IS_CUSTOM" +# Keyed on where the runtimes LAND, not on whether a master root was named: UNSLOTH_HOME set to +# the root an install already uses moves nothing, and calling that custom would demand an owner +# marker from a legacy source-built ~/.unsloth/llama.cpp that predates markers. +# +# Staging is excluded for the same reason, not as an exception to it: the placement below gives +# STAGE_ROOT precedence over the master root, so during a staged update the master root is not +# where anything lands and must not decide ownership. Assigned rather than raised, or a custom +# STUDIO_HOME kept the flag true while the runtimes went to the legacy root, which demanded +# markers from exactly the pre-marker trees this comparison exists to spare. +if [ -z "${STAGE_ROOT:-}" ] && [ -n "$_MASTER_ROOT" ]; then + # Canonicalised like _MASTER_ROOT, or a symlinked $HOME compares unequal to itself and the + # legacy root reads as custom after all. + _rrc_legacy="$HOME/.unsloth" + if [ -d "$_rrc_legacy" ]; then + _rrc_canon=$(CDPATH= cd -P -- "$_rrc_legacy" 2>/dev/null && pwd -P) || _rrc_canon="" + [ -z "$_rrc_canon" ] || _rrc_legacy="$_rrc_canon" + unset _rrc_canon + fi + if [ "$_MASTER_ROOT" = "$_rrc_legacy" ]; then + _RUNTIME_ROOT_IS_CUSTOM=false + else + _RUNTIME_ROOT_IS_CUSTOM=true + fi + unset _rrc_legacy +fi # Directory-local evidence Unsloth created "$1": only prebuilt-installer metadata # counts (UNSLOTH_PREBUILT_INFO.json for llama.cpp, UNSLOTH_NODE_PREBUILT_INFO.json # for Node, UNSLOTH_WHISPER_PREBUILT_INFO.json for whisper.cpp), all written only @@ -1527,11 +1584,37 @@ _report_denied_ancestor() { fi } +# What is at "$1", for the refusal message. A user told the path is "not an Unsloth install" +# when it is their own dangling symlink has no idea what to move aside. +_studio_path_shape() { + if [ -L "$1" ]; then + if [ -e "$1" ]; then printf 'a symlink'; else printf 'a dangling symlink'; fi + elif [ -f "$1" ]; then printf 'a regular file' + elif [ -d "$1" ]; then printf 'a directory' + else printf 'an existing path' + fi +} + +# $3 is the ownership flag: the runtime children pass _RUNTIME_ROOT_IS_CUSTOM, everything +# under the Studio home keeps _STUDIO_HOME_IS_CUSTOM. _assert_studio_owned_or_absent() { _aso_dir="$1" _aso_label="$2" - [ -d "$_aso_dir" ] || return 0 - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ ! -f "$_aso_dir/$_STUDIO_OWNED_MARKER" ]; then + _aso_custom="${3:-$_STUDIO_HOME_IS_CUSTOM}" + # -d alone read a dangling symlink or a regular file as "nothing is here", and the caller + # then rm -rf'd it or os.replace()d over it. A dangling link is the ordinary case: its + # volume is simply not mounted, and following it later installs onto somebody's other disk. + if [ ! -d "$_aso_dir" ] && [ ! -e "$_aso_dir" ] && [ ! -L "$_aso_dir" ]; then + return 0 + fi + if [ "$_aso_custom" = true ] && [ ! -f "$_aso_dir/$_STUDIO_OWNED_MARKER" ]; then + # Only a directory can carry the marker or the prebuilt metadata, so anything else here + # is unowned by construction and the adoption path below cannot apply to it. + if [ ! -d "$_aso_dir" ]; then + echo "ERROR: $_aso_dir already exists and is not an Unsloth-owned $_aso_label." >&2 + echo " It is $(_studio_path_shape "$_aso_dir"). Move it aside before re-running." >&2 + setup_fail 1 "$_aso_label path is not an Unsloth-owned install: $_aso_dir" + fi if _studio_owned_adoptable "$_aso_dir"; then : > "$_aso_dir/$_STUDIO_OWNED_MARKER" 2>/dev/null || true return 0 @@ -1641,6 +1724,8 @@ decide_node_source() { # Mirror the llama.cpp UNSLOTH_HOME derivation; the frontend build runs first. if [ -n "$STAGE_ROOT" ]; then _NODE_PARENT="$RUNTIME_ROOT" +elif [ -n "$_MASTER_ROOT" ]; then + _NODE_PARENT="$_MASTER_ROOT" elif [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then _NODE_PARENT="$STUDIO_HOME" else @@ -1659,8 +1744,8 @@ elif [ "$NODE_SOURCE" = bundled ]; then mkdir -p "$_NODE_PARENT" # install_node_prebuilt.py uses os.replace(); guard a custom-home dir so we # never displace a user-owned $UNSLOTH_STUDIO_HOME/node. - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then - _assert_studio_owned_or_absent "$NODE_DIR" "Node install" + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$NODE_DIR" "Node install" "$_RUNTIME_ROOT_IS_CUSTOM" fi substep "installing isolated Node (system Node/npm left untouched)..." # Runs before the venv is activated, so bare `python` may be absent; resolve @@ -1695,7 +1780,7 @@ elif [ "$NODE_SOURCE" = bundled ]; then fi grep -Fq "already matches" "$_NODE_LOG" && verbose_substep "isolated Node already up to date" rm -f "$_NODE_LOG" - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$NODE_DIR" ]; then + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ] && [ -d "$NODE_DIR" ]; then : > "$NODE_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true fi # Prepend the isolated bin (this process only) so node/npm/bun resolve here. @@ -3664,12 +3749,107 @@ fi # default keeps ~/.unsloth/llama.cpp so pre-PR builds are still discovered. if [ -n "$STAGE_ROOT" ]; then UNSLOTH_HOME="$RUNTIME_ROOT" +elif [ -n "$_MASTER_ROOT" ]; then + UNSLOTH_HOME="$_MASTER_ROOT" elif [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then UNSLOTH_HOME="$STUDIO_HOME" else UNSLOTH_HOME="$HOME/.unsloth" fi mkdir -p "$UNSLOTH_HOME" +# Record the master root inside the Studio tree, for the uninstaller: UNSLOTH_HOME can be set +# for a single command, and without the note the uninstaller removed the Studio tree and +# stranded multi-gigabyte llama.cpp, node and whisper.cpp trees beside it. +# +# Only for a master root (the other branches derive UNSLOTH_HOME from paths the uninstaller +# already knows), and only when a reader will honour it: both rejected shapes were being +# written, so the file on disk claimed a recorded root while every reader refused it. +_master_root_note_is_honoured() { + # Canonicalised here rather than trusting the value from the top of the script: that one was + # resolved before the mkdir above, so a root created by this run kept its uncanonical + # spelling, and the containment test below is textual. + _mrn_root=$(CDPATH= cd -P -- "$1" 2>/dev/null && pwd -P) || _mrn_root="" + [ -n "$_mrn_root" ] || return 1 + # The legacy default is the one root every reader finds without help, and both uninstallers + # refuse it outright. Writing it turned a one-command `UNSLOTH_HOME=$HOME/.unsloth` into a + # permanent portable install: every later bare launch read the note back and HF_HUB_CACHE + # moved off ~/.cache/huggingface, the one cache this change promises not to move. + _mrn_legacy=$(CDPATH= cd -P -- "${HOME:-}/.unsloth" 2>/dev/null && pwd -P) || _mrn_legacy="${HOME:-}/.unsloth" + [ "$_mrn_root" != "$_mrn_legacy" ] || return 1 + # Keyed on the TREE as well, because that is what the readers key on: storage_roots' + # _is_legacy_studio_tree and the CLI both decline ANY note found in ~/.unsloth/studio, + # whatever it records. Without this, UNSLOTH_HOME=$HOME passed containment (the legacy tree + # is inside $HOME) and the value is not the legacy root, so the note was written, warned + # about nothing, and then honoured by nobody. + _mrn_studio=$(CDPATH= cd -P -- "$STUDIO_HOME" 2>/dev/null && pwd -P) || _mrn_studio="$STUDIO_HOME" + _mrn_legacy_studio=$(CDPATH= cd -P -- "${HOME:-}/.unsloth/studio" 2>/dev/null && pwd -P) \ + || _mrn_legacy_studio="${HOME:-}/.unsloth/studio" + if [ "$_mrn_studio" = "$_mrn_legacy_studio" ]; then + echo "NOTE: the managed runtimes were installed under $_mrn_root, but Studio itself is the" >&2 + echo " default install at $_mrn_studio, where no reader honours a recorded root. That" >&2 + echo " root cannot be recorded, so a later launch or uninstall will not find them." >&2 + echo " Set UNSLOTH_STUDIO_HOME=$_mrn_root/studio, or re-export UNSLOTH_HOME whenever" >&2 + echo " you run Unsloth." >&2 + return 1 + fi + # A note must describe the tree it is written into: all four readers require the Studio + # directory to lie INSIDE the root it names, so a tree copied between master roots cannot + # aim a removal at the original install. install.sh and install.ps1 do not read UNSLOTH_HOME + # yet, so `UNSLOTH_HOME=/mnt/portable unsloth studio update` leaves Studio at + # ~/.unsloth/studio with the runtimes at /mnt/portable, which is exactly that shape: warn, + # rather than leave a file that reads as a recorded root and behaves as none. + case "$STUDIO_HOME" in + "$_mrn_root"|"$_mrn_root"/*) return 0 ;; + esac + echo "NOTE: the managed runtimes were installed under $_mrn_root, but Studio itself lives at" >&2 + echo " $STUDIO_HOME, which is outside it. That root cannot be recorded, so a later launch" >&2 + echo " or uninstall will not find them. Set UNSLOTH_STUDIO_HOME=$_mrn_root/studio, or" >&2 + echo " re-export UNSLOTH_HOME whenever you run Unsloth." >&2 + return 1 +} +if [ -n "$_MASTER_ROOT" ] && [ -z "$STAGE_ROOT" ] && _master_root_note_is_honoured "$UNSLOTH_HOME"; then + if mkdir -p "$STUDIO_HOME/share" 2>/dev/null; then + # Staged then renamed: a reader catching a half-written note would name a truncated + # path, and this note licenses deletions. mktemp rather than "$$", as + # _uv_cache_root_is_writable states above: a redirection into a predictable name follows + # a symlink anyone who can write share/ could have precreated there. + if ! _mrn_tmp=$(mktemp "$STUDIO_HOME/share/.unsloth-master-root.XXXXXX" 2>/dev/null); then + _mrn_tmp="" + fi + if [ -n "$_mrn_tmp" ] && printf '%s\n' "$UNSLOTH_HOME" > "$_mrn_tmp" 2>/dev/null; then + mv -f "$_mrn_tmp" "$STUDIO_HOME/share/.unsloth-master-root" 2>/dev/null \ + || rm -f "$_mrn_tmp" 2>/dev/null || true + elif [ -n "$_mrn_tmp" ]; then + rm -f "$_mrn_tmp" 2>/dev/null || true + fi + unset _mrn_tmp + fi +fi +# Clear a note an earlier run left naming the legacy default root, whatever this run was asked +# to do. The gate above stops new ones, but an install that already has one keeps reading it +# back on every bare launch -- portable mode on, the Hugging Face hub cache moved off +# ~/.cache/huggingface -- and no later `unsloth studio update` would ever pass through the write +# block to correct it, since that block only runs when UNSLOTH_HOME is set again. Narrow on +# purpose: this is the single value both uninstallers already refuse, so removing it takes +# nothing any reader was entitled to act on. A note naming any other root is left alone, since +# an unreadable or unexpected one may simply describe an install this run cannot see. +if [ -f "$STUDIO_HOME/share/.unsloth-master-root" ]; then + _mrn_old=$(head -n 1 "$STUDIO_HOME/share/.unsloth-master-root" 2>/dev/null \ + | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') || _mrn_old="" + if [ -n "$_mrn_old" ]; then + case "$_mrn_old" in + "~") _mrn_old="${HOME:-}" ;; + "~/"*) _mrn_old="${HOME:-}/${_mrn_old#'~/'}" ;; + esac + _mrn_old_canon=$(CDPATH= cd -P -- "$_mrn_old" 2>/dev/null && pwd -P) || _mrn_old_canon="" + [ -n "$_mrn_old_canon" ] && _mrn_old="$_mrn_old_canon" + _mrn_legacy=$(CDPATH= cd -P -- "${HOME:-}/.unsloth" 2>/dev/null && pwd -P) || _mrn_legacy="${HOME:-}/.unsloth" + if [ "$_mrn_old" = "$_mrn_legacy" ]; then + rm -f "$STUDIO_HOME/share/.unsloth-master-root" 2>/dev/null || true + fi + fi + unset _mrn_old _mrn_old_canon +fi LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp" LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server" _NEED_LLAMA_SOURCE_BUILD=false @@ -4014,8 +4194,8 @@ if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then # for a custom UNSLOTH_STUDIO_HOME (the assert would otherwise follow the # link into the user's dir and reject it as unowned). [ -L "$LLAMA_CPP_DIR" ] && rm -f "$LLAMA_CPP_DIR" - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then - _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" "$_RUNTIME_ROOT_IS_CUSTOM" fi rm -rf "$LLAMA_CPP_DIR" || true if [ -e "$LLAMA_CPP_DIR" ]; then @@ -4040,8 +4220,8 @@ fi # swap only reaches its own guards after the whole build, so check here instead. # Local-link paths are excluded: they already replaced or reused the tree above. if [ "$_LOCAL_LLAMA_CPP_LINKED" != true ]; then - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then - _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" "$_RUNTIME_ROOT_IS_CUSTOM" fi if _studio_dir_unreadable "$LLAMA_CPP_DIR"; then _path_access_denied "$LLAMA_CPP_DIR" "llama.cpp install" @@ -4071,8 +4251,8 @@ else # why: install_llama_prebuilt.py uses os.replace(), which would displace # an unrelated $UNSLOTH_STUDIO_HOME/llama.cpp before the source-build # ownership check below ever runs. - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then - _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" "$_RUNTIME_ROOT_IS_CUSTOM" fi # The ownership check above misses the default cache; stop before pathlib # turns an unreadable one into a traceback. @@ -4145,7 +4325,7 @@ else else step "llama.cpp" "prebuilt installed and validated" fi - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$LLAMA_CPP_DIR" ]; then + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ] && [ -d "$LLAMA_CPP_DIR" ]; then : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true fi print_installed_llama_prebuilt_release "$LLAMA_CPP_DIR" @@ -4223,7 +4403,7 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && \ [ -x "$LLAMA_CPP_DIR/build/bin/llama-quantize" ]; then step "llama.cpp" "existing source build found; skipping rebuild" ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize" - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ]; then : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true fi _NEED_LLAMA_SOURCE_BUILD=false @@ -4784,7 +4964,7 @@ else fi # Swap only after build succeeds -- preserves existing install on failure if [ "$BUILD_OK" = true ]; then - _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" "$_RUNTIME_ROOT_IS_CUSTOM" # || true: without it a raw rm error aborts under errexit to a bare exit # code, build stranded. Keep stderr: rm names the exact subpath, we cannot. rm -rf "$LLAMA_CPP_DIR" || true @@ -4862,7 +5042,7 @@ if [ "$_LLAMA_CPP_DEGRADED" = true ] \ fi if [ ! -L "$LLAMA_CPP_DIR" ] && { - [ "$_STUDIO_HOME_IS_CUSTOM" != true ] || + [ "$_RUNTIME_ROOT_IS_CUSTOM" != true ] || [ -f "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" ] || _studio_owned_adoptable "$LLAMA_CPP_DIR" }; then @@ -4881,8 +5061,8 @@ if [ -n "${WHISPER_SERVER_PATH:-}" ] || [ -n "${UNSLOTH_WHISPER_CPP_PATH:-}" ]; elif [ "${UNSLOTH_SKIP_WHISPER_INSTALL:-0}" = "1" ]; then verbose_substep "whisper.cpp: install skipped (UNSLOTH_SKIP_WHISPER_INSTALL=1)" else - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then - _assert_studio_owned_or_absent "$WHISPER_CPP_DIR" "whisper.cpp install" + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$WHISPER_CPP_DIR" "whisper.cpp install" "$_RUNTIME_ROOT_IS_CUSTOM" fi _WHISPER_CMD=(python "$SCRIPT_DIR/install_whisper_prebuilt.py" --install-dir "$WHISPER_CPP_DIR") if [ -n "${UNSLOTH_WHISPER_RELEASE_TAG:-}" ]; then @@ -4914,7 +5094,7 @@ else else step "whisper.cpp" "prebuilt installed" fi - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$WHISPER_CPP_DIR" ]; then + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ] && [ -d "$WHISPER_CPP_DIR" ]; then : > "$WHISPER_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true fi rm -f "$_WHISPER_LOG" @@ -4940,7 +5120,7 @@ else env UNSLOTH_HOME="$UNSLOTH_HOME" sh "$_WHISPER_BUILD"; then _WHISPER_RECOVERED=true step "whisper.cpp" "source build installed" - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$WHISPER_CPP_DIR" ]; then + if [ "$_RUNTIME_ROOT_IS_CUSTOM" = true ] && [ -d "$WHISPER_CPP_DIR" ]; then : > "$WHISPER_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true fi else diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index e90b3d639b..ece3c7cbf5 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -580,8 +580,12 @@ fn spawn_script( // Tauri only does default-root installs; install.sh / install.ps1 reject // these under --tauri. Scrub so an inherited value can't trip the guard. - cmd.env_remove("UNSLOTH_STUDIO_HOME"); - cmd.env_remove("STUDIO_HOME"); + // Applied by hand here, the one managed spawn outside + // apply_managed_cli_context: the Python setup.sh starts reads whatever is + // exported, even though setup.sh assigns UNSLOTH_HOME itself. + for name in crate::process::MANAGED_CHILD_SCRUBBED_ENV { + cmd.env_remove(name); + } cmd.env( "UNSLOTH_DESKTOP_BACKEND_VERSION", crate::preflight::expected_backend_version(), diff --git a/studio/src-tauri/src/process.rs b/studio/src-tauri/src/process.rs index b617cdad35..fa5ce8b1d8 100644 --- a/studio/src-tauri/src/process.rs +++ b/studio/src-tauri/src/process.rs @@ -1964,6 +1964,7 @@ fn windows_roots_from( /// PATH are absent: they are not single paths. Mirrors `_RELATIVE_PATH_ENV` in /// unsloth_cli/_system_dir_guard.py, held identical by a parity test. pub(crate) const RELATIVE_PATH_ENV: &[&str] = &[ + "UNSLOTH_HOME", "UNSLOTH_STUDIO_HOME", "STUDIO_HOME", "UNSLOTH_STUDIO_DOCUMENTS_HOME", @@ -1997,6 +1998,18 @@ pub(crate) const RELATIVE_PATH_ENV: &[&str] = &[ "UNSLOTH_DG_SHIM", "UNSLOTH_COMPILE_LOCATION", "TORCHINDUCTOR_CACHE_DIR", + // storage_roots.py fills these only when blank, so a relative value the user set is kept as + // written and would name a different folder after the move. + "TORCH_EXTENSIONS_DIR", + "TORCH_HOME", + "TRITON_HOME", + "TRITON_CACHE_DIR", + "TRITON_DUMP_DIR", + "CUDA_CACHE_PATH", + "MPLCONFIGDIR", + "NUMBA_CACHE_DIR", + "DATA_DESIGNER_HOME", + "DATA_DESIGNER_MANAGED_ASSETS_PATH", "UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR", "UNSLOTH_DIFFUSION_COND_CACHE_DIR", "HF_HOME", @@ -2005,6 +2018,9 @@ pub(crate) const RELATIVE_PATH_ENV: &[&str] = &[ "HF_XET_CACHE", "HF_DATASETS_CACHE", "HF_ASSETS_CACHE", + // transformers appends this to sys.path, so a relative value would import a + // different generated module after the move. + "HF_MODULES_CACHE", // huggingface_hub resolves the credential file from here; a relative value // would follow the child and silently lose access to gated repos. "HF_TOKEN_PATH", @@ -2315,7 +2331,17 @@ fn names_a_path(name: &str, value: &str) -> bool { /// Names every managed spawn removes before starting the child: Tauri uses the /// legacy Unsloth root whatever the environment says. Resolving one can only /// invent a failure for a value the child is never going to see. -const MANAGED_CHILD_SCRUBBED_ENV: &[&str] = &["UNSLOTH_STUDIO_HOME", "STUDIO_HOME"]; +/// +/// UNSLOTH_HOME moves the databases, assets and caches exactly as +/// UNSLOTH_STUDIO_HOME does, plus the managed llama.cpp, node and whisper.cpp +/// dirs. UNSLOTH_PORTABLE names no root but portable_mode() reads it on its own, +/// repointing the Hugging Face and torch caches away from the shared user ones. +pub(crate) const MANAGED_CHILD_SCRUBBED_ENV: &[&str] = &[ + "UNSLOTH_HOME", + "UNSLOTH_STUDIO_HOME", + "STUDIO_HOME", + "UNSLOTH_PORTABLE", +]; /// Read only by the update and installer path (install_python_stack.py), so a /// stale value must not be able to fail a probe, a backend start or an auth @@ -3566,11 +3592,13 @@ pub fn start_backend( #[cfg(target_os = "linux")] scrub_appimage_python_env(&mut cmd); - // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME; - // scrub so the spawned Python backend can't diverge. UNSLOTH_LLAMA_CPP_PATH - // is a pre-existing user-controlled llama.cpp dir override; keep it. - cmd.env_remove("UNSLOTH_STUDIO_HOME"); - cmd.env_remove("STUDIO_HOME"); + // Tauri uses the legacy root whatever the environment says; scrub so the + // spawned Python backend can't diverge. Off the shared list, so a name added + // there cannot be honoured by the backend and missed here. + // UNSLOTH_LLAMA_CPP_PATH is a pre-existing user-controlled override; keep it. + for name in MANAGED_CHILD_SCRUBBED_ENV { + cmd.env_remove(name); + } // read_output_stream decodes as UTF-8; without these, Python encodes its // redirected streams with the locale code page and non-ASCII lands as U+FFFD. diff --git a/tests/sh/test_portable_isolated_venv.sh b/tests/sh/test_portable_isolated_venv.sh new file mode 100755 index 0000000000..36bd485c6b --- /dev/null +++ b/tests/sh/test_portable_isolated_venv.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +# The path resolver must work in a clean interpreter with only structlog, and +# seeding must create no studio.db (get_app_setting CREATES one). Needs uv. +set -u +REPO="${1:-$(CDPATH= cd -P -- "$(dirname "$0")/../.." && pwd -P)}" +if ! command -v uv >/dev/null 2>&1; then + echo "SKIP: uv not on PATH, cannot build an isolated venv" + exit 0 +fi +WORK="$(mktemp -d)" +trap 'rm -rf "$WORK"' EXIT +fails=0 +check() { if [ "$2" = 0 ]; then echo " PASS $1"; else echo " FAIL $1 ${3:-}"; fails=$((fails+1)); fi; } + +echo "[1] build an isolated venv (uv, no site packages from the workspace)" +uv venv --python 3.13 "$WORK/venv" >/dev/null 2>&1 || { echo "FAIL: uv venv"; exit 1; } +PY="$WORK/venv/bin/python" +# structlog is the only thing storage_roots pulls in via loggers/path_utils. +uv pip install --python "$PY" structlog >/dev/null 2>&1 || { echo "FAIL: pip install"; exit 1; } +echo " venv: $("$PY" -V), packages: $(uv pip list --python "$PY" 2>/dev/null | wc -l)" + +echo +echo "[2] resolver imports and runs with only structlog installed" +HOME_A="$WORK/home_a"; mkdir -p "$HOME_A" +out="$(env -i HOME="$HOME_A" PATH="/usr/bin:/bin" UNSLOTH_HOME="$WORK/root_a" \ + "$PY" -c " +import sys, os, json +sys.path.insert(0, '$REPO/studio/backend') +from utils.paths import storage_roots as sr +sr.setup_cache_env() +print(json.dumps({ + 'studio_root': str(sr.studio_root()), + 'portable': sr.portable_mode(), + 'compile': os.environ.get('UNSLOTH_COMPILE_LOCATION'), + 'triton_home': os.environ.get('TRITON_HOME'), +})) +" 2>&1)" +case "$out" in + *studio_root*) check "resolver runs on a bare interpreter" 0; echo " $out" ;; + *) check "resolver runs on a bare interpreter" 1 "$out" ;; +esac + +echo +echo "[3] no torch, transformers or other heavy import is required" +heavy="$(env -i HOME="$HOME_A" PATH="/usr/bin:/bin" "$PY" -c " +import sys +sys.path.insert(0, '$REPO/studio/backend') +from utils.paths import storage_roots as sr +sr.setup_cache_env() +bad = [m for m in ('torch','transformers','numpy','huggingface_hub','fastapi') if m in sys.modules] +print(','.join(bad) if bad else 'NONE') +" 2>&1 | tail -1)" +[ "$heavy" = "NONE" ] && check "seeding pulls in no heavy modules" 0 \ + || check "seeding pulls in no heavy modules" 1 "imported: $heavy" + +echo +echo "[4] CLI seeding side effects: must not create a studio DB or leak outside the root" +HOME_B="$WORK/home_b"; mkdir -p "$HOME_B" +ROOT_B="$WORK/root_b" +env -i HOME="$HOME_B" PATH="/usr/bin:/bin" UNSLOTH_HOME="$ROOT_B" "$PY" -c " +import sys +sys.path.insert(0, '$REPO/studio/backend') +from utils.paths.storage_roots import setup_cache_env +setup_cache_env() +" >/dev/null 2>&1 +dbs="$(find "$ROOT_B" "$HOME_B" -name '*.db' 2>/dev/null | wc -l)" +[ "$dbs" = 0 ] && check "seeding creates no database" 0 || check "seeding creates no database" 1 "found $dbs" +outside="$(find "$HOME_B" -mindepth 1 2>/dev/null | wc -l)" +[ "$outside" = 0 ] && check "seeding writes nothing to HOME" 0 \ + || check "seeding writes nothing to HOME" 1 "$(find "$HOME_B" -mindepth 1 | head -5 | tr '\n' ' ')" + +echo +echo "[5] seeding is idempotent and cheap on repeat" +t0=$(date +%s%N) +env -i HOME="$HOME_B" PATH="/usr/bin:/bin" UNSLOTH_HOME="$ROOT_B" "$PY" -c " +import sys +sys.path.insert(0, '$REPO/studio/backend') +from utils.paths.storage_roots import setup_cache_env +for _ in range(50): + setup_cache_env() +" >/dev/null 2>&1 +t1=$(date +%s%N) +ms=$(( (t1 - t0) / 1000000 )) +[ "$ms" -lt 20000 ] && check "50 repeat seedings in ${ms}ms" 0 || check "50 repeat seedings too slow (${ms}ms)" 1 + +echo +echo "[6] resolver survives a missing HOME entirely" +out6="$(env -i PATH="/usr/bin:/bin" UNSLOTH_HOME="$WORK/root_c" "$PY" -c " +import sys +sys.path.insert(0, '$REPO/studio/backend') +from utils.paths import storage_roots as sr +print(sr.studio_root()) +" 2>&1 | tail -1)" +case "$out6" in + "$WORK/root_c/studio") check "no HOME set still resolves the portable root" 0 ;; + *) check "no HOME set still resolves the portable root" 1 "$out6" ;; +esac + +echo +if [ "$fails" -ne 0 ]; then echo "$fails check(s) failed"; exit 1; fi +echo "All isolated-venv checks passed" diff --git a/tests/sh/test_setup_staged_root.sh b/tests/sh/test_setup_staged_root.sh index bb8edb5451..3a89216a10 100755 --- a/tests/sh/test_setup_staged_root.sh +++ b/tests/sh/test_setup_staged_root.sh @@ -98,8 +98,18 @@ check "setup.ps1 does not install Git while staging" \ "$(has "$SETUP_PS1" 'Background staging cannot install Git; retry with the foreground updater.')" check "setup.ps1 preserves foreground Git bootstrap" \ "$(has "$SETUP_PS1" 'if ($gitNeeded -or -not $StageRoot) {')" +# The branch that decides TORCHINDUCTOR_CACHE_DIR, by what it does rather than by its old +# spelling: a staged run must take the stage root and must be asked FIRST, ahead of the Studio +# cache and the short drive-root fallback that a normal install chooses between. +_tcd="$(awk '/^\$TorchCacheDir = \$null$/{g=1} g{print} g && /^\$env:TORCHINDUCTOR_CACHE_DIR/{exit}' "$SETUP_PS1")" check "setup.ps1 keeps the staging compiler cache under the stage root" \ - "$(has "$SETUP_PS1" 'if ($StageRoot -or $LongPathsEnabled) {')" + "$(printf '%s' "$_tcd" | grep -qF 'if ($StageRoot) {' \ + && printf '%s' "$_tcd" | grep -qF '$TorchCacheDir = Join-Path $RuntimeRoot "TORCHINDUCTOR_CACHE_DIR"' \ + && echo 0 || echo 1)" +check "setup.ps1 asks about staging before anything else" \ + "$([ "$(printf '%s' "$_tcd" | grep -n 'if ($StageRoot) {' | cut -d: -f1)" -lt \ + "$(printf '%s' "$_tcd" | grep -n 'LongPathsEnabled' | head -1 | cut -d: -f1)" ] \ + && echo 0 || echo 1)" # ── the staged activation cannot dot-source a copied Activate script ── # A venv copied out of $STUDIO_HOME still names the original root in its activate diff --git a/tests/sh/test_uninstall_master_root.sh b/tests/sh/test_uninstall_master_root.sh new file mode 100755 index 0000000000..e642818f4f --- /dev/null +++ b/tests/sh/test_uninstall_master_root.sh @@ -0,0 +1,320 @@ +#!/bin/bash +# 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 test for UNSLOTH_HOME master-root removal in scripts/uninstall.sh. +# +# setup.sh and setup.ps1 install llama.cpp, node and whisper.cpp as CHILDREN of the master +# root, beside studio/, so removing the Studio root alone strands multi-gigabyte trees. The +# root is user-chosen, so only a tree carrying the Unsloth owner marker may be deleted. +# Hermetic, following test_uninstall_sd_cpp_custom_root.sh: the real helpers and the real +# removal block are extracted from uninstall.sh and run against per-test fixtures. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +UNINSTALL_SH="$SCRIPT_DIR/../../scripts/uninstall.sh" +PASS=0 +FAIL=0 + +_TMP_ROOT=$(mktemp -d) +trap 'rm -rf "$_TMP_ROOT"' EXIT +HOME="$_TMP_ROOT/home" +mkdir -p "$HOME" + +assert_nodir() { _l="$1"; [ -e "$2" ] && { echo " FAIL: $_l (still present: $2)"; FAIL=$((FAIL+1)); } || { echo " PASS: $_l"; PASS=$((PASS+1)); }; } +assert_dir() { _l="$1"; [ -e "$2" ] && { echo " PASS: $_l"; PASS=$((PASS+1)); } || { echo " FAIL: $_l (missing $2)"; FAIL=$((FAIL+1)); }; } +assert_link() { _l="$1"; [ -L "$2" ] && { echo " PASS: $_l"; PASS=$((PASS+1)); } || { echo " FAIL: $_l (link missing $2)"; FAIL=$((FAIL+1)); }; } +# -e is false for a dangling symlink, which is exactly the shape one case below is about. +assert_present() { _l="$1"; { [ -e "$2" ] || [ -L "$2" ]; } && { echo " PASS: $_l"; PASS=$((PASS+1)); } || { echo " FAIL: $_l (missing $2)"; FAIL=$((FAIL+1)); }; } +assert_eq() { _l="$1"; [ "$2" = "$3" ] && { echo " PASS: $_l"; PASS=$((PASS+1)); } || { echo " FAIL: $_l (got '$2', want '$3')"; FAIL=$((FAIL+1)); }; } + +HELPERS_FILE=$(mktemp -p "$_TMP_ROOT") +{ + sed -n '/^_remove_path() {/,/^}/p' "$UNINSTALL_SH" + sed -n '/^_remove_lock_file() {/,/^}/p' "$UNINSTALL_SH" + sed -n '/^_is_unsafe_root() {/,/^}/p' "$UNINSTALL_SH" + sed -n '/^_master_root() {/,/^}/p' "$UNINSTALL_SH" + sed -n '/^_set_marker() {/,/^}/p' "$UNINSTALL_SH" +} > "$HELPERS_FILE" +# Every helper the block calls, by name. A helper that stops being extracted goes "command not +# found" inside a subshell the harness already swallows, which reads as "nothing was removed" +# and passes the keep-assertions while failing the remove-assertions for the wrong reason. +for _h in _remove_path _remove_lock_file _is_unsafe_root _master_root _set_marker; do + grep -q "^$_h() {" "$HELPERS_FILE" || { echo "FAIL: helpers missing $_h"; exit 1; } +done + +# The real block, anchored on its own first line so a rename fails loudly instead of vacuously. +BLOCK_FILE=$(mktemp -p "$_TMP_ROOT") +sed -n '/^[[:space:]]*_mr_root="\$_MASTER_ROOT_SAVED"/,/^[[:space:]]*# end master-root children$/p' "$UNINSTALL_SH" > "$BLOCK_FILE" +[ -s "$BLOCK_FILE" ] || { echo "FAIL: master-root removal block not extracted"; exit 1; } +grep -q 'unsloth-studio-owned' "$BLOCK_FILE" || { echo "FAIL: extracted block lost the marker gate"; exit 1; } + +run_block() { + # The real script resolves the master root ONCE before anything is removed, because the note + # it can read lives inside a tree the custom-root loop deletes. The block takes that saved + # value, so the harness has to seed it the same way. + ( set -e; HOME="$1"; UNSLOTH_HOME="$2"; export HOME UNSLOTH_HOME + . "$HELPERS_FILE"; _MASTER_ROOT_SAVED="$(_master_root)"; . "$BLOCK_FILE" ) || true +} + +echo "== an owned master root loses its runtime children ==" +MR="$_TMP_ROOT/portable" +mkdir -p "$MR"/{studio,llama.cpp,node,whisper.cpp} +for d in llama.cpp node whisper.cpp; do : > "$MR/$d/.unsloth-studio-owned"; done +: > "$MR/.llama.cpp.install.lock" +: > "$MR/.node.install.lock" +run_block "$HOME" "$MR" +assert_nodir "llama.cpp removed" "$MR/llama.cpp" +assert_nodir "node removed" "$MR/node" +assert_nodir "whisper.cpp removed" "$MR/whisper.cpp" +assert_nodir "llama lock removed" "$MR/.llama.cpp.install.lock" +assert_nodir "node lock removed" "$MR/.node.install.lock" +assert_dir "studio root left to the Studio removal" "$MR/studio" + +echo "== only a FILE is removed from an install-lock path ==" +# prebuilt_core.install_lock creates the lock with os.open(O_CREAT | O_EXCL), so it is always a +# regular file. The names are fixed, so in a user-chosen root a directory here is the user's: +# rm -rf on it would take a tree that never carried the owner marker the children above need. +MRL="$_TMP_ROOT/lockshapes" +mkdir -p "$MRL/studio" "$MRL/.node.install.lock/keep" +: > "$MRL/.node.install.lock/keep/user-file" +: > "$MRL/.llama.cpp.install.lock" +# A link at a lock name is the user's too: O_CREAT | O_EXCL never makes one, and +# uninstall.ps1 already excluded reparse points, so unlinking it here was the two halves of +# one rule disagreeing. Two shapes, since a link to a file is -f and a dangling one is not. +ln -s "$MRL/.llama.cpp.install.lock" "$MRL/.whisper.cpp.install.lock" +ln -s "$MRL/nowhere" "$MRL/.sd.cpp.install.lock" +: > "$MRL/.sd.cpp.install.lock.stale.4242" +run_block "$HOME" "$MRL" +assert_nodir "a real lock file is removed" "$MRL/.llama.cpp.install.lock" +assert_nodir "a stale lock file is removed" "$MRL/.sd.cpp.install.lock.stale.4242" +assert_link "a linked lock is kept" "$MRL/.whisper.cpp.install.lock" +assert_link "a dangling linked lock is kept too" "$MRL/.sd.cpp.install.lock" +assert_dir "a directory at a lock path is kept" "$MRL/.node.install.lock" +assert_dir "and so is everything under it" "$MRL/.node.install.lock/keep/user-file" + +echo "== an unmarked tree is somebody else's and is kept ==" +MR2="$_TMP_ROOT/mixed" +mkdir -p "$MR2"/{llama.cpp,node} +: > "$MR2/node/.unsloth-studio-owned" +: > "$MR2/llama.cpp/my-own-build" +run_block "$HOME" "$MR2" +assert_dir "unmarked llama.cpp kept" "$MR2/llama.cpp" +assert_dir "its contents kept" "$MR2/llama.cpp/my-own-build" +assert_nodir "marked node removed" "$MR2/node" + +echo "== an emptied master root is pruned, a used one is not ==" +MR3="$_TMP_ROOT/empty-after" +mkdir -p "$MR3/node" +: > "$MR3/node/.unsloth-studio-owned" +run_block "$HOME" "$MR3" +assert_nodir "empty master root pruned" "$MR3" +MR4="$_TMP_ROOT/still-used" +mkdir -p "$MR4/node" "$MR4/my-datasets" +: > "$MR4/node/.unsloth-studio-owned" +run_block "$HOME" "$MR4" +assert_dir "master root with user files kept" "$MR4" +assert_dir "the user's files kept" "$MR4/my-datasets" + +echo "== the value is stripped and tilde-expanded, like the installers ==" +MR5="$HOME/tilde-root" +mkdir -p "$MR5/node" +: > "$MR5/node/.unsloth-studio-owned" +run_block "$HOME" " ~/tilde-root " +assert_nodir "padded tilde value resolved" "$MR5/node" + +echo "== a shared .staging keeps whatever is not ours ==" +MR6="$_TMP_ROOT/staged" +mkdir -p "$MR6/node" "$MR6/.staging/somebody-elses" +: > "$MR6/node/.unsloth-studio-owned" +run_block "$HOME" "$MR6" +assert_dir "a non-empty .staging is kept whole" "$MR6/.staging/somebody-elses" +MR7="$_TMP_ROOT/staged-empty" +mkdir -p "$MR7/node" "$MR7/.staging" +: > "$MR7/node/.unsloth-studio-owned" +run_block "$HOME" "$MR7" +assert_nodir "an empty .staging is pruned" "$MR7/.staging" + +echo "== an unset or default root selects nothing ==" +got=$( ( HOME="$HOME"; export HOME; unset UNSLOTH_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "unset UNSLOTH_HOME yields nothing" "$got" "" +got=$( ( HOME="$HOME"; UNSLOTH_HOME=" "; export HOME UNSLOTH_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "blank UNSLOTH_HOME yields nothing" "$got" "" +mkdir -p "$HOME/.unsloth" +got=$( ( HOME="$HOME"; UNSLOTH_HOME="$HOME/.unsloth"; export HOME UNSLOTH_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "the default root is left to the blocks that own it" "$got" "" + +echo "== a deny-listed root is refused ==" +got=$( ( HOME="$HOME"; UNSLOTH_HOME="/"; export HOME UNSLOTH_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "/ yields nothing" "$got" "" + +echo "== a root that is gone does not end the uninstall ==" +# Under set -e the canonicalizing `cd` fails whenever the root is absent (a second uninstall, a +# portable install on an unplugged drive) and a bare assignment takes that exit status. bash +# hides it by clearing errexit inside command substitution, so this runs the shells the +# advertised `| sh` one-liner reaches. REACHED_END is the assertion, not the value returned. +for _shell in dash "bash --posix" sh; do + command -v ${_shell%% *} > /dev/null 2>&1 || continue + # || true: this harness runs under set -e too, so without it the failing child takes the + # whole test file down and the check below never reports which shell broke. + out=$( env -i HOME="$HOME" PATH="$PATH" UNSLOTH_HOME="$_TMP_ROOT/never-existed" \ + $_shell -c "set -e; . \"$HELPERS_FILE\"; r=\$(_master_root); echo REACHED_END" 2>&1 \ + || true ) + assert_eq "$_shell keeps going past a missing root" "$out" "REACHED_END" +done + +echo "== a dangling runtime symlink is not ours to unlink ==" +# -e is false for a dangling symlink, so an unmarked one named llama.cpp fell past the marker +# gate into _remove_path, which treats -L as present and removes the link. Its target volume +# being unmounted is the ordinary reason, and the link is the user's. +MRD="$_TMP_ROOT/dangling" +mkdir -p "$MRD/studio" "$MRD/node" +: > "$MRD/node/.unsloth-studio-owned" +ln -s "$_TMP_ROOT/no-such-volume/llama.cpp" "$MRD/llama.cpp" +run_block "$HOME" "$MRD" +assert_present "an unmarked dangling link is kept" "$MRD/llama.cpp" +assert_nodir "a marked tree beside it still goes" "$MRD/node" + +echo "== a whitespace-only Studio override does not suppress the master root ==" +# storage_roots.studio_root() trims these, so " " is unset to every resolver. A bare -n test +# called it present, skipped the master-root branch, and then discarded the whitespace path: +# the runtime siblings went and /studio stayed. +ROOTS_FILE=$(mktemp -p "$_TMP_ROOT") +sed -n '/^_custom_studio_roots() {/,/^}/p' "$UNINSTALL_SH" > "$ROOTS_FILE" 2>/dev/null || : > "$ROOTS_FILE" +if grep -q '_custom_studio_roots() {' "$ROOTS_FILE"; then + MRW="$_TMP_ROOT/blankoverride" + mkdir -p "$MRW/studio" + got=$( ( HOME="$HOME"; UNSLOTH_HOME="$MRW"; UNSLOTH_STUDIO_HOME=" " + export HOME UNSLOTH_HOME UNSLOTH_STUDIO_HOME + . "$HELPERS_FILE" + _emit() { printf '%s\n' "$1"; } + _from_conf() { :; } + . "$ROOTS_FILE" + _custom_studio_roots 2>/dev/null ) | grep -c "$MRW/studio" ) || got=0 + assert_eq "a blank override still reaches the master root" "$got" "1" +else + echo " FAIL: _custom_studio_roots could not be extracted"; FAIL=$((FAIL+1)) +fi + +echo "== the master root is remembered when the environment does not carry it ==" +# `UNSLOTH_HOME=/mnt/portable unsloth studio update` installs the runtimes there and leaves +# nothing behind in the environment. setup.sh writes a note in the Studio tree; without it an +# uninstall later removed the Studio tree and stranded the runtimes. +# The real layout: Studio lives at /studio, and with a bare environment the only thing +# that names that directory is the default-mode studio.conf install.sh wrote. +NOTED="$_TMP_ROOT/noted" +mkdir -p "$HOME/.unsloth/studio/share" "$NOTED/studio/share" "$HOME/.local/share/unsloth" +printf "UNSLOTH_EXE='%s'\n" "$NOTED/studio/unsloth_studio/bin/unsloth" \ + > "$HOME/.local/share/unsloth/studio.conf" +printf '%s\n' "$NOTED" > "$NOTED/studio/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; export HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "the note names the master root" "$got" "$NOTED" +# The environment still outranks it, and the deny list still applies to whatever it names. +got=$( ( HOME="$HOME"; UNSLOTH_HOME="$_TMP_ROOT/from-env"; export HOME UNSLOTH_HOME + . "$HELPERS_FILE"; _master_root ) ) +assert_eq "the environment outranks the note" "$got" "$_TMP_ROOT/from-env" +printf '%s\n' "$HOME/.unsloth" > "$NOTED/studio/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; export HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a note naming the default root is still refused" "$got" "" +# A tree copied from one master root to another keeps a note naming the ORIGINAL, whose +# runtimes carry the same owner markers the copy's would, so accepting it would authorise +# deleting the original install's llama.cpp, node and whisper.cpp. +COPIED="$_TMP_ROOT/copied" +mkdir -p "$COPIED/studio/share" +printf '%s\n' "$NOTED" > "$COPIED/studio/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; UNSLOTH_STUDIO_HOME="$COPIED/studio" + export HOME UNSLOTH_STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a note carried in from another master root is refused" "$got" "" +rm -f "$NOTED/studio/share/.unsloth-master-root" "$HOME/.local/share/unsloth/studio.conf" + +echo "== a named Studio root does not borrow another install's note ==" +# Two installs on one box, the legacy tree carrying a note and the named one not: probing the +# legacy path first let its note win here while the removal worked on the named tree, so the run +# deleted one Studio and the OTHER install's runtime children. +BORROWED="$_TMP_ROOT/borrowed" +NAMED_MASTER="$_TMP_ROOT/named-master" +NAMED="$NAMED_MASTER/studio" +mkdir -p "$HOME/.unsloth/studio/share" "$BORROWED/studio" "$NAMED/share" +printf '%s\n' "$BORROWED" > "$HOME/.unsloth/studio/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; UNSLOTH_STUDIO_HOME="$NAMED" + export HOME UNSLOTH_STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "an unrelated install's note is not adopted" "$got" "" +# The alias is the same rule, and the named root's OWN note is still read. +printf '%s\n' "$NAMED_MASTER" > "$NAMED/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME UNSLOTH_STUDIO_HOME; STUDIO_HOME="$NAMED" + export HOME STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "the named root's own note is still read" "$got" "$NAMED_MASTER" +rm -f "$HOME/.unsloth/studio/share/.unsloth-master-root" "$NAMED/share/.unsloth-master-root" + +echo "== a padded Studio override still finds the note setup.sh wrote ==" +# setup.sh trims UNSLOTH_STUDIO_HOME before choosing where to write, so appending to the padded +# value looked somewhere that does not exist: the note was never found, while the removal loop +# trimmed the same override and stranded the master-root runtimes. +PADDED_MASTER="$_TMP_ROOT/padded-master" +PADDED="$PADDED_MASTER/studio" +mkdir -p "$PADDED/share" +printf '%s\n' "$PADDED_MASTER" > "$PADDED/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; UNSLOTH_STUDIO_HOME=" $PADDED " + export HOME UNSLOTH_STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a padded override finds the note" "$got" "$PADDED_MASTER" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME UNSLOTH_STUDIO_HOME; STUDIO_HOME=" $PADDED " + export HOME STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "the padded alias finds it too" "$got" "$PADDED_MASTER" + +echo "== only the installer's own stale locks are swept ==" +# prebuilt_core.py moves a lock aside as .stale.. A bare .*.install.lock.stale.* glob +# also matched an unrelated hidden file, and in a root the user chose that file is theirs. +STALE="$_TMP_ROOT/stale-root" +mkdir -p "$STALE" +: > "$STALE/.node.install.lock.stale.4242" +: > "$STALE/.backup.install.lock.stale.copy" +: > "$STALE/.llama.cpp.install.lock.stale.notapid" +run_block "$HOME" "$STALE" +assert_nodir "the installer's own stale lock goes" "$STALE/.node.install.lock.stale.4242" +assert_present "an unrelated hidden file stays" "$STALE/.backup.install.lock.stale.copy" +assert_present "a non-numeric suffix stays" "$STALE/.llama.cpp.install.lock.stale.notapid" + +echo "== a note with a second line is refused, as the comment beside it promises ==" +# The four readers have to agree, and this is the one that authorises a delete: the two Python +# readers strip the WHOLE file, so a two-line note is not a directory to them, while taking +# line 1 here licensed removing /llama.cpp, node, whisper.cpp and sd.cpp. +MULTI_MASTER="$_TMP_ROOT/multiline-master" +MULTI="$MULTI_MASTER/studio" +mkdir -p "$MULTI/share" "$_TMP_ROOT/elsewhere" +printf '%s\n%s\n' "$MULTI_MASTER" "$_TMP_ROOT/elsewhere" > "$MULTI/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; UNSLOTH_STUDIO_HOME="$MULTI" + export HOME UNSLOTH_STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a two-line note names no master root" "$got" "" + +# ... and a blank second line is just the trailing newline of an ordinary note. +printf '%s\n\n' "$MULTI_MASTER" > "$MULTI/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; UNSLOTH_STUDIO_HOME="$MULTI" + export HOME UNSLOTH_STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a trailing blank line is still one line" "$got" "$MULTI_MASTER" + +# The ordinary note must keep working, or the check above proves only that nothing is read. +printf '%s\n' "$MULTI_MASTER" > "$MULTI/share/.unsloth-master-root" +got=$( ( HOME="$HOME"; unset UNSLOTH_HOME; UNSLOTH_STUDIO_HOME="$MULTI" + export HOME UNSLOTH_STUDIO_HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a one-line note is still honoured" "$got" "$MULTI_MASTER" + +echo "== a note in the legacy tree cannot aim the uninstall at a second install ==" +# The legacy install carries a note naming $HOME, which setup used to accept: the profile +# CONTAINS ~/.unsloth/studio, so containment passed and the value is not the legacy root. +# _custom_studio_roots then emitted "$HOME/studio", and a SECOND marked install kept there was +# accepted as a Studio root to remove, taking its database and outputs. The deny list does not +# catch it: it refuses $HOME as a master root, not the Studio candidate derived from it. +# storage_roots._is_legacy_studio_tree and the CLI decline that note outright; so must this. +LEGACY_HOME="$_TMP_ROOT/legacy-note-home" +mkdir -p "$LEGACY_HOME/.unsloth/studio/share" +printf '%s\n' "$LEGACY_HOME" > "$LEGACY_HOME/.unsloth/studio/share/.unsloth-master-root" +mkdir -p "$LEGACY_HOME/studio/share" "$LEGACY_HOME/studio/outputs" +: > "$LEGACY_HOME/studio/.unsloth-studio-owned" +: > "$LEGACY_HOME/studio/studio.db" +got=$( ( HOME="$LEGACY_HOME"; unset UNSLOTH_HOME UNSLOTH_STUDIO_HOME STUDIO_HOME + export HOME; . "$HELPERS_FILE"; _master_root ) ) +assert_eq "a legacy-tree note names no master root" "$got" "" +assert_dir "the second install is untouched" "$LEGACY_HOME/studio/studio.db" + +echo +echo "PASS=$PASS FAIL=$FAIL" +[ "$FAIL" -eq 0 ] diff --git a/tests/sh/test_uninstall_sd_cpp_custom_root.sh b/tests/sh/test_uninstall_sd_cpp_custom_root.sh index 4b06cfb5a3..f740f1cfde 100755 --- a/tests/sh/test_uninstall_sd_cpp_custom_root.sh +++ b/tests/sh/test_uninstall_sd_cpp_custom_root.sh @@ -33,6 +33,10 @@ HELPERS_FILE=$(mktemp -p "$_TMP_ROOT") sed -n '/^_is_venv_dir() {/,/^}/p' "$UNINSTALL_SH" sed -n '/^_is_studio_root() {/,/^}/p' "$UNINSTALL_SH" sed -n '/^_is_unsafe_root() {/,/^}/p' "$UNINSTALL_SH" + # The loop now asks whether the Studio root is also the master root before removing it + # whole, so this helper has to come across or the block dies on "command not found" and + # every keep-assertion passes for the wrong reason. + sed -n '/^_master_root() {/,/^}/p' "$UNINSTALL_SH" # The loop removes roots through this wrapper; without it and its marker helper the # fragment dies with "command not found" and every assertion below is vacuous. sed -n '/^_set_marker() {/,/^}/p' "$UNINSTALL_SH" @@ -45,7 +49,7 @@ HELPERS_FILE=$(mktemp -p "$_TMP_ROOT") } > "$HELPERS_FILE" grep -q '_owned_sd_cpp_roots' "$HELPERS_FILE" || { echo "FAIL: helpers missing _owned_sd_cpp_roots"; exit 1; } # _is_studio_root's own helpers, or every root reads as foreign and nothing below is removed. -for _needed in _is_owner_marker _is_venv_dir _restore_owner_marker; do +for _needed in _is_owner_marker _is_venv_dir _restore_owner_marker _master_root; do grep -q "^$_needed() {" "$HELPERS_FILE" || { echo "FAIL: helpers missing $_needed"; exit 1; } done grep -q '_sd_cpp_sibling_bases() {' "$HELPERS_FILE" || { echo "FAIL: helpers missing _sd_cpp_sibling_bases"; exit 1; } diff --git a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh index 2d87b7f541..0e8fac1c63 100644 --- a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh +++ b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh @@ -41,6 +41,11 @@ setup_fail() { exit "$1"; } _assert_studio_owned_or_absent() { :; } C_ERR="" _STUDIO_HOME_IS_CUSTOM=false +# The runtime children took their own ownership flag when UNSLOTH_HOME arrived: with no master +# root it carries _STUDIO_HOME_IS_CUSTOM, which is what this fixture is about. Unseeded it is an +# unbound variable under set -u, and the block dies before printing anything, which reads here as +# every state assertion getting an empty value rather than as a broken harness. +_RUNTIME_ROOT_IS_CUSTOM=false _NEED_LLAMA_SOURCE_BUILD=UNSET _SKIP_PREBUILT_INSTALL=UNSET ' @@ -54,10 +59,23 @@ SNIP="$PREAMBLE"$'\n'"$block"$'\n'"$EPILOGUE" # run_link -> prints state lines; RC in $RC run_link() { + # stderr kept on failure: the block runs under set -u, so a variable this preamble does not + # seed kills it silently and every val() below returns empty. That looked like seven wrong + # answers rather than one missing seed, so say which it is. OUT="$(env -i PATH="$PATH" HOME="$T" \ UNSLOTH_LOCAL_LLAMA_CPP_DIR="$1" LLAMA_CPP_DIR="$2" \ - bash -c "$SNIP" 2>/dev/null)" + bash -c "$SNIP" 2>"$T/.stderr")" RC=$? + # Only when the block claims it succeeded: some fixtures below expect setup_fail, and those + # legitimately exit before the epilogue. + if [ "$RC" -eq 0 ]; then + case "$OUT" in + *LINKED=*) : ;; + *) echo "FAIL: the block exited 0 and printed no state; stderr was:" >&2 + sed 's/^/ /' "$T/.stderr" >&2 + exit 1 ;; + esac + fi } val() { printf '%s\n' "$OUT" | grep "^$1=" | head -1 | cut -d= -f2-; } diff --git a/tests/studio/install/test_denied_llama_cpp_preflight.py b/tests/studio/install/test_denied_llama_cpp_preflight.py index 0d3769c1c4..9d7f278cb2 100644 --- a/tests/studio/install/test_denied_llama_cpp_preflight.py +++ b/tests/studio/install/test_denied_llama_cpp_preflight.py @@ -34,6 +34,7 @@ SHARED_FUNCTIONS = ( "Write-PathAccessDenied", "Get-CanonicalDir", "Test-StudioHomeIsCustom", + "Get-MasterRootOverride", "Get-ManagedLlamaCppDir", "Invoke-ManagedLlamaCppPreflight", ) @@ -216,11 +217,12 @@ def test_a_custom_studio_home_is_never_called_a_cache_we_own() -> None: """Do not call an unreadable custom Unsloth home a managed cache.""" body = _function_source(INSTALL_PS1, "Invoke-ManagedLlamaCppPreflight") # Use the same predicate for path selection and ownership wording. - assert "$homeIsCustom = Test-StudioHomeIsCustom" in body + # A master root counts as custom too: it moves llama.cpp out of the default location. + assert "$homeIsCustom = (Test-StudioHomeIsCustom) -or [bool](Get-MasterRootOverride)" in body assert "-OwnershipUnverified:$homeIsCustom" in body assert ( 'Exit-PathAccessDenied -Path $LlamaCppDir -Label "llama.cpp install"' - " -OwnershipUnverified:$StudioHomeIsCustom" in SETUP_PS1 + " -OwnershipUnverified:$RuntimeRootIsCustom" in SETUP_PS1 ) diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 0d69b451a9..9bebda1f0b 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -4408,6 +4408,7 @@ _LLAMA_CPP_NO_SPACE=false _LLAMA_CPP_DEGRADED=false _explicit_llama_backend="" _STUDIO_HOME_IS_CUSTOM=false +_RUNTIME_ROOT_IS_CUSTOM=false _STUDIO_OWNED_MARKER=".unsloth-owned" step() { echo "step: $2"; } substep() { echo "substep: $1"; } diff --git a/tests/studio/install/test_pr10648_offline_and_status.py b/tests/studio/install/test_pr10648_offline_and_status.py index ae244b0a89..11f8122605 100644 --- a/tests/studio/install/test_pr10648_offline_and_status.py +++ b/tests/studio/install/test_pr10648_offline_and_status.py @@ -720,6 +720,9 @@ _LLAMA_CPP_NO_SPACE=false _LLAMA_CPP_DEGRADED=false _explicit_llama_backend="" _STUDIO_HOME_IS_CUSTOM=false +# The flag the runtime children switched to when UNSLOTH_HOME arrived: the guards inside the +# extracted block read it, and setup.sh derives it in the section this harness stands in for. +_RUNTIME_ROOT_IS_CUSTOM=false _STUDIO_OWNED_MARKER=".unsloth-owned" _WHISPER_RECOVERED=false step() { echo "step|$1|$2"; } @@ -737,6 +740,8 @@ $ErrorActionPreference = "Stop" $NeedLlamaSourceBuild = $false $script:LlamaCppDegraded = $false $StudioHomeIsCustom = $false +# Same reason as _RUNTIME_ROOT_IS_CUSTOM in the shell harness above. +$RuntimeRootIsCustom = $false function step { param($a, $b, $c) Write-Output "step|$a|$b" } function substep { param($a, $b) Write-Output "substep|$a" } function Write-LlamaFailureLog { param($Output) } diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index faa1af5a8e..2b848b4b2b 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -94,10 +94,19 @@ def load_studio_run_module(monkeypatch): startup_banner.stdout_supports_color = lambda: False monkeypatch.setitem(sys.modules, "startup_banner", startup_banner) + # run.py exports UNSLOTH_STUDIO_HOME and UNSLOTH_LLAMA_CPP_PATH at module scope and this + # helper execs it in-process, so without monkeypatch they outlive the test. + for _leaked in ("UNSLOTH_STUDIO_HOME", "UNSLOTH_LLAMA_CPP_PATH"): + monkeypatch.setenv(_leaked, os.environ.get(_leaked, "")) + monkeypatch.delenv(_leaked) + paths = types.ModuleType("utils.paths") paths.__path__ = [] storage_roots = types.ModuleType("utils.paths.storage_roots") storage_roots.studio_root = lambda: PACKAGE_ROOT / ".studio-test-root" + # Imported at module scope too, so the stub carries it. None keeps these tests on the + # legacy-default path they were written for. + storage_roots.unsloth_home = lambda: None paths.storage_roots = storage_roots monkeypatch.setitem(sys.modules, "utils.paths", paths) monkeypatch.setitem(sys.modules, "utils.paths.storage_roots", storage_roots) diff --git a/tests/studio/install/test_setup_denied_install_tree.py b/tests/studio/install/test_setup_denied_install_tree.py index e39f1ca64b..a4d1d8300b 100644 --- a/tests/studio/install/test_setup_denied_install_tree.py +++ b/tests/studio/install/test_setup_denied_install_tree.py @@ -75,7 +75,9 @@ def test_every_denial_route_reports_instead_of_proceeding(): assert '$llamaGitState = Get-PathState -Path (Join-Path $LlamaCppDir ".git")' in SETUP_PS1 assert '$llamaGitState -eq "Denied"' in SETUP_PS1 assert "$pathState = Get-PathState -Path $Path -PathType Container" in SETUP_PS1 - assert '$StudioHomeIsCustom -and $pathState -eq "Denied"' in SETUP_PS1 + # $isCustomRoot is the parameter form of $StudioHomeIsCustom: the runtime children pass + # $RuntimeRootIsCustom, which a master root widens, and it defaults to the other flag. + assert '$isCustomRoot -and $pathState -eq "Denied"' in SETUP_PS1 # The junction path replaces this destination, so it needs its own stop. assert "$destState = Get-PathState -Path $LlamaCppDir" in SETUP_PS1 assert '$destState -eq "Denied"' in SETUP_PS1 @@ -129,7 +131,8 @@ def test_ownership_guard_distinguishes_denied_from_unowned(): # genuinely-unowned case only. assert "is not marked as an Unsloth-owned $Label" in guard # Both stops stay gated, so default-home installs behave exactly as before. - assert guard.count("$StudioHomeIsCustom -and") >= 3 + assert "$isCustomRoot = $StudioHomeIsCustom" in guard + assert guard.count("$isCustomRoot -and") >= 3 def test_no_bare_test_path_probes_inside_the_llama_install_tree(): @@ -321,7 +324,7 @@ def test_the_source_build_denial_never_advises_deleting_an_unproven_tree(): block = _slice("$llamaBinState = ", "$WillBuildLlamaFromSource") denials = [ln.strip() for ln in block.splitlines() if "Exit-PathAccessDenied" in ln] assert len(denials) >= 2, denials - assert all(d.endswith("-OwnershipUnverified:$StudioHomeIsCustom") for d in denials), denials + assert all(d.endswith("-OwnershipUnverified:$RuntimeRootIsCustom") for d in denials), denials assert all(d.startswith("Exit-PathAccessDenied -Path $LlamaCppDir ") for d in denials), denials @@ -368,10 +371,13 @@ def test_the_whisper_phase_survives_an_unreadable_whisper_tree(): # Only the denial is handed back; an unowned tree must still stop. assert 'Exit-SetupFailure "$Label path is not an Unsloth-owned install' in guard whisper = _whisper_phase() - assert '-Label "whisper.cpp install" -NonFatal) -eq "Denied"' in whisper, whisper + assert ( + '-Label "whisper.cpp install" -NonFatal -IsCustom $RuntimeRootIsCustom) -eq "Denied"' + in whisper + ), whisper # Scoped to the new branch: both phrases occur elsewhere in the phase, so a # phase-wide match proves nothing about this branch. - marker = '-NonFatal) -eq "Denied") {' + marker = '-NonFatal -IsCustom $RuntimeRootIsCustom) -eq "Denied") {' assert marker in whisper, whisper denial = whisper.split(marker, 1)[1].split("\n} elseif", 1)[0] assert re.search(r'^\s*step "whisper\.cpp" ', denial, re.M), denial @@ -389,7 +395,7 @@ def test_the_whisper_phase_survives_an_unreadable_whisper_tree(): def test_the_whisper_skip_stays_behind_the_installer_gate(): """The guard used to live inside the installer branch, so a tree without the installer was a no-op. Hoisting it must not make that case fatal.""" - marker = "-NonFatal) -eq" + marker = "-NonFatal -IsCustom $RuntimeRootIsCustom) -eq" whisper = _whisper_phase() assert marker in whisper, whisper head = whisper.split(marker, 1)[0] diff --git a/tests/studio/test_path_probe_access_denied.ps1 b/tests/studio/test_path_probe_access_denied.ps1 index 9f49add2d7..e1605db011 100644 --- a/tests/studio/test_path_probe_access_denied.ps1 +++ b/tests/studio/test_path_probe_access_denied.ps1 @@ -52,7 +52,14 @@ Check "the source-build .git probe stops on a denied checkout" ( $setupText -match '\$llamaGitState -eq "Denied"') Check "the ownership guard stops on a denied root instead of returning" ( $setupText -match '\$pathState = Get-PathState -Path \$Path -PathType Container' -and - $setupText -match '\$StudioHomeIsCustom -and \$pathState -eq "Denied"') + $setupText -match '\$isCustomRoot -and \$pathState -eq "Denied"') +# The guard now takes the flag as a parameter, because the runtime children beside studio\ under +# a master root are owned even when the Studio home itself is the legacy one. With no override +# it must still be the script-level flag this file's other checks are about, or the guard would +# be reading something nothing sets. +Check "the guard's flag defaults to the script-level one" ( + $setupText -match '\$isCustomRoot = \$StudioHomeIsCustom' -and + $setupText -match 'if \(\$null -ne \$IsCustom\) \{ \$isCustomRoot = \[bool\]\$IsCustom \}') Check "guidance says an app reinstall does not reset the folder" ( $setupText -match 'reinstalling Unsloth Studio, to any drive, reuses it' -and $setupText -match 'Reinstalling the app does not reset it\.') @@ -445,6 +452,28 @@ if ($assertSrc -and $markSrc) { try { $null = Assert-StudioOwnedOrAbsent -Path $nfUnowned -Label "whisper.cpp install" -NonFatal } catch { $threw = $true; $thrownBy = $_.Exception.Message } Check "-NonFatal does not excuse an unowned tree" ($threw -and $thrownBy -eq "EXIT-SETUP") + + # A non-directory at a runtime path under a user-chosen root. install_llama_prebuilt's + # activate_install_tree moves aside whatever Path.exists() finds, so a Container-only + # probe let a user's file be displaced by the install that followed. + $nfFile = Join-Path $nfRoot "llama.cpp" + Set-Content -LiteralPath $nfFile -Value "mine" + $nfLink = Join-Path $nfRoot "node" + New-Item -ItemType SymbolicLink -Path $nfLink -Target (Join-Path $nfRoot "gone") -ErrorAction SilentlyContinue | Out-Null + foreach ($shape in @(@("a regular file", $nfFile), @("a dangling link", $nfLink))) { + if (-not (Get-Item -LiteralPath $shape[1] -Force -ErrorAction SilentlyContinue)) { continue } + $threw = $false + $thrownBy = $null + try { $null = Assert-StudioOwnedOrAbsent -Path $shape[1] -Label "llama.cpp install" -NonFatal } + catch { $threw = $true; $thrownBy = $_.Exception.Message } + Check "$($shape[0]) at a custom runtime path stops setup" ($threw -and $thrownBy -eq "EXIT-SETUP") + # The legacy default home keeps the behaviour it had: only a chosen root is the + # user's directory, and setup.sh draws the line in the same place. + $threw = $false + try { $null = Assert-StudioOwnedOrAbsent -Path $shape[1] -Label "llama.cpp install" -IsCustom $false } + catch { $threw = $true } + Check "$($shape[0]) at a default runtime path is left to the caller" (-not $threw) + } } finally { Set-NfDenied $false Remove-Item -Recurse -Force -LiteralPath $nfRoot -ErrorAction SilentlyContinue @@ -458,7 +487,7 @@ $installPath = [System.IO.Path]::Combine($repoRoot, "install.ps1") $preflightFns = @("Test-AccessDeniedError", "Get-PathState", "Get-LlamaCppInstallReadState", "Get-PathDenialDetail", "Get-SecuritySoftwareNote", "Write-PathAccessDenied", "Get-CanonicalDir", - "Test-StudioHomeIsCustom", "Get-ManagedLlamaCppDir", + "Test-StudioHomeIsCustom", "Get-MasterRootOverride", "Get-ManagedLlamaCppDir", "Invoke-ManagedLlamaCppPreflight") $preflightSrc = @() foreach ($fn in $preflightFns) { @@ -466,6 +495,31 @@ foreach ($fn in $preflightFns) { Check "install.ps1 defines $fn" ($null -ne $src) if ($src) { $preflightSrc += $src } } +# Anything the lifted bodies call that install.ps1 defines and the lift then leaves behind, +# other than the stubs the harness supplies. A missing helper does not announce itself: the +# child dies on the first call, every Write-Host after it is lost, and the assertions below +# fail as though the preflight had resolved the wrong directory. This is how +# Get-MasterRootOverride was missed. The test is "install.ps1 defines it", not "this session +# can resolve it": Get-SecuritySoftwareNote reaches Get-CimInstance and Get-MpPreference, which +# exist on the Windows host that runs the preflight and on no Linux runner, so resolvability +# would report those two as unlifted helpers here and nothing at all there. +$harnessStubs = @("step", "substep", "Write-StudioLine") +$preflightCalls = @() +foreach ($src in $preflightSrc) { + $tokens = [System.Management.Automation.PSParser]::Tokenize($src, [ref]$null) + foreach ($t in $tokens) { + if ($t.Type -eq "Command" -and $t.Content -match "^(Get|Test|Invoke|Write|Set|New)-[A-Za-z]+$") { + $preflightCalls += $t.Content + } + } +} +$unlifted = @($preflightCalls | Sort-Object -Unique | Where-Object { + $_ -notin $preflightFns -and $_ -notin $harnessStubs -and + $null -ne (Get-FunctionSource -Path $installPath -Name $_) +}) +Check ("the preflight lifts every install.ps1 helper it calls" + + $(if ($unlifted.Count) { " (not lifted: " + ($unlifted -join ", ") + ")" } else { "" })) ` + ($unlifted.Count -eq 0) if ($preflightSrc.Count -eq $preflightFns.Count) { $preflightHarness = @" `$ErrorActionPreference = "Stop" diff --git a/tests/test_installer_system32_guard.py b/tests/test_installer_system32_guard.py index bf6f9b9137..34714ac45b 100644 --- a/tests/test_installer_system32_guard.py +++ b/tests/test_installer_system32_guard.py @@ -1036,6 +1036,8 @@ def test_cli_guard_pins_every_storage_root_override_studio_reads(): source = storage_roots.read_text(encoding = "utf-8") overrides = set(re.findall(r'environ\.get\(\s*"(UNSLOTH_[A-Z_]*(?:HOME|PATH|DIR))"', source)) assert overrides, "no storage root overrides found: has storage_roots.py moved?" + # Same case: _setup_cache_env fills each only when blank, so a user's relative value survives. + overrides |= set(re.findall(r'^\s*"([A-Z][A-Z0-9_]*)":\s*str\(', source, re.MULTILINE)) missing = sorted(overrides - set(_system_dir_guard._RELATIVE_PATH_ENV)) assert not missing, f"relative values of {missing} would be retargeted by the move" diff --git a/tests/test_managed_tools_master_root.py b/tests/test_managed_tools_master_root.py new file mode 100644 index 0000000000..6abdccdc34 --- /dev/null +++ b/tests/test_managed_tools_master_root.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""UNSLOTH_HOME names the tree, not the Studio directory inside it, and the native runtimes are +siblings of studio/, the spelling studio/setup.sh and scripts/build_whisper_cpp.sh already use. A +resolver deriving them from studio_root() would look in /studio/ for what the +installer put at /, so managed Node and whisper.cpp go missing and run.py pins the +wrong llama.cpp path into every worker. + +Run in a subprocess per case: these modules read the environment at import time. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parents[1] +BACKEND = REPO / "studio" / "backend" + +PROBE = """ +import json, os, sys +sys.path.insert(0, os.environ["_BACKEND"]) +from utils.paths import storage_roots as sr +from utils.node_runtime import managed_node_dir +from core.inference.stt_ggml_sidecar import _managed_whisper_cpp_dir + +# studio_root() is called constantly, so a warning it emits for a supported +# layout is not one line, it is a flooded log. +_warnings = [] +sr.logger.warning = lambda msg, *a, **k: _warnings.append(msg % a if a else msg) + +print(json.dumps({ + "studio": str(sr.studio_root()), + "master": None if sr.unsloth_home() is None else str(sr.unsloth_home()), + "node": str(managed_node_dir()), + "whisper": str(_managed_whisper_cpp_dir()), + "warnings": _warnings, +})) +""" + + +def _resolve(env_overrides: dict[str, str], home: Path) -> dict[str, str]: + env = { + "PATH": os.environ.get("PATH", "/usr/bin:/bin"), + "HOME": str(home), + "USERPROFILE": str(home), + "_BACKEND": str(BACKEND), + } + # A hand-built environment has to carry what the interpreter needs: Windows python exits 1 + # with no usable message when SYSTEMROOT is absent, which read as "the resolver answered + # wrongly" on every Windows runner. + for name in ("SYSTEMROOT", "SystemRoot", "COMSPEC", "PATHEXT", "TEMP", "TMP", "WINDIR"): + value = os.environ.get(name) + if value: + env.setdefault(name, value) + env.update(env_overrides) + out = subprocess.run([sys.executable, "-c", PROBE], env = env, capture_output = True, text = True) + # Not check = True: the child's stderr is the only thing that says why, and swallowing it is + # how a dead interpreter passes for a wrong answer. + assert out.returncode == 0, f"probe failed ({out.returncode}): {out.stderr.strip()}" + return json.loads(out.stdout.strip().splitlines()[-1]) + + +def test_portable_root_puts_the_tools_beside_studio(tmp_path): + home = tmp_path / "home" + home.mkdir() + root = tmp_path / "portable" + r = _resolve({"UNSLOTH_HOME": str(root)}, home) + assert r["studio"] == str(root / "studio") + assert r["master"] == str(root) + assert r["node"] == str(root / "node") + assert r["whisper"] == str(root / "whisper.cpp") + + +def test_a_default_install_is_untouched(tmp_path): + home = tmp_path / "home" + home.mkdir() + r = _resolve({}, home) + assert r["master"] is None + assert r["studio"] == str(home / ".unsloth" / "studio") + assert r["node"] == str(home / ".unsloth" / "node") + assert r["whisper"] == str(home / ".unsloth" / "whisper.cpp") + + +def test_a_plain_custom_studio_home_is_untouched(tmp_path): + # No UNSLOTH_HOME: the tools stay children of the Studio root, as before. + home = tmp_path / "home" + home.mkdir() + custom = tmp_path / "custom" + r = _resolve({"UNSLOTH_STUDIO_HOME": str(custom)}, home) + assert r["master"] is None + assert r["studio"] == str(custom) + assert r["node"] == str(custom / "node") + assert r["whisper"] == str(custom / "whisper.cpp") + + +def test_a_flat_root_keeps_the_tools_at_that_root(tmp_path): + # UNSLOTH_HOME == UNSLOTH_STUDIO_HOME, so "beside studio/" and "inside it" are one directory. + home = tmp_path / "home" + home.mkdir() + root = tmp_path / "flat" + r = _resolve({"UNSLOTH_HOME": str(root), "UNSLOTH_STUDIO_HOME": str(root)}, home) + assert r["studio"] == str(root) + assert r["node"] == str(root / "node") + assert r["whisper"] == str(root / "whisper.cpp") + # Path.parents excludes the path itself, so the equality check is what keeps this warning off. + assert r["warnings"] == [] + + +def test_a_studio_home_outside_the_master_root_still_warns(tmp_path): + home = tmp_path / "home" + home.mkdir() + r = _resolve( + { + "UNSLOTH_HOME": str(tmp_path / "portable"), + "UNSLOTH_STUDIO_HOME": str(tmp_path / "elsewhere"), + }, + home, + ) + assert any("not self-contained" in w for w in r["warnings"]) + + +def test_the_builder_and_the_resolver_agree_on_the_same_directory(tmp_path): + # build_whisper_cpp.sh installs under UNSLOTH_HOME; a resolver one level off reports + # dictation unavailable with whisper-server sitting right there. + home = tmp_path / "home" + home.mkdir() + root = tmp_path / "portable" + built = root / "whisper.cpp" + built.mkdir(parents = True) + r = _resolve({"UNSLOTH_HOME": str(root)}, home) + assert r["whisper"] == str(built) + + +def _record_note(studio: Path, master: Path) -> None: + """What setup.sh writes at the end of a master-root install.""" + (studio / "share").mkdir(parents = True, exist_ok = True) + (studio / "share" / ".unsloth-master-root").write_text(f"{master}\n", encoding = "utf-8") + + +def test_a_recorded_master_root_outlives_the_command_that_set_it(tmp_path): + # UNSLOTH_HOME is settable for one command, which installs node and whisper.cpp BESIDE + # studio/, and nothing persists it: the next launch resolved both one level down, at + # /node and /whisper.cpp, with the real trees untouched next door. + home = tmp_path / "home" + home.mkdir() + root = tmp_path / "portable" + studio = root / "studio" + studio.mkdir(parents = True) + _record_note(studio, root) + # UNSLOTH_STUDIO_HOME alone: what the installer's launcher actually persists. + r = _resolve({"UNSLOTH_STUDIO_HOME": str(studio)}, home) + assert r["master"] == str(root) + assert r["studio"] == str(studio) + assert r["node"] == str(root / "node") + assert r["whisper"] == str(root / "whisper.cpp") + assert r["warnings"] == [] + + +def test_a_note_whose_root_has_since_moved_is_ignored(tmp_path): + # A note licenses this process to adopt a root for caches and runtimes both. One naming a + # tree that is no longer there must not win over the layout in front of it. + home = tmp_path / "home" + home.mkdir() + studio = tmp_path / "custom" + studio.mkdir() + _record_note(studio, tmp_path / "gone") + r = _resolve({"UNSLOTH_STUDIO_HOME": str(studio)}, home) + assert r["master"] is None + assert r["node"] == str(studio / "node") + + +def test_a_note_copied_into_an_unrelated_install_is_ignored(tmp_path): + # The recorded root exists and has a studio/ child, but it is not THIS install's studio + # directory, so the note travelled rather than described. Checking only is_dir() would + # redirect this install's runtimes into someone else's tree. + home = tmp_path / "home" + home.mkdir() + other = tmp_path / "other" + (other / "studio").mkdir(parents = True) + studio = tmp_path / "custom" + studio.mkdir() + _record_note(studio, other) + r = _resolve({"UNSLOTH_STUDIO_HOME": str(studio)}, home) + assert r["master"] is None + assert r["node"] == str(studio / "node") + + +def test_an_empty_note_is_not_a_root(tmp_path): + # A truncated or zero-length note must read as "no record", not as the current directory. + home = tmp_path / "home" + home.mkdir() + studio = tmp_path / "custom" + (studio / "share").mkdir(parents = True) + (studio / "share" / ".unsloth-master-root").write_text("\n", encoding = "utf-8") + r = _resolve({"UNSLOTH_STUDIO_HOME": str(studio)}, home) + assert r["master"] is None + assert r["node"] == str(studio / "node") + + +def test_an_explicit_studio_home_without_a_note_does_not_borrow_anothers(tmp_path): + # Two installs on one box, only the legacy tree carrying a note: studio_root() stays on the + # named tree, so reading past it would send the runtimes and the portable caches to the + # OTHER install while Studio ran from here. + home = tmp_path / "home" + legacy_master = home / ".unsloth" + (legacy_master / "studio").mkdir(parents = True) + _record_note(legacy_master / "studio", legacy_master) + named = tmp_path / "named" + named.mkdir() + r = _resolve({"UNSLOTH_STUDIO_HOME": str(named)}, home) + assert r["master"] is None + assert r["studio"] == str(named) + assert r["node"] == str(named / "node") + assert r["whisper"] == str(named / "whisper.cpp") + + +def test_a_flat_recorded_root_is_still_honoured(tmp_path): + # The flat layout is supported, and the note then sits in the root it names rather than in a + # studio/ child: an exact /studio match would refuse it, containment satisfies it. + home = tmp_path / "home" + home.mkdir() + flat = tmp_path / "flat" + (flat / "share").mkdir(parents = True) + (flat / "share" / ".unsloth-master-root").write_text(f"{flat}\n", encoding = "utf-8") + r = _resolve({"UNSLOTH_STUDIO_HOME": str(flat)}, home) + assert r["master"] == str(flat) + assert r["node"] == str(flat / "node") + assert r["warnings"] == [] + + +def test_a_default_install_reads_no_note(tmp_path): + # Nothing writes the note for a default install, so the legacy tree must not acquire a + # master root by accident: this is the path every existing user is on. + home = tmp_path / "home" + (home / ".unsloth" / "studio" / "share").mkdir(parents = True) + r = _resolve({}, home) + assert r["master"] is None + assert r["node"] == str(home / ".unsloth" / "node") + + +_DISCOVERY_PROBE = """ +import json, os, sys +sys.path.insert(0, os.environ["_BACKEND"]) +from pathlib import Path +from utils.paths.storage_roots import studio_root, unsloth_home +from utils.llama_cpp_path_settings import mark_managed_llama_cpp_path + +# Replay run.py's module-level block, which is what a real server start does +# before anything asks where llama-server is. +resolved = studio_root().resolve() +if resolved != (Path.home() / ".unsloth" / "studio"): + os.environ.setdefault("UNSLOTH_STUDIO_HOME", str(resolved)) + managed = (unsloth_home() or resolved) / "llama.cpp" + os.environ.setdefault("UNSLOTH_LLAMA_CPP_PATH", str(managed)) + mark_managed_llama_cpp_path(managed) + +from core.inference.llama_cpp import LlamaCppBackend + +print(json.dumps({ + "exported": os.environ.get("UNSLOTH_LLAMA_CPP_PATH"), + "found": LlamaCppBackend._find_llama_server_binary(), +})) +""" + + +def _install_llama_server(directory: Path) -> Path: + # The name _find_llama_server_binary looks for on THIS platform: it appends .exe on Windows, + # so a fixture that only ever writes the POSIX name made discovery correctly answer None and + # all three discovery tests fail on a Windows runner for a reason in the fixture. + name = "llama-server.exe" if sys.platform == "win32" else "llama-server" + binary = directory / "build" / "bin" / name + binary.parent.mkdir(parents = True, exist_ok = True) + binary.write_text("#!/bin/sh\nexit 0\n", encoding = "utf-8") + binary.chmod(0o755) + return binary + + +def _discover(env_overrides: dict[str, str], home: Path) -> dict[str, str]: + env = { + "PATH": os.environ.get("PATH", "/usr/bin:/bin"), + "HOME": str(home), + "USERPROFILE": str(home), + "_BACKEND": str(BACKEND), + } + # Same reason as _resolve: the interpreter's own requirements travel with it. + for name in ("SYSTEMROOT", "SystemRoot", "COMSPEC", "PATHEXT", "TEMP", "TMP", "WINDIR"): + value = os.environ.get(name) + if value: + env.setdefault(name, value) + env.update(env_overrides) + out = subprocess.run( + [sys.executable, "-c", _DISCOVERY_PROBE], + env = env, + capture_output = True, + text = True, + ) + assert out.returncode == 0, f"discovery probe failed ({out.returncode}): {out.stderr.strip()}" + return json.loads(out.stdout.strip().splitlines()[-1]) + + +def test_discovery_finds_the_llama_server_the_master_root_holds(tmp_path): + # The managed marker makes discovery SKIP the env var for its own derivation, so the two + # must name one directory or every GGUF model reports no runtime. + home = tmp_path / "home" + home.mkdir() + root = tmp_path / "portable" + (root / "studio").mkdir(parents = True) + binary = _install_llama_server(root / "llama.cpp") + + result = _discover({"UNSLOTH_HOME": str(root)}, home) + + assert result["exported"] == str(root / "llama.cpp") + assert result["found"] == str(binary) + + +def test_discovery_still_prefers_a_plain_custom_studio_root(tmp_path): + home = tmp_path / "home" + home.mkdir() + custom = tmp_path / "custom" + custom.mkdir() + binary = _install_llama_server(custom / "llama.cpp") + + result = _discover({"UNSLOTH_STUDIO_HOME": str(custom)}, home) + + assert result["found"] == str(binary) + + +def test_discovery_still_finds_a_legacy_install(tmp_path): + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + binary = _install_llama_server(home / ".unsloth" / "llama.cpp") + + result = _discover({}, home) + + assert result["found"] == str(binary) diff --git a/tests/test_setup_master_root.py b/tests/test_setup_master_root.py new file mode 100644 index 0000000000..c090f219cb --- /dev/null +++ b/tests/test_setup_master_root.py @@ -0,0 +1,1473 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""setup.sh / setup.ps1 must install the managed runtimes at the same root the resolvers read. + +`llama.cpp`, `node` and `whisper.cpp` are siblings of `studio/` under `UNSLOTH_HOME`, and +`storage_roots.unsloth_home()` sends every runtime resolver there. The CLI exports +`UNSLOTH_STUDIO_HOME=/studio` before running setup, so a setup that derived the runtime +parent from the Studio home alone would install them one level too deep and no GGUF model, +managed Node or dictation engine would be found. Companion to +tests/test_managed_tools_master_root.py, which holds the resolvers to each other. +""" + +from __future__ import annotations + +import os +import pathlib +import shutil +import subprocess +from pathlib import Path + +import pytest + +# tests/conftest.py puts tests/_shared on sys.path; see unsloth_pwsh_runner for why a bare +# subprocess call to pwsh is not enough under xdist. +from unsloth_pwsh_runner import PWSH, run_pwsh + +REPO_ROOT = Path(__file__).resolve().parents[1] +SETUP_SH = REPO_ROOT / "studio" / "setup.sh" +SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1" + + +# These execute blocks of the POSIX installer, so they need bash AND a POSIX filesystem. Git +# Bash puts bash on PATH on a Windows runner, where the blocks run against Windows path +# semantics and hand back UTF-16 (`assert '\x00' == 'ALLOWED'`). setup.ps1 is covered separately +# here and in tests/studio/, so skipping loses no coverage. +NEEDS_POSIX_BASH = pytest.mark.skipif( + shutil.which("bash") is None or os.name == "nt", + reason = "runs studio/setup.sh blocks: needs bash on a POSIX filesystem", +) + + +def _ps_quote(value: str) -> str: + """A PowerShell single-quoted literal: only the quote itself needs escaping.""" + return "'" + value.replace("'", "''") + "'" + + +def _slice(src: str, start: str, end: str) -> str: + begin = src.index(start) + return src[begin : src.index(end, begin)] + + +def _runtime_parent(env: dict[str, str]) -> tuple[str, str]: + """Run the shipped setup.sh derivations and report (node parent, llama.cpp parent). + + The two blocks are executed rather than pattern-matched, so a later edit that keeps the + words and changes the order still fails here. + """ + src = SETUP_SH.read_text(encoding = "utf-8") + master = _slice(src, "# Stripped before anything else", "# Directory-local evidence") + node = _slice(src, "# Mirror the llama.cpp UNSLOTH_HOME derivation", "NODE_DIR=") + llama = _slice(src, 'if [ -n "$STAGE_ROOT" ]; then\n UNSLOTH_HOME=', "LLAMA_CPP_DIR=") + script = "\n".join( + ( + "set -u", + master, + node, + llama, + 'printf "%s\\n%s\\n" "$_NODE_PARENT" "$UNSLOTH_HOME"', + ) + ) + completed = subprocess.run( + ["bash", "-c", script], + env = env, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + node_parent, llama_parent = completed.stdout.splitlines() + return node_parent, llama_parent + + +def _env(home: Path, **overrides: str) -> dict[str, str]: + env = { + "HOME": str(home), + "PATH": "/usr/bin:/bin", + "STAGE_ROOT": "", + "RUNTIME_ROOT": "", + "STUDIO_HOME": str(home / ".unsloth" / "studio"), + "_STUDIO_HOME_IS_CUSTOM": "false", + } + env.update(overrides) + return env + + +@NEEDS_POSIX_BASH +def _master_root_answer(tmp_path: Path, environment: str) -> str: + """What uninstall.ps1's own _MasterRoot answers, with *environment* run before it. + + The functions are Invoke-Expression'd straight out of the shipped script, so a rewrite that + keeps the words and loses the behaviour fails here instead of passing against a copy. + """ + script = tmp_path / "probe.ps1" + script.write_text( + f"""$txt = Get-Content -Raw "{UNINSTALL_PS1}" +foreach ($n in @("_ExpandTilde", "_MasterRoot")) {{ + $m = [regex]::Match($txt, "(?ms)^ function $n \\{{.*?^ \\}}") + if (-not $m.Success) {{ Write-Output "EXTRACT-FAILED:$n"; exit 1 }} + Invoke-Expression $m.Value +}} +{environment} +""", + encoding = "utf-8", + ) + out = run_pwsh( + [PWSH, "-NoProfile", "-File", str(script)], + capture_output = True, + text = True, + check = True, + ).stdout.strip() + assert "EXTRACT-FAILED" not in out, out + return out + + +def test_a_master_root_puts_the_runtimes_beside_studio(tmp_path): + root = tmp_path / "portable" + (root / "studio").mkdir(parents = True) + env = _env( + tmp_path / "home", + UNSLOTH_HOME = str(root), + # What unsloth_cli/commands/studio.py exports before it runs setup. + STUDIO_HOME = str(root / "studio"), + _STUDIO_HOME_IS_CUSTOM = "true", + ) + node_parent, llama_parent = _runtime_parent(env) + assert Path(node_parent) == root + assert Path(llama_parent) == root + + +@NEEDS_POSIX_BASH +def test_no_master_root_keeps_the_studio_home_layout(tmp_path): + """The pre-existing custom-root behaviour, which the new branch must not disturb.""" + studio = tmp_path / "elsewhere" / "studio" + studio.mkdir(parents = True) + env = _env( + tmp_path / "home", + STUDIO_HOME = str(studio), + _STUDIO_HOME_IS_CUSTOM = "true", + ) + node_parent, llama_parent = _runtime_parent(env) + assert Path(node_parent) == studio + assert Path(llama_parent) == studio + + +@NEEDS_POSIX_BASH +def test_a_default_install_still_uses_the_legacy_root(tmp_path): + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + node_parent, llama_parent = _runtime_parent(_env(home)) + assert Path(node_parent) == home / ".unsloth" + assert Path(llama_parent) == home / ".unsloth" + + +@NEEDS_POSIX_BASH +def test_a_staging_root_still_outranks_the_master_root(tmp_path): + stage = tmp_path / "stage" + stage.mkdir() + env = _env( + tmp_path / "home", + UNSLOTH_HOME = str(tmp_path / "portable"), + STAGE_ROOT = str(stage), + RUNTIME_ROOT = str(stage), + ) + node_parent, llama_parent = _runtime_parent(env) + assert Path(node_parent) == stage + assert Path(llama_parent) == stage + + +def test_setup_ps1_derives_both_runtimes_from_the_master_root(): + """No PowerShell on the Linux runners, so the Windows half is held structurally: both + derivations must go through the one helper, and it must read UNSLOTH_HOME.""" + src = SETUP_PS1.read_text(encoding = "utf-8") + assert "function Get-MasterRootOverride" in src + helper = _slice(src, "function Get-MasterRootOverride", "function Get-ManagedLlamaCppDir") + assert "$env:UNSLOTH_HOME" in helper + llama = _slice(src, "function Get-ManagedLlamaCppDir", "\n# Failure reason when the managed") + assert "Get-MasterRootOverride" in llama + # The master root has to lose to an explicit staging root, as it does in setup.sh. + assert llama.index("$StagingRoot") < llama.index("Get-MasterRootOverride") + node = _slice(src, " $_masterRoot = Get-MasterRootOverride", " $NodeDir = Join-Path") + assert "$NodeParent = $_masterRoot" in node + assert node.index("$StageRoot") < node.index("$NodeParent = $_masterRoot") + + +@NEEDS_POSIX_BASH +def test_a_padded_master_root_names_the_same_directory(tmp_path): + """storage_roots.unsloth_home() and the CLI both .strip(), so setup has to as well: an + unstripped value installs under a directory whose name carries the whitespace.""" + root = tmp_path / "portable" + (root / "studio").mkdir(parents = True) + env = _env( + tmp_path / "home", + UNSLOTH_HOME = f" {root} ", + STUDIO_HOME = str(root / "studio"), + _STUDIO_HOME_IS_CUSTOM = "true", + ) + node_parent, llama_parent = _runtime_parent(env) + assert Path(node_parent) == root + assert Path(llama_parent) == root + + +@NEEDS_POSIX_BASH +def test_a_blank_master_root_counts_as_unset(tmp_path): + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + node_parent, llama_parent = _runtime_parent(_env(home, UNSLOTH_HOME = " ")) + assert Path(node_parent) == home / ".unsloth" + assert Path(llama_parent) == home / ".unsloth" + + +BUILD_WHISPER = REPO_ROOT / "scripts" / "build_whisper_cpp.sh" + + +def _whisper_root_block() -> str: + """The shipped root selection, from the normalizer down to the INSTALL_DIR it decides.""" + src = BUILD_WHISPER.read_text(encoding = "utf-8") + return _slice(src, "_root_value() {", "STUDIO_OWNED_MARKER=") + + +@pytest.fixture +def whisper_path(tmp_path_factory): + """PATH for the extracted whisper block, with git and cmake satisfied by stubs. + + The slice above is the shipped file verbatim, so it carries the builder's `command -v` + preflight along with the root selection these tests are about. On a runner without cmake the + block exits 1 before choosing anything and four tests fail for a reason that has nothing to + do with what they assert. Stubbed rather than excised: cutting the preflight out of the slice + would mean the tests no longer run the shipped text, which is the whole point of lifting it. + """ + stub_bin = tmp_path_factory.mktemp("stubbin") + for tool in ("git", "cmake"): + path = stub_bin / tool + path.write_text("#!/bin/sh\nexit 0\n", encoding = "utf-8") + path.chmod(0o755) + return f"{stub_bin}:/usr/bin:/bin" + + +@NEEDS_POSIX_BASH +def test_the_whisper_builder_installs_under_the_master_root(tmp_path, whisper_path): + """setup.sh runs this with UNSLOTH_STUDIO_HOME=/studio still inherited, so a builder + that preferred it would install a level below _managed_whisper_cpp_dir().""" + root = tmp_path / "portable" + src = BUILD_WHISPER.read_text(encoding = "utf-8") + block = _whisper_root_block() + script = "\n".join(("set -eu", block, 'printf "%s\\n" "$INSTALL_DIR"')) + completed = subprocess.run( + ["bash", "-c", script], + env = { + "HOME": str(tmp_path / "home"), + "PATH": whisper_path, + "UNSLOTH_HOME": str(root), + "UNSLOTH_STUDIO_HOME": str(root / "studio"), + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert Path(completed.stdout.strip()) == root / "whisper.cpp" + + +@NEEDS_POSIX_BASH +def test_the_whisper_builder_still_honours_a_studio_home_alone(tmp_path, whisper_path): + studio = tmp_path / "elsewhere" / "studio" + src = BUILD_WHISPER.read_text(encoding = "utf-8") + block = _whisper_root_block() + script = "\n".join(("set -eu", block, 'printf "%s\\n" "$INSTALL_DIR"')) + completed = subprocess.run( + ["bash", "-c", script], + env = { + "HOME": str(tmp_path / "home"), + "PATH": whisper_path, + "UNSLOTH_STUDIO_HOME": str(studio), + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert Path(completed.stdout.strip()) == studio / "whisper.cpp" + + +@NEEDS_POSIX_BASH +def test_a_blank_master_root_does_not_outrank_the_whisper_studio_home(tmp_path, whisper_path): + """${VAR:-} only treats the EMPTY string as unset, so an unstripped whitespace value would + win the new precedence and name a relative " /whisper.cpp".""" + studio = tmp_path / "elsewhere" / "studio" + completed = subprocess.run( + [ + "bash", + "-c", + "\n".join( + ( + "set -eu", + _whisper_root_block(), + 'printf "%s\\n" "$INSTALL_DIR"', + ) + ), + ], + env = { + "HOME": str(tmp_path / "home"), + "PATH": whisper_path, + "UNSLOTH_HOME": " ", + "UNSLOTH_STUDIO_HOME": str(studio), + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert Path(completed.stdout.strip()) == studio / "whisper.cpp" + + +@NEEDS_POSIX_BASH +def test_a_tilde_master_root_expands_for_the_whisper_builder(tmp_path, whisper_path): + home = tmp_path / "home" + completed = subprocess.run( + [ + "bash", + "-c", + "\n".join( + ( + "set -eu", + _whisper_root_block(), + 'printf "%s\\n" "$INSTALL_DIR"', + ) + ), + ], + env = {"HOME": str(home), "PATH": whisper_path, "UNSLOTH_HOME": "~/portable"}, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert Path(completed.stdout.strip()) == home / "portable" / "whisper.cpp" + + +@NEEDS_POSIX_BASH +def test_a_relative_master_root_resolves_against_the_caller(tmp_path): + """Node is chosen before setup.sh's first `cd "$SCRIPT_DIR"` and llama.cpp after it, so a + value left relative names two different directories and the backend's neither.""" + caller = tmp_path / "caller" + caller.mkdir() + src = SETUP_SH.read_text(encoding = "utf-8") + script = "\n".join( + ( + "set -u", + _slice(src, "# Stripped before anything else", "# Directory-local evidence"), + 'printf "%s\\n" "$_MASTER_ROOT"', + ) + ) + completed = subprocess.run( + ["bash", "-c", script], + cwd = str(caller), + env = { + "HOME": str(tmp_path / "home"), + "PATH": "/usr/bin:/bin", + "PWD": str(caller), + "UNSLOTH_HOME": "not-created-yet", + "_STUDIO_HOME_IS_CUSTOM": "false", + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert Path(completed.stdout.strip()) == caller / "not-created-yet" + + +@NEEDS_POSIX_BASH +def test_a_master_root_alone_still_asserts_ownership_of_the_runtimes(tmp_path): + """With only UNSLOTH_HOME set, STUDIO_HOME stays legacy and _STUDIO_HOME_IS_CUSTOM stays + false, which is what licenses the installers to os.replace() and rm -rf without checking + the Unsloth-owned marker. The runtimes moved, so the guard has to follow them.""" + root = tmp_path / "portable" + root.mkdir() + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + src = SETUP_SH.read_text(encoding = "utf-8") + script = "\n".join( + ( + "set -u", + _slice(src, "# Stripped before anything else", "# Directory-local evidence"), + 'printf "%s %s\\n" "$_STUDIO_HOME_IS_CUSTOM" "$_RUNTIME_ROOT_IS_CUSTOM"', + ) + ) + completed = subprocess.run( + ["bash", "-c", script], + env = { + "HOME": str(home), + "PATH": "/usr/bin:/bin", + "UNSLOTH_HOME": str(root), + "_STUDIO_HOME_IS_CUSTOM": "false", + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.split() == ["false", "true"] + + +@NEEDS_POSIX_BASH +@pytest.mark.parametrize( + "environment, expected", + [ + # Staging: the placement below gives STAGE_ROOT precedence over the master root, so the + # master root is not where anything lands and must not decide ownership. Raised anyway, + # it demanded an owner marker from a markerless tree the old updater had staged, and an + # update that worked on the merge base exited instead. + pytest.param( + {"UNSLOTH_STUDIO_STAGE_ROOT": "STAGE", "_STUDIO_HOME_IS_CUSTOM": "false"}, + ["false", "false"], + id = "a staged update ignores the master root", + ), + # And the flag has to be ASSIGNED, not only raised: a custom STUDIO_HOME whose runtimes + # land in the legacy root kept it true and demanded markers from exactly the pre-marker + # ~/.unsloth/llama.cpp this comparison exists to spare. + pytest.param( + {"UNSLOTH_HOME": "LEGACY", "_STUDIO_HOME_IS_CUSTOM": "true"}, + ["true", "false"], + id = "a custom studio home whose runtimes stay legacy", + ), + ], +) +def test_the_ownership_flag_follows_where_the_runtimes_land(tmp_path, environment, expected): + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + stage = tmp_path / "stage" + stage.mkdir() + resolved = { + key: str(stage) + if value == "STAGE" + else str(home / ".unsloth") + if value == "LEGACY" + else value + for key, value in environment.items() + } + if "UNSLOTH_HOME" not in resolved: + resolved["UNSLOTH_HOME"] = str(tmp_path / "portable") + (tmp_path / "portable").mkdir() + src = SETUP_SH.read_text(encoding = "utf-8") + script = "\n".join( + ( + "set -u", + f'STAGE_ROOT="{resolved.pop("UNSLOTH_STUDIO_STAGE_ROOT", "")}"', + _slice(src, "# Stripped before anything else", "# Directory-local evidence"), + 'printf "%s %s\\n" "$_STUDIO_HOME_IS_CUSTOM" "$_RUNTIME_ROOT_IS_CUSTOM"', + ) + ) + completed = subprocess.run( + ["bash", "-c", script], + env = {"HOME": str(home), "PATH": "/usr/bin:/bin", **resolved}, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.split() == expected + + +@NEEDS_POSIX_BASH +def test_no_master_root_leaves_the_ownership_flag_alone(tmp_path): + src = SETUP_SH.read_text(encoding = "utf-8") + script = "\n".join( + ( + "set -u", + _slice(src, "# Stripped before anything else", "# Directory-local evidence"), + 'printf "%s %s\\n" "$_STUDIO_HOME_IS_CUSTOM" "$_RUNTIME_ROOT_IS_CUSTOM"', + ) + ) + for flag in ("false", "true"): + completed = subprocess.run( + ["bash", "-c", script], + env = { + "HOME": str(tmp_path / "home"), + "PATH": "/usr/bin:/bin", + "_STUDIO_HOME_IS_CUSTOM": flag, + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.split() == [flag, flag] + + +def test_every_runtime_ownership_guard_uses_the_runtime_flag(): + """The Studio home and its venvs keep _STUDIO_HOME_IS_CUSTOM; every guard that names a + runtime child has to move, or a master-root install loses the marker check on it.""" + src = SETUP_SH.read_text(encoding = "utf-8") + for line in src.splitlines(): + if "_STUDIO_HOME_IS_CUSTOM" not in line: + continue + assert not any( + name in line for name in ("$NODE_DIR", "$LLAMA_CPP_DIR", "$WHISPER_CPP_DIR") + ), line + ps = SETUP_PS1.read_text(encoding = "utf-8") + # Seeded from the Studio flag and raised only when the master root is elsewhere, as setup.sh + # does; test_the_windows_legacy_root_named_explicitly_is_not_custom runs the divergence. + assert "$RuntimeRootIsCustom = $StudioHomeIsCustom\n" in ps + assert "$_masterRootForOwnership -ine $_legacyRuntimeRoot" in ps + for line in ps.splitlines(): + if "$StudioHomeIsCustom" not in line: + continue + assert not any( + name in line for name in ("$LlamaCppDir", "$WhisperCppDir", "$NodeDir") + ), line + + +@NEEDS_POSIX_BASH +def test_the_legacy_root_named_explicitly_is_not_custom(tmp_path): + """UNSLOTH_HOME=$HOME/.unsloth names the root a default install already uses. + + Nothing moves, so demanding an owner marker there rejects an update that used to reuse a + legacy source-built ~/.unsloth/llama.cpp, which predates the marker entirely. Customness has + to come from where the runtimes LAND, not from whether the variable was set. + """ + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + src = SETUP_SH.read_text(encoding = "utf-8") + block = _slice(src, "# Stripped before anything else", "# Directory-local evidence") + script = "\n".join( + ( + "set -u", + "_STUDIO_HOME_IS_CUSTOM=false", + block, + 'printf "%s\\n" "$_RUNTIME_ROOT_IS_CUSTOM"', + ) + ) + + def flag(master: str) -> str: + done = subprocess.run( + ["bash", "-c", script], + env = {"HOME": str(home), "PATH": "/usr/bin:/bin", "UNSLOTH_HOME": master}, + capture_output = True, + text = True, + timeout = 60, + ) + assert done.returncode == 0, done.stderr + return done.stdout.strip() + + assert flag(str(home / ".unsloth")) == "false" + # Non-vacuity: a root that really is elsewhere still takes the strict path. + assert flag(str(tmp_path / "portable")) == "true" + assert flag("") == "false" + + +@pytest.mark.skipif(PWSH is None, reason = "needs pwsh") +def test_the_windows_legacy_root_named_explicitly_is_not_custom(tmp_path): + """The Windows half of the test above, run rather than pattern-matched. + + setup.ps1 read `$StudioHomeIsCustom -or [bool](Get-MasterRootOverride)`, so + UNSLOTH_HOME=%USERPROFILE%\\.unsloth turned the ownership guard on for a root a default + install is already using. Nothing moves there, and the guard then demands an owner marker + from a markerless source-built .unsloth\\llama.cpp, so Assert-StudioOwnedOrAbsent exits an + update that used to reuse that build. setup.sh adopts it, which made the two halves answer + differently for one environment. + """ + profile = tmp_path / "profile" + (profile / ".unsloth" / "studio").mkdir(parents = True) + ps = SETUP_PS1.read_text(encoding = "utf-8") + # The derivation itself, from the flag it seeds to the line after it. + block = _slice(ps, "$RuntimeRootIsCustom = $StudioHomeIsCustom", "$LlamaCppDir = ") + script = tmp_path / "probe.ps1" + script.write_text( + "\n".join( + ( + f"$txt = Get-Content -Raw '{SETUP_PS1}'", + # The shipped helpers, so a rewrite that keeps the words and changes the + # canonicalisation still fails here. + 'foreach ($n in @("Get-PathState", "Test-AccessDeniedError", "Get-CanonicalDir",' + ' "Get-MasterRootOverride")) {', + ' $m = [regex]::Match($txt, "(?ms)^function $n \\{.*?^\\}")', + ' if (-not $m.Success) { Write-Output "EXTRACT-FAILED:$n"; exit 1 }', + " Invoke-Expression $m.Value", + "}", + # Test-StudioHomeIsCustom closes over $StudioHome, so it is stated here rather + # than extracted: the point of this probe is the master-root comparison. + "$StudioHomeIsCustom = $false", + block, + "Write-Output $RuntimeRootIsCustom", + ) + ), + encoding = "utf-8", + ) + + def flag(master: str) -> str: + out = run_pwsh( + [PWSH, "-NoProfile", "-File", str(script)], + env = { + "PATH": os.environ.get("PATH", ""), + "HOME": str(profile), + "USERPROFILE": str(profile), + "UNSLOTH_HOME": master, + }, + capture_output = True, + text = True, + check = True, + ).stdout.strip() + assert "EXTRACT-FAILED" not in out, out + return out.splitlines()[-1].strip() + + assert flag(str(profile / ".unsloth")) == "False" + # Non-vacuity, and the same two controls the bash test uses. + assert flag(str(tmp_path / "portable")) == "True" + assert flag("") == "False" + + +@pytest.mark.skipif(PWSH is None, reason = "needs pwsh") +def test_the_windows_uninstaller_only_removes_a_root_that_is_really_empty(tmp_path): + """The master root and its .staging go only when empty, and by an operation that cannot do + anything else. + + Windows asked with Get-ChildItem -ErrorAction SilentlyContinue and then called _RemovePath, + which is Remove-Item -Recurse -Force, in the one directory this branch teaches users to + point at their own files. Two separate calls, so a file arriving between the two was taken + by the -Recurse; and a swallowed enumeration error read as an empty directory. + uninstall.sh uses rmdir for both paths, which is one call that refuses a non-empty + directory, and this file's own %TEMP% prune already asks with -ErrorAction Stop. + + The behavioural halves below are what can be shown on a POSIX host: a recursive delete needs + to enumerate too, so an unlistable directory survives either version here and the ACL shape + that separates DELETE from LIST is Windows-only. What is checked instead is the mechanism, + which is what removes the race as well. + """ + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + assert "function _RemoveDirIfEmpty" in ps + helper = _slice(ps, " function _RemoveDirIfEmpty {", "\n # The exact shape") + # One atomic call: no -Recurse anywhere in it, and no emptiness question asked separately. + assert "[System.IO.Directory]::Delete($Path, $false)" in helper + assert "-Recurse" not in helper, helper + assert "Get-ChildItem" not in helper, helper + # And both master-root sites go through it rather than keeping their own version. + master_block = _slice(ps, "$masterRoot = $masterRootToStop", "if ($defaultLlamaCpp)") + assert master_block.count("_RemoveDirIfEmpty") == 2, master_block + assert "-ErrorAction SilentlyContinue)) {\n _RemovePath" not in master_block + + script = tmp_path / "probe.ps1" + script.write_text( + "\n".join( + ( + f"$txt = Get-Content -Raw '{UNINSTALL_PS1}'", + '$m = [regex]::Match($txt, "(?ms)^ function _RemoveDirIfEmpty \\{.*?^ \\}")', + 'if (-not $m.Success) { Write-Output "EXTRACT-FAILED"; exit 1 }', + "Invoke-Expression $m.Value", + "function _Substep { param($a, $b) }", + "_RemoveDirIfEmpty $args[0]", + 'Write-Output ("EXISTS=" + (Test-Path -LiteralPath $args[0]))', + ) + ), + encoding = "utf-8", + ) + + def still_there(path: Path) -> bool: + out = run_pwsh( + [PWSH, "-NoProfile", "-File", str(script), str(path)], + env = {"PATH": os.environ.get("PATH", ""), "HOME": str(tmp_path)}, + capture_output = True, + text = True, + check = True, + ).stdout + assert "EXTRACT-FAILED" not in out, out + assert "EXISTS=" in out, out + return out.strip().splitlines()[-1] == "EXISTS=True" + + empty = tmp_path / "empty" + empty.mkdir() + assert not still_there(empty), "an empty root is the one case that may be removed" + + occupied = tmp_path / "occupied" + occupied.mkdir() + (occupied / "theirs.txt").write_text("mine", encoding = "utf-8") + assert still_there(occupied) + assert (occupied / "theirs.txt").is_file(), "a non-empty root must keep its contents" + + # A directory holding only a subdirectory: rmdir refuses this too, and a -Recurse would not. + nested = tmp_path / "nested" + (nested / "child").mkdir(parents = True) + assert still_there(nested) + assert (nested / "child").is_dir() + + +@NEEDS_POSIX_BASH +def test_only_a_directory_can_be_adopted_at_a_runtime_path(tmp_path): + """A dangling symlink and a regular file are both things the user put there. + + -d alone read either as "nothing is here", and the caller then rm -rf'd the path or let + install_*_prebuilt.py os.replace() over it. A dangling link is the ordinary case: its target + volume is simply not mounted, and resolving it later installs onto somebody's other disk. + """ + src = SETUP_SH.read_text(encoding = "utf-8") + block = _slice(src, "_studio_path_shape() {", "\n_packaged_frontend_available") + script = "\n".join( + ( + "set -u", + "_STUDIO_OWNED_MARKER=.unsloth-studio-owned", + "_studio_owned_adoptable() { return 1; }", + "_studio_dir_unsearchable() { return 1; }", + "_path_access_denied() { echo DENIED; exit 9; }", + "setup_fail() { echo REFUSED; exit 1; }", + block, + '_assert_studio_owned_or_absent "$1" llama.cpp true', + "echo ALLOWED", + ) + ) + + def verdict(path: pathlib.Path) -> str: + done = subprocess.run( + ["bash", "-c", script, "_", str(path)], + env = {"HOME": str(tmp_path), "PATH": "/usr/bin:/bin"}, + capture_output = True, + text = True, + timeout = 60, + ) + return done.stdout.strip().splitlines()[-1] if done.stdout.strip() else done.stderr[-120:] + + assert verdict(tmp_path / "absent") == "ALLOWED" + + dangling = tmp_path / "dangling" + dangling.symlink_to(tmp_path / "no-such-volume" / "llama.cpp") + assert verdict(dangling) == "REFUSED" + + plain = tmp_path / "afile" + plain.write_text("mine") + assert verdict(plain) == "REFUSED" + + owned = tmp_path / "owned" + owned.mkdir() + (owned / ".unsloth-studio-owned").touch() + assert verdict(owned) == "ALLOWED" + + +def test_the_windows_inductor_cache_agrees_with_the_resolver(): + """setup.ps1 persists TORCHINDUCTOR_CACHE_DIR to the USER environment. + + Every later Studio process inherits it, so _setup_cache_env's fill-if-unset default never + applies on Windows and the containment this branch is for does not happen there. It has to + name the directory the resolver would have chosen. Two things still outrank that, and both + are recorded here so a later edit cannot quietly drop them: long paths off keeps the short + drive-root directory for MAX_PATH headroom, and a path containing a space is refused for the + same reason storage_roots does, since the C++ builders paste it in unquoted. + """ + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice(ps, "$TorchCacheDir = $null", "$env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir") + assert 'Join-Path (Join-Path $StudioHome "cache") "torchinductor"' in block + assert "$LongPathsEnabled" in block + assert "'\\s'" in block or '"\\s"' in block, "the whitespace refusal is gone" + assert '"C:\\tc"' in block + + roots = (REPO_ROOT / "studio" / "backend" / "utils" / "paths" / "storage_roots.py").read_text( + encoding = "utf-8" + ) + # The same key, named by both sides, so the two cannot drift apart silently. + assert '"TORCHINDUCTOR_CACHE_DIR",' in roots + assert 'str(root / "torchinductor")' in roots + + +def test_the_windows_node_guard_covers_a_master_root(): + """setup.ps1's Node ownership guard and its marker both hung off $NodeOverride alone. + + $NodeOverride is set only in the UNSLOTH_STUDIO_HOME / STUDIO_HOME branch. The master-root + branch sets $NodeParent and leaves it null, so \\node reached the whole-directory + os.replace() in install_node_prebuilt.py with no ownership evidence at all, and the tree the + run then created stayed unmarked, which makes the uninstaller decline to remove it later. + setup.sh had already moved these two sites to _RUNTIME_ROOT_IS_CUSTOM. + + The sibling test above only rejects $StudioHomeIsCustom beside $NodeDir, which this bug + never wrote: it named a third variable. So the rule here is positive, not a denial. + """ + ps = SETUP_PS1.read_text(encoding = "utf-8") + # Every site in the Node install section that decides whether the Node path is the user's. + # Keyed on $NodeOverride, the variable the bug named, not $NodeDir: the guard's probe moved + # to its own line, so a $NodeDir test would pass by seeing less. The section start drops the + # parsing block far above, which rightly says nothing about the runtime root. + lines = ps.splitlines() + start = next(i for i, line in enumerate(lines) if "$nodeEntry = if (" in line) + guards = [line for line in lines[start:] if "$NodeOverride" in line] + # The guard's probe, and the marker's `if`. Fewer means the block was restructured and this + # test would otherwise pass by finding nothing. + assert len(guards) >= 2, guards + for line in guards: + assert "$RuntimeRootIsCustom" in line, line + + +UNINSTALL_SH = REPO_ROOT / "scripts" / "uninstall.sh" +UNINSTALL_PS1 = REPO_ROOT / "scripts" / "uninstall.ps1" + + +@pytest.mark.skipif(shutil.which("pwsh") is None, reason = "pwsh not available") +def test_the_windows_uninstaller_resolves_a_relative_root_like_setup(tmp_path): + """setup.ps1 resolves through PowerShell's own location; uninstall.ps1 used + [IO.Path]::GetFullPath, which anchors at [Environment]::CurrentDirectory. + + PowerShell does not keep those two in step, so after a Set-Location a relative UNSLOTH_HOME + named one install to setup and a different one to the uninstaller. The owner marker does not + save you there: it spares trees that are not Unsloth's, and the wrongly resolved path is + another Unsloth install, marker and all. + + Runs the shipped function rather than matching its text, so a rewrite that keeps the words + and loses the behaviour still fails. + """ + initial = tmp_path / "initial" + chosen = tmp_path / "chosen" + (chosen / "portable").mkdir(parents = True) + initial.mkdir() + out = _master_root_answer( + tmp_path, + f"""[System.Environment]::CurrentDirectory = "{initial}" +Set-Location "{chosen}" +$env:UNSLOTH_HOME = "portable" +$env:USERPROFILE = "{tmp_path}/profile" +Write-Output (_MasterRoot) +""", + ) + assert out == str(chosen / "portable"), out + + +def test_neither_uninstaller_recurses_into_an_install_lock_path(): + """A lock is always a regular file, so a directory at one of those fixed names is the user's. + + prebuilt_core.install_lock creates it with os.open(O_CREAT | O_EXCL). Both uninstallers + reached the lock names through their recursive remover, which in a user-chosen master root + deletes a whole tree with none of the owner-marker proof the runtime children require. + Behaviour is covered by tests/sh/test_uninstall_master_root.sh for the POSIX half; this + holds the PowerShell twin, which has no runner here. + """ + sh = UNINSTALL_SH.read_text(encoding = "utf-8") + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + + assert "_remove_lock_file() {" in sh + assert "function _RemoveLockFile" in ps + + # The call sites do not all spell the lock out (PowerShell iterates $lockName, both sweep a + # $stale), so matching only "install.lock" let the master-root loop keep the recursive + # remover and still pass. + lockish = ("install.lock", "lockName", "$stale", "_stale", "_mr_lock") + for text, remover, shape in ( + (sh, "_remove_path ", "_remove_lock_file"), + (ps, "_RemovePath ", "_RemoveLockFile"), + ): + hits = 0 + for line in text.splitlines(): + if line.lstrip().startswith("#"): + continue + if remover not in line and shape not in line: + continue + if not any(token in line for token in lockish): + continue + hits += 1 + assert shape in line, line + # Counts the call sites found, not the ones left wrong: once they are all correct the + # remover no longer appears beside a lock at all, and a count of zero would mean the + # tokens above had stopped matching and this loop proved nothing. + assert hits >= 4, (remover, hits) + + # The rename keeps the leading dot, and neither a glob without one nor the dot alone is + # enough, so the sweep matches the installer's full shape; held by + # test_neither_uninstaller_sweeps_a_stale_lock_name_it_did_not_make below. + assert "install.lock.stale" in ps + assert '"*.install.lock.stale.*"' not in ps + + +def test_both_uninstallers_clear_the_master_root_children(): + """setup installs llama.cpp, node and whisper.cpp as children of the master root, so an + uninstaller that only knows the legacy siblings and the Studio root strands them. Behaviour + is covered by tests/sh/test_uninstall_master_root.sh; this holds the PowerShell twin, which + the Linux runners cannot execute, and pins the marker gate on both.""" + sh = UNINSTALL_SH.read_text(encoding = "utf-8") + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + assert "_master_root() {" in sh + assert "function _MasterRoot" in ps + for src, marker in ((sh, ".unsloth-studio-owned"), (ps, ".unsloth-studio-owned")): + block = _slice(src, "master root's own children", "llama.cpp build + cache") + assert marker in block, block + for child in ("llama.cpp", "node", "whisper.cpp"): + assert child in block, (child, block) + # A user-chosen root reaches the deny list on both sides before anything is removed. + assert '_is_unsafe_root "$_mr_root"' in sh + assert "_IsUnsafeRoot $masterRoot" in ps + # And it only contributes its studio child when neither exact override is set, as in + # storage_roots.studio_root(). + assert '_emit "$(_master_root)/studio"' in sh + assert '$envRoot = (Join-Path $master "studio")' in ps + + +def test_the_stop_pass_covers_the_master_root_runtimes(): + """Windows locks a loaded executable, so a runtime still running under the master root has + to be stopped before its tree is removed or the delete exhausts its retries. _MasterRoot is + resolved before the stop pass, and only marker-owned children join it, so an unmarked + neighbour's process is never killed.""" + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + stop_line = next(l for l in ps.splitlines() if l.strip().startswith("$stopRoots = ")) + assert "$masterChildrenToStop" in stop_line, stop_line + block = _slice(ps, "$masterRootToStop = _MasterRoot", "$stopRoots = ") + assert ".unsloth-studio-owned" in block, block + assert "_IsUnsafeRoot $masterRootToStop" in block, block + assert ps.index("$masterRootToStop = _MasterRoot") < ps.index( + "_StopProcessesLockingRoots -Roots" + ) + + +def test_a_shared_staging_directory_is_pruned_not_deleted(): + """The prebuilt installers share /.staging and prune it only when empty, so anything + left in a user-chosen root is not ours to delete recursively.""" + sh = UNINSTALL_SH.read_text(encoding = "utf-8") + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + assert 'rmdir "$_mr_root/.staging"' in sh + assert '_remove_path "$_mr_root/.staging"' not in sh + # The Windows equivalent of rmdir, for the same reason. Asking with Get-ChildItem and then + # calling _RemovePath is not it: see + # test_the_windows_uninstaller_only_removes_a_root_that_is_really_empty. + staging = _slice( + ps, '_RemoveDirIfEmpty (Join-Path $masterRoot ".staging")', "# Shared llama.cpp build" + ) + assert "_RemovePath" not in staging, staging + + +def test_the_windows_uninstaller_refuses_a_note_carried_in_from_elsewhere(tmp_path): + """A Studio tree copied from master root A to B keeps a note naming A. + + A's llama.cpp, node, whisper.cpp and sd.cpp carry exactly the owner markers B's would, so + the marker gates cannot tell them apart: accepting the copied note has the uninstall of B + delete the ORIGINAL install's runtimes. The note has to describe the tree it was found in, + which means the Studio directory it was read from lying inside the root it names. + """ + original = tmp_path / "original" + (original / "studio").mkdir(parents = True) + copied = tmp_path / "copied" + (copied / "studio" / "share").mkdir(parents = True) + (copied / "studio" / "share" / ".unsloth-master-root").write_text( + f"{original}\n", encoding = "utf-8" + ) + profile = tmp_path / "profile" + (profile / ".unsloth" / "studio" / "share").mkdir(parents = True) + + script = tmp_path / "probe.ps1" + script.write_text( + f"""$txt = Get-Content -Raw "{UNINSTALL_PS1}" +foreach ($n in @("_ExpandTilde", "_RootFromConf", "_MasterRoot")) {{ + $m = [regex]::Match($txt, "(?ms)^ function $n \\{{.*?^ \\}}") + if (-not $m.Success) {{ Write-Output "EXTRACT-FAILED:$n"; exit 1 }} + Invoke-Expression $m.Value +}} +$env:UNSLOTH_HOME = "" +$env:STUDIO_HOME = "" +$env:UNSLOTH_STUDIO_HOME = "{copied}/studio" +$env:USERPROFILE = "{profile}" +$answer = _MasterRoot +Write-Output "ANSWER:$answer" +""", + encoding = "utf-8", + ) + out = run_pwsh( + [PWSH, "-NoProfile", "-File", str(script)], + capture_output = True, + text = True, + check = True, + ).stdout.strip() + assert "EXTRACT-FAILED" not in out, out + assert out == "ANSWER:", out + + +def test_the_windows_uninstaller_does_not_borrow_another_installs_note(tmp_path): + """Two installs on one box: the legacy tree carries a note, the named one does not. + + _MasterRoot probed the legacy path first, so its note won here while the removal still + worked on the tree UNSLOTH_STUDIO_HOME names. The run deleted the named Studio and then + followed the OTHER install's master root and took its marked runtime children with it. + Runs the shipped function, as the test below does. + """ + profile = tmp_path / "profile" + (profile / ".unsloth" / "studio" / "share").mkdir(parents = True) + borrowed = tmp_path / "borrowed" + (borrowed / "studio").mkdir(parents = True) + (profile / ".unsloth" / "studio" / "share" / ".unsloth-master-root").write_text( + f"{borrowed}\n", encoding = "utf-8" + ) + named = tmp_path / "named" + (named / "share").mkdir(parents = True) + + out = _master_root_answer( + tmp_path, + f"""$env:UNSLOTH_HOME = "" +$env:STUDIO_HOME = "" +$env:UNSLOTH_STUDIO_HOME = "{named}" +$env:USERPROFILE = "{profile}" +$answer = _MasterRoot +Write-Output "ANSWER:$answer" +""", + ) + assert out == "ANSWER:", out + + +def test_the_windows_uninstaller_finds_a_master_root_from_the_note(tmp_path): + """`$env:UNSLOTH_HOME = 'D:\\portable'; unsloth studio update` names the root for one command. + + setup.ps1 puts node\\, llama.cpp\\ and whisper.cpp\\ under it and marks them, and a later + uninstall run from an ordinary shell has no UNSLOTH_HOME at all. Reading only the current + environment, _MasterRoot answered null there, so the run removed \\studio and left + multi-gigabyte runtimes beside it. setup.sh already wrote the note; setup.ps1 did not, and + the PowerShell uninstaller did not read it. + + Runs the shipped function, so a rewrite that keeps the words and loses the behaviour fails. + """ + master = tmp_path / "portable" + (master / "studio" / "share").mkdir(parents = True) + (master / "studio" / "share" / ".unsloth-master-root").write_text( + f"{master}\n", encoding = "utf-8" + ) + profile = tmp_path / "profile" + (profile / ".unsloth" / "studio" / "share").mkdir(parents = True) + + out = _master_root_answer( + tmp_path, + f"""$env:UNSLOTH_HOME = "" +$env:STUDIO_HOME = "" +$env:UNSLOTH_STUDIO_HOME = "{master}/studio" +$env:USERPROFILE = "{profile}" +Write-Output (_MasterRoot) +""", + ) + assert out == str(master), out + + +def test_the_windows_setup_records_the_master_root_for_the_uninstaller(): + """The note the test above reads has to be written, and only where it is true. + + A staging run installs to a throwaway root, and every non-master branch derives the root + from paths the uninstaller already knows, so a note there would only ever be able to go + stale. This is the same rule setup.sh applies. + """ + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice( + ps, "# Record the master root inside the Studio tree", "$WithLlamaCppDir = $null" + ) + assert "(Get-MasterRootOverride)" in block + assert "-not $StageRoot" in block + assert '".unsloth-master-root"' in block + # Staged then renamed: a reader catching a half-written note would name a truncated path, + # and this note licenses deletions. + assert "$noteTmp" in block and "Move-Item" in block + # The 3-argument overwrite overload is .NET Core only, and setup.ps1 runs under 5.1. + assert "[System.IO.File]::Move(" not in block + + +def test_the_windows_uninstaller_clears_the_inductor_path_it_persisted(): + """setup.ps1 writes TORCHINDUCTOR_CACHE_DIR to the USER environment, so it outlives the + install. Every later PyTorch process on the account inherits it, including ones unrelated to + Unsloth, and they compile into the deleted tree and rebuild part of it. + + Only a value inside a root this run owned: a directory the user chose is theirs, and the + shared C:\\tc fallback is not install specific and is not deleted here either. + """ + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + block = _slice( + ps, "# Clear the persisted Inductor cache path", "# Remove HKCU\\Software\\Unsloth" + ) + assert "GetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', 'User')" in block + assert "[NullString]::Value, 'User'" in block + # Scoped to what this run owns. $knownRoots includes roots the gates refused to delete. + assert "$ownedRoots" in block and "$knownRoots" not in block + # Comments stripped first: this block explains why C:\tc is spared, and a comment saying so + # is not the same thing as code naming it. + code = "\n".join(line for line in block.splitlines() if not line.lstrip().startswith("#")) + assert "C:\\tc" not in code + + +def test_the_windows_node_guard_treats_a_file_as_occupied(): + """install_node_prebuilt's _swap_into_place renames whatever it finds out of the way. + + Guarding on -PathType Container meant a regular file or a symlink named `node` under a + user-selected master root was invisible to the check and got displaced. Ownership evidence + lives inside a directory, so a non-directory can never carry it and is refused outright. + setup.sh's _assert_studio_owned_or_absent takes the same view of -d against -e and -L. + """ + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice( + ps, + " $nodeEntry = if ($NodeOverride -or $RuntimeRootIsCustom) {", + "install_node_prebuilt.py", + ) + code = "\n".join(line for line in block.splitlines() if not line.lstrip().startswith("#")) + # Get-Item -Force, not Test-Path: under 5.1 Test-Path answers for the link TARGET, so a + # dangling `node` link read as absent and install_node_prebuilt.py installed through it. + assert "Get-Item -LiteralPath $NodeDir -Force" in code, code + assert "Test-Path -LiteralPath $NodeDir" not in code, code + # A reparse point is not a directory here either, however it resolves. + assert "$nodeEntry.PSIsContainer" in code and "ReparsePoint" in code, code + # And a non-directory must fail rather than fall through to the marker questions. + assert "$nodeIsDir" in code and "-not $nodeIsDir -or" in code + + +def test_the_windows_runtime_guard_treats_a_file_as_occupied(): + """The Node guard above is one of three: llama.cpp and whisper.cpp share + Assert-StudioOwnedOrAbsent, which was still container-only. + + install_llama_prebuilt's activate_install_tree moves aside whatever Path.exists() finds, so + a regular file named `llama.cpp` under a master root was renamed to a rollback name and + replaced. The behaviour is run for real in tests/studio/test_path_probe_access_denied.ps1, + which the Linux runners cannot execute; this holds the shape there too. + """ + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice(ps, "function Assert-StudioOwnedOrAbsent", "function Mark-StudioOwned") + code = "\n".join(line for line in block.splitlines() if not line.lstrip().startswith("#")) + # Get-Item -Force, not another Test-Path: 5.1 answers false for a dangling link. + assert "Get-Item -LiteralPath $Path -Force" in code + # A non-directory is refused, and only under a root the user chose. + assert "if (-not $isCustomRoot) { return }" in code + assert "already exists and is not a directory" in block + # Fatal even in -NonFatal mode: that mode rescues a denial, not somebody else's file. + refusal = code[ + code.index("Get-Item -LiteralPath $Path -Force") : code.index("Exit-SetupFailure") + ] + assert "$NonFatal" not in refusal, refusal + + +def test_neither_uninstaller_takes_a_studio_root_that_is_also_the_master_root(): + """The flat layout, UNSLOTH_HOME and UNSLOTH_STUDIO_HOME naming one directory. + + A Studio root is removed WHOLE once it carries the ownership marker, so in the flat case + anything else the user keeps in that directory goes with the install. Every other + master-root child is individually marker-gated precisely so a user-chosen root is never + removed wholesale; this was the one hole in that rule. Kept rather than pruned: data left + behind is recoverable and printed, a deleted file is not. + """ + sh = UNINSTALL_SH.read_text(encoding = "utf-8") + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + + sh_block = _slice(sh, "_crf_canon=", '_remove_root_recording_db "$_custom_root"') + sh_code = "\n".join(l for l in sh_block.splitlines() if not l.lstrip().startswith("#")) + assert "_MASTER_ROOT_SAVED" in sh_code and "continue" in sh_code + # Canonicalised on both sides, or a symlinked path compares unequal to itself and the guard + # never fires on the very layout it is for. + assert "cd -P --" in sh_code + + ps_block = _slice(ps, "$flatMaster = $masterRootToStop", "_RemoveRootRecordingDb $r") + ps_code = "\n".join(l for l in ps_block.splitlines() if not l.lstrip().startswith("#")) + assert "$masterRootToStop" in ps_code and "continue" in ps_code + # Case-insensitive: Windows paths differing only in case are the same directory. + assert "-ieq" in ps_code + + +def test_neither_uninstaller_re_resolves_the_master_root_after_deleting_it(): + """_master_root can read its answer from a note INSIDE a Studio tree. + + The custom-root loop removes that tree, so a second call after it returns nothing and the + marked llama.cpp, Node and whisper.cpp siblings are stranded: the very failure the note was + added to prevent, reintroduced by asking too late. Both scripts resolve it once, before + anything is deleted, and every later use takes that value. + """ + sh = UNINSTALL_SH.read_text(encoding = "utf-8") + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + + sh_code = "\n".join(l for l in sh.splitlines() if not l.lstrip().startswith("#")) + saved = sh_code.index('_MASTER_ROOT_SAVED="$(_master_root)"') + removal = sh_code.index("_custom_studio_roots | while") + # Resolved before the first deletion, and nothing asks again after it. Calls BEFORE the loop + # are fine: those enumerate while every tree is still on disk. + assert saved < removal + after = sh_code[removal:] + assert "_master_root" not in after, after[: after.index("_master_root") + 200] + # Both consumers take the saved value: the flat-layout guard and the children block. + assert after.count('"$_MASTER_ROOT_SAVED"') >= 2 + + ps_code = "\n".join(l for l in ps.splitlines() if not l.lstrip().startswith("#")) + ps_saved = ps_code.index("$masterRootToStop = _MasterRoot") + ps_removal = ps_code.index("_RemoveRootRecordingDb $r") + assert ps_saved < ps_removal + ps_after = ps_code[ps_removal:] + assert "_MasterRoot" not in ps_after, ps_after[: ps_after.index("_MasterRoot") + 200] + assert ps_after.count("$masterRootToStop") >= 1 + + +def test_neither_uninstaller_sweeps_a_stale_lock_name_it_did_not_make(): + """prebuilt_core.py leaves .stale.. A dotted glob also matched a user's own + ".backup.install.lock.stale.copy", and in a root the user chose that file is theirs. + Behaviour for the POSIX half is in tests/sh/test_uninstall_master_root.sh.""" + sh = UNINSTALL_SH.read_text(encoding = "utf-8") + ps = UNINSTALL_PS1.read_text(encoding = "utf-8") + + sh_code = "\n".join(l for l in sh.splitlines() if not l.lstrip().startswith("#")) + assert ".*.install.lock.stale.*" not in sh_code + # A numeric pid, checked at both sweeps: the master root and the default home. + assert sh_code.count("*[!0-9]*) continue ;;") >= 2 + + ps_code = "\n".join(l for l in ps.splitlines() if not l.lstrip().startswith("#")) + assert '-like ".*.install.lock.stale.*"' not in ps_code + assert "$script:StaleLockPattern" in ps_code + assert "[0-9]+$" in ps_code + + +@NEEDS_POSIX_BASH +def test_the_master_root_note_does_not_write_through_a_planted_link(tmp_path): + """The note is staged under an unpredictable name, so a link left in share/ is not followed. + + The staging name used to be "$notePath.$$", and a shell redirection into an existing symlink + truncates its target. Anyone who could write share/ -- a portable volume mounted for two + accounts is the shape that matters, since a master root is what this note records -- could + therefore point that name at a file of theirs and have setup empty it. The attacker's + capability is simulated exactly rather than approximated: the link is created by the same + shell that then runs the shipped block, so "$$" matches the way a guessed pid would. + + setup.sh already states this rule for the uv cache probe ("mktemp (O_EXCL) not $$"); the note + writer is now held to it too. + """ + src = SETUP_SH.read_text(encoding = "utf-8") + block = _slice(src, "_master_root_note_is_honoured() {", "LLAMA_CPP_DIR=") + + # Contained, so the honoured-note gate passes and the writer below actually runs. The gate + # itself is covered by test_the_master_root_note_is_only_written_when_a_reader_honours_it. + master = tmp_path / "portable" + studio_home = master / "studio" + (studio_home / "share").mkdir(parents = True) + victim = tmp_path / "victim.txt" + victim.write_text("do not truncate me\n", encoding = "utf-8") + + script = "\n".join( + ( + "set -u", + # The predictable name the old block would have opened. + 'ln -s "$VICTIM" "$STUDIO_HOME/share/.unsloth-master-root.$$"', + block, + ) + ) + completed = subprocess.run( + ["bash", "-c", script], + env = { + "PATH": os.environ.get("PATH", "/usr/bin:/bin"), + "STUDIO_HOME": str(studio_home), + "UNSLOTH_HOME": str(master), + "_MASTER_ROOT": str(master), + "STAGE_ROOT": "", + "VICTIM": str(victim), + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + + assert victim.read_text(encoding = "utf-8") == "do not truncate me\n" + note = studio_home / "share" / ".unsloth-master-root" + assert note.read_text(encoding = "utf-8").strip() == str(master) + # The rename replaces a link at the final path rather than writing through it. + assert not note.is_symlink() + # No staging file survives the run, whatever name it was given. + leftovers = sorted( + p.name + for p in (studio_home / "share").iterdir() + if p.name.startswith(".unsloth-master-root.") and not p.is_symlink() + ) + assert leftovers == [], leftovers + + +def test_the_windows_note_writer_creates_its_staging_file_exclusively(): + """The setup.ps1 half of the test above. Held structurally: the Linux runners have no + Windows filesystem to plant a link on, and CreateNew is the thing that must not regress.""" + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice(ps, "if ((Get-MasterRootOverride) -and -not $StageRoot -and", "\n# ") + # Comments stripped, as the other structural checks here do: the block explains the old + # staging name in prose, and the name is only a finding when something executes it. + block = "\n".join(l for l in block.splitlines() if not l.lstrip().startswith("#")) + + assert "$notePath.$PID" not in block, "predictable staging name is back" + assert "GetRandomFileName()" in block + assert "[System.IO.FileMode]::CreateNew" in block + # WriteAllText opens an existing path rather than failing on it. + assert "WriteAllText($noteTmp" not in block + + +def _run_note_block( + tmp_path, + studio_home, + master, + *, + existing_note = None, + home = None, +): + """Run the shipped note gate + writer + legacy sweep against a fixture, and report the note. + + The whole region is executed, not pattern-matched: the gate, the writer it guards and the + sweep below it are one decision, and slicing them apart is how a gate that never runs passes. + """ + src = SETUP_SH.read_text(encoding = "utf-8") + block = _slice(src, "_master_root_note_is_honoured() {", "LLAMA_CPP_DIR=") + note = studio_home / "share" / ".unsloth-master-root" + note.parent.mkdir(parents = True, exist_ok = True) + if existing_note is not None: + note.write_text(str(existing_note) + "\n", encoding = "utf-8") + completed = subprocess.run( + ["bash", "-c", "set -u\n" + block], + env = { + "PATH": os.environ.get("PATH", "/usr/bin:/bin"), + "HOME": str(home if home is not None else tmp_path / "home"), + "STUDIO_HOME": str(studio_home), + "UNSLOTH_HOME": str(master), + "_MASTER_ROOT": str(master), + "STAGE_ROOT": "", + }, + capture_output = True, + text = True, + timeout = 60, + ) + assert completed.returncode == 0, completed.stderr + value = note.read_text(encoding = "utf-8").strip() if note.exists() else None + return value, completed.stderr + + +@NEEDS_POSIX_BASH +def test_the_master_root_note_is_only_written_when_a_reader_honours_it(tmp_path): + """A note whose Studio tree is not INSIDE the root it names is refused by every reader -- + storage_roots.py, the CLI and both uninstallers all require containment -- so writing one + records nothing and hides the fact that nothing was recorded. + + This is not a corner. install.sh does not read UNSLOTH_HOME yet, so + `UNSLOTH_HOME=/mnt/portable unsloth studio update` on an ordinary install is exactly this + shape: Studio stays at ~/.unsloth/studio while llama.cpp, node and whisper.cpp go to + /mnt/portable. That is the case the note was added for, and the case in which it was written + and then thrown away by every reader. Warn instead. + """ + home = tmp_path / "home" + studio_home = home / ".unsloth" / "studio" + master = tmp_path / "portable" + master.mkdir() + value, stderr = _run_note_block(tmp_path, studio_home, master, home = home) + assert value is None, f"wrote a note no reader will honour: {value}" + assert "cannot be recorded" in stderr, stderr + + # The contained layout, which every reader does honour, still records. + contained = master / "studio" + value, _ = _run_note_block(tmp_path, contained, master, home = home) + assert value == str(master) + + +@NEEDS_POSIX_BASH +def test_the_legacy_default_root_is_never_recorded_and_an_old_note_is_cleared(tmp_path): + """`UNSLOTH_HOME=$HOME/.unsloth` names the directory the install is already in. Both + uninstallers refuse that value outright, so a note carrying it licenses nothing -- but + storage_roots honours it, and honouring it flips portable_mode() on for every later BARE + launch, which moves HF_HUB_CACHE off ~/.cache/huggingface. One command that named the default + root therefore turned a shared model cache into a private one, permanently, with no variable + set at the time. + + The write is refused, and an existing note carrying that value is cleared on the next run -- + without the sweep the damage survives the fix, since the writer only runs when UNSLOTH_HOME + is set again and the user who hit this set it once. + """ + home = tmp_path / "home" + studio_home = home / ".unsloth" / "studio" + legacy_master = home / ".unsloth" + studio_home.mkdir(parents = True) + + # Contained (studio IS inside ~/.unsloth), so only the legacy-root rule can refuse it. + value, _ = _run_note_block(tmp_path, studio_home, legacy_master, home = home) + assert value is None, f"recorded the legacy default root: {value}" + + # An install that already carries one from an earlier build is repaired in place. + value, _ = _run_note_block( + tmp_path, + studio_home, + legacy_master, + existing_note = legacy_master, + home = home, + ) + assert value is None, "a stale legacy-root note survived" + + # A note naming any other root is left alone: it may describe an install this run cannot see. + other = tmp_path / "portable" + other.mkdir() + value, _ = _run_note_block( + tmp_path, + studio_home, + legacy_master, + existing_note = other, + home = home, + ) + assert value == str(other) + + +def test_the_windows_note_gate_holds_the_same_two_rules(): + """The setup.ps1 half of the two tests above. Structural: no PowerShell-on-Linux run can + exercise %USERPROFILE% canonicalisation the way a real Windows profile does.""" + ps = SETUP_PS1.read_text(encoding = "utf-8") + gate = _slice(ps, "function Test-MasterRootNoteIsHonoured {", "\nif ((Get-MasterRootOverride)") + assert "$legacy = Get-CanonicalDir" in gate, "the legacy default root is not refused" + assert "StartsWith($norm + $sep" in gate, "containment is not checked" + # The gate has to be wired into the writer, not merely defined beside it. + assert "Test-MasterRootNoteIsHonoured -Root $UnslothHome -StudioRoot $StudioHome" in ps + sweep = _slice(ps, "$staleNote = Join-Path", "} catch {") + assert "Remove-Item -LiteralPath $staleNote" in sweep + + +def test_the_windows_note_writer_is_a_no_op_when_the_note_already_says_this(): + """Move-Item -Force is not an atomic replace: PowerShell's provider deletes the destination + and then moves the source, so an interruption in between leaves NO note, and a backend that + looks next caches "no master root" from it. Re-entering that window on every update to + rewrite bytes that are already there is the avoidable half, and the staging file has to go + whether the move threw or not. + + Run under pwsh rather than asserted from source text: the comparison is the point, and a + structural check would pass on a version that compares the wrong two things. + """ + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice( + ps, ' $notePath = Join-Path $noteDir ".unsloth-master-root"', "\n } catch {" + ) + assert "Remove-Item -LiteralPath $noteTmp" in block, "the staging file is never cleaned up" + assert "$noteSame" in block, "the note is rewritten even when it already says this" + + +@pytest.mark.skipif(PWSH is None, reason = "needs pwsh") +def test_the_windows_note_writer_rewrites_only_when_the_value_changed(tmp_path): + """The behaviour, measured: same value twice leaves the file untouched; a different value + replaces it; and no staging file survives either way.""" + ps = SETUP_PS1.read_text(encoding = "utf-8") + block = _slice( + ps, ' $notePath = Join-Path $noteDir ".unsloth-master-root"', "\n } catch {" + ) + note_dir = tmp_path / "share" + note_dir.mkdir() + script = tmp_path / "write_note.ps1" + script.write_text( + f"$noteDir = {_ps_quote(str(note_dir))}\n" + f"$UnslothHome = {_ps_quote(str(tmp_path / 'master'))}\n" + + block + + "\n$final = Join-Path $noteDir '.unsloth-master-root'\n" + "$stamp = (Get-Item -LiteralPath $final).LastWriteTimeUtc.Ticks\n" + "Write-Output ((Get-Content -LiteralPath $final -Raw) + '|' + $stamp)\n", + encoding = "utf-8", + ) + argv = [PWSH, "-NoProfile", "-File", str(script)] + first = run_pwsh(argv, capture_output = True, text = True, check = True) + note = note_dir / ".unsloth-master-root" + + # Back-dated so a rewrite is unmistakable: two runs a millisecond apart land on one + # filesystem tick, which is how the first version of this test passed against the unfixed + # writer. + old_stamp = 1_000_000_000 + os.utime(note, (old_stamp, old_stamp)) + + second = run_pwsh(argv, capture_output = True, text = True, check = True) + + body_1, _ = first.stdout.strip().rsplit("|", 1) + body_2, _ = second.stdout.strip().rsplit("|", 1) + assert body_1 == body_2 == str(tmp_path / "master") + "\n" + assert ( + int(note.stat().st_mtime) == old_stamp + ), "an unchanged note was rewritten, re-entering the delete-then-move window" + leftovers = [p.name for p in note_dir.iterdir() if p.name != ".unsloth-master-root"] + assert leftovers == [], leftovers + + # And a CHANGED value must still be written, or the check above passes by doing nothing. + script.write_text( + script.read_text(encoding = "utf-8").replace( + _ps_quote(str(tmp_path / "master")), _ps_quote(str(tmp_path / "moved")) + ), + encoding = "utf-8", + ) + run_pwsh(argv, capture_output = True, text = True, check = True) + assert note.read_text(encoding = "utf-8").strip() == str(tmp_path / "moved") + assert int(note.stat().st_mtime) != old_stamp diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py index 6bbfe72771..45ffb5b3eb 100644 --- a/tests/test_studio_install_workspace_guard.py +++ b/tests/test_studio_install_workspace_guard.py @@ -697,15 +697,30 @@ def test_setup_helpers_gate_on_canonical_custom_root(): """setup.sh/setup.ps1 ownership guards must gate on a canonical custom-vs-legacy root comparison.""" sh_src = SETUP_SH.read_text(encoding = "utf-8") sh_idx = sh_src.index("_assert_studio_owned_or_absent() {") - sh_func = sh_src[sh_idx : sh_idx + 600] + # To the end of the function, not a fixed width, for the reason the PowerShell half below + # already gives: a new parameter or comment pushes the assertions out of a fixed window and + # the test fails while the guard it is about is intact. + sh_func = sh_src[sh_idx:].split("\n}\n", 1)[0] + # The caller may name the flag (the runtime children pass _RUNTIME_ROOT_IS_CUSTOM), but the + # default has to stay the canonical one. assert ( - '"$_STUDIO_HOME_IS_CUSTOM" = true' in sh_func - ), "setup.sh _assert_studio_owned_or_absent must gate on _STUDIO_HOME_IS_CUSTOM" + '_aso_custom="${3:-$_STUDIO_HOME_IS_CUSTOM}"' in sh_func + and '"$_aso_custom" = true' in sh_func + ), "setup.sh _assert_studio_owned_or_absent must gate on the canonical custom-root flag" assert ( "_LEGACY_STUDIO_HOME=" in sh_src and "_studio_home_canon=" in sh_src and "_STUDIO_HOME_IS_CUSTOM=" in sh_src ), "setup.sh must compute the canonical custom-root flag" + # A master root moves the runtime children, so the flag has to widen to cover them. Not the + # exact one-liner this used to name: UNSLOTH_HOME can legitimately name ~/.unsloth, the root a + # default install already uses, and treating that as custom made an update reject a legacy + # source-built llama.cpp that carries no owner marker. tests/test_setup_master_root.py runs + # the shipped derivation for that case; the rule here is the two halves being present. + assert '_RUNTIME_ROOT_IS_CUSTOM="$_STUDIO_HOME_IS_CUSTOM"' in sh_src + assert ( + "_RUNTIME_ROOT_IS_CUSTOM=true" in sh_src and "$_MASTER_ROOT" in sh_src + ), "setup.sh must widen the flag to a master root, which moves the runtime children" ps_src = SETUP_PS1.read_text(encoding = "utf-8") ps_idx = ps_src.index("function Assert-StudioOwnedOrAbsent") @@ -713,8 +728,8 @@ def test_setup_helpers_gate_on_canonical_custom_root(): # out of. ps_func = ps_src[ps_idx:].split("\nfunction ", 1)[0] assert ( - "$StudioHomeIsCustom -and" in ps_func - ), "setup.ps1 Assert-StudioOwnedOrAbsent must gate on $StudioHomeIsCustom" + "$isCustomRoot = $StudioHomeIsCustom" in ps_func and "$isCustomRoot -and" in ps_func + ), "setup.ps1 Assert-StudioOwnedOrAbsent must gate on the canonical custom-root flag" assert ( "$StudioOwnedMarker) -PathType Leaf" in ps_func ), "setup.ps1 marker check must use -PathType Leaf so a directory cannot satisfy it" @@ -730,8 +745,8 @@ def test_setup_ps1_inplace_git_sync_marks_studio_owned(): "Mark-StudioOwned -Path $LlamaCppDir" in inplace_block ), "in-place git-sync branch must call Mark-StudioOwned on success" assert ( - "$StudioHomeIsCustom" in inplace_block - ), "in-place Mark-StudioOwned call should be gated on $StudioHomeIsCustom" + "$RuntimeRootIsCustom" in inplace_block + ), "in-place Mark-StudioOwned call should be gated on $RuntimeRootIsCustom" def test_setup_ps1_inplace_git_sync_asserts_studio_owned_before_mutation(): @@ -926,6 +941,99 @@ def test_tauri_preflight_scrubs_studio_home_env(): assert ( 'cmd.env_remove("STUDIO_HOME")' in commands ), "commands.rs check_install_status must scrub STUDIO_HOME" + # Belt and braces: the managed context is what removes the whole list at these sites. + assert ( + preflight.count("apply_managed_cli_context_tokio(&mut cmd)") >= 2 + ), "preflight probes must build the managed context, which is what applies the scrub list" + assert ( + "apply_managed_cli_context_tokio(&mut cmd)" in commands + ), "commands.rs check_install_status must build the managed context" + + +# The three storage_roots.py resolvers that CHOOSE a data root, as opposed to the ones that +# join a subdirectory onto whatever they chose. Every environment variable read inside them +# moves the desktop's databases, assets and caches, so every one has to be scrubbed. +# _env_unsloth_home is the env-reading half unsloth_home splits into once it grows an on-disk +# fallback. Absent resolvers are skipped, not failed, so this holds for either spelling. +_ROOT_CHOOSING_RESOLVERS = ( + "unsloth_home", + "_env_unsloth_home", + "portable_mode", + "studio_root", +) + + +def _root_moving_env_names() -> set[str]: + """The env vars storage_roots.py picks a root from, read out of the shipped source. + + Derived rather than written down, so a fourth root variable fails this test on the commit + that adds it instead of on the bug report that follows it. + """ + source = (REPO_ROOT / "studio" / "backend" / "utils" / "paths" / "storage_roots.py").read_text( + encoding = "utf-8" + ) + names: set[str] = set() + seen = 0 + for resolver in _ROOT_CHOOSING_RESOLVERS: + marker = f"\ndef {resolver}(" + if marker not in source: + continue + seen += 1 + start = source.index(marker) + end = source.index("\ndef ", start + 1) + names |= set(re.findall(r'os\.environ\.get\(\s*"([A-Z][A-Z0-9_]*)"', source[start:end])) + assert seen, "none of the root-choosing resolvers exist under these names any more" + assert names, "the root-choosing resolvers read no environment variable any more" + return names + + +def _rust_string_list(source: str, const_name: str) -> list[str]: + """The string literals in a `const NAME: &[&str] = &[...];` declaration.""" + start = source.index(f"const {const_name}:") + return re.findall(r'"([^"]*)"', source[start : source.index(";", start)]) + + +def test_tauri_managed_children_scrub_every_root_moving_env(): + """A shell-level Unsloth root must not reach the packaged desktop's Python children. + + The desktop pins the legacy ~/.unsloth root and hardcodes it in Rust, so any variable + storage_roots.py would resolve a different root from has to be removed before the spawn, + or the databases, assets and caches move while the Rust half keeps reading the old place. + """ + src_root = REPO_ROOT / "studio" / "src-tauri" / "src" + process = (src_root / "process.rs").read_text(encoding = "utf-8") + install = (src_root / "install.rs").read_text(encoding = "utf-8") + + scrubbed = _rust_string_list(process, "MANAGED_CHILD_SCRUBBED_ENV") + assert len(scrubbed) == len(set(scrubbed)), f"duplicate names in the scrub list: {scrubbed}" + missing = sorted(_root_moving_env_names() - set(scrubbed)) + assert not missing, ( + f"{missing} would redirect the managed backend's data root away from the legacy one " + "the Tauri shell hardcodes" + ) + + # Second scrub path: start_backend removes the list again after the managed context, so + # the skip is a fact about the child. The whole list, not a subset that can drift. + start = process.index("pub fn start_backend(") + end = min( + offset + for offset in (process.find("\npub fn ", start + 1), process.find("\nfn ", start + 1)) + if offset != -1 + ) + backend_start = process[start:end] + assert ( + "for name in MANAGED_CHILD_SCRUBBED_ENV" in backend_start + ), "start_backend must scrub the whole list, not the names that were in it when it was written" + + # Third spawn path: the installer runs install.sh / install.ps1 directly, never reaching + # apply_managed_cli_context, so it applies the list itself. + assert ( + "for name in crate::process::MANAGED_CHILD_SCRUBBED_ENV" in install + ), "install.rs spawns outside the managed context and must scrub the whole list" + for name in scrubbed: + assert ( + f'cmd.env_remove("{name}")' not in install + ), f"install.rs scrubs {name} by hand as well as by list; the hand-written copy can rot" def test_install_sh_shim_uses_atomic_replace(): diff --git a/tests/test_studio_root_resilience.py b/tests/test_studio_root_resilience.py index edd820c8b2..d19fcd27d2 100644 --- a/tests/test_studio_root_resilience.py +++ b/tests/test_studio_root_resilience.py @@ -115,7 +115,10 @@ def test_kill_orphan_catches_oserror_from_studio_root(): def _exec_search_roots_block( - home: Path, studio_root_value: Path, resolve_raises: bool + home: Path, + studio_root_value: Path, + resolve_raises: bool, + master: Path | None = None, ) -> list[Path]: """Run _find_llama_server_binary's search_roots derivation -- plus the shared _resolved_studio_root_and_is_legacy() classifier it delegates to -- with a @@ -138,6 +141,7 @@ def _exec_search_roots_block( block = textwrap.dedent(" " * 8 + src[block_start:block_end]) fake_module = type(sys)("fake_storage_roots") fake_module.studio_root = lambda: studio_root_value + fake_module.unsloth_home = lambda: master sys.modules["utils.paths.storage_roots"] = fake_module try: original_resolve = Path.resolve @@ -183,6 +187,20 @@ def test_search_roots_keeps_custom_when_resolve_fails(tmp_path): ) not in roots, f"legacy llama path must not appear in custom-mode search_roots: {roots}" +def test_search_roots_follow_the_master_root(tmp_path): + # llama.cpp is a sibling of studio/ there; a failing resolve() must not drop the master root. + home = tmp_path / "home" + home.mkdir() + master = tmp_path / "portable" + roots = _exec_search_roots_block( + home = home, + studio_root_value = master / "studio", + resolve_raises = True, + master = master, + ) + assert roots == [master / "llama.cpp"] + + def test_search_roots_default_mode_uses_legacy_only(tmp_path): home = tmp_path / "home" home.mkdir() diff --git a/tests/test_unsloth_home_root_agreement.py b/tests/test_unsloth_home_root_agreement.py new file mode 100644 index 0000000000..882e261fc8 --- /dev/null +++ b/tests/test_unsloth_home_root_agreement.py @@ -0,0 +1,333 @@ +# Unsloth Zoo - Utilities for Unsloth +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""UNSLOTH_HOME has to mean the same install to the CLI and to the backend: if the two disagree +the backend writes studio.db, auth and the pid file under /studio while the CLI +reads ~/.unsloth/studio. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] + +_PROBE = """ +import json, sys +sys.path.insert(0, {backend!r}) +import os +from unsloth_cli.commands import studio as cli +from utils.paths.storage_roots import studio_root, unsloth_home +cli._ensure_studio_env_exported() +master = unsloth_home() +print(json.dumps({{ + "cli": str(cli.STUDIO_HOME), + "cli_is_custom": bool(cli._STUDIO_HOME_IS_CUSTOM), + "backend": str(studio_root()), + "exported_studio_home": os.environ.get("UNSLOTH_STUDIO_HOME"), + "cli_llama": os.environ.get("UNSLOTH_LLAMA_CPP_PATH"), + "exported_master": os.environ.get("UNSLOTH_HOME"), + "backend_llama": str((master or studio_root()) / "llama.cpp"), +}})) +""" + + +def _probe(env_overrides: dict) -> dict: + env = dict(os.environ) + # UNSLOTH_LLAMA_CPP_PATH too, as _main_probe does: an inherited value measures the runner. + for key in ("UNSLOTH_HOME", "UNSLOTH_STUDIO_HOME", "STUDIO_HOME", "UNSLOTH_LLAMA_CPP_PATH"): + env.pop(key, None) + env.update(env_overrides) + env["PYTHONPATH"] = str(REPO_ROOT) + source = _PROBE.format(backend = str(REPO_ROOT / "studio" / "backend")) + out = subprocess.run( + [sys.executable, "-c", source], + capture_output = True, + text = True, + cwd = str(REPO_ROOT), + env = env, + check = True, + ) + return json.loads(out.stdout.strip().splitlines()[-1]) + + +def test_unsloth_home_resolves_to_one_studio_root(tmp_path): + master = tmp_path / "portable" + result = _probe({"UNSLOTH_HOME": str(master)}) + + assert result["backend"] == str(master / "studio") + assert result["cli"] == result["backend"] + # Custom, so `unsloth studio ...` re-exports UNSLOTH_STUDIO_HOME. + assert result["cli_is_custom"] is True + + +def test_studio_home_still_outranks_unsloth_home(tmp_path): + explicit = tmp_path / "explicit" + result = _probe( + {"UNSLOTH_HOME": str(tmp_path / "portable"), "UNSLOTH_STUDIO_HOME": str(explicit)} + ) + + assert result["backend"] == str(explicit) + assert result["cli"] == str(explicit) + + +def test_no_unsloth_home_keeps_the_legacy_default(tmp_path): + result = _probe({"HOME": str(tmp_path / "home"), "USERPROFILE": str(tmp_path / "home")}) + + legacy = str(tmp_path / "home" / ".unsloth" / "studio") + assert result["cli"] == legacy + assert result["backend"] == legacy + assert result["cli_is_custom"] is False + + +def test_the_cli_exports_the_llama_cpp_path_the_backend_will_use(tmp_path): + # run.py keeps a non-blank UNSLOTH_LLAMA_CPP_PATH, so a CLI exporting + # /studio/llama.cpp pins that wrong path for every worker. + master = tmp_path / "portable" + result = _probe({"UNSLOTH_HOME": str(master)}) + + assert result["cli_llama"] == str(master / "llama.cpp") + assert result["cli_llama"] == result["backend_llama"] + + +def test_the_cli_recovers_a_recorded_master_root_for_the_setup_subprocess(tmp_path): + # A one-command UNSLOTH_HOME leaves nothing in a later environment, and the backend recovers + # the root from the note, so a CLI that did not would have setup refresh the runtimes at + # /studio/ while the backend launched the stale ones at /. Exported rather + # than read, since setup is a subprocess. + master = tmp_path / "portable" + studio = master / "studio" + (studio / "share").mkdir(parents = True) + (studio / "share" / ".unsloth-master-root").write_text(f"{master}\n", encoding = "utf-8") + result = _probe({"UNSLOTH_STUDIO_HOME": str(studio)}) + + assert result["exported_master"] == str(master) + assert result["cli_llama"] == str(master / "llama.cpp") + assert result["cli_llama"] == result["backend_llama"] + + +def test_the_cli_and_the_backend_accept_the_same_nested_studio_root(tmp_path): + # A real pairing: setup records , and the Studio directory is inside it without being + # its studio/ child. A stricter rule in the CLI would decline a note the backend accepts. + master = tmp_path / "root" + studio = master / "custom" / "studio" + (studio / "share").mkdir(parents = True) + (studio / "share" / ".unsloth-master-root").write_text(f"{master}\n", encoding = "utf-8") + result = _probe({"UNSLOTH_STUDIO_HOME": str(studio)}) + + assert result["exported_master"] == str(master) + assert result["cli"] == str(studio) + assert result["cli_llama"] == str(master / "llama.cpp") + assert result["cli_llama"] == result["backend_llama"] + + +def test_the_cli_refuses_a_note_carried_in_from_another_master_root(tmp_path): + # A Studio tree copied from master root A to B keeps a note naming A. Exporting it would + # point an update at the original install; the Studio directory has to lie inside the root + # its note names, which is the rule storage_roots and both uninstallers apply. + original = tmp_path / "original" + (original / "studio").mkdir(parents = True) + copied = tmp_path / "copied" / "studio" + (copied / "share").mkdir(parents = True) + (copied / "share" / ".unsloth-master-root").write_text(f"{original}\n", encoding = "utf-8") + result = _probe({"UNSLOTH_STUDIO_HOME": str(copied)}) + + assert result["exported_master"] is None + assert result["cli_llama"] == str(copied / "llama.cpp") + assert result["cli_llama"] == result["backend_llama"] + + +def test_a_plain_custom_root_still_keeps_llama_cpp_inside_it(tmp_path): + explicit = tmp_path / "explicit" + result = _probe({"UNSLOTH_STUDIO_HOME": str(explicit)}) + + assert result["cli_llama"] == str(explicit / "llama.cpp") + assert result["cli_llama"] == result["backend_llama"] + + +def test_a_master_root_the_cli_declined_is_still_told_to_the_backend(tmp_path): + """The one root the CLI picks that the backend would not pick for itself. + + install.sh and install.ps1 do not read UNSLOTH_HOME yet, so _resolve_studio_home keeps an + exported-but-uninstalled master root off the legacy install that actually exists. That + fallback is not custom, and returning early on non-custom exported nothing at all, so + studio_root() in the backend went on honouring UNSLOTH_HOME with no install check: the CLI + ran the legacy venv while the backend under it wrote studio.db, auth and the pid file to + /studio. The other tests here point the master root at an empty directory AND leave + the legacy root uninstalled, so the fallback is never entered by them. + """ + home = tmp_path / "home" + # The sentinel _looks_like_installer_managed_studio_home reads. + conf = home / ".unsloth" / "studio" / "share" / "studio.conf" + conf.parent.mkdir(parents = True) + conf.write_text("installed\n", encoding = "utf-8") + + result = _probe( + { + "HOME": str(home), + "USERPROFILE": str(home), + "UNSLOTH_HOME": str(tmp_path / "portable"), + } + ) + + legacy = str(home / ".unsloth" / "studio") + assert result["cli"] == legacy + # Not custom: the root is the Unsloth-owned legacy one, and false is what lets the + # installers replace the tree without an owner marker. + assert result["cli_is_custom"] is False + assert result["exported_studio_home"] == legacy + assert result["backend"] == legacy + + +def test_a_whitespace_studio_home_does_not_defeat_the_export(tmp_path): + """The same fallback as above, with an inherited UNSLOTH_STUDIO_HOME of " ". + + Every resolver strips before deciding, so " " means unset to _resolve_studio_home and to + studio_root alike. The export guard read the raw value, where " " is truthy, so it left the + whitespace in place: the CLI ran the legacy install and the backend under it resolved + /studio, which is the split this file exists to prevent. Blank was already covered by + the guard; whitespace-only was the shape that got through. + """ + home = tmp_path / "home" + conf = home / ".unsloth" / "studio" / "share" / "studio.conf" + conf.parent.mkdir(parents = True) + conf.write_text("installed\n", encoding = "utf-8") + + legacy = str(home / ".unsloth" / "studio") + for inherited in (" ", "", "\t\n"): + result = _probe( + { + "HOME": str(home), + "USERPROFILE": str(home), + "UNSLOTH_HOME": str(tmp_path / "portable"), + "UNSLOTH_STUDIO_HOME": inherited, + "UNSLOTH_LLAMA_CPP_PATH": inherited, + } + ) + assert result["cli"] == legacy, inherited + assert result["exported_studio_home"] == legacy, inherited + assert result["backend"] == legacy, inherited + # The runtimes live beside studio/ at the master root, and a whitespace value left in + # place would be scanned as a directory rather than replaced. + assert result["cli_llama"] == str(tmp_path / "portable" / "llama.cpp"), inherited + + +def test_a_legacy_install_still_keeps_llama_cpp_at_the_legacy_path(tmp_path): + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + result = _probe({"HOME": str(home), "USERPROFILE": str(home)}) + + # Not custom, so nothing is exported at all and the backend default stands. + assert result["cli_llama"] is None + assert result["backend_llama"] == str(home / ".unsloth" / "studio" / "llama.cpp") + + +_MAIN_PROBE = """ +import json, sys, types +sys.path.insert(0, {backend!r}) +# main.py pulls in the whole app on import, so exec only its module-level +# llama.cpp block, which is what decides the managed path and marks it. +import os, re +from pathlib import Path +src = Path({backend!r}, "main.py").read_text(encoding = "utf-8") +# Anchored on the master-root lookup rather than the `if`, because the guard +# reads _MASTER_ROOT and a slice starting at the `if` would not define it. +start = src.index('from utils.paths.storage_roots import unsloth_home as _unsloth_home') +end = src.index('# The studio bundles unsloth_zoo', start) +if 'if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT' not in src[start:end]: + raise AssertionError("the extracted slice no longer contains the export guard") +from utils.paths.storage_roots import studio_root as _studio_root +_LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve() +_STUDIO_ROOT_RESOLVED = _studio_root().resolve() +marked = [] +mod = types.ModuleType("utils.llama_cpp_path_settings") +mod.mark_managed_llama_cpp_path = lambda p: marked.append(str(p)) +sys.modules["utils.llama_cpp_path_settings"] = mod +exec(compile(src[start:end], "main.py", "exec"), globals()) +print(json.dumps({{ + "exported": os.environ.get("UNSLOTH_LLAMA_CPP_PATH"), + "marked": marked, +}})) +""" + + +def _main_probe(env_overrides: dict) -> dict: + env = dict(os.environ) + for key in ("UNSLOTH_HOME", "UNSLOTH_STUDIO_HOME", "STUDIO_HOME", "UNSLOTH_LLAMA_CPP_PATH"): + env.pop(key, None) + env.update(env_overrides) + env["PYTHONPATH"] = str(REPO_ROOT) + source = _MAIN_PROBE.format(backend = str(REPO_ROOT / "studio" / "backend")) + out = subprocess.run( + [sys.executable, "-c", source], + capture_output = True, + text = True, + cwd = str(REPO_ROOT), + env = env, + check = True, + ) + return json.loads(out.stdout.strip().splitlines()[-1]) + + +def test_main_marks_the_same_llama_cpp_path_run_py_exports(tmp_path): + # A direct `uvicorn main:app` exports this outright; under run.py main.py would instead + # mark the wrong path managed, making the bundled one look like an override. + master = tmp_path / "portable" + result = _main_probe({"UNSLOTH_HOME": str(master)}) + + assert result["exported"] == str(master / "llama.cpp") + assert result["marked"] == [str(master / "llama.cpp")] + + +def test_main_leaves_a_plain_custom_root_alone(tmp_path): + explicit = tmp_path / "explicit" + result = _main_probe({"UNSLOTH_STUDIO_HOME": str(explicit)}) + + assert result["exported"] == str(explicit / "llama.cpp") + assert result["marked"] == [str(explicit / "llama.cpp")] + + +def test_main_exports_for_a_master_root_that_is_the_legacy_path(tmp_path): + # A portable install pointed at the legacy Studio path still owns /llama.cpp, so the + # legacy equality alone skipped the export and left unsloth_zoo on ~/.unsloth/llama.cpp. + home = tmp_path / "home" + legacy = home / ".unsloth" / "studio" + legacy.mkdir(parents = True) + result = _main_probe( + { + "HOME": str(home), + "USERPROFILE": str(home), + "UNSLOTH_HOME": str(legacy), + "UNSLOTH_STUDIO_HOME": str(legacy), + } + ) + + assert result["exported"] == str(legacy / "llama.cpp") + assert result["marked"] == [str(legacy / "llama.cpp")] + + +def test_main_still_exports_nothing_without_a_master_root(tmp_path): + # The other half: no master root, so the guard stays closed rather than pinning the default. + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + result = _main_probe({"HOME": str(home), "USERPROFILE": str(home)}) + + assert result["exported"] is None + assert result["marked"] == [] diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index 10ad1cf1c3..76e576e087 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -132,10 +132,41 @@ def read_json_checking_deferred_error(url: str, response): return require_completed_padded_body(url, raise_for_deferred_error(url, body)) -def ensure_studio_backend_path() -> None: +_cache_env_seeded = False + + +def _seed_cache_env() -> None: + """Pin the cache locations the backend pins, for in-process CLI commands. + + Otherwise they inherit unsloth_zoo's relative UNSLOTH_COMPILE_LOCATION, resolved against the + working directory, leaving an unsloth_compiled_cache cache_cleanup will not remove (#8865). + """ + global _cache_env_seeded + if _cache_env_seeded: + return + _cache_env_seeded = True + try: + from utils.paths.storage_roots import setup_cache_env + setup_cache_env() + except Exception: # noqa: BLE001 - never fail a command over cache placement + pass + + +def ensure_studio_backend_path(*, seed_cache_env: bool = True) -> None: + """Put studio/backend on sys.path, and by default pin the cache locations too. + + `seed_cache_env = False` is for callers that only want to IMPORT a path helper. + setup_cache_env() creates every cache directory it pins, so seeding it from + `unsloth start`'s Node discovery turned a read-only lookup into 18 mkdirs under a home + that may have nothing to do with the command being run -- including one against a remote + server. Seeding stays on for the ML entry points, where the pins are the point. + """ backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend") if backend_dir not in sys.path: sys.path.insert(0, backend_dir) + if seed_cache_env: + # After the path insert, before the caller's backend import pulls in unsloth_zoo.compiler. + _seed_cache_env() def configure_quiet_logging() -> None: diff --git a/unsloth_cli/_studio_deps.py b/unsloth_cli/_studio_deps.py index 516070e7d8..b5db9d270c 100644 --- a/unsloth_cli/_studio_deps.py +++ b/unsloth_cli/_studio_deps.py @@ -681,7 +681,15 @@ def studio_backend_imports(feature: str = "This command", *, studio_only: bool = Only ModuleNotFoundError is intercepted; any other ImportError from the backend is a real bug and keeps its traceback. + + Every CLI backend import goes through here, so the cache env is pinned here too: `export` and + `list-checkpoints` reach studio.backend.core.export without calling + ensure_studio_backend_path(), and would import unsloth_zoo.compiler with + UNSLOTH_COMPILE_LOCATION unset, writing unsloth_compiled_cache into the working dir (#8865). """ + from unsloth_cli._inference import ensure_studio_backend_path + + ensure_studio_backend_path() try: yield except ModuleNotFoundError as exc: diff --git a/unsloth_cli/_system_dir_guard.py b/unsloth_cli/_system_dir_guard.py index 39afbddd85..790dd15c0e 100644 --- a/unsloth_cli/_system_dir_guard.py +++ b/unsloth_cli/_system_dir_guard.py @@ -281,6 +281,8 @@ def is_relocatable_invocation(argv, environ): # Path overrides written relative to the folder being left: Unsloth resolves them with # Path.resolve(), so moving first would silently retarget them. _RELATIVE_PATH_ENV = ( + # Unsloth roots: storage_roots.py. + "UNSLOTH_HOME", "UNSLOTH_STUDIO_HOME", "STUDIO_HOME", "UNSLOTH_STUDIO_DOCUMENTS_HOME", @@ -321,6 +323,18 @@ _RELATIVE_PATH_ENV = ( "UNSLOTH_DG_SHIM", "UNSLOTH_COMPILE_LOCATION", "TORCHINDUCTOR_CACHE_DIR", + # storage_roots.py fills these only when blank, so a relative value the user set is kept as + # written and would name a different folder after the move. + "TORCH_EXTENSIONS_DIR", + "TORCH_HOME", + "TRITON_HOME", + "TRITON_CACHE_DIR", + "TRITON_DUMP_DIR", + "CUDA_CACHE_PATH", + "MPLCONFIGDIR", + "NUMBA_CACHE_DIR", + "DATA_DESIGNER_HOME", + "DATA_DESIGNER_MANAGED_ASSETS_PATH", "UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR", "UNSLOTH_DIFFUSION_COND_CACHE_DIR", "HF_HOME", @@ -329,6 +343,9 @@ _RELATIVE_PATH_ENV = ( "HF_XET_CACHE", "HF_DATASETS_CACHE", "HF_ASSETS_CACHE", + # transformers appends this to sys.path, so a relative value would import a different + # generated module after the move. + "HF_MODULES_CACHE", # The credential file: a relative value would follow the child and lose access to gated repos. "HF_TOKEN_PATH", # Authoritative when non-blank (storage_roots.py), so `unsloth studio update` would install diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index d85255fe7f..bdfc819f39 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -3534,7 +3534,9 @@ def _refresh_windows_path() -> None: def _managed_node_tools() -> Optional[tuple[Path, Path, bool]]: # Best-effort: any failure here means "no managed Node", never a broken launch. try: - ensure_studio_backend_path() + # Discovery only: this answers "is there a managed Node", including for a launch aimed + # at a remote server, so it must not create the cache tree on the way past. + ensure_studio_backend_path(seed_cache_env = False) from utils.node_runtime import managed_node_binary, resolve_node_executable node = Path(managed_node_binary()) except (ImportError, OSError, RuntimeError, TypeError, ValueError): diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index a89b81ddf8..c06b229ed8 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -40,7 +40,9 @@ def _enable_verbose_access_logs() -> None: os.environ["UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS"] = "0" -# Root order: UNSLOTH_STUDIO_HOME, STUDIO_HOME, sys.prefix, legacy ~/.unsloth/studio. Markers mirror install.ps1 / uninstall.ps1 and are matched as bytes. +# Root order: UNSLOTH_STUDIO_HOME, STUDIO_HOME, UNSLOTH_HOME/studio, sys.prefix, +# legacy ~/.unsloth/studio. Keep this aligned with storage_roots.studio_root(). +# Shim markers mirror install.ps1 / uninstall.ps1 and are matched as bytes. _CMD_SHIM_MARKERS = (b"unsloth-studio-managed-launcher", b"from unsloth_cli import app") _CMD_SHIM_MAX_BYTES = 8192 @@ -76,6 +78,25 @@ def _resolve_studio_home() -> tuple[Path, bool]: return Path(override).expanduser().resolve(), True except (OSError, ValueError): return Path(override).expanduser(), True + # Keeps the CLI on the same root as storage_roots.py; see test_unsloth_home_root_agreement.py. + master = (os.environ.get("UNSLOTH_HOME") or "").strip() + if master: + try: + candidate = Path(master).expanduser().resolve() / "studio" + except (OSError, ValueError): + candidate = Path(master).expanduser() / "studio" + try: + legacy = (Path.home() / ".unsloth" / "studio").resolve() + except (OSError, ValueError): + legacy = Path.home() / ".unsloth" / "studio" + # install.sh and install.ps1 do not read UNSLOTH_HOME yet, so they leave the venv and the + # launcher at the legacy root while setup puts the runtimes under the master root: + # preferring /studio unconditionally reported "Unsloth Studio not set up" for an + # install that is right there. The master root wins only when it HAS an install. + if candidate != legacy and not _looks_like_installer_managed_studio_home(candidate): + if _looks_like_installer_managed_studio_home(legacy): + return legacy, False + return candidate, candidate != legacy try: prefix = Path(sys.prefix).resolve() if prefix.name == "unsloth_studio": @@ -90,25 +111,93 @@ def _resolve_studio_home() -> tuple[Path, bool]: STUDIO_HOME, _STUDIO_HOME_IS_CUSTOM = _resolve_studio_home() +MASTER_ROOT_NOTE = ".unsloth-master-root" + + +def _recorded_master_root() -> Optional[Path]: + """The master root setup recorded in this Studio tree, or None. + + Keep this aligned with storage_roots._recorded_master_root(); see + test_unsloth_home_root_agreement.py. `UNSLOTH_HOME=/mnt/portable unsloth studio update` puts + node, llama.cpp and whisper.cpp BESIDE studio/ and leaves nothing in a later environment. The + backend recovers that from the note, so a `unsloth studio update` that did not would hand + setup a plain Studio root, have it refresh the runtimes one level down at /studio/, + and leave the backend still launching the stale trees at /. + + The recorded root must exist and THIS Studio directory must lie inside it, so a tree copied + from one master root to another does not send the update, or a removal, into the original + install. Containment, not an exact /studio match: the flat layout names one directory + for both, and UNSLOTH_HOME=/root with UNSLOTH_STUDIO_HOME=/root/custom/studio is a root that + genuinely contains its Studio somewhere other than the default child. The backend and both + uninstallers apply exactly this rule, and a stricter one here would have the CLI decline a + note the backend accepts, which is the same split this function exists to close. + """ + try: + recorded = (STUDIO_HOME / "share" / MASTER_ROOT_NOTE).read_text(encoding = "utf-8").strip() + except (OSError, ValueError, UnicodeDecodeError): + return None + if not recorded: + return None + try: + master = Path(recorded).expanduser().resolve() + here = STUDIO_HOME.resolve() + if not (master.is_dir() and (here == master or master in here.parents)): + return None + # A legacy-rooted install has no master root, and both uninstallers already refuse what + # that shape records. storage_roots._is_legacy_studio_tree declines it too, and + # test_unsloth_home_root_agreement.py holds the two together: without this the CLI would + # export UNSLOTH_HOME for a note the backend has declined. Keyed on the TREE, so a note + # naming $HOME or any other ancestor goes with it. + try: + if here == (Path.home() / ".unsloth" / "studio").resolve(): + return None + except (OSError, RuntimeError, ValueError): + pass + return master + except (OSError, ValueError): + return None + return None + def _ensure_studio_env_exported() -> None: - """Re-export UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH for custom roots only, per - subcommand rather than at import, so unrelated importers see no env changes.""" - if not _STUDIO_HOME_IS_CUSTOM: + """Re-export UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH for custom roots, and for a master + root the resolver above declined, per subcommand rather than at import, so unrelated + importers see no env changes.""" + # storage_roots.studio_root() honours UNSLOTH_HOME with no install check, so the fallback + # keeping this CLI on an installed legacy root must be told to the backend too: unexported, + # `unsloth studio` runs the legacy venv while the backend inside it writes studio.db, auth + # and the pid file under /studio. Exporting the root does not make it custom, since + # the value equals the legacy path both installers compare against. + # The note, when this run has no UNSLOTH_HOME of its own. Exported rather than merely read, + # because setup runs as a subprocess and would otherwise refresh the runtimes at + # /studio/ while the backend kept launching the ones at /. + if not (os.environ.get("UNSLOTH_HOME") or "").strip(): + _recorded = _recorded_master_root() + if _recorded is not None: + os.environ["UNSLOTH_HOME"] = str(_recorded) + if not _STUDIO_HOME_IS_CUSTOM and not (os.environ.get("UNSLOTH_HOME") or "").strip(): return # Truthy-check, not setdefault: a blank UNSLOTH_STUDIO_HOME= must not win. - if not os.environ.get("UNSLOTH_STUDIO_HOME"): + if not (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip(): os.environ["UNSLOTH_STUDIO_HOME"] = str(STUDIO_HOME) try: _legacy_studio = (Path.home() / ".unsloth" / "studio").resolve() _is_legacy = STUDIO_HOME.resolve() == _legacy_studio except (OSError, ValueError): _is_legacy = STUDIO_HOME == (Path.home() / ".unsloth" / "studio") - if _is_legacy: + # The runtimes are siblings of studio/, at the master root, so STUDIO_HOME/llama.cpp is one + # level too deep. run.py keeps a non-blank value, so a wrong export here wins everywhere. + _master = (os.environ.get("UNSLOTH_HOME") or "").strip() + if _master: + try: + _llama_dir = Path(_master).expanduser().resolve() / "llama.cpp" + except (OSError, ValueError): + _llama_dir = Path(_master).expanduser() / "llama.cpp" + elif _is_legacy: _llama_dir = Path.home() / ".unsloth" / "llama.cpp" else: _llama_dir = STUDIO_HOME / "llama.cpp" - if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + if not (os.environ.get("UNSLOTH_LLAMA_CPP_PATH") or "").strip(): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_llama_dir) diff --git a/unsloth_cli/tests/test_export_cache_env.py b/unsloth_cli/tests/test_export_cache_env.py new file mode 100644 index 0000000000..b41a6ae1bb --- /dev/null +++ b/unsloth_cli/tests/test_export_cache_env.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth export` / `list-checkpoints` must pin the cache env before importing the export +backend. + +Neither calls ensure_studio_backend_path(), so before the seeding moved into +studio_backend_imports() they reached studio.backend.core.export with UNSLOTH_COMPILE_LOCATION +unset and unsloth_zoo resolved its relative default against the shell's cwd (#8865). +""" + +import json +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] + +# Runs in its own interpreter: the seeding is process-global (os.environ plus a one-shot flag), +# so an in-process assertion would depend on test ordering. +PROBE = textwrap.dedent( + """ + import json, os, sys, types + + import typer + from typer.testing import CliRunner + + command_name = sys.argv[1] + recorded = {} + + # Stand in for the heavy backend (unsloth, torch) so the probe stays + # import-light. The CLI's own import statement, and therefore everything that + # has to happen before it, is untouched. + class ExportBackend: + def __init__(self): + recorded["compile_location"] = os.environ.get("UNSLOTH_COMPILE_LOCATION") + + def load_checkpoint(self, **kwargs): + return True, "loaded" + + def export_merged_model(self, **kwargs): + return True, "exported", kwargs["save_directory"] + + def scan_checkpoints(self, **kwargs): + return [] + + # Each stub keeps the REAL directory on __path__ instead of []. The point is to skip + # the heavy studio.backend.core.export import, not to make the whole studio tree + # unimportable: unsloth_cli/__init__ eagerly imports commands.start, which pulls + # studio.backend.utils.coding_agents, and an empty __path__ turned that into + # ModuleNotFoundError. Submodules still load from disk; the package __init__ files + # are the thing that never runs, which is what kept this probe light to begin with. + repo_root = sys.argv[2] + for name, relative in ( + ("studio", "studio"), + ("studio.backend", "studio/backend"), + ("studio.backend.core", "studio/backend/core"), + ): + module = types.ModuleType(name) + module.__path__ = [os.path.join(repo_root, *relative.split("/"))] + sys.modules[name] = module + export_module = types.ModuleType("studio.backend.core.export") + export_module.ExportBackend = ExportBackend + sys.modules["studio.backend.core.export"] = export_module + + from unsloth_cli.commands import export as export_command + + app = typer.Typer() + if command_name == "export": + app.command()(export_command.export) + args = ["checkpoint", "exported"] + else: + app.command()(export_command.list_checkpoints) + args = [] + result = CliRunner().invoke(app, args) + recorded["exit_code"] = result.exit_code + print("PROBE " + json.dumps(recorded)) + """ +) + + +@pytest.mark.parametrize("command_name", ["export", "list-checkpoints"]) +def test_export_commands_seed_compile_location(tmp_path, command_name): + unsloth_home = tmp_path / "unsloth" + workdir = tmp_path / "cwd" + workdir.mkdir() + + env = dict(os.environ) + env.pop("UNSLOTH_COMPILE_LOCATION", None) + # Both outrank UNSLOTH_HOME in the resolver, and conftest.py sets a session-scoped + # UNSLOTH_STUDIO_HOME, so an inherited one would send the probe somewhere unnamed. + for key in ("UNSLOTH_STUDIO_HOME", "STUDIO_HOME", "UNSLOTH_PORTABLE"): + env.pop(key, None) + env.update( + HOME = str(tmp_path / "home"), + UNSLOTH_HOME = str(unsloth_home), + PYTHONPATH = str(REPO_ROOT), + ) + completed = subprocess.run( + [sys.executable, "-c", PROBE, command_name, str(REPO_ROOT)], + cwd = str(workdir), + env = env, + capture_output = True, + text = True, + timeout = 300, + ) + assert completed.returncode == 0, completed.stderr + line = next(l for l in completed.stdout.splitlines() if l.startswith("PROBE ")) + recorded = json.loads(line[len("PROBE ") :]) + + assert recorded["exit_code"] == 0, recorded + location = recorded["compile_location"] + assert location, "UNSLOTH_COMPILE_LOCATION was unset when the export backend was imported" + # Absolute and inside the install root, not resolved against the CWD. + assert Path(location).is_absolute() + assert Path(location) == unsloth_home / "studio" / "compiled_cache" + assert not (workdir / "unsloth_compiled_cache").exists() + + +def test_node_discovery_does_not_build_the_cache_tree(tmp_path): + """`unsloth start` asks whether a managed Node exists; that is a read, not an install. + + ensure_studio_backend_path() seeds the cache environment, and setup_cache_env() CREATES + every directory it pins, so routing Node discovery through it turned the lookup into 18 + mkdirs under a home that may have nothing to do with the command -- including a launch + aimed at a remote server. The merge base created nothing here. + """ + home = tmp_path / "home" + (home / ".unsloth" / "studio").mkdir(parents = True) + probe = ( + "import json, os, sys\n" + f"sys.path.insert(0, {str(REPO_ROOT)!r})\n" + "from pathlib import Path\n" + "home = Path(os.environ['HOME'])\n" + "before = {str(p) for p in home.rglob('*') if p.is_dir()}\n" + "from unsloth_cli.commands import start\n" + "start._managed_node_tools()\n" + "after = {str(p) for p in home.rglob('*') if p.is_dir()}\n" + "print('PROBE ' + json.dumps(sorted(after - before)))\n" + ) + completed = subprocess.run( + [sys.executable, "-c", probe], + cwd = str(tmp_path), + env = {"HOME": str(home), "PATH": os.environ.get("PATH", "/usr/bin:/bin")}, + capture_output = True, + text = True, + timeout = 300, + ) + assert completed.returncode == 0, completed.stderr + line = next(l for l in completed.stdout.splitlines() if l.startswith("PROBE ")) + assert json.loads(line[len("PROBE ") :]) == []