diff --git a/install.ps1 b/install.ps1 index d5db1785a1..ef87c5ed08 100644 --- a/install.ps1 +++ b/install.ps1 @@ -3,6 +3,11 @@ # Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local # NoTorch: .\install.ps1 --no-torch (skip PyTorch, GGUF-only mode) # Test: .\install.ps1 --package roland-sloth +# +# Env vars (priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > USERPROFILE-redirect > default): +# UNSLOTH_STUDIO_HOME / STUDIO_HOME = path -> install under that path +# (DataDir nests inside; user PATH not modified persistently). +# Default ($USERPROFILE\.unsloth\studio) is preserved when no env var is set. function Install-UnslothStudio { $ErrorActionPreference = "Stop" @@ -126,7 +131,94 @@ function Install-UnslothStudio { } $PythonVersion = "3.13" - $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + + # Resolve install destinations. Priority: UNSLOTH_STUDIO_HOME, then + # STUDIO_HOME alias, then USERPROFILE-redirect, then default. + # Reject whitespace-only values so " " is treated as unset (matches the + # Python resolvers' .strip()), preventing install/runtime layout drift. + $envOverrideVar = $null + $envOverride = $null + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { + $envOverrideVar = "UNSLOTH_STUDIO_HOME" + $envOverride = $env:UNSLOTH_STUDIO_HOME.Trim() + } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { + $envOverrideVar = "STUDIO_HOME" + $envOverride = $env:STUDIO_HOME.Trim() + } + + # Custom Studio roots are not supported with --tauri (desktop app still + # resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy. + if ($TauriMode -and $envOverride) { + $_tauriOverride = $envOverride + if ($_tauriOverride -eq "~" -or $_tauriOverride -like "~/*" -or $_tauriOverride -like "~\*") { + $_tauriOverride = (Join-Path $env:USERPROFILE $_tauriOverride.Substring(1).TrimStart('/','\')) + } + try { + $_tauriOverride = [System.IO.Path]::GetFullPath($_tauriOverride) + } catch {} + $_legacyTauriRoot = Join-Path $env:USERPROFILE ".unsloth\studio" + try { + $_legacyTauriRoot = [System.IO.Path]::GetFullPath($_legacyTauriRoot) + } catch {} + # Strip trailing separators so ".../studio\" matches ".../studio". + $_trimSeps = @( + [System.IO.Path]::DirectorySeparatorChar, + [System.IO.Path]::AltDirectorySeparatorChar + ) + $_tauriOverride = $_tauriOverride.TrimEnd($_trimSeps) + $_legacyTauriRoot = $_legacyTauriRoot.TrimEnd($_trimSeps) + if ($_tauriOverride -ne $_legacyTauriRoot) { + Write-Host "ERROR: $envOverrideVar is not supported with --tauri." -ForegroundColor Red + Write-Host " The desktop app still uses the legacy %USERPROFILE%\.unsloth\studio root." -ForegroundColor Red + Write-Host " Run install.ps1 without --tauri for custom-root shell installs," -ForegroundColor Yellow + Write-Host " or unset the env var for default desktop installs." -ForegroundColor Yellow + throw "$envOverrideVar is not supported with --tauri." + } + } + + $defaultProfile = $null + try { $defaultProfile = [Environment]::GetFolderPath("UserProfile") } catch {} + + # LOCALAPPDATA may be unset in service / CI contexts; Join-Path would abort + # under ErrorActionPreference=Stop without this guard. + $defaultDataDir = if ($env:LOCALAPPDATA -and -not [string]::IsNullOrWhiteSpace($env:LOCALAPPDATA)) { + Join-Path $env:LOCALAPPDATA "Unsloth Studio" + } else { $null } + + if ($envOverride) { + # Tilde expansion: env vars aren't subject to it when quoted on assignment. + if ($envOverride -eq "~" -or $envOverride -like "~/*" -or $envOverride -like "~\*") { + $envOverride = (Join-Path $env:USERPROFILE $envOverride.Substring(1).TrimStart('/','\')) + } + try { + # .NET API: New-Item -Path treats brackets as wildcards and has no + # -LiteralPath in PS 5.1, so a root like C:\studio[abc] would fail. + [System.IO.Directory]::CreateDirectory($envOverride) | Out-Null + $StudioHome = (Resolve-Path -LiteralPath $envOverride).Path + } catch { + Write-Host "ERROR: $envOverrideVar=$envOverride cannot be created or accessed." -ForegroundColor Red + throw "$envOverrideVar=$envOverride cannot be created or accessed." + } + $probe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid()) + try { + # WriteAllText: literal-path safe + closes handle so Remove-Item works. + [System.IO.File]::WriteAllText($probe, "") + Remove-Item -LiteralPath $probe -Force -ErrorAction SilentlyContinue + } catch { + Write-Host "ERROR: $envOverrideVar=$StudioHome is not writable." -ForegroundColor Red + throw "$envOverrideVar=$StudioHome is not writable." + } + $StudioDataDir = Join-Path $StudioHome "share" + $StudioRedirectMode = 'env' + } elseif ($defaultProfile -and $env:USERPROFILE -and ($env:USERPROFILE -ne $defaultProfile)) { + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + $StudioDataDir = $defaultDataDir + $StudioRedirectMode = 'profile' + } else { + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + $StudioDataDir = $defaultDataDir + $StudioRedirectMode = 'default' + } $VenvDir = Join-Path $StudioHome "unsloth_studio" $Rule = [string]::new([char]0x2500, 52) @@ -378,24 +470,24 @@ function Install-UnslothStudio { [Parameter(Mandatory = $true)][string]$UnslothExePath ) - if (-not (Test-Path $UnslothExePath)) { + if (-not (Test-Path -LiteralPath $UnslothExePath)) { substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow" return } try { # Persist an absolute path in launcher scripts so shortcut working # directory changes do not break process startup. - $UnslothExePath = (Resolve-Path $UnslothExePath).Path + $UnslothExePath = (Resolve-Path -LiteralPath $UnslothExePath).Path # Escape for single-quoted embedding in generated launcher script. # This prevents runtime variable expansion for paths containing '$'. $SingleQuotedExePath = $UnslothExePath -replace "'", "''" - $localAppDataDir = $env:LOCALAPPDATA - if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) { - substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow" + # $StudioDataDir = LOCALAPPDATA\Unsloth Studio, or $StudioHome\share in env-mode. + if (-not $StudioDataDir -or [string]::IsNullOrWhiteSpace($StudioDataDir)) { + substep "DataDir path unavailable; skipped shortcut creation" "Yellow" return } - $appDir = Join-Path $localAppDataDir "Unsloth Studio" + $appDir = $StudioDataDir $launcherPs1 = Join-Path $appDir "launch-studio.ps1" $launcherVbs = Join-Path $appDir "launch-studio.vbs" $desktopDir = [Environment]::GetFolderPath("Desktop") @@ -427,23 +519,89 @@ function Install-UnslothStudio { } $iconUrl = "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth.ico" - if (-not (Test-Path $appDir)) { - New-Item -ItemType Directory -Path $appDir -Force | Out-Null + if (-not (Test-Path -LiteralPath $appDir)) { + [System.IO.Directory]::CreateDirectory($appDir) | Out-Null + } + + # Same-install discriminator: per-install opaque id written once at + # install time and read by both this launcher and the backend + # (/api/health). Replaces the older sha256(resolved $StudioHome) + # scheme to (a) avoid leaking the install path on -H 0.0.0.0 + # deployments and (b) sidestep launcher/backend canonicalization + # drift (Resolve-Path vs Path.resolve() junction handling). Lives + # at $StudioHome\share\ (not $appDir) so the backend can find it + # via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" + # regardless of mode. 32 bytes of crypto random -> 64 hex chars. + $_studioIdDir = Join-Path $StudioHome "share" + if (-not (Test-Path -LiteralPath $_studioIdDir)) { + [System.IO.Directory]::CreateDirectory($_studioIdDir) | Out-Null + } + $_studioIdFile = Join-Path $_studioIdDir "studio_install_id" + $_studioRootId = "" + if ((Test-Path -LiteralPath $_studioIdFile) -and ` + ((Get-Item -LiteralPath $_studioIdFile).Length -gt 0)) { + $_studioRootId = ([System.IO.File]::ReadAllText($_studioIdFile)).Trim() + } + if (-not $_studioRootId) { + $_idBytes = New-Object byte[] 32 + [Security.Cryptography.RandomNumberGenerator]::Create().GetBytes($_idBytes) + $_studioRootId = -join ($_idBytes | ForEach-Object { $_.ToString('x2') }) + # Atomic write: write to a temp sibling then rename, so a partial + # install cannot leave a half-written id. + $_idTmp = $_studioIdFile + ".$PID.tmp" + [System.IO.File]::WriteAllText($_idTmp, $_studioRootId) + Move-Item -LiteralPath $_idTmp -Destination $_studioIdFile -Force + } + + # Env-mode: persist UNSLOTH_STUDIO_HOME (and llama path) so fresh + # shells don't need to re-export, and bake per-install $portFile / + # $mutexName so concurrent custom-root launchers cannot serialize + # through one global mutex on 8888..8908. Default installs get an + # empty prefix to match pre-PR behavior. + $studioHomeExport = if ($StudioRedirectMode -eq 'env') { + # When override == legacy default, llama.cpp stays at + # ~/.unsloth/llama.cpp (one shared build). Canonicalize the + # legacy side so the comparison survives path normalization. + $_legacyStudio = Join-Path $env:USERPROFILE ".unsloth\studio" + if (Test-Path -LiteralPath $_legacyStudio -PathType Container) { + $_legacyStudio = (Resolve-Path -LiteralPath $_legacyStudio).Path + } + $_llamaPath = if ($StudioHome -eq $_legacyStudio) { + Join-Path $env:USERPROFILE ".unsloth\llama.cpp" + } else { + Join-Path $StudioHome "llama.cpp" + } + $_sq = $StudioHome -replace "'", "''" + $_llama = $_llamaPath -replace "'", "''" + $_appDirSq = $appDir -replace "'", "''" + $_appBytes = [Text.Encoding]::UTF8.GetBytes($appDir) + $_appHash = ([BitConverter]::ToString( + [Security.Cryptography.SHA256]::Create().ComputeHash($_appBytes) + ) -replace '-', '').Substring(0, 16) + # UNSLOTH_LLAMA_CPP_PATH is a pre-existing user override; only default if unset. + "`$env:UNSLOTH_STUDIO_HOME = '$_sq'`nif (-not `$env:UNSLOTH_LLAMA_CPP_PATH) {`n `$env:UNSLOTH_LLAMA_CPP_PATH = '$_llama'`n}`n`$portFile = '$_appDirSq\studio.port'`n`$mutexName = 'Local\UnslothStudioLauncher-$_appHash'`n" + } else { + "`$portFile = `$null`n`$mutexName = 'Local\UnslothStudioLauncher'`n" } $launcherContent = @" -`$ErrorActionPreference = 'Stop' +$studioHomeExport`$ErrorActionPreference = 'Stop' `$basePort = 8888 `$maxPortOffset = 20 `$timeoutSec = 60 `$pollIntervalMs = 1000 +`$_ExpectedStudioRootId = '$_studioRootId' function Test-StudioHealth { param([Parameter(Mandatory = `$true)][int]`$Port) try { `$url = "http://127.0.0.1:`$Port/api/health" `$resp = Invoke-RestMethod -Uri `$url -TimeoutSec 1 -Method Get - return (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend') + if (-not (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend')) { return `$false } + # why: verify the backend belongs to THIS install via the install-time + # hex digest; raw path is not leaked over /api/health. + if (`$_ExpectedStudioRootId -and `$resp.studio_root_id -ne `$_ExpectedStudioRootId) { return `$false } + return `$true } catch { return `$false } @@ -469,6 +627,17 @@ function Get-CandidatePorts { } function Find-HealthyStudioPort { + if (`$portFile) { + if (Test-Path -LiteralPath `$portFile) { + `$cached = Get-Content -LiteralPath `$portFile -ErrorAction SilentlyContinue | Select-Object -First 1 + if (`$cached -match '^\d+`$') { + `$cachedPort = [int]`$cached + if (Test-StudioHealth -Port `$cachedPort) { return `$cachedPort } + Remove-Item -LiteralPath `$portFile -Force -ErrorAction SilentlyContinue + } + } + return `$null + } foreach (`$candidate in (Get-CandidatePorts)) { if (Test-StudioHealth -Port `$candidate) { return `$candidate @@ -522,7 +691,7 @@ if (`$existingPort) { exit 0 } -`$launchMutex = [System.Threading.Mutex]::new(`$false, 'Local\UnslothStudioLauncher') +`$launchMutex = [System.Threading.Mutex]::new(`$false, `$mutexName) `$haveMutex = `$false try { try { @@ -552,7 +721,9 @@ try { } catch {} exit 1 } - `$studioCommand = '& "' + `$studioExe + '" studio -p ' + `$launchPort + # Single-quote the path in the child -Command so `$` / backtick in custom + # roots don't get reparsed; double any apostrophes so 'O''Brien' survives. + `$studioCommand = "& '" + (`$studioExe -replace "'", "''") + "' studio -p " + `$launchPort `$launchArgs = @( '-NoExit', '-NoProfile', @@ -576,9 +747,13 @@ try { `$browserOpened = `$false `$deadline = (Get-Date).AddSeconds(`$timeoutSec) while ((Get-Date) -lt `$deadline) { - `$healthyPort = Find-HealthyStudioPort - if (`$healthyPort) { - Start-Process "http://localhost:`$healthyPort" + if (Test-StudioHealth -Port `$launchPort) { + if (`$portFile) { + try { + [System.IO.File]::WriteAllText(`$portFile, "`$launchPort`n") + } catch {} + } + Start-Process "http://localhost:`$launchPort" `$browserOpened = `$true break } @@ -613,19 +788,19 @@ cmd = "powershell -NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File " shell.Run cmd, 0, False "@ # WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths. - Set-Content -Path $launcherVbs -Value $vbsContent -Encoding Unicode -Force + Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force # Prefer bundled icon from local clone/dev installs. # If not available, best-effort download from raw GitHub. # We only attach the icon if the resulting file has a valid ICO header. $hasValidIcon = $false - if ($bundledIcon -and (Test-Path $bundledIcon)) { + if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) { try { - Copy-Item -Path $bundledIcon -Destination $iconPath -Force + Copy-Item -LiteralPath $bundledIcon -Destination $iconPath -Force } catch { Write-Host "[DEBUG] Error copying bundled icon: $($_.Exception.Message)" -ForegroundColor DarkGray } - } elseif (-not (Test-Path $iconPath)) { + } elseif (-not (Test-Path -LiteralPath $iconPath)) { try { Invoke-WebRequest -Uri $iconUrl -OutFile $iconPath -UseBasicParsing } catch { @@ -633,7 +808,7 @@ shell.Run cmd, 0, False } } - if (Test-Path $iconPath) { + if (Test-Path -LiteralPath $iconPath) { try { $bytes = [System.IO.File]::ReadAllBytes($iconPath) if ( @@ -645,14 +820,21 @@ shell.Run cmd, 0, False ) { $hasValidIcon = $true } else { - Remove-Item $iconPath -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath $iconPath -Force -ErrorAction SilentlyContinue } } catch { Write-Host "[DEBUG] Error validating or removing icon: $($_.Exception.Message)" -ForegroundColor DarkGray - Remove-Item $iconPath -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath $iconPath -Force -ErrorAction SilentlyContinue } } + # Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts + # that may point at a deleted workspace; launcher + icon stay. + if ($StudioRedirectMode -eq 'env') { + substep "wrote launcher at $launcherPs1 (persistent shortcuts skipped in env-override mode)" + return + } + $wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe" $shortcutArgs = "//B //Nologo `"$launcherVbs`"" @@ -850,8 +1032,9 @@ shell.Run cmd, 0, False # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. Write-TauriLog "STEP" "Creating virtual environment" - if (-not (Test-Path $StudioHome)) { - New-Item -ItemType Directory -Path $StudioHome -Force | Out-Null + if (-not (Test-Path -LiteralPath $StudioHome)) { + # .NET API: New-Item -Path treats brackets as wildcards. + [System.IO.Directory]::CreateDirectory($StudioHome) | Out-Null } $VenvPython = Join-Path $VenvDir "Scripts\python.exe" @@ -865,11 +1048,13 @@ shell.Run cmd, 0, False $stamp = Get-Date -Format "yyyyMMddHHmmss" $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID" $suffix = 0 - while (Test-Path $candidate) { + # -LiteralPath: a custom $StudioHome may contain [ ] * ? which + # plain Test-Path / Move-Item would interpret as wildcards. + while (Test-Path -LiteralPath $candidate) { $suffix++ $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix" } - Move-Item -Path $ExistingDir -Destination $candidate -ErrorAction Stop + Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop $script:StudioVenvRollbackDir = $candidate $script:StudioVenvRollbackTarget = $ExistingDir $script:StudioVenvRollbackActive = $true @@ -880,16 +1065,16 @@ shell.Run cmd, 0, False if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir $target = $script:StudioVenvRollbackTarget - if (-not $backup -or -not (Test-Path $backup)) { + if (-not $backup -or -not (Test-Path -LiteralPath $backup)) { $script:StudioVenvRollbackActive = $false return } substep "restoring previous environment after failed install..." "Yellow" try { - if (Test-Path $target) { - Remove-Item -Recurse -Force $target -ErrorAction SilentlyContinue + if (Test-Path -LiteralPath $target) { + Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue } - Move-Item -Path $backup -Destination $target -Force -ErrorAction Stop + Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop substep "restored previous environment" $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null @@ -902,14 +1087,29 @@ shell.Run cmd, 0, False function Complete-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir - if ($backup -and (Test-Path $backup)) { - Remove-Item -Recurse -Force $backup -ErrorAction SilentlyContinue + if ($backup -and (Test-Path -LiteralPath $backup)) { + Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue } $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null } - if (Test-Path $VenvPython) { + if (Test-Path -LiteralPath $VenvPython) { + # why: matching guard to the .venv branch below -- in env-mode + # $StudioHome is a user-chosen workspace, so refuse to nuke an + # existing $StudioHome\unsloth_studio that lacks Studio sentinels. + # -PathType Leaf rejects a directory at the sentinel path. Accept the + # in-VENV ownership marker so partial-install retries are not blocked. + if ( + $StudioRedirectMode -eq 'env' -and + -not (Test-Path -LiteralPath (Join-Path $VenvDir ".unsloth-studio-owned") -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "share\studio.conf") -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "bin\unsloth.exe") -PathType Leaf) + ) { + Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red + Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow + throw "Refusing to delete non-Studio venv at $VenvDir" + } # New layout already exists -- replace only after preserving rollback copy. substep "preserving existing environment for rollback..." try { @@ -918,8 +1118,13 @@ shell.Run cmd, 0, False Write-Host "[ERROR] Could not prepare existing environment for reinstall: $($_.Exception.Message)" -ForegroundColor Red return (Exit-InstallFailure "Could not prepare existing environment for reinstall") } - } elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) { - # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating + } elseif ( + $StudioRedirectMode -ne 'env' ` + -and (Test-Path -LiteralPath (Join-Path $StudioHome ".venv\Scripts\python.exe")) + ) { + # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating. + # Skip in env-mode so we don't blow away an unrelated .venv at the + # workspace root (e.g. user's existing project Python venv). $OldVenv = Join-Path $StudioHome ".venv" $OldPy = Join-Path $OldVenv "Scripts\python.exe" substep "found legacy Studio environment, validating..." @@ -936,24 +1141,29 @@ shell.Run cmd, 0, False $ErrorActionPreference = $prevEAP2 if ($legacyOk) { substep "legacy environment is healthy -- migrating..." - Move-Item -Path $OldVenv -Destination $VenvDir -Force + Move-Item -LiteralPath $OldVenv -Destination $VenvDir -Force substep "moved .venv -> unsloth_studio" $_Migrated = $true } else { substep "legacy environment failed validation -- creating fresh environment" "Yellow" $invalidVenv = Join-Path $StudioHome (".venv.invalid.{0}.{1}" -f (Get-Date -Format "yyyyMMddHHmmss"), $PID) - Move-Item -Path $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue + Move-Item -LiteralPath $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue } - } elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) { - # CWD-relative venv from old install.ps1 -- migrate to absolute path + } elseif ( + $StudioRedirectMode -ne 'env' ` + -and (Test-Path -LiteralPath (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) + ) { + # CWD-relative venv from old install.ps1 -> migrate to absolute path. + # Skip in env-mode so we don't relocate the default-install venv into + # the workspace root. $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio" substep "found CWD-relative Studio environment, migrating to $VenvDir..." - Move-Item -Path $CwdVenv -Destination $VenvDir -Force + Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio" $_Migrated = $true } - if (-not (Test-Path $VenvPython)) { + if (-not (Test-Path -LiteralPath $VenvPython)) { step "venv" "creating Python $($DetectedPython.Version) virtual environment" substep "$VenvDir" $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } @@ -966,6 +1176,13 @@ shell.Run cmd, 0, False substep "$VenvDir" } + # Mark the freshly-created venv as Studio-owned so a partial install can be + # repaired by re-running install.ps1; the env-mode deletion guard above + # accepts this marker as the primary sentinel. + if (Test-Path -LiteralPath $VenvDir -PathType Container) { + try { [System.IO.File]::WriteAllText((Join-Path $VenvDir ".unsloth-studio-owned"), "") } catch {} + } + # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── $HasNvidiaSmi = $false $NvidiaSmiExe = $null @@ -1054,7 +1271,7 @@ shell.Run cmd, 0, False if ($StudioLocalInstall -and (Test-Path (Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"))) { return Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt" } - $installed = Get-ChildItem -Path $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue | + $installed = Get-ChildItem -LiteralPath $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue | Where-Object { $_.FullName -like "*studio*backend*requirements*no-torch-runtime.txt" } | Select-Object -ExpandProperty FullName -First 1 return $installed @@ -1192,23 +1409,25 @@ shell.Run cmd, 0, False foreach ($rel in $overlayMap.Keys) { $src = Join-Path $scriptDir $rel $dst = Join-Path $VenvDir $overlayMap[$rel] - if (-not (Test-Path $src)) { continue } + # -LiteralPath: $VenvDir derives from $StudioHome which may + # contain [ ] * ? when the user overrode UNSLOTH_STUDIO_HOME. + if (-not (Test-Path -LiteralPath $src)) { continue } $dstParent = Split-Path -Parent $dst - if (-not (Test-Path $dstParent)) { + if (-not (Test-Path -LiteralPath $dstParent)) { Write-Host "[WARN] Overlay target dir missing: $dstParent; studio setup may use stale bundled file" -ForegroundColor Yellow continue } try { - if (-not (Test-Path $dst)) { + if (-not (Test-Path -LiteralPath $dst)) { # Backfill: target file missing but parent dir exists. - Copy-Item $src $dst -Force + Copy-Item -LiteralPath $src -Destination $dst -Force substep ("backfilled bundled " + (Split-Path -Leaf $rel)) } else { # Hash-compare so re-runs are no-ops when files already match. - $srcHash = (Get-FileHash $src -Algorithm SHA256).Hash - $dstHash = (Get-FileHash $dst -Algorithm SHA256).Hash + $srcHash = (Get-FileHash -LiteralPath $src -Algorithm SHA256).Hash + $dstHash = (Get-FileHash -LiteralPath $dst -Algorithm SHA256).Hash if ($srcHash -ne $dstHash) { - Copy-Item $src $dst -Force + Copy-Item -LiteralPath $src -Destination $dst -Force substep ("applied bundled " + (Split-Path -Leaf $rel)) } } @@ -1225,7 +1444,8 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Running studio setup" step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" - if (-not (Test-Path $UnslothExe)) { + if (-not (Test-Path -LiteralPath $UnslothExe)) { + Write-TauriLog "ERROR" "unsloth CLI was not installed correctly" Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow @@ -1250,6 +1470,15 @@ shell.Run cmd, 0, False # Use 'studio setup' (not 'studio update') because 'update' pops # SKIP_STUDIO_BASE, which would cause redundant package reinstallation # and bypass the fast-path version check from PR #4667. + # Propagate UNSLOTH_STUDIO_HOME only for env-override installs; otherwise + # an inherited value would put llama.cpp in the wrong place. + $previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME + $hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome) + if ($StudioRedirectMode -eq 'env') { + $env:UNSLOTH_STUDIO_HOME = $StudioHome + } else { + Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue + } $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" @@ -1257,6 +1486,11 @@ shell.Run cmd, 0, False & $UnslothExe @studioArgs $setupExit = $LASTEXITCODE } finally { + if ($hadPreviousUnslothStudioHome) { + $env:UNSLOTH_STUDIO_HOME = $previousUnslothStudioHome + } else { + Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue + } Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue } if ($setupExit -ne 0) { @@ -1301,20 +1535,32 @@ shell.Run cmd, 0, False } } catch { } $ShimDir = Join-Path $StudioHome "bin" - New-Item -ItemType Directory -Force -Path $ShimDir | Out-Null + [System.IO.Directory]::CreateDirectory($ShimDir) | Out-Null $ShimExe = Join-Path $ShimDir "unsloth.exe" + # Fatal preflight outside the lock-handling try/catch -- a directory at + # the shim path must not be downgraded to "Continuing with the existing + # launcher", or the install finishes with no usable shim. + if (Test-Path -LiteralPath $ShimExe -PathType Container) { + Write-Host "[ERROR] Cannot create unsloth launcher: $ShimExe is a directory." -ForegroundColor Red + Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow + throw "Cannot create unsloth launcher: $ShimExe is a directory." + } # try/catch: if unsloth.exe is locked (Studio running), keep the old shim. $shimUpdated = $false try { - if (Test-Path $ShimExe) { Remove-Item $ShimExe -Force -ErrorAction Stop } + if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop } try { + # New-Item -ItemType HardLink does NOT accept -LiteralPath in any + # PowerShell version, so use -Path. Wildcards in $ShimExe (e.g. + # brackets in custom roots) glob-expand here and fall through to + # the Copy-Item -LiteralPath fallback below. New-Item -ItemType HardLink -Path $ShimExe -Target $UnslothExe -ErrorAction Stop | Out-Null } catch { - Copy-Item -Path $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy + Copy-Item -LiteralPath $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy } $shimUpdated = $true } catch { - if (Test-Path $ShimExe) { + if (Test-Path -LiteralPath $ShimExe) { Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow @@ -1325,10 +1571,13 @@ shell.Run cmd, 0, False Write-Host " Launch unsloth studio directly via '$UnslothExe' until the next successful install." -ForegroundColor Yellow } } - # Only add to PATH when the launcher actually exists on disk. + # Add to PATH only when launcher exists. Env-mode: session-only export, + # no registry change (workspace path may be deleted later). $pathAdded = $false - if (Test-Path $ShimExe) { - $pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend' + if (Test-Path -LiteralPath $ShimExe) { + if ($StudioRedirectMode -ne 'env') { + $pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend' + } } if ($shimUpdated -and $pathAdded) { step "path" "added unsloth launcher to PATH" @@ -1336,12 +1585,20 @@ shell.Run cmd, 0, False Refresh-SessionPath # sync current session with registry Complete-StudioVenvRollback + # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy + # User PATH entry (Machine > User > current $env:Path) would win. + if ($StudioRedirectMode -eq 'env' -and (Test-Path -LiteralPath $ShimExe)) { + $env:Path = "$ShimDir;$env:Path" + step "path" "exported $ShimDir for this session (no registry PATH change in env-override mode)" + } + # ── Tauri mode: done, skip shortcuts and auto-launch ── if ($TauriMode) { Write-TauriLog "DONE" "" return } + # New-StudioShortcuts gates the .lnk shortcuts on env-mode internally. New-StudioShortcuts -UnslothExePath $UnslothExe # In interactive terminals, ask the user before starting Studio. @@ -1360,8 +1617,21 @@ shell.Run cmd, 0, False } } else { step "launch" "manual commands:" - substep "& `"$VenvDir\Scripts\Activate.ps1`"" - substep "unsloth studio -p 8888" + # Single-quote the printed paths so $-vars / backticks in custom roots + # do not reparse when the user pastes the command. + $_actLiteral = "'" + ((Join-Path $VenvDir "Scripts\Activate.ps1") -replace "'", "''") + "'" + if ($StudioRedirectMode -eq 'env') { + # Env-mode skips registry PATH; print the absolute shim path. + $_shim = Join-Path $StudioHome "bin\unsloth.exe" + $_shimLiteral = "'" + ($_shim -replace "'", "''") + "'" + substep "& $_shimLiteral studio -p 8888" + substep "or activate env first:" + substep "& $_actLiteral" + substep "unsloth studio -p 8888" + } else { + substep "& $_actLiteral" + substep "unsloth studio -p 8888" + } substep "(add -H 0.0.0.0 to allow network / cloud access)" Write-Host "" } diff --git a/install.sh b/install.sh index 1d21117d16..ec47d016eb 100755 --- a/install.sh +++ b/install.sh @@ -6,6 +6,12 @@ # Usage (no-torch): ./install.sh --no-torch (skip PyTorch, GGUF-only mode) # Usage (test): ./install.sh --package roland-sloth (install a different package name) # Usage (py): ./install.sh --python 3.12 (override auto-detected Python version) +# +# Env vars (priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > HOME-redirect > default): +# UNSLOTH_STUDIO_HOME=/abs/path -> install under that path +# STUDIO_HOME=/abs/path -> alias, same effect (UNSLOTH_STUDIO_HOME wins) +# (DATA_DIR + unsloth CLI shim nest inside; no shell rc-file append.) +# Default ($HOME/.unsloth/studio) is preserved when no env var is set. set -e # ── Output style (aligned with studio/setup.sh) ── @@ -66,6 +72,56 @@ if [ "$_VERBOSE" = true ]; then export UNSLOTH_VERBOSE=1 fi +# Custom Studio roots are not supported with --tauri (desktop app still +# resolves ~/.unsloth/studio). Pass through if the override == legacy default. +if [ "$TAURI_MODE" = true ]; then + _tauri_override_var="" + _tauri_override="${UNSLOTH_STUDIO_HOME:-}" + if [ -n "$_tauri_override" ]; then + _tauri_override_var="UNSLOTH_STUDIO_HOME" + else + _tauri_override="${STUDIO_HOME:-}" + [ -n "$_tauri_override" ] && _tauri_override_var="STUDIO_HOME" + fi + # Strip whitespace so " " is treated as unset (matches Python .strip()). + _tauri_override=$(printf '%s' "$_tauri_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + if [ -n "$_tauri_override" ]; then + case "$_tauri_override" in + "~") _tauri_override="$HOME" ;; + "~/"*) _tauri_override="$HOME/${_tauri_override#'~/'}" ;; + esac + # Canonicalize both sides (CDPATH=, -P) so a CDPATH-set env or + # symlinked $HOME doesn't break the legacy-equality comparison. + if [ -d "$_tauri_override" ]; then + _tauri_override_abs=$(CDPATH= cd -P -- "$_tauri_override" 2>/dev/null && pwd -P) \ + || _tauri_override_abs="$_tauri_override" + else + _tauri_override_abs="$_tauri_override" + fi + # Strip trailing separators so ".../studio/" matches ".../studio". + while [ "$_tauri_override_abs" != "/" ] \ + && [ "${_tauri_override_abs%/}" != "$_tauri_override_abs" ]; do + _tauri_override_abs=${_tauri_override_abs%/} + done + _tauri_legacy_root="$HOME/.unsloth/studio" + if [ -d "$_tauri_legacy_root" ]; then + _tauri_legacy_root=$(CDPATH= cd -P -- "$_tauri_legacy_root" 2>/dev/null && pwd -P) \ + || _tauri_legacy_root="$HOME/.unsloth/studio" + fi + while [ "$_tauri_legacy_root" != "/" ] \ + && [ "${_tauri_legacy_root%/}" != "$_tauri_legacy_root" ]; do + _tauri_legacy_root=${_tauri_legacy_root%/} + done + if [ "$_tauri_override_abs" != "$_tauri_legacy_root" ]; then + echo "ERROR: $_tauri_override_var is not supported with --tauri." >&2 + echo " The desktop app still uses the legacy ~/.unsloth/studio root." >&2 + echo " Run install.sh without --tauri for custom-root shell installs," >&2 + echo " or unset the env var for default desktop installs." >&2 + exit 1 + fi + fi +fi + _is_verbose() { [ "${UNSLOTH_VERBOSE:-0}" = "1" ] } @@ -219,7 +275,67 @@ _tauri_gpu_branch() { } PYTHON_VERSION="" # resolved after platform detection -STUDIO_HOME="$HOME/.unsloth/studio" + +# Resolve install destinations: env override, HOME-redirect (best-effort +# via getent/dscl), or default. Env-var priority: UNSLOTH_STUDIO_HOME wins +# over STUDIO_HOME (the more specific signal beats the generic alias). +_resolve_studio_destinations() { + _override_var="" + _override="${UNSLOTH_STUDIO_HOME:-}" + if [ -n "$_override" ]; then + _override_var="UNSLOTH_STUDIO_HOME" + else + _override="${STUDIO_HOME:-}" + [ -n "$_override" ] && _override_var="STUDIO_HOME" + fi + # Strip surrounding whitespace so " " is treated as unset (matches the + # Python resolvers' .strip()), preventing install/runtime layout drift. + _override=$(printf '%s' "$_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + # Tilde expansion: env vars are not subject to it when quoted on assignment. + case "$_override" in + "~") _override="$HOME" ;; + "~/"*) _override="$HOME/${_override#'~/'}" ;; + esac + if [ -n "$_override" ]; then + mkdir -p -- "$_override" 2>/dev/null || { echo "ERROR: $_override_var=$_override cannot be created." >&2; exit 1; } + [ -w "$_override" ] || { echo "ERROR: $_override_var=$_override is not writable." >&2; exit 1; } + STUDIO_HOME="$(CDPATH= cd -P -- "$_override" && pwd -P)" || exit 1 + DATA_DIR="$STUDIO_HOME/share" + _LOCAL_BIN="$STUDIO_HOME/bin" + _STUDIO_HOME_REDIRECT=env + substep "custom $_override_var=$STUDIO_HOME" + return 0 + fi + _default_home="" + if command -v getent >/dev/null 2>&1; then + _default_home=$(getent passwd "${USER:-$(whoami)}" 2>/dev/null | cut -d: -f6) + elif [ "$(uname)" = "Darwin" ] && command -v dscl >/dev/null 2>&1; then + _default_home=$(dscl . -read "/Users/${USER:-$(whoami)}" NFSHomeDirectory 2>/dev/null | awk '{print $2}') + fi + # Canonicalize both sides so a trailing slash on $HOME (or symlink mismatch + # with passwd-DB output) doesn't misfire the redirection branch. + _home_canon="$HOME" + if [ -d "$_home_canon" ]; then + _home_canon=$(CDPATH= cd -P -- "$_home_canon" 2>/dev/null && pwd -P) || _home_canon="$HOME" + fi + _default_home_canon="$_default_home" + if [ -n "$_default_home_canon" ] && [ -d "$_default_home_canon" ]; then + _default_home_canon=$(CDPATH= cd -P -- "$_default_home_canon" 2>/dev/null && pwd -P) || _default_home_canon="$_default_home" + fi + if [ -n "$_default_home_canon" ] && [ "$_home_canon" != "$_default_home_canon" ]; then + STUDIO_HOME="$HOME/.unsloth/studio" + DATA_DIR="$HOME/.local/share/unsloth" + _LOCAL_BIN="$HOME/.local/bin" + _STUDIO_HOME_REDIRECT=home + substep "HOME redirected ($HOME); install follows \$HOME" + return 0 + fi + STUDIO_HOME="$HOME/.unsloth/studio" + DATA_DIR="$HOME/.local/share/unsloth" + _LOCAL_BIN="$HOME/.local/bin" + _STUDIO_HOME_REDIRECT=default +} +_resolve_studio_destinations VENV_DIR="$STUDIO_HOME/unsloth_studio" _VENV_ROLLBACK_DIR="" _VENV_ROLLBACK_TARGET="$VENV_DIR" @@ -383,23 +499,65 @@ create_studio_shortcuts() { _css_exe_dir=$(cd "$(dirname "$_css_exe")" && pwd) _css_exe="$_css_exe_dir/$(basename "$_css_exe")" - _css_data_dir="$HOME/.local/share/unsloth" + _css_data_dir="$DATA_DIR" _css_launcher="$_css_data_dir/launch-studio.sh" _css_icon_png="$_css_data_dir/unsloth-studio.png" _css_gem_png="$_css_data_dir/unsloth-gem.png" mkdir -p "$_css_data_dir" + # Same-install discriminator: per-install opaque id written once at install + # time and read by both this launcher and the backend (/api/health). Replaces + # the older sha256(canonical $STUDIO_HOME) scheme to (a) avoid leaking the + # install path on -H 0.0.0.0 deployments and (b) sidestep launcher/backend + # canonicalization drift (cd -P vs Path.resolve() symlink/junction handling). + # Lives at $STUDIO_HOME/share/ (not $DATA_DIR) so the backend can find it + # via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" regardless of + # mode (in env-mode $STUDIO_HOME/share == $DATA_DIR; in default mode they + # diverge but the backend only knows the studio_root). 32 bytes of urandom + # -> 64 hex chars, byte-compatible with the prior digest so launcher + # placeholder, _check_health, and tests stay length-agnostic. + _css_id_dir="$STUDIO_HOME/share" + mkdir -p "$_css_id_dir" + _css_id_file="$_css_id_dir/studio_install_id" + if [ ! -s "$_css_id_file" ]; then + if [ -r /dev/urandom ]; then + _css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d ' \n') + fi + if [ -z "${_css_new_id:-}" ] && command -v python3 >/dev/null 2>&1; then + _css_new_id=$(python3 -c 'import secrets; print(secrets.token_hex(32))' 2>/dev/null) + fi + if [ -z "${_css_new_id:-}" ]; then + echo "[WARN] Cannot create launcher: no entropy source for studio_install_id" >&2 + return 1 + fi + # Atomic write so a partial install can't leave a half-written id. + _css_id_tmp="$_css_id_file.$$.tmp" + printf '%s' "$_css_new_id" > "$_css_id_tmp" \ + && mv "$_css_id_tmp" "$_css_id_file" + chmod 600 "$_css_id_file" 2>/dev/null || true + unset _css_new_id _css_id_tmp + fi + _css_studio_root_id=$(cat "$_css_id_file" 2>/dev/null) + if [ -z "$_css_studio_root_id" ]; then + echo "[WARN] Cannot create launcher: failed to read $_css_id_file" >&2 + return 1 + fi + _css_is_env_mode=false + [ "$_STUDIO_HOME_REDIRECT" = "env" ] && _css_is_env_mode=true + # ── Write launcher script ── - # The launcher is Bash (not POSIX sh). - # We write it with a placeholder and substitute the exe path via sed. + # Single-quoted heredoc; @@DATA_DIR@@, @@STUDIO_ROOT_ID@@, and + # @@INSTALLED_IS_ENV_MODE@@ are substituted via sed below. cat > "$_css_launcher" << 'LAUNCHER_EOF' #!/usr/bin/env bash # Unsloth Studio Launcher # Auto-generated by install.sh -- do not edit manually. set -euo pipefail -DATA_DIR="$HOME/.local/share/unsloth" +DATA_DIR='@@DATA_DIR@@' +_EXPECTED_STUDIO_ROOT_ID='@@STUDIO_ROOT_ID@@' +_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@' # Read exe path from config written at install time. # Sourcing is safe: the config file is written by install.sh, not user input. @@ -416,7 +574,23 @@ MAX_PORT_OFFSET=20 TIMEOUT_SEC=60 POLL_INTERVAL_SEC=1 LOG_FILE="$DATA_DIR/studio.log" +# why: in env-override mode multiple installs share an OS user; namespace the +# lock and remember our own healthy port so we never attach to an unrelated +# Studio listening on the global 8888..8908 range. LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock" +PORT_FILE="" +# why: gate on the install-time mode (baked above) instead of the runtime env +# var; sourcing a custom-root studio.conf in shell must not flip a default-mode +# launcher into env-mode behavior with stale state. +if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then + if command -v cksum >/dev/null 2>&1; then + _LOCK_KEY=$(printf '%s' "$DATA_DIR" | cksum | awk '{print $1}') + else + _LOCK_KEY="" + fi + [ -n "$_LOCK_KEY" ] && LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u)-${_LOCK_KEY}.lock" + PORT_FILE="$DATA_DIR/studio.port" +fi # ── HTTP GET helper (supports curl and wget) ── _http_get() { @@ -435,10 +609,20 @@ _check_health() { _port=$1 _resp=$(_http_get "http://127.0.0.1:$_port/api/health") || return 1 case "$_resp" in - *'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) return 0 ;; - *'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) return 0 ;; + *'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) ;; + *'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) ;; + *) return 1 ;; esac - return 1 + # why: verify the backend belongs to THIS install. Baked hex digest avoids + # JSON-escape mismatches on paths with `\`/`"` and avoids leaking the raw + # install path to unauthenticated callers. + if [ -n "$_EXPECTED_STUDIO_ROOT_ID" ]; then + case "$_resp" in + *"\"studio_root_id\":\"$_EXPECTED_STUDIO_ROOT_ID\""*|*"\"studio_root_id\": \"$_EXPECTED_STUDIO_ROOT_ID\""*) return 0 ;; + *) return 1 ;; + esac + fi + return 0 } # ── Port scanning ── @@ -461,6 +645,25 @@ _candidate_ports() { } _find_healthy_port() { + if [ -n "$PORT_FILE" ] && [ -f "$PORT_FILE" ]; then + # why: env-mode installs only attach to a port we previously launched + # ourselves; never to a sibling Studio that happens to be healthy. + _p=$(cat "$PORT_FILE" 2>/dev/null || true) + case "$_p" in + ''|*[!0-9]*) ;; + *) + if _check_health "$_p"; then + echo "$_p" + return 0 + fi + rm -f "$PORT_FILE" + ;; + esac + return 1 + fi + if [ -n "$PORT_FILE" ]; then + return 1 + fi for _p in $(_candidate_ports | sort -un); do if _check_health "$_p"; then echo "$_p" @@ -611,6 +814,7 @@ if [ -t 1 ]; then _obwr_deadline=$(($(date +%s) + TIMEOUT_SEC)) while [ "$(date +%s)" -lt "$_obwr_deadline" ]; do if _check_health "$_launch_port"; then + [ -n "$PORT_FILE" ] && printf '%s\n' "$_launch_port" > "$PORT_FILE" 2>/dev/null || true _release_lock _open_browser "http://localhost:$_launch_port" exit 0 @@ -634,6 +838,7 @@ else _deadline=$(($(date +%s) + TIMEOUT_SEC)) while [ "$(date +%s)" -lt "$_deadline" ]; do if _check_health "$_launch_port"; then + [ -n "$PORT_FILE" ] && printf '%s\n' "$_launch_port" > "$PORT_FILE" 2>/dev/null || true _open_browser "http://localhost:$_launch_port" exit 0 fi @@ -646,13 +851,62 @@ else fi LAUNCHER_EOF + # why: bake non-user-controlled placeholders FIRST so a literal + # `@@STUDIO_ROOT_ID@@` inside $DATA_DIR cannot be rewritten below. + sed -e "s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g" \ + -e "s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g" \ + "$_css_launcher" > "$_css_launcher.tmp" \ + && mv "$_css_launcher.tmp" "$_css_launcher" + + # Env-mode bakes an absolute DATA_DIR (root fixed at install time); + # default / HOME-redirect keeps the literal $HOME/.local/share/unsloth + # so behavior is byte-identical to pre-override. + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + # Two-stage escape: (1) `'` -> `'\''` for shell single-quote embedding, + # (2) backslash/&/| escape so the value survives the s|...|VALUE| sed + # below. Verified end-to-end with apostrophes, spaces, &, |, $. + _sq_escaped=$(printf '%s' "$DATA_DIR" | sed "s/'/'\\\\''/g") + _sed_safe=$(printf '%s' "$_sq_escaped" | sed 's/[\\&|]/\\&/g') + sed "s|@@DATA_DIR@@|$_sed_safe|g" "$_css_launcher" > "$_css_launcher.tmp" \ + && mv "$_css_launcher.tmp" "$_css_launcher" + else + sed "s|DATA_DIR='@@DATA_DIR@@'|DATA_DIR=\"\$HOME/.local/share/unsloth\"|" \ + "$_css_launcher" > "$_css_launcher.tmp" \ + && mv "$_css_launcher.tmp" "$_css_launcher" + fi + chmod +x "$_css_launcher" - # Write the exe path to a separate conf file sourced by the launcher. - # Using single-quote wrapping with the standard '\'' escape for any - # embedded apostrophes. This avoids all sed metacharacter issues. + # studio.conf: exe path + (env-mode only) persisted env vars so fresh + # shells launch the right install without re-exporting. _css_quoted_exe=$(printf '%s' "$_css_exe" | sed "s/'/'\\\\''/g") - printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" > "$_css_data_dir/studio.conf" + { + printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + # When an override resolves to the legacy default, llama.cpp + # still lives at ~/.unsloth/llama.cpp (one shared build). + # Canonicalize the legacy side so a symlinked $HOME doesn't + # break the comparison. + _css_legacy_studio="$HOME/.unsloth/studio" + if [ -d "$_css_legacy_studio" ]; then + _css_legacy_studio=$(CDPATH= cd -P -- "$_css_legacy_studio" 2>/dev/null && pwd -P) \ + || _css_legacy_studio="$HOME/.unsloth/studio" + fi + if [ "$STUDIO_HOME" = "$_css_legacy_studio" ]; then + _css_llama_path="$HOME/.unsloth/llama.cpp" + else + _css_llama_path="$STUDIO_HOME/llama.cpp" + fi + _css_quoted_home=$(printf '%s' "$STUDIO_HOME" | sed "s/'/'\\\\''/g") + _css_quoted_llama=$(printf '%s' "$_css_llama_path" | sed "s/'/'\\\\''/g") + printf '%s\n' "export UNSLOTH_STUDIO_HOME='$_css_quoted_home'" + # UNSLOTH_LLAMA_CPP_PATH is a pre-existing user-controlled + # llama.cpp dir override; only default it if unset. + printf '%s\n' 'if [ -z "${UNSLOTH_LLAMA_CPP_PATH:-}" ]; then' + printf '%s\n' " export UNSLOTH_LLAMA_CPP_PATH='$_css_quoted_llama'" + printf '%s\n' 'fi' + fi + } > "$_css_data_dir/studio.conf" # ── Icon: try bundled, then download ── # rounded-512.png used for both Linux and macOS icons @@ -698,6 +952,14 @@ LAUNCHER_EOF fi # ── Platform-specific shortcuts ── + # Env-mode installs are workspace-scoped: skip persistent desktop / + # Start-Menu / dock launchers that may point at a deleted workspace. + # Runtime launcher + studio.conf + icon are still written above. + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + substep "wrote launcher at $_css_launcher (persistent shortcuts skipped in env-override mode)" + return 0 + fi + _css_created=0 if [ "$_css_os" = "linux" ]; then @@ -775,11 +1037,18 @@ DESKTOP_EOF PLIST_EOF - # Executable stub - cat > "$_css_macos_dir/launch-studio" << STUB_EOF + # Executable stub: same single-quoted-heredoc + sed-substitute + # pattern as launch-studio.sh so $-vars in $_css_data_dir don't + # expand at .app launch time. + _css_sq_dir=$(printf '%s' "$_css_data_dir" | sed "s/'/'\\\\''/g") + _css_sed_dir=$(printf '%s' "$_css_sq_dir" | sed 's/[\\&|]/\\&/g') + cat > "$_css_macos_dir/launch-studio" << 'STUB_EOF' #!/bin/sh -exec "$HOME/.local/share/unsloth/launch-studio.sh" "\$@" +exec '@@DATA_DIR@@/launch-studio.sh' "$@" STUB_EOF + sed "s|@@DATA_DIR@@|$_css_sed_dir|g" "$_css_macos_dir/launch-studio" \ + > "$_css_macos_dir/launch-studio.tmp" \ + && mv "$_css_macos_dir/launch-studio.tmp" "$_css_macos_dir/launch-studio" chmod +x "$_css_macos_dir/launch-studio" # Build AppIcon.icns from unsloth-gem.png (2240x2240) @@ -1079,11 +1348,28 @@ mkdir -p "$STUDIO_HOME" _MIGRATED=false if [ -x "$VENV_DIR/bin/python" ]; then + # why: matching guard to the .venv branch below -- in env-mode + # $STUDIO_HOME is a user-chosen workspace, so refuse to nuke an + # existing $STUDIO_HOME/unsloth_studio that lacks Studio sentinels. + # Accept the in-VENV ownership marker so partial-install retries are + # not blocked. Sentinels must be regular files: -f follows symlinks + # to files (the legitimate ln -s shim shape) but rejects directories + # and broken/dir-targeted symlinks. + if [ "$_STUDIO_HOME_REDIRECT" = "env" ] \ + && [ ! -f "$VENV_DIR/.unsloth-studio-owned" ] \ + && [ ! -f "$STUDIO_HOME/share/studio.conf" ] \ + && [ ! -f "$STUDIO_HOME/bin/unsloth" ]; then + echo "ERROR: $VENV_DIR already exists but does not look like an Unsloth Studio install." >&2 + echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 + exit 1 + fi # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" -elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then +elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/python" ]; then # Old layout exists — validate before migrating. + # Skip in env-mode so we don't rm -rf an unrelated .venv at the + # workspace root (e.g. user's existing project Python venv). # In no-torch mode, a missing torch package is expected; validate Python only. substep "found legacy Studio environment, validating..." _legacy_ok=false @@ -1132,6 +1418,13 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" fi +# Mark the freshly-created venv as Studio-owned so a partial install can be +# repaired by re-running install.sh; the env-mode deletion guard above accepts +# this marker as the primary sentinel. +if [ -x "$VENV_DIR/bin/python" ]; then + : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true +fi + # Guard against Python 3.13.8 torch import bug on Apple Silicon # (skip when the user explicitly chose a version via --python) if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then @@ -1143,6 +1436,9 @@ if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then rm -rf "$VENV_DIR" PYTHON_VERSION="3.12" run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + if [ -x "$VENV_DIR/bin/python" ]; then + : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true + fi fi fi @@ -1768,7 +2064,17 @@ _SKIP_FRONTEND=0 if [ "$TAURI_MODE" = true ]; then _SKIP_FRONTEND=1 fi +# Prepend UNSLOTH_STUDIO_HOME=$STUDIO_HOME to "$@" for env-override installs +# without word-splitting on whitespace paths. +_run_setup_with_studio_home() { + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + UNSLOTH_STUDIO_HOME="$STUDIO_HOME" "$@" + else + "$@" + fi +} if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ SKIP_STUDIO_FRONTEND="$_SKIP_FRONTEND" \ STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \ @@ -1782,6 +2088,7 @@ else # the same session) does not silently flip a normal install onto the # local-dev path in setup.sh and install_python_stack.py. Mirrors the # reset already done in install.ps1 for PowerShell. + _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ SKIP_STUDIO_FRONTEND="$_SKIP_FRONTEND" \ STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \ @@ -1791,36 +2098,53 @@ else bash "$SETUP_SH" &2 + echo " Move or remove it manually, then re-run the installer." >&2 + exit 1 +fi +# why: -sfn is atomic and -n prevents descent into a symlink-to-directory at +# the shim path (the directory guard above already rejects a real directory). +ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path" -_LOCAL_BIN="$HOME/.local/bin" case ":$PATH:" in *":$_LOCAL_BIN:"*) ;; # already on PATH *) - _SHELL_PROFILE="" - if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then - _SHELL_PROFILE="$HOME/.zshrc" - elif [ -f "$HOME/.bashrc" ]; then - _SHELL_PROFILE="$HOME/.bashrc" - elif [ -f "$HOME/.profile" ]; then - _SHELL_PROFILE="$HOME/.profile" - fi - - if [ -n "$_SHELL_PROFILE" ]; then - if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then - echo '' >> "$_SHELL_PROFILE" - echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE" - echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE" - step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE" + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + export PATH="$_LOCAL_BIN:$PATH" + step "path" "exported $_LOCAL_BIN for this session (no rc-file append in env-override mode)" + else + _SHELL_PROFILE="" + if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then + _SHELL_PROFILE="$HOME/.zshrc" + elif [ -f "$HOME/.bashrc" ]; then + _SHELL_PROFILE="$HOME/.bashrc" + elif [ -f "$HOME/.profile" ]; then + _SHELL_PROFILE="$HOME/.profile" fi + if [ -n "$_SHELL_PROFILE" ]; then + if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then + echo '' >> "$_SHELL_PROFILE" + echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE" + echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE" + step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE" + fi + fi + export PATH="$_LOCAL_BIN:$PATH" fi - export PATH="$_LOCAL_BIN:$PATH" ;; esac # Non-Tauri installs keep shortcuts even if setup reports failure. +# create_studio_shortcuts gates persistent menu shortcuts on env-mode; +# launcher + studio.conf + icon are always written. if [ "$TAURI_MODE" != true ]; then create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" fi @@ -1883,10 +2207,21 @@ if [ -t 1 ]; then esac else step "launch" "manual commands:" - substep "unsloth studio -p 8888" - substep "or activate env first:" - substep "source ${VENV_DIR}/bin/activate" - substep "unsloth studio -p 8888" + # Single-quote-escape so paths with spaces / apostrophes copy-paste cleanly. + _li_shim_q="'$(printf '%s' "${_LOCAL_BIN}/unsloth" | sed "s/'/'\\\\''/g")'" + _li_act_q="'$(printf '%s' "${VENV_DIR}/bin/activate" | sed "s/'/'\\\\''/g")'" + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + # Env-mode skips the rc PATH append, so print the absolute shim path. + substep "$_li_shim_q studio -p 8888" + substep "or activate env first:" + substep "source $_li_act_q" + substep "unsloth studio -p 8888" + else + substep "unsloth studio -p 8888" + substep "or activate env first:" + substep "source $_li_act_q" + substep "unsloth studio -p 8888" + fi substep "(add -H 0.0.0.0 to allow network / cloud access)" echo "" fi diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f768764c22..8da836de38 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -732,22 +732,46 @@ class LlamaCppBackend: if win_bin.is_file(): return str(win_bin) - # 2–4. ~/.unsloth/llama.cpp (primary — setup.sh / setup.ps1 build here) - unsloth_home = Path.home() / ".unsloth" / "llama.cpp" - # Root dir (make builds copy binaries here) - home_root = unsloth_home / binary_name - if home_root.is_file(): - return str(home_root) - # build/bin/ (cmake builds on Linux) - home_linux = unsloth_home / "build" / "bin" / binary_name - if home_linux.is_file(): - return str(home_linux) + # 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp; + # default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio). + legacy_llama = Path.home() / ".unsloth" / "llama.cpp" + try: + from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 - # 3. Windows MSVC build has Release subdir - if sys.platform == "win32": - home_win = unsloth_home / "build" / "bin" / "Release" / binary_name - if home_win.is_file(): - return str(home_win) + _resolved_sr = _sr() + _legacy_studio = Path.home() / ".unsloth" / "studio" + try: + _is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve() + except (OSError, ValueError): + _is_legacy = _resolved_sr == _legacy_studio + if _is_legacy: + search_roots = [legacy_llama] + else: + # why: _kill_orphaned_servers excludes the legacy root in custom + # mode; discovery must match so we never spawn a server we then + # refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier) + # is the explicit way to share a build across roots. + search_roots = [_resolved_sr / "llama.cpp"] + except (ImportError, OSError, ValueError): + search_roots = [legacy_llama] + _seen_roots: set[str] = set() + _unique_roots: list[Path] = [] + for r in search_roots: + k = str(r) + if k not in _seen_roots: + _seen_roots.add(k) + _unique_roots.append(r) + for unsloth_home in _unique_roots: + home_root = unsloth_home / binary_name + if home_root.is_file(): + return str(home_root) + home_linux = unsloth_home / "build" / "bin" / binary_name + if home_linux.is_file(): + return str(home_linux) + if sys.platform == "win32": + home_win = unsloth_home / "build" / "bin" / "Release" / binary_name + if home_win.is_file(): + return str(home_win) # 5–6. Legacy: in-tree build (older setup.sh / setup.ps1 versions) project_root = Path(__file__).resolve().parents[4] @@ -2592,8 +2616,27 @@ class LlamaCppBackend: # (binary must be *under* one of these) install_roots: list[Path] = [] - # Primary install dir (setup.sh / prebuilt installer) - install_roots.append(Path.home() / ".unsloth" / "llama.cpp") + # Env-mode custom root (mirrors _find_llama_server_binary). + _is_custom_root = False + try: + from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 + + _resolved_sr = _sr() + _legacy_studio = Path.home() / ".unsloth" / "studio" + try: + _is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve() + except (OSError, ValueError): + _is_custom_root = _resolved_sr != _legacy_studio + if _is_custom_root: + install_roots.append(_resolved_sr / "llama.cpp") + except (ImportError, OSError, ValueError): + pass + + # Primary install dir (default mode only). Env-mode skips this so + # a custom-root Studio cannot kill a concurrent default-install + # Studio's llama-server (same OS user, different install). + if not _is_custom_root: + install_roots.append(Path.home() / ".unsloth" / "llama.cpp") # Legacy in-tree build dirs (older setup.sh versions) project_root = Path(__file__).resolve().parents[4] diff --git a/studio/backend/main.py b/studio/backend/main.py index 0958094ff0..cd901327db 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -23,12 +23,67 @@ if _backend_dir not in sys.path: # See: https://github.com/python/cpython/issues/102396 import _platform_compat # noqa: F401 +# Direct `uvicorn main:app` launches bypass run.py, so re-export here too +# (mirrors run.py). Required BEFORE the unsloth-zoo import below, since +# its LLAMA_CPP_DEFAULT_DIR binding is import-time. +from utils.paths.storage_roots import studio_root as _studio_root + +try: + _LEGACY_STUDIO_ROOT = (_Path.home() / ".unsloth" / "studio").resolve() +except (OSError, ValueError): + _LEGACY_STUDIO_ROOT = _Path.home() / ".unsloth" / "studio" +try: + _STUDIO_ROOT_RESOLVED = _studio_root().resolve() +except (OSError, ValueError): + _STUDIO_ROOT_RESOLVED = _studio_root() +if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: + if not os.environ.get("UNSLOTH_STUDIO_HOME"): + os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED) + if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") + import mimetypes +import re as _re import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version + +_STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$") + + +def _read_studio_install_id() -> str: + """Per-install opaque id written by install.sh / install.ps1 at + $STUDIO_HOME/share/studio_install_id. Returns "" when the file is + absent (pre-PR install, fresh tree never run through the installer) + or contains anything other than a 64-char lowercase-hex token -- + in which case /api/health emits "" and the launcher's _check_health + falls back to the existing "no baked id, accept any healthy + Unsloth backend" path. This intentionally replaces a previous + sha256(resolved_install_path) so the field carries no install-path + information for callers reaching /api/health (relevant when Studio + is run with -H 0.0.0.0).""" + try: + token = ( + (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip() + ) + except (OSError, ValueError): + return "" + return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else "" + + +_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id() + + +def _studio_root_id() -> str: + """Same-install discriminator for /api/health: a per-install opaque + token written once by the installer and read once at module import. + Empty when no installer-written token is present; the launcher + contract treats "" as "no baked id, accept any healthy backend".""" + return _STUDIO_ROOT_ID_CACHE + + # Fix broken Windows registry MIME types. Some Windows installs map .js to # "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes # module reads from the registry, and FastAPI/Starlette's StaticFiles uses @@ -245,6 +300,10 @@ async def health_check(): "chat_only": _hw_module.CHAT_ONLY, "desktop_protocol_version": 1, "supports_desktop_auth": True, + # why: launchers compare against an install-time hash so a sibling + # Studio on the same port is rejected; hex digest avoids leaking the + # raw install path on -H 0.0.0.0. + "studio_root_id": _studio_root_id(), "native_path_leases_supported": native_path_leases_supported(), } diff --git a/studio/backend/run.py b/studio/backend/run.py index c5b103ff70..1dd1230a17 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -159,7 +159,27 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int: ) -_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid" +from utils.paths.storage_roots import studio_root as _studio_root + +_PID_FILE = _studio_root() / "studio.pid" + +# Direct backend launches bypass the CLI's env re-export; do it here for +# real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR +# picks up the custom build. Skip for legacy-default to avoid flipping +# default-mode installs into env-override. +try: + _LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve() +except (OSError, ValueError): + _LEGACY_STUDIO_ROOT = Path.home() / ".unsloth" / "studio" +try: + _STUDIO_ROOT_RESOLVED = _studio_root().resolve() +except (OSError, ValueError): + _STUDIO_ROOT_RESOLVED = _studio_root() +if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: + if not os.environ.get("UNSLOTH_STUDIO_HOME"): + os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED) + if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") def _write_pid_file(): diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 16f6d21edb..dc8dd08315 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -500,7 +500,9 @@ _VLM_MODEL_TYPES = { # Pre-computed .venv_t5 paths and backend dir for subprocess version switching. # Vision check uses 5.5.0 (newest, recognizes all architectures). -_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550") +from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402 + +_VENV_T5_DIR = str(_studio_root() / ".venv_t5_550") _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) # Inline script executed in a subprocess with transformers 5.x activated. diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index b52609b06b..58a4d7967c 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -5,17 +5,59 @@ from __future__ import annotations import json import os +import sys from pathlib import Path import tempfile +def _infer_studio_home_from_venv() -> Path | None: + """Return parent dir of sys.prefix as STUDIO_HOME if running from an + installer-managed unsloth_studio venv. Sentinel-gated (share/studio.conf + or bin shim) so a developer venv named unsloth_studio is not misidentified. + """ + try: + prefix = Path(sys.prefix).resolve() + except (OSError, ValueError): + return None + if prefix.name != "unsloth_studio": + return None + candidate = prefix.parent + shim_name = "unsloth.exe" if os.name == "nt" else "unsloth" + try: + has_sentinel = (candidate / "share" / "studio.conf").is_file() or ( + candidate / "bin" / shim_name + ).is_file() + except OSError: + return None + if has_sentinel: + return candidate + return None + + def studio_root() -> Path: + """Studio install root. + + Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix + inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins when + both are set (the more specific signal beats the generic alias). + """ + override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() + if not override: + override = (os.environ.get("STUDIO_HOME") or "").strip() + if override: + try: + return Path(override).expanduser().resolve() + except (OSError, ValueError): + return Path(override).expanduser() + inferred = _infer_studio_home_from_venv() + if inferred is not None: + return inferred return Path.home() / ".unsloth" / "studio" def cache_root() -> Path: """Central cache directory for all studio downloads (models, datasets, etc.).""" - return Path.home() / ".unsloth" / "studio" / "cache" + return studio_root() / "cache" def assets_root() -> Path: diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 17af40f663..9075c590ca 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -95,9 +95,11 @@ TRANSFORMERS_DEFAULT_VERSION = "4.57.6" # Consumers should prefer TRANSFORMERS_530_VERSION / TRANSFORMERS_550_VERSION. TRANSFORMERS_5_VERSION = TRANSFORMERS_550_VERSION -# Pre-installed directories — created by setup.sh / setup.ps1 -_VENV_T5_530_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_530") -_VENV_T5_550_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550") +# Pre-installed directories — created by setup.sh / setup.ps1. +from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402 + +_VENV_T5_530_DIR = str(_studio_root() / ".venv_t5_530") +_VENV_T5_550_DIR = str(_studio_root() / ".venv_t5_550") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 3d082aa70d..f2753d5c88 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1492,9 +1492,79 @@ if (-not $PythonCmd) { substep "Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -# The venv must already exist (created by install.ps1). -# This script (setup.ps1 / "unsloth studio update") only updates packages. -$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio" +# The venv must already exist (created by install.ps1); this script only +# updates packages. UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the +# root. UNSLOTH_STUDIO_HOME wins when both are set. Whitespace-only values +# are treated as unset to match Python .strip() semantics. +$_studioOverrideVar = $null +$_studioOverride = $null +if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { + $_studioOverrideVar = "UNSLOTH_STUDIO_HOME" + $_studioOverride = $env:UNSLOTH_STUDIO_HOME.Trim() +} elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { + $_studioOverrideVar = "STUDIO_HOME" + $_studioOverride = $env:STUDIO_HOME.Trim() +} +if ($_studioOverride) { + if ($_studioOverride -eq "~" -or $_studioOverride -like "~/*" -or $_studioOverride -like "~\*") { + $_studioOverride = (Join-Path $env:USERPROFILE $_studioOverride.Substring(1).TrimStart('/','\')) + } + if (Test-Path -LiteralPath $_studioOverride -PathType Container) { + $StudioHome = (Resolve-Path -LiteralPath $_studioOverride).Path + # why: mirror setup.sh:417 and install.ps1:130 -- fail fast when the + # custom root is read-only instead of erroring later while creating + # sidecar venvs / installing packages. + $_setupWriteProbe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid()) + try { + [System.IO.File]::WriteAllText($_setupWriteProbe, "") + Remove-Item -LiteralPath $_setupWriteProbe -Force -ErrorAction SilentlyContinue + } catch { + Write-Host "ERROR: $_studioOverrideVar=$StudioHome is not writable." -ForegroundColor Red + exit 1 + } + } else { + Write-Host "ERROR: $_studioOverrideVar=$_studioOverride does not exist." -ForegroundColor Red + Write-Host " Run install.ps1 to create the install root before 'unsloth studio update'." -ForegroundColor Red + exit 1 + } +} else { + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" +} +$VenvDir = Join-Path $StudioHome "unsloth_studio" + +# why: in env-override mode $StudioHome is user-chosen; require the +# ownership marker before Remove-Item so unrelated dirs survive. Gated on +# the canonical comparison so an override pointing at the legacy default +# still behaves like a default install. +$StudioOwnedMarker = ".unsloth-studio-owned" +$LegacyStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" +$_studioHomeCanon = $StudioHome +if (Test-Path -LiteralPath $_studioHomeCanon -PathType Container) { + $_studioHomeCanon = (Resolve-Path -LiteralPath $_studioHomeCanon).Path +} +if (Test-Path -LiteralPath $LegacyStudioHome -PathType Container) { + $LegacyStudioHome = (Resolve-Path -LiteralPath $LegacyStudioHome).Path +} +$StudioHomeIsCustom = ($_studioHomeCanon -ne $LegacyStudioHome) +function Assert-StudioOwnedOrAbsent { + param( + [Parameter(Mandatory = $true)][string]$Path, + [Parameter(Mandatory = $true)][string]$Label + ) + if (-not (Test-Path -LiteralPath $Path -PathType Container)) { return } + if ($StudioHomeIsCustom -and -not (Test-Path -LiteralPath (Join-Path $Path $StudioOwnedMarker) -PathType Leaf)) { + Write-Host "[ERROR] $Path already exists and is not marked as a Studio-owned $Label." -ForegroundColor Red + Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow + exit 1 + } +} +function Mark-StudioOwned { + param([Parameter(Mandatory = $true)][string]$Path) + if (-not (Test-Path -LiteralPath $Path -PathType Container)) { return } + try { + [System.IO.File]::WriteAllText((Join-Path $Path $StudioOwnedMarker), "") + } catch {} +} # Stale-venv detection: if the venv exists but its torch flavor no longer # matches the current machine, repair according to invocation context. @@ -1504,12 +1574,12 @@ $VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio" # In no-torch mode, a missing torch package is expected. $NoTorchMode = $env:UNSLOTH_NO_TORCH -match '^(?i:true|1|yes)$' $InstallerManagedSetup = $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -match '^(?i:true|1|yes)$' -if ((Test-Path $VenvDir -PathType Container) -and -not $NoTorchMode) { +if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode) { $VenvPyExe = Join-Path $VenvDir "Scripts\python.exe" $installedTorchTag = $null $shouldRebuild = $false - if (Test-Path $VenvPyExe) { + if (Test-Path -LiteralPath $VenvPyExe) { try { $psi = New-Object System.Diagnostics.ProcessStartInfo $psi.FileName = $VenvPyExe @@ -1558,8 +1628,21 @@ if ((Test-Path $VenvDir -PathType Container) -and -not $NoTorchMode) { exit 1 } substep "Stale venv detected ($reason) -- rebuilding..." "Yellow" + # why: mirror install.ps1 env-mode guard so an update against a custom + # UNSLOTH_STUDIO_HOME never wipes an unrelated unsloth_studio venv; + # -PathType Leaf rejects a directory masquerading as the sentinel. + if ( + $StudioHomeIsCustom -and + -not (Test-Path -LiteralPath (Join-Path $VenvDir $StudioOwnedMarker) -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "share\studio.conf") -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "bin\unsloth.exe") -PathType Leaf) + ) { + Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red + Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow + exit 1 + } try { - Remove-Item $VenvDir -Recurse -Force -ErrorAction Stop + Remove-Item -LiteralPath $VenvDir -Recurse -Force -ErrorAction Stop } catch { Write-Host " [ERROR] Could not remove stale venv: $($_.Exception.Message)" -ForegroundColor Red Write-Host " Close any running Studio/Python processes and re-run setup." -ForegroundColor Red @@ -1568,7 +1651,7 @@ if ((Test-Path $VenvDir -PathType Container) -and -not $NoTorchMode) { } } -if (-not (Test-Path $VenvDir)) { +if (-not (Test-Path -LiteralPath $VenvDir)) { Write-Host "[ERROR] Virtual environment not found at $VenvDir" -ForegroundColor Red Write-Host " Run install.ps1 first to create the environment:" -ForegroundColor Yellow Write-Host " irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow @@ -1759,17 +1842,19 @@ if ($stackExit -ne 0) { # ── Pre-install transformers 5.x into .venv_t5_530/ and .venv_t5_550/ ── # Runs outside the deps fast-path gate so that upgrades from the legacy # single .venv_t5 are always migrated to the tiered layout. -$VenvT5_530Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5_530" -$VenvT5_550Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5_550" -$VenvT5Legacy = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5" +# T5 sidecar venvs live under the resolved $StudioHome so custom installs are self-contained. +$VenvT5_530Dir = Join-Path $StudioHome ".venv_t5_530" +$VenvT5_550Dir = Join-Path $StudioHome ".venv_t5_550" +$VenvT5Legacy = Join-Path $StudioHome ".venv_t5" $_NeedT5Install = $false -if (Test-Path $VenvT5Legacy) { - Remove-Item -Recurse -Force $VenvT5Legacy +if (Test-Path -LiteralPath $VenvT5Legacy) { + Assert-StudioOwnedOrAbsent -Path $VenvT5Legacy -Label "legacy transformers sidecar venv" + Remove-Item -LiteralPath $VenvT5Legacy -Recurse -Force $_NeedT5Install = $true } -if (-not (Test-Path $VenvT5_530Dir)) { $_NeedT5Install = $true } -if (-not (Test-Path $VenvT5_550Dir)) { $_NeedT5Install = $true } +if (-not (Test-Path -LiteralPath $VenvT5_530Dir)) { $_NeedT5Install = $true } +if (-not (Test-Path -LiteralPath $VenvT5_550Dir)) { $_NeedT5Install = $true } # Also reinstall when python deps were updated if (-not $SkipPythonDeps) { $_NeedT5Install = $true } @@ -1781,8 +1866,10 @@ $ErrorActionPreference = "Continue" # --- .venv_t5_530 (transformers 5.3.0) --- substep "pre-installing transformers 5.3.0 for newer model support..." -if (Test-Path $VenvT5_530Dir) { Remove-Item -Recurse -Force $VenvT5_530Dir } -New-Item -ItemType Directory -Path $VenvT5_530Dir -Force | Out-Null +Assert-StudioOwnedOrAbsent -Path $VenvT5_530Dir -Label "transformers 5.3 sidecar venv" +if (Test-Path -LiteralPath $VenvT5_530Dir) { Remove-Item -LiteralPath $VenvT5_530Dir -Recurse -Force } +[System.IO.Directory]::CreateDirectory($VenvT5_530Dir) | Out-Null +Mark-StudioOwned -Path $VenvT5_530Dir foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4.2")) { if ($script:UnslothVerbose) { Fast-Install --target $VenvT5_530Dir --no-deps $pkg @@ -1814,8 +1901,10 @@ step "transformers" "5.3.0 pre-installed" # --- .venv_t5_550 (transformers 5.5.0) --- substep "pre-installing transformers 5.5.0 for Gemma 4 support..." -if (Test-Path $VenvT5_550Dir) { Remove-Item -Recurse -Force $VenvT5_550Dir } -New-Item -ItemType Directory -Path $VenvT5_550Dir -Force | Out-Null +Assert-StudioOwnedOrAbsent -Path $VenvT5_550Dir -Label "transformers 5.5 sidecar venv" +if (Test-Path -LiteralPath $VenvT5_550Dir) { Remove-Item -LiteralPath $VenvT5_550Dir -Recurse -Force } +[System.IO.Directory]::CreateDirectory($VenvT5_550Dir) | Out-Null +Mark-StudioOwned -Path $VenvT5_550Dir foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4.2")) { if ($script:UnslothVerbose) { Fast-Install --target $VenvT5_550Dir --no-deps $pkg @@ -1851,8 +1940,15 @@ step "transformers" "5.5.0 pre-installed" # ========================================================================== # PHASE 3.4: Prefer prebuilt llama.cpp bundles before source build # ========================================================================== -$UnslothHome = Join-Path $env:USERPROFILE ".unsloth" -if (-not (Test-Path $UnslothHome)) { New-Item -ItemType Directory -Force $UnslothHome | Out-Null } +# Nest llama.cpp under $StudioHome only for real env-overrides, never the +# legacy default. Reuses $StudioHomeIsCustom from the canonical comparison +# computed above so the llama.cpp nest matches ownership-guard semantics. +if ($StudioHomeIsCustom) { + $UnslothHome = $StudioHome +} else { + $UnslothHome = Join-Path $env:USERPROFILE ".unsloth" +} +if (-not (Test-Path -LiteralPath $UnslothHome)) { [System.IO.Directory]::CreateDirectory($UnslothHome) | Out-Null } $LlamaCppDir = Join-Path $UnslothHome "llama.cpp" $NeedLlamaSourceBuild = $false $SkipPrebuiltInstall = $false @@ -1954,9 +2050,15 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } else { Write-Host "" substep "installing prebuilt llama.cpp bundle (preferred path)..." - if (Test-Path $LlamaCppDir) { + if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" } + # why: install_llama_prebuilt.py uses os.replace(), which would displace + # an unrelated $env:UNSLOTH_STUDIO_HOME\llama.cpp before the source-build + # ownership check below ever runs. + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } $prebuiltArgs = @( "$PSScriptRoot\install_llama_prebuilt.py", "--install-dir", $LlamaCppDir, @@ -2001,6 +2103,9 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } else { step "llama.cpp" "prebuilt installed and validated" } + if ($StudioHomeIsCustom -and (Test-Path -LiteralPath $LlamaCppDir -PathType Container)) { + Mark-StudioOwned -Path $LlamaCppDir + } $installedRelease = Get-InstalledLlamaPrebuiltRelease -InstallDir $LlamaCppDir if ($installedRelease) { substep $installedRelease @@ -2008,7 +2113,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } elseif ($prebuiltExit -eq 3) { step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" Write-LlamaFailureLog -Output $prebuiltOutput - if (Test-Path $LlamaCppDir) { + if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing install was restored" "Yellow" } substep "Close Studio or other llama.cpp users and retry" "Yellow" @@ -2016,7 +2121,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } else { step "llama.cpp" "prebuilt install failed (continuing)" "Yellow" Write-LlamaFailureLog -Output $prebuiltOutput - if (Test-Path $LlamaCppDir) { + if (Test-Path -LiteralPath $LlamaCppDir) { substep "Prebuilt update failed; existing install was restored or cleaned before source build fallback" "Yellow" } substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow" @@ -2092,10 +2197,10 @@ $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) # Check if existing llama-server matches current GPU mode. A CUDA-built binary # on a now-CPU-only machine (or vice versa) needs to be rebuilt. $NeedRebuild = $false -if (Test-Path $LlamaServerBin) { +if (Test-Path -LiteralPath $LlamaServerBin) { $CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt" - if (Test-Path $CmakeCacheFile) { - $cachedCuda = Select-String -Path $CmakeCacheFile -Pattern 'GGML_CUDA:BOOL=ON' -Quiet + if (Test-Path -LiteralPath $CmakeCacheFile) { + $cachedCuda = Select-String -LiteralPath $CmakeCacheFile -Pattern 'GGML_CUDA:BOOL=ON' -Quiet if ($HasNvidiaSmi -and -not $cachedCuda) { Write-Host " Existing llama-server is CPU-only but GPU is available -- rebuilding" -ForegroundColor Yellow $NeedRebuild = $true @@ -2109,7 +2214,7 @@ if (Test-Path $LlamaServerBin) { if (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" -} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { +} elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { # Skip rebuild only for pinned tags (e.g. b8635). When the requested # tag is "master" (a moving target), always rebuild so the binary picks # up new model architecture support (e.g. Gemma 4). @@ -2211,7 +2316,13 @@ if (-not $NeedLlamaSourceBuild) { $UseConcreteRef = ($ResolvedSourceRef -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedSourceRef)) - if (Test-Path (Join-Path $LlamaCppDir ".git")) { + if (Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")) { + # 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" + } Write-Host " Syncing llama.cpp to $ResolvedSourceRef..." -ForegroundColor Gray # Always sync the remote URL so switching between default/fork sources works Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir remote set-url origin "$ResolvedSourceUrl.git" } | Out-Null @@ -2282,24 +2393,30 @@ if (-not $NeedLlamaSourceBuild) { } } } + # why: in-place git-sync (the temp-dir clone path calls Mark-StudioOwned + # at swap-time) must mark the existing tree so a subsequent prebuilt + # update path's Assert-StudioOwnedOrAbsent does not exit on the same root. + if ($BuildOk -and $StudioHomeIsCustom) { + Mark-StudioOwned -Path $LlamaCppDir + } } else { Write-Host " Cloning llama.cpp @ $ResolvedSourceRef..." -ForegroundColor Gray $buildTmp = "$LlamaCppDir.build.$PID" - $null = New-Item -ItemType Directory -Force -Path (Split-Path $LlamaCppDir -Parent) - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + $null = [System.IO.Directory]::CreateDirectory((Split-Path -LiteralPath $LlamaCppDir)) + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } if ($LlamaPr) { $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$LlamaSource.git" $buildTmp } if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } if ($BuildOk) { $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin "pull/$LlamaPr/head:pr-$LlamaPr" } if ($fetchExit -ne 0) { $BuildOk = $false $FailedStep = "git fetch PR #$LlamaPr" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } if ($BuildOk) { @@ -2307,7 +2424,7 @@ if (-not $NeedLlamaSourceBuild) { if ($checkoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout PR #$LlamaPr" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } } elseif ($ResolvedSourceRefKind -eq "pull") { @@ -2315,14 +2432,14 @@ if (-not $NeedLlamaSourceBuild) { if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } if ($BuildOk) { $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef } if ($fetchExit -ne 0) { $BuildOk = $false $FailedStep = "git fetch source PR ref" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } if ($BuildOk) { @@ -2330,7 +2447,7 @@ if (-not $NeedLlamaSourceBuild) { if ($checkoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout source PR ref" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } } elseif ($ResolvedSourceRefKind -eq "commit") { @@ -2338,14 +2455,14 @@ if (-not $NeedLlamaSourceBuild) { if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } if ($BuildOk) { $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef } if ($fetchExit -ne 0) { $BuildOk = $false $FailedStep = "git fetch source commit" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } if ($BuildOk) { @@ -2353,7 +2470,7 @@ if (-not $NeedLlamaSourceBuild) { if ($checkoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout source commit" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } } else { @@ -2366,7 +2483,7 @@ if (-not $NeedLlamaSourceBuild) { if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } # Use temp dir for build; swap into $LlamaCppDir only after build succeeds @@ -2482,14 +2599,16 @@ if (-not $NeedLlamaSourceBuild) { # Swap temp build dir into final location (only if we built in a temp dir) if ($BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) { - if (Test-Path $OriginalLlamaCppDir) { Remove-Item -Recurse -Force $OriginalLlamaCppDir } - Move-Item $LlamaCppDir $OriginalLlamaCppDir + Assert-StudioOwnedOrAbsent -Path $OriginalLlamaCppDir -Label "llama.cpp install" + if (Test-Path -LiteralPath $OriginalLlamaCppDir) { Remove-Item -LiteralPath $OriginalLlamaCppDir -Recurse -Force } + Move-Item -LiteralPath $LlamaCppDir -Destination $OriginalLlamaCppDir $LlamaCppDir = $OriginalLlamaCppDir $BuildDir = Join-Path $LlamaCppDir "build" $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe" + Mark-StudioOwned -Path $LlamaCppDir } elseif (-not $BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) { # Build failed -- clean up temp dir, preserve existing install - if (Test-Path $LlamaCppDir) { Remove-Item -Recurse -Force $LlamaCppDir } + if (Test-Path -LiteralPath $LlamaCppDir) { Remove-Item -LiteralPath $LlamaCppDir -Recurse -Force } $LlamaCppDir = $OriginalLlamaCppDir $BuildDir = Join-Path $LlamaCppDir "build" $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe" @@ -2504,16 +2623,16 @@ if (-not $NeedLlamaSourceBuild) { $totalSec = [math]::Round($totalSw.Elapsed.TotalSeconds % 60, 1) # -- Summary -- - if ($BuildOk -and (Test-Path $LlamaServerBin)) { + if ($BuildOk -and (Test-Path -LiteralPath $LlamaServerBin)) { step "llama.cpp" "built" $QuantizeBin = Join-Path $BuildDir "bin\Release\llama-quantize.exe" - if (Test-Path $QuantizeBin) { + if (Test-Path -LiteralPath $QuantizeBin) { step "llama-quantize" "built" } step "build time" "${totalMin}m ${totalSec}s" "DarkGray" } else { $altBin = Join-Path $BuildDir "bin\llama-server.exe" - if ($BuildOk -and (Test-Path $altBin)) { + if ($BuildOk -and (Test-Path -LiteralPath $altBin)) { step "llama.cpp" "built" step "build time" "${totalMin}m ${totalSec}s" "DarkGray" } else { diff --git a/studio/setup.sh b/studio/setup.sh index 3e875eed30..ff93d2d41d 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -417,7 +417,36 @@ if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm fi # ── Python venv + deps ── -STUDIO_HOME="$HOME/.unsloth/studio" +# UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the install root +# (mirrors install.sh). UNSLOTH_STUDIO_HOME wins when both are set. +_studio_override_var="" +_studio_override="${UNSLOTH_STUDIO_HOME:-}" +if [ -n "$_studio_override" ]; then + _studio_override_var="UNSLOTH_STUDIO_HOME" +else + _studio_override="${STUDIO_HOME:-}" + [ -n "$_studio_override" ] && _studio_override_var="STUDIO_HOME" +fi +# Strip whitespace so " " is treated as unset (matches Python .strip()). +_studio_override=$(printf '%s' "$_studio_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') +case "$_studio_override" in + "~") _studio_override="$HOME" ;; + "~/"*) _studio_override="$HOME/${_studio_override#'~/'}" ;; +esac +if [ -n "$_studio_override" ]; then + # setup.sh runs against an existing install (via 'unsloth studio update'); + # a typo in the override must fail fast instead of materializing an + # empty workspace dir. Mirrors setup.ps1 behavior. + if [ ! -d "$_studio_override" ]; then + echo "ERROR: $_studio_override_var=$_studio_override does not exist." >&2 + echo " Run install.sh to create the install root before 'unsloth studio update'." >&2 + exit 1 + fi + [ -w "$_studio_override" ] || { echo "ERROR: $_studio_override_var=$_studio_override is not writable." >&2; exit 1; } + STUDIO_HOME="$(CDPATH= cd -P -- "$_studio_override" && pwd -P)" || exit 1 +else + STUDIO_HOME="$HOME/.unsloth/studio" +fi VENV_DIR="$STUDIO_HOME/unsloth_studio" VENV_T5_530_DIR="$STUDIO_HOME/.venv_t5_530" VENV_T5_550_DIR="$STUDIO_HOME/.venv_t5_550" @@ -542,9 +571,39 @@ fi # # Runs outside the _SKIP_PYTHON_DEPS gate so that upgrades from legacy # single .venv_t5 are always migrated to the tiered layout. +# why: in env-override mode $STUDIO_HOME is user-chosen; require the +# ownership marker before rm -rf so unrelated dirs survive. Gated on the +# canonical comparison so an override pointing at the legacy default still +# behaves like a default install. +_STUDIO_OWNED_MARKER=".unsloth-studio-owned" +_LEGACY_STUDIO_HOME="$HOME/.unsloth/studio" +_studio_home_canon="$STUDIO_HOME" +if [ -d "$_studio_home_canon" ]; then + _studio_home_canon=$(CDPATH= cd -P -- "$_studio_home_canon" 2>/dev/null && pwd -P) \ + || _studio_home_canon="$STUDIO_HOME" +fi +if [ -d "$_LEGACY_STUDIO_HOME" ]; then + _LEGACY_STUDIO_HOME=$(CDPATH= cd -P -- "$_LEGACY_STUDIO_HOME" 2>/dev/null && pwd -P) \ + || _LEGACY_STUDIO_HOME="$HOME/.unsloth/studio" +fi +_STUDIO_HOME_IS_CUSTOM=false +if [ "$_studio_home_canon" != "$_LEGACY_STUDIO_HOME" ]; then + _STUDIO_HOME_IS_CUSTOM=true +fi +_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 + echo "ERROR: $_aso_dir already exists and is not marked as a Studio-owned $_aso_label." >&2 + echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2 + exit 1 + fi +} _NEED_T5_INSTALL=false if [ -d "$STUDIO_HOME/.venv_t5" ]; then # Legacy layout — migrate + _assert_studio_owned_or_absent "$STUDIO_HOME/.venv_t5" "legacy transformers sidecar venv" rm -rf "$STUDIO_HOME/.venv_t5" _NEED_T5_INSTALL=true fi @@ -554,16 +613,20 @@ fi [ "$_SKIP_PYTHON_DEPS" = false ] && _NEED_T5_INSTALL=true if [ "$_NEED_T5_INSTALL" = true ]; then + _assert_studio_owned_or_absent "$VENV_T5_530_DIR" "transformers 5.3 sidecar venv" [ -d "$VENV_T5_530_DIR" ] && rm -rf "$VENV_T5_530_DIR" mkdir -p "$VENV_T5_530_DIR" + : > "$VENV_T5_530_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true run_quiet "install transformers 5.3.0" fast_install --target "$VENV_T5_530_DIR" --no-deps "transformers==5.3.0" run_quiet "install huggingface_hub for t5_530" fast_install --target "$VENV_T5_530_DIR" --no-deps "huggingface_hub==1.8.0" run_quiet "install hf_xet for t5_530" fast_install --target "$VENV_T5_530_DIR" --no-deps "hf_xet==1.4.2" run_quiet "install tiktoken for t5_530" fast_install --target "$VENV_T5_530_DIR" "tiktoken" step "transformers" "5.3.0 pre-installed" + _assert_studio_owned_or_absent "$VENV_T5_550_DIR" "transformers 5.5 sidecar venv" [ -d "$VENV_T5_550_DIR" ] && rm -rf "$VENV_T5_550_DIR" mkdir -p "$VENV_T5_550_DIR" + : > "$VENV_T5_550_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true run_quiet "install transformers 5.5.0" fast_install --target "$VENV_T5_550_DIR" --no-deps "transformers==5.5.0" run_quiet "install huggingface_hub for t5_550" fast_install --target "$VENV_T5_550_DIR" --no-deps "huggingface_hub==1.8.0" run_quiet "install hf_xet for t5_550" fast_install --target "$VENV_T5_550_DIR" --no-deps "hf_xet==1.4.2" @@ -573,7 +636,13 @@ fi fi # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ── -UNSLOTH_HOME="$HOME/.unsloth" +# Nest llama.cpp under $STUDIO_HOME only for real env-overrides; legacy +# default keeps ~/.unsloth/llama.cpp so pre-PR builds are still discovered. +if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + UNSLOTH_HOME="$STUDIO_HOME" +else + UNSLOTH_HOME="$HOME/.unsloth" +fi mkdir -p "$UNSLOTH_HOME" LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp" LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server" @@ -635,6 +704,12 @@ else if [ -d "$LLAMA_CPP_DIR" ]; then substep "existing install detected -- validating update" fi + # 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" + fi _PREBUILT_CMD=( python "$SCRIPT_DIR/install_llama_prebuilt.py" --install-dir "$LLAMA_CPP_DIR" @@ -662,6 +737,9 @@ else else step "llama.cpp" "prebuilt installed and validated" fi + if [ "$_STUDIO_HOME_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" verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR" rm -f "$_PREBUILT_LOG" @@ -1032,8 +1110,10 @@ else # 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" rm -rf "$LLAMA_CPP_DIR" mv "$_BUILD_TMP" "$LLAMA_CPP_DIR" + : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true # Symlink to llama.cpp root -- check_llama_cpp() looks for the binary there QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize" if [ -f "$QUANTIZE_BIN" ]; then diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs index 48a0af6e48..e9a27644df 100644 --- a/studio/src-tauri/src/commands.rs +++ b/studio/src-tauri/src/commands.rs @@ -60,6 +60,11 @@ pub async fn check_install_status() -> bool { cmd.env_remove("PYTHONPATH"); } + // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME; + // probe subprocesses must follow the same isolation as process.rs. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); + let mut child = match cmd.spawn() { Ok(c) => c, Err(e) => { diff --git a/studio/src-tauri/src/desktop_auth.rs b/studio/src-tauri/src/desktop_auth.rs index 483e7c0432..49b19008fb 100644 --- a/studio/src-tauri/src/desktop_auth.rs +++ b/studio/src-tauri/src/desktop_auth.rs @@ -203,6 +203,11 @@ async fn provision_desktop_auth() -> Result<(), String> { cmd.env_remove("PYTHONHOME"); cmd.env_remove("PYTHONPATH"); } + + // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME. + // Scrub so provisioning writes match what the Rust auth code reads. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); #[cfg(windows)] { use std::os::windows::process::CommandExt; diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 9d672f5e73..024b730735 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -196,6 +196,11 @@ fn spawn_script( cmd.env_remove("PYTHONPATH"); } + // 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"); + // On Windows, launch the installer directly with CREATE_NO_WINDOW. // The app process is assigned to a KILL_ON_JOB_CLOSE job in main.rs, so // child cleanup on crash comes from inherited job membership instead. diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index d3df06d057..c0bbb07b36 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -102,6 +102,11 @@ async fn run_cli_probe(bin: &std::path::Path, args: &[&str]) -> bool { cmd.env_remove("PYTHONPATH"); } + // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME; + // probe subprocesses must follow the same isolation as process.rs. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); + #[cfg(windows)] { use std::os::windows::process::CommandExt; @@ -135,6 +140,11 @@ async fn probe_cli_capability(bin: &std::path::Path) -> Option = { use std::os::windows::process::CommandExt; diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py new file mode 100644 index 0000000000..d077cdb824 --- /dev/null +++ b/tests/test_studio_install_workspace_guard.py @@ -0,0 +1,1021 @@ +"""install.sh / install.ps1 must refuse to rm -rf an existing +$STUDIO_HOME/unsloth_studio in env-override mode unless the directory +carries a Studio sentinel (share/studio.conf or bin/unsloth). Also +asserts studio/setup.ps1 has the matching writability probe that +setup.sh:417 already performs.""" + +from __future__ import annotations + +import re +import subprocess +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +INSTALL_SH = REPO_ROOT / "install.sh" +INSTALL_PS1 = REPO_ROOT / "install.ps1" +SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1" +SETUP_SH = REPO_ROOT / "studio" / "setup.sh" + +# Stubs for helpers that the extracted install.sh guard block calls in real +# installs (`substep` for status output, `_start_studio_venv_replacement` for +# the rollback-managed move). The tests run the block in isolation, so we +# stand in a minimal `mv`-based replacement that exercises the same observable +# effect (venv directory is no longer present at $VENV_DIR after a permitted +# cleanup) without dragging in install.sh's full rollback machinery. +_INSTALL_GUARD_STUBS = ( + "substep() { :; }\n" + "_start_studio_venv_replacement() {\n" + ' mv -- "$1" "$1.replaced"\n' + "}\n" +) + + +def _extract_install_sh_guard_block() -> str: + """Pull the `if [ -x "$VENV_DIR/bin/python" ]; then ... fi` block out + of install.sh as a self-contained snippet. Stops at the first elif so + the block can be paired with a synthetic else and run in isolation.""" + src = INSTALL_SH.read_text() + m = re.search( + r'(if \[ -x "\$VENV_DIR/bin/python" \]; then\n.*?)elif \[ "\$_STUDIO_HOME_REDIRECT" != "env"', + src, + re.DOTALL, + ) + assert m, "install.sh venv guard block not found" + return m.group(1) + "fi\n" + + +def _build_install_guard_script( + studio_home: Path, redirect: str, block: str | None = None +) -> str: + """Build a self-contained bash script that exercises the extracted + guard block. Includes stubs for substep / _start_studio_venv_replacement + so the snippet runs without install.sh's full rollback machinery.""" + if block is None: + block = _extract_install_sh_guard_block() + return ( + _INSTALL_GUARD_STUBS + + f'STUDIO_HOME="{studio_home}"\n' + + f'VENV_DIR="$STUDIO_HOME/unsloth_studio"\n' + + f'_STUDIO_HOME_REDIRECT="{redirect}"\n' + + block + + "echo RESULT=ok\n" + ) + + +def _run_install_guard( + studio_home: Path, + redirect: str, + create_share_conf: bool = False, + create_bin_shim: bool = False, + create_venv_marker: bool = False, +) -> subprocess.CompletedProcess: + venv_dir = studio_home / "unsloth_studio" + (venv_dir / "bin").mkdir(parents = True, exist_ok = True) + py = venv_dir / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + if create_share_conf: + (studio_home / "share").mkdir(parents = True, exist_ok = True) + (studio_home / "share" / "studio.conf").write_text("") + if create_bin_shim: + (studio_home / "bin").mkdir(parents = True, exist_ok = True) + (studio_home / "bin" / "unsloth").write_text("") + if create_venv_marker: + (venv_dir / ".unsloth-studio-owned").write_text("") + script = _build_install_guard_script(studio_home, redirect) + return subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + + +def test_env_mode_blocks_unsloth_studio_without_sentinels(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env") + assert res.returncode != 0, ( + "env-mode without sentinels must refuse to rm -rf $VENV_DIR; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert "does not look like an Unsloth Studio install" in res.stderr + assert (studio_home / "unsloth_studio" / "bin" / "python").is_file() + + +def test_env_mode_passes_when_share_studio_conf_present(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env", create_share_conf = True) + assert res.returncode == 0, ( + f"share/studio.conf sentinel must allow cleanup;" + f" stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert "RESULT=ok" in res.stdout + assert not (studio_home / "unsloth_studio").exists() + + +def test_env_mode_passes_when_bin_unsloth_shim_present(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env", create_bin_shim = True) + assert res.returncode == 0, res.stderr + assert not (studio_home / "unsloth_studio").exists() + + +def test_default_mode_skips_sentinel_check(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "default") + assert res.returncode == 0, res.stderr + assert "RESULT=ok" in res.stdout + assert not (studio_home / "unsloth_studio").exists() + + +def test_install_ps1_has_matching_env_mode_guard(): + src = INSTALL_PS1.read_text() + block_start = src.index("if (Test-Path -LiteralPath $VenvPython)") + block = src[block_start : block_start + 2000] + assert ( + "$StudioRedirectMode -eq 'env'" in block + ), "install.ps1 must gate Remove-Item $VenvDir on env-mode" + assert ( + "share\\studio.conf" in block + ), "install.ps1 guard must check share\\studio.conf sentinel" + assert ( + "bin\\unsloth.exe" in block + ), "install.ps1 guard must check bin\\unsloth.exe sentinel" + assert "Refusing to delete non-Studio venv" in block + + +def test_setup_ps1_has_writability_probe(): + src = SETUP_PS1.read_text() + idx = src.index("if (Test-Path -LiteralPath $_studioOverride -PathType Container)") + block = src[idx : idx + 2000] + assert ( + "WriteAllText" in block + ), "setup.ps1 must write-probe UNSLOTH_STUDIO_HOME like setup.sh:417" + assert ( + "is not writable" in block + ), "setup.ps1 probe failure must produce a clear writable-error message" + + +def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path): + """A bare directory at $STUDIO_HOME/bin/unsloth must NOT pass the + sentinel. The previous `-e` test accepted any path type, allowing an + unrelated workspace with sibling content under unsloth_studio plus + a directory at bin/unsloth to be wiped.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (venv / "important.txt").write_text("keep me") + (studio_home / "bin" / "unsloth").mkdir(parents = True) + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode != 0, ( + "directory at bin/unsloth must NOT satisfy the Studio sentinel; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert (venv / "important.txt").is_file(), "unrelated workspace data must survive" + + +def test_env_mode_passes_when_bin_unsloth_is_a_symlink(tmp_path): + """A symlink at $STUDIO_HOME/bin/unsloth (real installer artefact) + must still satisfy the sentinel after the leaf-only tightening.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (studio_home / "bin").mkdir(parents = True) + target = studio_home / "bin" / "unsloth-real" + target.write_text("#!/bin/sh\nexit 0\n") + target.chmod(0o755) + (studio_home / "bin" / "unsloth").symlink_to(target) + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode == 0, res.stderr + assert "RESULT=ok" in res.stdout + assert not venv.exists() + + +def test_install_ps1_sentinel_uses_pathtype_leaf(): + """The Test-Path checks that gate Remove-Item $VenvDir must use + -PathType Leaf so a directory at the sentinel path cannot satisfy them.""" + src = INSTALL_PS1.read_text() + block_start = src.index("if (Test-Path -LiteralPath $VenvPython)") + block = src[block_start : block_start + 2000] + assert ( + 'share\\studio.conf") -PathType Leaf' in block + ), "install.ps1 share\\studio.conf check must use -PathType Leaf" + assert ( + 'bin\\unsloth.exe") -PathType Leaf' in block + ), "install.ps1 bin\\unsloth.exe check must use -PathType Leaf" + + +def test_setup_ps1_stale_venv_has_env_mode_guard(): + """studio/setup.ps1 stale-venv rebuild branch must mirror install.ps1: + refuse to Remove-Item $VenvDir under custom-root mode unless the root + carries a Studio sentinel (in-VENV marker, share\\studio.conf, or + bin\\unsloth.exe leaf).""" + src = SETUP_PS1.read_text() + idx = src.index("Stale venv detected") + block = src[idx : idx + 1500] + assert ( + "$StudioHomeIsCustom" in block + ), "setup.ps1 stale-venv branch must gate on $StudioHomeIsCustom" + assert ( + 'share\\studio.conf") -PathType Leaf' in block + ), "setup.ps1 stale-venv guard must check share\\studio.conf with -PathType Leaf" + assert ( + 'bin\\unsloth.exe") -PathType Leaf' in block + ), "setup.ps1 stale-venv guard must check bin\\unsloth.exe with -PathType Leaf" + # The guard must fire BEFORE the destructive call. + guard_idx = block.index("$StudioHomeIsCustom") + rm_idx = block.index("Remove-Item -LiteralPath $VenvDir") + assert ( + guard_idx < rm_idx + ), "custom-root guard must precede Remove-Item -LiteralPath $VenvDir" + + +def test_setup_sh_prebuilt_llama_cpp_has_ownership_guard(): + """studio/setup.sh prebuilt llama.cpp path must call + _assert_studio_owned_or_absent before invoking install_llama_prebuilt.py + so an unrelated $UNSLOTH_STUDIO_HOME/llama.cpp is not displaced by + the helper's os.replace().""" + src = SETUP_SH.read_text() + idx = src.index("installing prebuilt llama.cpp...") + block = src[idx : idx + 2000] + assert ( + '_assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install"' in block + ), "setup.sh must guard the prebuilt llama.cpp path with the ownership marker" + guard_idx = block.index('_assert_studio_owned_or_absent "$LLAMA_CPP_DIR"') + # Anchor on the actual command-array entry, not the why-comment mention. + helper_idx = block.index('python "$SCRIPT_DIR/install_llama_prebuilt.py"') + assert ( + guard_idx < helper_idx + ), "ownership guard must precede the install_llama_prebuilt.py call" + + +def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard(): + """Mirror check for studio/setup.ps1: prebuilt llama.cpp path must + call Assert-StudioOwnedOrAbsent before invoking install_llama_prebuilt.py.""" + src = SETUP_PS1.read_text() + idx = src.index("installing prebuilt llama.cpp bundle (preferred path)") + block = src[idx : idx + 2000] + assert ( + 'Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install"' + in block + ), "setup.ps1 must guard the prebuilt llama.cpp path with Assert-StudioOwnedOrAbsent" + guard_idx = block.index("Assert-StudioOwnedOrAbsent -Path $LlamaCppDir") + # Anchor on the actual command-array entry, not the why-comment mention. + helper_idx = block.index('"$PSScriptRoot\\install_llama_prebuilt.py"') + assert ( + guard_idx < helper_idx + ), "Assert-StudioOwnedOrAbsent must precede the install_llama_prebuilt.py call" + + +def test_env_mode_passes_when_venv_marker_present(tmp_path): + """install.sh env-mode guard must accept the in-VENV + .unsloth-studio-owned marker as a primary sentinel so a partial + install (uv venv created, sentinels not yet written) is recoverable + by re-running install.sh.""" + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env", create_venv_marker = True) + assert res.returncode == 0, ( + f"in-VENV marker must allow cleanup; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert "RESULT=ok" in res.stdout + assert not (studio_home / "unsloth_studio").exists() + + +def test_env_mode_blocks_when_bin_unsloth_is_symlink_to_directory(tmp_path): + """install.sh env-mode guard must NOT accept a symlink-to-directory at + bin/unsloth as a Studio sentinel. Iter1's standalone -L test let any + symlink (including symlinks to dirs and broken symlinks) bypass the + guard; iter2 dropped that test so only -f (file or symlink-to-file) + counts.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (venv / "important.txt").write_text("keep me") + (studio_home / "bin").mkdir(parents = True) + target_dir = studio_home / "bin" / "unsloth-target-dir" + target_dir.mkdir() + (studio_home / "bin" / "unsloth").symlink_to(target_dir) + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode != 0, ( + "symlink-to-directory at bin/unsloth must NOT pass; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert (venv / "important.txt").is_file(), "unrelated workspace data must survive" + + +def test_env_mode_blocks_when_bin_unsloth_is_broken_symlink(tmp_path): + """install.sh guard must reject a broken symlink at bin/unsloth.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (venv / "important.txt").write_text("keep me") + (studio_home / "bin").mkdir(parents = True) + (studio_home / "bin" / "unsloth").symlink_to(studio_home / "bin" / "does-not-exist") + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode != 0, ( + "broken symlink at bin/unsloth must NOT pass; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert (venv / "important.txt").is_file() + + +def test_install_sh_writes_venv_marker_after_uv_venv(): + """install.sh must write the .unsloth-studio-owned marker into + $VENV_DIR right after `uv venv` succeeds so the env-mode deletion + guard accepts it on the next install run.""" + src = INSTALL_SH.read_text() + create_idx = src.index('run_install_cmd "create venv" uv venv "$VENV_DIR"') + tail = src[create_idx : create_idx + 600] + assert ( + ".unsloth-studio-owned" in tail + ), "install.sh must write .unsloth-studio-owned after uv venv create" + + +def test_install_ps1_writes_venv_marker_after_uv_venv(): + """install.ps1 must write the .unsloth-studio-owned marker into + $VenvDir after `uv venv` succeeds.""" + src = INSTALL_PS1.read_text() + venv_create = src.index("uv venv $VenvDir --python") + tail = src[venv_create : venv_create + 1500] + assert ( + ".unsloth-studio-owned" in tail + ), "install.ps1 must write .unsloth-studio-owned after uv venv create" + + +def test_install_ps1_guard_accepts_venv_marker(): + """install.ps1 env-mode guard must accept the in-VENV + .unsloth-studio-owned marker as a primary sentinel.""" + src = INSTALL_PS1.read_text() + block_start = src.index("if (Test-Path -LiteralPath $VenvPython)") + block = src[block_start : block_start + 2000] + assert ( + '$VenvDir ".unsloth-studio-owned") -PathType Leaf' in block + ), "install.ps1 guard must check the in-VENV marker with -PathType Leaf" + + +def test_setup_helpers_gate_on_canonical_custom_root(): + """Both _assert_studio_owned_or_absent (setup.sh) and + Assert-StudioOwnedOrAbsent (setup.ps1) must gate on a canonical + custom-vs-legacy comparison so an explicit override that resolves + to the legacy default does not trip the guard for pre-PR T5 + sidecar venvs or llama.cpp dirs.""" + sh_src = SETUP_SH.read_text() + sh_idx = sh_src.index("_assert_studio_owned_or_absent() {") + sh_func = sh_src[sh_idx : sh_idx + 600] + assert ( + '"$_STUDIO_HOME_IS_CUSTOM" = true' in sh_func + ), "setup.sh _assert_studio_owned_or_absent must gate on _STUDIO_HOME_IS_CUSTOM" + 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" + + ps_src = SETUP_PS1.read_text() + ps_idx = ps_src.index("function Assert-StudioOwnedOrAbsent") + ps_func = ps_src[ps_idx : ps_idx + 800] + assert ( + "$StudioHomeIsCustom -and" in ps_func + ), "setup.ps1 Assert-StudioOwnedOrAbsent must gate on $StudioHomeIsCustom" + assert ( + "$StudioOwnedMarker) -PathType Leaf" in ps_func + ), "setup.ps1 marker check must use -PathType Leaf so a directory cannot satisfy it" + + +def test_setup_ps1_inplace_git_sync_marks_studio_owned(): + """setup.ps1 in-place git-sync branch (when $LlamaCppDir/.git exists) + must call Mark-StudioOwned after a successful sync so a later prebuilt + update path's Assert-StudioOwnedOrAbsent does not exit.""" + src = SETUP_PS1.read_text() + inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")') + # The in-place branch ends just before the temp-dir clone branch. + clone_idx = src.index("Cloning llama.cpp @", inplace_idx) + inplace_block = src[inplace_idx:clone_idx] + assert ( + "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" + + +def test_setup_ps1_inplace_git_sync_asserts_studio_owned_before_mutation(): + """setup.ps1 in-place git-sync branch must call Assert-StudioOwnedOrAbsent + BEFORE any destructive git operation (remote set-url, checkout -B, clean + -fdx). Asymmetric to the prebuilt path and the temp-dir-swap path which + both guard.""" + src = SETUP_PS1.read_text() + inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")') + clone_idx = src.index("Cloning llama.cpp @", inplace_idx) + inplace_block = src[inplace_idx:clone_idx] + assert ( + "Assert-StudioOwnedOrAbsent -Path $LlamaCppDir" in inplace_block + ), "in-place git-sync must Assert-StudioOwnedOrAbsent before mutating $LlamaCppDir" + guard_idx = inplace_block.index("Assert-StudioOwnedOrAbsent -Path $LlamaCppDir") + git_idx = inplace_block.index("git -C $LlamaCppDir remote set-url") + assert ( + guard_idx < git_idx + ), "Assert-StudioOwnedOrAbsent must precede the first git mutation" + + +def _extract_check_health_function() -> str: + src = INSTALL_SH.read_text() + fn_start = src.index("_check_health() {") + fn_end = src.index("\n}\n", fn_start) + 2 + return src[fn_start:fn_end] + + +def _run_check_health(expected_root_id: str, response_json: str) -> int: + fn = _extract_check_health_function() + script = ( + f"_EXPECTED_STUDIO_ROOT_ID={expected_root_id!r}\n" + "_http_get() { printf '%s' \"$1\"; }\n" + + fn.replace( + '_resp=$(_http_get "http://127.0.0.1:$_port/api/health") || return 1', + f"_resp={response_json!r}", + ) + + "\n_check_health 8888\n" + "echo rc=$?\n" + ) + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + rc_lines = [l for l in res.stdout.splitlines() if l.startswith("rc=")] + return int(rc_lines[0].split("=")[1]) if rc_lines else res.returncode + + +def test_check_health_accepts_matching_studio_root_id(): + """Hex digest baked at install time matches the backend's + /api/health studio_root_id -- launcher attaches to its own backend.""" + expected_id = "a" * 64 + rc = _run_check_health( + expected_id, + f'{{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"{expected_id}"}}', + ) + assert rc == 0, f"matching studio_root_id must allow attach (rc={rc})" + + +def test_check_health_rejects_mismatched_studio_root_id(): + """Different install root → different sha256 → reject. Workspace + isolation: launcher A must not open Studio B running on the same port.""" + expected_id = "a" * 64 + other_id = "b" * 64 + rc = _run_check_health( + expected_id, + f'{{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"{other_id}"}}', + ) + assert rc != 0, "mismatched studio_root_id must reject attach (workspace isolation)" + + +def test_check_health_rejects_missing_studio_root_id_field(): + """A backend that omits studio_root_id (older or non-conforming) must + not be attached to when an expected id is baked into the launcher.""" + expected_id = "a" * 64 + rc = _run_check_health( + expected_id, + '{"status":"healthy","service":"Unsloth UI Backend"}', + ) + assert rc != 0, "missing studio_root_id field must reject attach" + + +def test_check_health_no_baked_id_accepts_any_healthy_backend(): + """If _EXPECTED_STUDIO_ROOT_ID is empty (e.g. install-time hash failed + to compute), the launcher falls back to the legacy contract and accepts + any healthy Unsloth backend.""" + rc = _run_check_health( + "", + '{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"deadbeef"}', + ) + assert rc == 0, "no baked id → accept any healthy Unsloth backend" + + +def test_check_health_rejects_non_unsloth_service(): + rc = _run_check_health( + "", + '{"status":"healthy","service":"Other UI Backend"}', + ) + assert rc != 0, "non-Unsloth service must be rejected" + + +def test_check_health_handles_arbitrary_id_token(): + """Iter3 used a raw shell match against the JSON-escaped studio_root, + which failed for paths containing `\\` or `"` (FastAPI emits `\\\\` and + `\\\"`). The per-install id token is hex-only by construction, so its + JSON form has no escapes regardless of where the install lives or what + the path contains. This test pins the round-trip on a fully arbitrary + 64-char hex token.""" + expected_id = "f0" + ("ed" * 31) # 64 hex chars, not derived from any path + rc = _run_check_health( + expected_id, + f'{{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"{expected_id}"}}', + ) + assert ( + rc == 0 + ), "arbitrary 64-hex install id must round-trip cleanly (no JSON escape issue)" + + +def test_install_ps1_test_studio_health_verifies_studio_root_id(): + """install.ps1 Test-StudioHealth must compare studio_root_id against + the install-time-baked $_ExpectedStudioRootId, not the runtime env var.""" + src = INSTALL_PS1.read_text() + fn_start = src.index("function Test-StudioHealth") + fn_end = src.index("\n}\n", fn_start) + 2 + fn = src[fn_start:fn_end] + assert ( + "studio_root_id" in fn + ), "Test-StudioHealth must inspect the studio_root_id field" + assert ( + "$_ExpectedStudioRootId" in fn + ), "Test-StudioHealth must compare against the install-time baked $_ExpectedStudioRootId" + + +def test_install_ps1_bakes_studio_root_id_into_launcher(): + """install.ps1 must persist a per-install opaque id at + $StudioHome\\share\\studio_install_id and bake the value into the + generated launcher as $_ExpectedStudioRootId so the launcher can + verify the backend belongs to THIS install. The id is generated + via a CSPRNG so /api/health does not leak the install path.""" + src = INSTALL_PS1.read_text() + assert ( + "$_studioRootId" in src + ), "install.ps1 must compute $_studioRootId for the launcher" + assert ( + '"share"' in src and "studio_install_id" in src + ), "install.ps1 must persist the id at $StudioHome\\share\\studio_install_id" + assert ( + "RandomNumberGenerator" in src + ), "install.ps1 must seed the id from a CSPRNG (RandomNumberGenerator)" + assert ( + "$_ExpectedStudioRootId" in src + ), "install.ps1 must bake $_ExpectedStudioRootId into the launcher" + + +def test_health_endpoint_exposes_studio_root_id_not_raw_path(): + """studio/backend/main.py /api/health must expose studio_root_id (a + hex digest) and NOT the raw studio_root path. Studio supports + `-H 0.0.0.0`; an unauthenticated /api/health that returns the raw + install path leaks username, home dir, workspace name, etc.""" + main_py = REPO_ROOT / "studio" / "backend" / "main.py" + src = main_py.read_text() + health_idx = src.index('@app.get("/api/health")') + health_block = src[health_idx : health_idx + 1500] + assert ( + '"studio_root_id"' in health_block + ), "/api/health must expose studio_root_id (hex digest)" + assert ( + '"studio_root":' not in health_block + ), "/api/health must NOT expose the raw studio_root path (information disclosure)" + assert ( + "_studio_root_id()" in health_block + ), "/api/health must call the _studio_root_id helper" + + +def test_install_sh_bakes_studio_root_id_into_launcher(): + """install.sh must persist a per-install opaque id at + $STUDIO_HOME/share/studio_install_id and substitute its content into + the launcher heredoc placeholder for ALL modes (env / home / default), + so the launcher's _check_health rejects sibling Studios on the same + port. The id is seeded from /dev/urandom (or python3 secrets fallback) + so /api/health does not leak the install path.""" + src = INSTALL_SH.read_text() + assert ( + "_css_studio_root_id" in src + ), "install.sh must compute _css_studio_root_id for the launcher" + assert ( + '_css_id_file="$_css_id_dir/studio_install_id"' in src + ), "install.sh must persist the id at $STUDIO_HOME/share/studio_install_id" + assert ( + "od -An -N32 -tx1 /dev/urandom" in src + ), "install.sh must seed new ids from /dev/urandom (CSPRNG)" + assert ( + "@@STUDIO_ROOT_ID@@" in src + ), "install.sh must use @@STUDIO_ROOT_ID@@ placeholder in the launcher heredoc" + assert ( + "s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g" in src + ), "install.sh must sed-substitute @@STUDIO_ROOT_ID@@ unconditionally (not just env-mode)" + + +def test_tauri_preflight_scrubs_studio_home_env(): + """All three Tauri CLI-spawn sites that lacked the scrub must now + env_remove UNSLOTH_STUDIO_HOME and STUDIO_HOME, mirroring + process.rs / install.rs / desktop_auth.rs / update.rs.""" + preflight = ( + REPO_ROOT / "studio" / "src-tauri" / "src" / "preflight.rs" + ).read_text() + commands = (REPO_ROOT / "studio" / "src-tauri" / "src" / "commands.rs").read_text() + # Both functions in preflight.rs (run_cli_probe + probe_cli_capability) + # must scrub. Count occurrences -- expect 2 in preflight, 1 in commands. + assert ( + preflight.count('cmd.env_remove("UNSLOTH_STUDIO_HOME")') >= 2 + ), "preflight.rs must scrub UNSLOTH_STUDIO_HOME in both run_cli_probe and probe_cli_capability" + assert ( + preflight.count('cmd.env_remove("STUDIO_HOME")') >= 2 + ), "preflight.rs must scrub STUDIO_HOME in both run_cli_probe and probe_cli_capability" + assert ( + 'cmd.env_remove("UNSLOTH_STUDIO_HOME")' in commands + ), "commands.rs check_install_status must scrub UNSLOTH_STUDIO_HOME" + assert ( + 'cmd.env_remove("STUDIO_HOME")' in commands + ), "commands.rs check_install_status must scrub STUDIO_HOME" + + +def test_install_sh_shim_uses_atomic_replace(): + """install.sh shim install must use ln -sfn for atomic replace; the + older `rm -f ...; ln -s ...` left a window where the shim was missing.""" + src = INSTALL_SH.read_text() + shim_idx = src.index('_shim_path="$_LOCAL_BIN/unsloth"') + block = src[shim_idx : shim_idx + 1500] + assert ( + 'ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path"' in block + ), "install.sh must use ln -sfn for atomic shim replacement" + assert ( + 'rm -f -- "$_shim_path"' not in block + ), "the explicit rm + ln pair must be replaced by atomic ln -sfn" + + +def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback( + tmp_path, +): + """_create_shortcuts must seed new ids from /dev/urandom first (no + interpreter spawn cost on the install hot path) and fall back to + `python3 -c 'secrets.token_hex(32)'` only when urandom is unreadable. + Re-running the function with an existing id file must not regenerate + the id (otherwise re-runs would invalidate previously-baked launchers).""" + src = INSTALL_SH.read_text() + fn_start = src.index('_css_data_dir="$DATA_DIR"') + block = src[fn_start : fn_start + 3000] + urandom_idx = block.index("od -An -N32 -tx1 /dev/urandom") + py_fallback_idx = block.index("python3 -c 'import secrets;", urandom_idx) + assert ( + urandom_idx < py_fallback_idx + ), "/dev/urandom must be tried before the python3 secrets fallback" + # The id file is checked for non-empty content before we generate; this is + # what makes re-runs idempotent. + assert ( + 'if [ ! -s "$_css_id_file" ]; then' in block + ), "install.sh must skip id generation when the file already has content" + + # Behavioral check: extract the generation block and run it in isolation + # twice to confirm idempotence. + studio_home = tmp_path / "studio" + (studio_home / "share").mkdir(parents = True) + gen_script = ( + f'STUDIO_HOME="{studio_home}"\n' + '_css_id_dir="$STUDIO_HOME/share"\n' + '_css_id_file="$_css_id_dir/studio_install_id"\n' + # Replicate the generation block (kept narrowly so the test fails loud + # if install.sh changes the surrounding contract). + "gen() {\n" + ' if [ ! -s "$_css_id_file" ]; then\n' + ' _css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d " \\n")\n' + ' printf "%s" "$_css_new_id" > "$_css_id_file.$$.tmp"\n' + ' mv "$_css_id_file.$$.tmp" "$_css_id_file"\n' + " fi\n" + ' cat "$_css_id_file"\n' + "}\n" + "a=$(gen); b=$(gen)\n" + '[ "$a" = "$b" ] || { echo MISMATCH; exit 1; }\n' + 'echo "ID=$a"\n' + 'echo "LEN=${#a}"\n' + ) + res = subprocess.run(["bash", "-c", gen_script], text = True, capture_output = True) + assert res.returncode == 0, res.stderr + out = dict( + line.split("=", 1) for line in res.stdout.strip().splitlines() if "=" in line + ) + assert ( + out.get("LEN") == "64" + ), f"id must be 64 hex chars, got LEN={out.get('LEN')!r}" + assert all( + c in "0123456789abcdef" for c in out.get("ID", "") + ), f"id must be lowercase hex, got {out.get('ID')!r}" + + +def test_install_sh_create_shortcuts_fails_fast_when_no_entropy(): + """If neither /dev/urandom nor python3 is available, _create_shortcuts + must `return 1` instead of silently baking an empty studio_root_id + (which would disable the launcher's same-install discriminator).""" + src = INSTALL_SH.read_text() + fn_start = src.index('_css_data_dir="$DATA_DIR"') + block = src[fn_start : fn_start + 3000] + assert ( + "[WARN] Cannot create launcher: no entropy source for studio_install_id" + in block + ), "install.sh must warn when neither urandom nor python3 is available" + assert ( + "[WARN] Cannot create launcher: failed to read" in block + ), "install.sh must warn when the id file read produces no content" + assert ( + block.count("return 1") >= 2 + ), "both the no-entropy branch and the empty-read branch must `return 1`" + + +def test_install_sh_bakes_installed_is_env_mode_flag_in_launcher(): + """install.sh must bake the install-time mode (env vs default/home) into + the generated launcher so PORT_FILE / namespaced LOCK_DIR cannot be + flipped on by a sourced custom-root studio.conf in the user's shell.""" + src = INSTALL_SH.read_text() + assert ( + "_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'" in src + ), "launcher heredoc must declare _INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'" + assert ( + "_css_is_env_mode=false" in src + ), "install.sh must default _css_is_env_mode to false" + assert ( + '[ "$_STUDIO_HOME_REDIRECT" = "env" ] && _css_is_env_mode=true' in src + ), "install.sh must set _css_is_env_mode=true only when _STUDIO_HOME_REDIRECT=env" + assert ( + "s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g" in src + ), "install.sh sed pipeline must substitute @@INSTALLED_IS_ENV_MODE@@" + + +def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env(): + """The launcher's PORT_FILE / namespaced LOCK_DIR must be gated on the + baked $_INSTALLED_IS_ENV_MODE flag, not the runtime $UNSLOTH_STUDIO_HOME. + Sourcing a custom-root studio.conf in shell must not flip a default-mode + launcher into env-mode behavior.""" + src = INSTALL_SH.read_text() + heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'") + heredoc_end = src.index("LAUNCHER_EOF\n", heredoc_start) + heredoc = src[heredoc_start:heredoc_end] + assert ( + 'if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then' in heredoc + ), "launcher must gate PORT_FILE/LOCK_DIR on baked _INSTALLED_IS_ENV_MODE" + port_block_start = heredoc.index('if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then') + port_block_end = heredoc.index("\nfi\n", port_block_start) + len("\nfi\n") + port_block = heredoc[port_block_start:port_block_end] + assert 'PORT_FILE="$DATA_DIR/studio.port"' in port_block + assert ( + 'if [ -n "${UNSLOTH_STUDIO_HOME:-}" ]; then\n if command -v cksum' + not in heredoc + ), "launcher must NOT gate PORT_FILE on runtime UNSLOTH_STUDIO_HOME" + + def _run_launcher_gate(installed_flag: str, runtime_env: dict) -> str: + # Reproduce just the LOCK_DIR/PORT_FILE init block in isolation. + script = ( + f"_INSTALLED_IS_ENV_MODE={installed_flag!r}\n" + "DATA_DIR=/tmp/test_data_dir\n" + 'LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock"\n' + 'PORT_FILE=""\n' + port_block + '\necho "PORT_FILE=$PORT_FILE"\n' + ) + env = {"PATH": "/usr/bin:/bin"} + env.update(runtime_env) + res = subprocess.run( + ["bash", "-c", script], + text = True, + capture_output = True, + env = env, + ) + for line in res.stdout.splitlines(): + if line.startswith("PORT_FILE="): + return line[len("PORT_FILE=") :] + return "" + + # default-mode install should NEVER set PORT_FILE, even if UNSLOTH_STUDIO_HOME leaks in. + assert ( + _run_launcher_gate("false", {"UNSLOTH_STUDIO_HOME": "/tmp/leaked"}) == "" + ), "default-mode launcher must keep PORT_FILE empty even with UNSLOTH_STUDIO_HOME in env" + # env-mode install should set PORT_FILE regardless of runtime env. + assert ( + _run_launcher_gate("true", {}) == "/tmp/test_data_dir/studio.port" + ), "env-mode launcher must set PORT_FILE based on baked DATA_DIR" + + +def test_main_py_studio_root_id_caches_at_module_load(): + """_studio_root_id() is called on every /api/health poll; the id is + stable for the lifetime of the process so it must be read once at + module load and re-used (avoids a hot-path filesystem probe and + protects against transient FS errors during health polling).""" + main_py = (REPO_ROOT / "studio" / "backend" / "main.py").read_text() + assert ( + "_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()" in main_py + ), "main.py must populate _STUDIO_ROOT_ID_CACHE from _read_studio_install_id() at module load" + fn_idx = main_py.index("def _studio_root_id() -> str:") + next_def_idx = main_py.index("\ndef ", fn_idx + 1) + fn_block = main_py[fn_idx:next_def_idx] + assert ( + "return _STUDIO_ROOT_ID_CACHE" in fn_block + ), "_studio_root_id() body must return the cached value" + assert ( + "read_text(" not in fn_block and "hashlib" not in fn_block + ), "_studio_root_id() must NOT do filesystem or hash work on every call" + + +def test_main_py_read_studio_install_id_validates_hex_and_handles_missing( + tmp_path, monkeypatch +): + """_read_studio_install_id reads $STUDIO_HOME/share/studio_install_id and + returns "" when the file is absent, empty, contains non-hex content, or + is the wrong length. "" triggers the launcher's "no baked id, accept any + healthy backend" fallback path (see test_check_health_no_baked_id_*). + Behavioral check: spin up a stub _STUDIO_ROOT_RESOLVED and exercise + _read_studio_install_id directly without importing main.py (which + pulls in heavy deps). Test the rejection rules verbatim.""" + import re + + pattern = re.compile(r"^[0-9a-f]{64}$") + + def _read(root: Path) -> str: + # Mirror the implementation; this test pins the exact contract so a + # future refactor can't silently widen what's accepted. + try: + token = (root / "share" / "studio_install_id").read_text().strip() + except (OSError, ValueError): + return "" + return token if pattern.fullmatch(token) else "" + + root = tmp_path / "studio" + (root / "share").mkdir(parents = True) + + # Missing file -> empty + assert _read(root) == "" + + id_file = root / "share" / "studio_install_id" + # Empty file -> empty + id_file.write_text("") + assert _read(root) == "" + # Non-hex content -> empty + id_file.write_text( + "not-a-hex-id-just-text-padded-to-64-chars-zzzzzzzzzzzzzzzzzzzzzz" + ) + assert _read(root) == "" + # Uppercase hex -> empty (must be lowercase) + id_file.write_text("F" * 64) + assert _read(root) == "" + # Wrong length -> empty (32 chars, not 64) + id_file.write_text("a" * 32) + assert _read(root) == "" + # Valid 64-char lowercase hex with surrounding whitespace -> stripped+accepted + valid = "0123456789abcdef" * 4 + id_file.write_text(f"\n {valid} \n") + assert _read(root) == valid + + +def test_llama_cpp_search_roots_handles_studio_root_oserror(): + """_find_llama_server_binary calls studio_root() which can raise + OSError or ValueError from Path.expanduser().resolve() (broken symlink, + null byte). The except clause must mirror sibling _kill_orphaned_servers + (which catches the same trio) so inference startup does not crash.""" + llama_cpp = ( + REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" + ).read_text() + find_block_start = llama_cpp.index("_find_llama_server_binary") + find_block = llama_cpp[find_block_start : find_block_start + 4000] + assert ( + "except (ImportError, OSError, ValueError):" in find_block + ), "_find_llama_server_binary must catch (ImportError, OSError, ValueError) from studio_root()" + kill_def_idx = llama_cpp.index("def _kill_orphaned_servers") + kill_block = llama_cpp[kill_def_idx : kill_def_idx + 4000] + assert ( + "except (ImportError, OSError, ValueError):" in kill_block + ), "sibling _kill_orphaned_servers must keep its (ImportError, OSError, ValueError) handler" + + +def test_install_sh_install_id_survives_symlinked_studio_home(tmp_path): + """End-to-end behavioral check: when $STUDIO_HOME is reached via a + symlinked parent (e.g. symlinked $HOME on Linux, junctioned %USERPROFILE% + on Windows), install.sh and the backend agree on the install id BY + CONSTRUCTION because the id is read from a file whose location resolves + the same way for both. The previous sha256(canonical_path) scheme + required `cd -P/pwd -P` and Path.resolve() to produce identical strings, + which broke under symlinks/junctions and required cycles 17-27 of the + PR's review history to fully canonicalize. This is the regression test + pinning that the new design has no such drift.""" + real = tmp_path / "realhome" + real.mkdir() + link = tmp_path / "linkhome" + link.symlink_to(real) + studio_home = real / ".unsloth" / "studio" + (studio_home / "share").mkdir(parents = True) + # Write a stub install id at the canonical location. + valid_id = "ab12" * 16 + (studio_home / "share" / "studio_install_id").write_text(valid_id) + # Read it back via both the canonical and the symlinked path; both must + # see the SAME content (which is what makes install.sh's cat and the + # backend's read_text agree without any canonicalization dance). + raw_via_link = link / ".unsloth" / "studio" / "share" / "studio_install_id" + raw_direct = studio_home / "share" / "studio_install_id" + assert raw_via_link.read_text() == valid_id + assert raw_direct.read_text() == valid_id + # And install.sh's `cat` would see the same. + import subprocess as _sp + + res = _sp.run(["cat", str(raw_via_link)], capture_output = True, text = True) + assert res.returncode == 0 + assert res.stdout == valid_id + + +def test_install_sh_substitutes_root_id_before_data_dir(): + """The two-stage sed substitution must bake @@STUDIO_ROOT_ID@@ / + @@INSTALLED_IS_ENV_MODE@@ first (non-user-controlled), then @@DATA_DIR@@ + (user-controlled). A custom $DATA_DIR containing the literal text + @@STUDIO_ROOT_ID@@ must not be mutated by the global root-id sed pass.""" + src = INSTALL_SH.read_text() + root_id_idx = src.index("s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g") + env_mode_idx = src.index("s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g") + data_dir_idx = src.index("s|@@DATA_DIR@@|$_sed_safe|g") + assert root_id_idx < data_dir_idx, ( + "@@STUDIO_ROOT_ID@@ substitution must happen BEFORE @@DATA_DIR@@ " + "(non-user-controlled placeholders first)" + ) + assert ( + env_mode_idx < data_dir_idx + ), "@@INSTALLED_IS_ENV_MODE@@ substitution must happen BEFORE @@DATA_DIR@@" + + +def test_install_sh_root_id_pass_does_not_mutate_user_data_dir(tmp_path): + """Behavioral subprocess test: a $DATA_DIR containing the literal text + `@@STUDIO_ROOT_ID@@` must not be mutated when the placeholder pass runs + first; only the actual placeholder occurrences in the launcher template + are replaced.""" + src = INSTALL_SH.read_text() + heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'") + heredoc_body_start = src.index("\n", heredoc_start) + 1 + heredoc_body_end = src.index("LAUNCHER_EOF\n", heredoc_start) + template = src[heredoc_body_start:heredoc_body_end] + launcher_path = tmp_path / "launch.sh" + launcher_path.write_text(template) + # Run the iter6 sed order: root-id first, then data-dir. + weird_data_dir = "/tmp/with-@@STUDIO_ROOT_ID@@/share" + root_id = "deadbeef" * 8 + is_env = "true" + script = f""" +sed -e "s|@@STUDIO_ROOT_ID@@|{root_id}|g" \\ + -e "s|@@INSTALLED_IS_ENV_MODE@@|{is_env}|g" \\ + "{launcher_path}" > "{launcher_path}.tmp" && mv "{launcher_path}.tmp" "{launcher_path}" +_sq_escaped=$(printf '%s' "{weird_data_dir}" | sed "s/'/'\\\\\\\\''/g") +_sed_safe=$(printf '%s' "$_sq_escaped" | sed 's/[\\\\&|]/\\\\&/g') +sed "s|@@DATA_DIR@@|$_sed_safe|g" "{launcher_path}" > "{launcher_path}.tmp" \\ + && mv "{launcher_path}.tmp" "{launcher_path}" +""" + subprocess.run(["bash", "-c", script], check = True) + final = launcher_path.read_text() + assert ( + f"DATA_DIR='{weird_data_dir}'" in final + ), f"DATA_DIR must be preserved verbatim (no @@STUDIO_ROOT_ID@@ mutation); got: {final[:500]}" + assert ( + f"_EXPECTED_STUDIO_ROOT_ID='{root_id}'" in final + ), "STUDIO_ROOT_ID placeholder must still be substituted in the launcher heredoc" + + +def test_install_ps1_install_id_file_layout_matches_backend_read_path(): + """install.ps1 must write the id at $StudioHome\\share\\studio_install_id + so the backend (studio/backend/main.py:_read_studio_install_id) can find + it via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" without + mode-specific path knowledge. Persistence-across-runs is enforced by the + pre-write Test-Path check.""" + src = INSTALL_PS1.read_text() + id_idx = src.index('$_studioIdDir = Join-Path $StudioHome "share"') + context = src[id_idx : id_idx + 1500] + assert ( + '$_studioIdFile = Join-Path $_studioIdDir "studio_install_id"' in context + ), "install.ps1 must persist the id at $StudioHome\\share\\studio_install_id" + assert ( + "Test-Path -LiteralPath $_studioIdFile" in context + ), "install.ps1 must skip id generation when the file already has content (re-run idempotence)" + assert ( + "RandomNumberGenerator" in context and "GetBytes($_idBytes)" in context + ), "install.ps1 must seed new ids from a CSPRNG (RandomNumberGenerator)" + assert ( + "Move-Item -LiteralPath $_idTmp" in context + ), "install.ps1 must atomic-rename the temp file into place to avoid half-written ids" diff --git a/tests/test_studio_root_resilience.py b/tests/test_studio_root_resilience.py new file mode 100644 index 0000000000..1ce0430dc4 --- /dev/null +++ b/tests/test_studio_root_resilience.py @@ -0,0 +1,154 @@ +"""Resilience checks for Studio install-root inference under hostile +filesystem conditions: +- _infer_studio_home_from_venv must NOT propagate PermissionError / + OSError out through studio_root() (it would crash module import in + run.py / main.py / transformers_version.py / model_config.py). +- _kill_orphaned_servers must catch (ImportError, OSError, ValueError) + on the studio_root() probe so a transient resolve / sentinel failure + cannot crash server startup. +- _find_llama_server_binary must keep the custom-root in search_roots + when the inner resolve() comparison itself fails.""" + +from __future__ import annotations + +import importlib.util +import re +import sys +import textwrap +from pathlib import Path +from unittest import mock + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +STORAGE_ROOTS = ( + REPO_ROOT / "studio" / "backend" / "utils" / "paths" / "storage_roots.py" +) +LLAMA_CPP = REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" + + +def _load(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +def test_infer_studio_home_swallows_permission_error(tmp_path, monkeypatch): + candidate = tmp_path / "fake_root" + venv = candidate / "unsloth_studio" + venv.mkdir(parents = True) + monkeypatch.setattr(sys, "prefix", str(venv)) + sys.modules.pop("sr_perm", None) + mod = _load("sr_perm", STORAGE_ROOTS) + with mock.patch.object(Path, "is_file", side_effect = PermissionError("denied")): + # Must NOT raise. + assert mod._infer_studio_home_from_venv() is None + + +def test_studio_root_does_not_crash_on_permission_error(tmp_path, monkeypatch): + """studio_root() must remain callable even when the venv inference + encounters a restricted filesystem; it should fall through to the + legacy default.""" + candidate = tmp_path / "fake_root" + venv = candidate / "unsloth_studio" + venv.mkdir(parents = True) + monkeypatch.setattr(sys, "prefix", str(venv)) + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + monkeypatch.delenv("STUDIO_HOME", raising = False) + sys.modules.pop("sr_studio_perm", None) + mod = _load("sr_studio_perm", STORAGE_ROOTS) + with mock.patch.object(Path, "is_file", side_effect = OSError("ebusy")): + result = mod.studio_root() + assert result == Path.home() / ".unsloth" / "studio" + + +def test_kill_orphan_catches_oserror_from_studio_root(): + """_kill_orphaned_servers must catch (ImportError, OSError, ValueError) + on the studio_root() probe specifically; the sister function + _find_llama_server_binary uses the same broader catch on its own probe.""" + src = LLAMA_CPP.read_text() + fn_start = src.index("def _kill_orphaned_servers") + fn_body = src[fn_start : fn_start + 4000] + # The studio_root() probe in this fn is the one that imports as `_sr` + # and assigns `_resolved_sr = _sr()`. Find the except that closes it. + probe_idx = fn_body.index("storage_roots import studio_root as _sr") + # The matching except is the next `except ...:` after the inner + # OSError/ValueError block that wraps resolve(). + after = fn_body[probe_idx:] + # Skip over the inner `except (OSError, ValueError):` that wraps resolve(). + inner_idx = after.index("except (OSError, ValueError):") + after_inner = after[inner_idx + len("except (OSError, ValueError):") :] + outer_match = re.search(r"except\s*\(?[^)]*?\)?:", after_inner) + assert outer_match, "outer except for studio_root probe missing" + clause = outer_match.group(0) + assert ( + "OSError" in clause and "ValueError" in clause + ), f"_kill_orphaned_servers studio_root probe catch too narrow: {clause!r}" + + +def _exec_search_roots_block( + home: Path, studio_root_value: Path, resolve_raises: bool +) -> list[Path]: + """Extract _find_llama_server_binary's env-mode search_roots block + and execute it with controlled inputs.""" + src = LLAMA_CPP.read_text() + block_start = src.index('legacy_llama = Path.home() / ".unsloth" / "llama.cpp"') + block_end = src.index("_seen_roots: set[str]", block_start) + raw = src[block_start:block_end] + indent = " " * 8 + block = textwrap.dedent(indent + raw) + fake_module = type(sys)("fake_storage_roots") + fake_module.studio_root = lambda: studio_root_value + sys.modules["utils.paths.storage_roots"] = fake_module + try: + original_resolve = Path.resolve + + def _resolve(self, *a, **k): + if resolve_raises: + raise OSError("ebusy") + return original_resolve(self, *a, **k) + + with ( + mock.patch.object(Path, "home", classmethod(lambda cls: home)), + mock.patch.object(Path, "resolve", _resolve), + ): + ns: dict = {"Path": Path} + exec(block, ns) # noqa: S102 + return ns["search_roots"] + finally: + sys.modules.pop("utils.paths.storage_roots", None) + + +def test_search_roots_keeps_custom_when_resolve_fails(tmp_path): + home = tmp_path / "home" + home.mkdir() + custom = tmp_path / "custom_studio" + custom.mkdir() + roots = _exec_search_roots_block( + home = home, studio_root_value = custom, resolve_raises = True + ) + # On resolve() failure, the inner except falls back to direct equality; + # custom != legacy_studio so the custom root must remain in search_roots. + assert ( + custom / "llama.cpp" in roots + ), f"custom root dropped on resolve() failure: {roots}" + # custom-mode discovery excludes the legacy tree to match _kill_orphaned_servers. + assert ( + (home / ".unsloth" / "llama.cpp") not in roots + ), f"legacy llama path must not appear in custom-mode search_roots: {roots}" + + +def test_search_roots_default_mode_uses_legacy_only(tmp_path): + home = tmp_path / "home" + home.mkdir() + legacy = home / ".unsloth" / "studio" + legacy.mkdir(parents = True) + roots = _exec_search_roots_block( + home = home, studio_root_value = legacy, resolve_raises = False + ) + # Default mode: only legacy_llama. + assert roots == [home / ".unsloth" / "llama.cpp"] diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 140940209f..76aac3dc15 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -20,7 +20,73 @@ import typer studio_app = typer.Typer(help = "Unsloth Studio commands.") -STUDIO_HOME = Path.home() / ".unsloth" / "studio" + +# Resolve install root: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then +# sys.prefix inference (so a direct call to /bin/unsloth resolves after +# the installer's env var has expired), then legacy ~/.unsloth/studio. +# UNSLOTH_STUDIO_HOME wins when both env vars are set. +def _looks_like_installer_managed_studio_home(candidate: Path) -> bool: + """Sentinel check (studio.conf or bin shim) so a dev venv named + unsloth_studio is not misidentified as a custom Studio root. + """ + shim_name = "unsloth.exe" if platform.system() == "Windows" else "unsloth" + return (candidate / "share" / "studio.conf").is_file() or ( + candidate / "bin" / shim_name + ).is_file() + + +def _resolve_studio_home() -> tuple[Path, bool]: + override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() + if not override: + override = (os.environ.get("STUDIO_HOME") or "").strip() + if override: + try: + return Path(override).expanduser().resolve(), True + except (OSError, ValueError): + return Path(override).expanduser(), True + try: + prefix = Path(sys.prefix).resolve() + if prefix.name == "unsloth_studio": + inferred = prefix.parent + legacy = (Path.home() / ".unsloth" / "studio").resolve() + if inferred != legacy and _looks_like_installer_managed_studio_home( + inferred + ): + return inferred, True + except (OSError, ValueError): + pass + return Path.home() / ".unsloth" / "studio", False + + +STUDIO_HOME, _STUDIO_HOME_IS_CUSTOM = _resolve_studio_home() + + +def _ensure_studio_env_exported() -> None: + """Re-export UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH only for real + custom roots so subprocesses inherit the right install. Called from each + studio subcommand entry rather than at import time, to avoid leaking env + state into unrelated importers (tests, --help, CLI introspection). + """ + if not _STUDIO_HOME_IS_CUSTOM: + return + # Truthy-check (not setdefault) so a blank UNSLOTH_STUDIO_HOME= does not + # suppress the inferred custom root. + if not os.environ.get("UNSLOTH_STUDIO_HOME"): + os.environ["UNSLOTH_STUDIO_HOME"] = str(STUDIO_HOME) + # When override == legacy default, llama.cpp stays at ~/.unsloth/llama.cpp. + 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: + _llama_dir = Path.home() / ".unsloth" / "llama.cpp" + else: + _llama_dir = STUDIO_HOME / "llama.cpp" + if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_llama_dir) + + BOOTSTRAP_PASSWORD_FILE = ".bootstrap_password" DESKTOP_SECRET_FILE = ".desktop_secret" DEFAULT_ADMIN_USERNAME = "unsloth" @@ -427,6 +493,8 @@ def studio_default( ), ): """Launch the Unsloth Studio server.""" + # Runs before any subcommand; covers run/setup/update/etc in one place. + _ensure_studio_env_exported() if ctx.invoked_subcommand is not None: return