diff --git a/install.ps1 b/install.ps1 index e57b3f670e..16337f4b9c 100644 --- a/install.ps1 +++ b/install.ps1 @@ -144,6 +144,10 @@ function Install-UnslothStudio { # UNSLOTH_PYTHON pins the version (mirrors install.sh --python); default 3.13. $PythonVersion = if ($env:UNSLOTH_PYTHON) { $env:UNSLOTH_PYTHON } else { "3.13" } + # python.org fallback patch, used only when winget is unavailable/broken AND + # the live python.org listing can't be fetched. The installer URL scheme is + # stable so an older patch still installs. Bump alongside $PythonVersion. + $PythonFallbackFullVersion = "3.13.13" # Resolve install destinations. Priority: UNSLOTH_STUDIO_HOME, then # STUDIO_HOME alias, then USERPROFILE-redirect, then default. @@ -854,6 +858,7 @@ shell.Run cmd, 0, False try { $wshell = New-Object -ComObject WScript.Shell $createdShortcutCount = 0 + $createdShortcutPaths = @() foreach ($linkPath in @($desktopLink, $startMenuLink)) { if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue } try { @@ -867,12 +872,46 @@ shell.Run cmd, 0, False } $shortcut.Save() $createdShortcutCount++ + $createdShortcutPaths += $linkPath } catch { substep "could not create shortcut at ${linkPath}: $($_.Exception.Message)" "Yellow" } } if ($createdShortcutCount -gt 0) { substep "Created Unsloth Studio shortcut" + # Force Explorer to re-read each new shortcut's icon so it renders + # immediately instead of a stale/generic entry (a same-name .lnk + # recreated across reinstalls keeps Explorer's cached per-item icon). + # The reliable, non-disruptive fix (no explorer restart) is a per-item + # SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW per .lnk; the global + # SHCNE_ASSOCCHANGED broadcast alone does NOT recover a stale item. + # Also clear the on-disk icon cache (covers heavier staleness). + try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {} + try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {} + try { + Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);' -ErrorAction SilentlyContinue + # SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut + foreach ($scPath in $createdShortcutPaths) { + try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, $scPath, [System.IntPtr]::Zero) } catch {} + } + # SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders) + [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero) + } catch {} + # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN + # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT + # invalidate, so a rewritten same-name shortcut shows the old tile + # until the host restarts. Drop only the render caches (NEVER + # start2.bin -- the pinned layout) and let the host rebuild. + # Best-effort; Win10 has no such host (Test-Path skips it). + try { + $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" + if (Test-Path -LiteralPath $smehTemp) { + Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue | + Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } + } catch {} } else { substep "no Unsloth Studio shortcuts were created" "Yellow" } @@ -985,6 +1024,79 @@ shell.Run cmd, 0, False return $null } + # ── Fallback: install CPython directly from python.org ── + # Used when winget is unavailable or fails (notably msstore cert-pinning error + # 0x8a15005e, which aborts `winget install` unless --source winget is given). + # Downloads the official installer and runs it silently as a per-user install + # (no UAC), putting python.exe + the py launcher on PATH. Mirrors the uv -> + # astral.sh fallback below. Returns @{ Version; Path } or $null. + function Install-PythonFromPythonOrg { + # python.org ships one installer per architecture. + $archSuffix = switch (Get-TauriDiagArch) { + "x86_64" { "-amd64" } + "arm64" { "-arm64" } + "x86" { "" } + default { $null } + } + if ($null -eq $archSuffix) { + substep "No python.org installer is available for this architecture." "Yellow" + return $null + } + + # Resolve the latest $PythonVersion.x patch from the python.org listing, + # falling back to a same-minor version if the listing cannot be fetched. + # Use the pinned full version only when it matches the requested minor so a + # non-default UNSLOTH_PYTHON (e.g. 3.12) doesn't silently install 3.13. + $full = if ($PythonFallbackFullVersion -like "$PythonVersion.*") { $PythonFallbackFullVersion } else { "$PythonVersion.0" } + try { + $listing = [string](Invoke-RestMethod -Uri "https://www.python.org/ftp/python/" -UseBasicParsing -TimeoutSec 20) + $patches = [regex]::Matches($listing, ([regex]::Escape($PythonVersion) + '\.(\d+)/')) | + ForEach-Object { [int]$_.Groups[1].Value } | Sort-Object -Descending + if ($patches.Count -gt 0) { $full = "$PythonVersion.$($patches[0])" } + } catch {} + + $file = "python-$full$archSuffix.exe" + $url = "https://www.python.org/ftp/python/$full/$file" + $dest = Join-Path ([System.IO.Path]::GetTempPath()) $file + substep "downloading Python $full from python.org..." "Yellow" + try { + Invoke-WebRequest -Uri $url -OutFile $dest -UseBasicParsing + } catch { + substep "python.org download failed: $($_.Exception.Message)" "Yellow" + return $null + } + + # Per-user install => no UAC. PrependPath puts python + py on PATH; + # Include_launcher installs py.exe (preferred by Find-CompatiblePython). + substep "installing Python $full (silent, per-user)..." + $installArgs = @( + "/quiet", + "InstallAllUsers=0", + "PrependPath=1", + "Include_launcher=1", + # Launcher per-user too: Include_launcher defaults InstallLauncherAllUsers=1, + # which needs admin and would break this non-admin per-user fallback. + "InstallLauncherAllUsers=0", + "Include_pip=1", + "AssociateFiles=0", + "Shortcuts=0" + ) + $rc = 1 + try { + $proc = Start-Process -FilePath $dest -ArgumentList $installArgs -Wait -PassThru + $rc = $proc.ExitCode + } catch { + substep "python.org installer failed to start: $($_.Exception.Message)" "Yellow" + } finally { + Remove-Item -LiteralPath $dest -Force -ErrorAction SilentlyContinue + } + if ($rc -ne 0) { + substep "python.org installer exited with code $rc." "Yellow" + } + Refresh-SessionPath + return (Find-CompatiblePython) + } + # ── Install Python if no compatible version (3.11-3.13) found ── # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. Write-TauriLog "STEP" "Installing Python" @@ -994,47 +1106,62 @@ shell.Run cmd, 0, False step "python" "Python $($DetectedPython.Version) already installed" } if (-not $DetectedPython) { - if (-not $script:WingetAvailable) { - Write-Host "[ERROR] No compatible Python (3.11-3.13) found and winget is unavailable on this host." -ForegroundColor Red - Write-Host " Install Python $PythonVersion from https://www.python.org/downloads/" -ForegroundColor Yellow - Write-Host " and re-run this installer (make sure 'Add Python to PATH' is checked)." -ForegroundColor Yellow - return (Exit-InstallFailure "winget required to install Python on this host") - } substep "installing Python ${PythonVersion}..." $pythonPackageId = "Python.Python.$PythonVersion" - # Temporarily lower ErrorActionPreference so that winget stderr - # (progress bars, warnings) does not become a terminating error - # on PowerShell 5.1 where native-command stderr is ErrorRecord. - $prevEAP = $ErrorActionPreference - $ErrorActionPreference = "Continue" - try { - winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements - $wingetExit = $LASTEXITCODE - } catch { $wingetExit = 1 } - $ErrorActionPreference = $prevEAP - Refresh-SessionPath + $wingetExit = $null - # Re-detect after install (PATH may have changed) - $DetectedPython = Find-CompatiblePython - - if (-not $DetectedPython) { - # Python still not functional after winget -- force reinstall. - # This handles both real failures AND "already installed" codes where - # winget thinks Python is present but it's not actually on PATH - # (e.g. user partially uninstalled, or installed via a different method). - substep "Python not found on PATH after winget. Retrying with --force..." "Yellow" + if ($script:WingetAvailable) { + # --source winget avoids the msstore source, which can fail with + # cert-pinning error 0x8a15005e and abort the whole `winget install` + # (winget then demands --source). Python and uv both live in the + # winget source, so pinning it is correct and faster. + # + # Lower ErrorActionPreference so winget stderr (progress/warnings) is + # not a terminating error on PS 5.1 (native stderr is ErrorRecord). + $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" try { - winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements --force + winget install -e --id $pythonPackageId --source winget --accept-package-agreements --accept-source-agreements $wingetExit = $LASTEXITCODE } catch { $wingetExit = 1 } $ErrorActionPreference = $prevEAP Refresh-SessionPath + + # Re-detect after install (PATH may have changed) $DetectedPython = Find-CompatiblePython + + if (-not $DetectedPython) { + # Python still not functional after winget -- force reinstall. + # This handles both real failures AND "already installed" codes where + # winget thinks Python is present but it's not actually on PATH + # (e.g. user partially uninstalled, or installed via a different method). + substep "Python not found on PATH after winget. Retrying with --force..." "Yellow" + $ErrorActionPreference = "Continue" + try { + winget install -e --id $pythonPackageId --source winget --accept-package-agreements --accept-source-agreements --force + $wingetExit = $LASTEXITCODE + } catch { $wingetExit = 1 } + $ErrorActionPreference = $prevEAP + Refresh-SessionPath + $DetectedPython = Find-CompatiblePython + } + } + + # Fall back to python.org if winget is unavailable OR couldn't install a + # working Python (missing/broken winget, msstore cert errors --source + # winget can't fix). Keeps the install automatic instead of failing out. + if (-not $DetectedPython) { + if ($script:WingetAvailable) { + substep "winget could not install Python -- falling back to python.org..." "Yellow" + } else { + substep "winget is unavailable -- installing Python from python.org..." "Yellow" + } + $DetectedPython = Install-PythonFromPythonOrg } if (-not $DetectedPython) { - Write-Host "[ERROR] Python installation failed (exit code $wingetExit)" -ForegroundColor Red + $exitNote = if ($null -ne $wingetExit) { " (winget exit code $wingetExit)" } else { "" } + Write-Host "[ERROR] Python installation failed$exitNote" -ForegroundColor Red Write-Host " Please install Python $PythonVersion manually from https://www.python.org/downloads/" -ForegroundColor Yellow Write-Host " Make sure to check 'Add Python to PATH' during installation." -ForegroundColor Yellow Write-Host " Then re-run this installer." -ForegroundColor Yellow @@ -1054,7 +1181,7 @@ shell.Run cmd, 0, False if ($script:WingetAvailable) { $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" - try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} + try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} $ErrorActionPreference = $prevEAP Refresh-SessionPath } @@ -1230,6 +1357,55 @@ shell.Run cmd, 0, False try { [System.IO.File]::WriteAllText((Join-Path $VenvDir ".unsloth-studio-owned"), "") } catch {} } + # ── Helper: run amd-smi without triggering a UAC elevation prompt ── + # amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing + # DiskPart UAC prompt mid-install (Studio backend amd.py hits the same). + # __COMPAT_LAYER=RunAsInvoker forces it (and helpers it spawns) to run + # un-elevated; on failure the WMI name -> gfx fallback still resolves the arch. + function Invoke-AmdSmiNoElevate { + param( + [Parameter(Mandatory = $true, Position = 0)][string]$Exe, + [Parameter(Position = 1)][string[]]$SmiArgs = @(), + [int]$TimeoutSec = 30 + ) + # RunAsInvoker blocks the auto-elevation/UAC prompt; the timeout bounds a + # flaky amd-smi that can otherwise spin for minutes (30s mirrors amd.py). + $prevCompat = [Environment]::GetEnvironmentVariable('__COMPAT_LAYER', 'Process') + $env:__COMPAT_LAYER = 'RunAsInvoker' + try { + # [Process]::Start, NOT Start-Process -PassThru: the latter leaves + # .ExitCode $null after WaitForExit on PS 5.1, so $LASTEXITCODE (checked + # by callers) reads non-zero and kills detection. Async reads drain the + # pipes (no deadlock); amd-smi args have no spaces so a plain join is safe. + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $Exe + $psi.Arguments = ($SmiArgs -join ' ') + $psi.UseShellExecute = $false + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + if (-not $proc.WaitForExit($TimeoutSec * 1000)) { + try { $proc.Kill() } catch {} + $global:LASTEXITCODE = 124 + return "" + } + $global:LASTEXITCODE = $proc.ExitCode + return ($outTask.Result + "`n" + $errTask.Result) + } catch { + $global:LASTEXITCODE = 1 + return "" + } finally { + if ($null -eq $prevCompat) { + Remove-Item Env:__COMPAT_LAYER -ErrorAction SilentlyContinue + } else { + $env:__COMPAT_LAYER = $prevCompat + } + } + } + # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── $HasNvidiaSmi = $false $NvidiaSmiExe = $null @@ -1303,11 +1479,21 @@ shell.Run cmd, 0, False } } catch {} } - if (-not $HasROCm) { + # On hosts without a working HIP runtime amd-smi elevates a child at runtime, + # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (manifest is + # asInvoker). So only probe when a HIP SDK is present (hipinfo found -> + # un-elevated) or the user opts in; else fall through to WMI name inference + # (enough to pick ROCm wheels + lemonade llama.cpp). + # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the + # HIP-SDK heuristic: a HIP SDK binary with a broken runtime can still pop the + # prompt, so $HipSdkInstalled must NOT silently re-enable it. + $amdSmiOptOut = $env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(0|false|no|off)$' + $amdSmiAllowed = (-not $amdSmiOptOut) -and ($HipSdkInstalled -or ($env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(1|true|yes|on)$')) + if (-not $HasROCm -and $amdSmiAllowed) { $amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue if ($amdSmiExe) { try { - $smiOut = & $amdSmiExe.Source list 2>&1 | Out-String + $smiOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('list') if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { $HasROCm = $true # Mirror the hipinfo path: collect all gfx tokens in enumeration @@ -1322,7 +1508,7 @@ shell.Run cmd, 0, False # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, # including the GFX target needed for wheel index selection. $smiAsicOut = "" - try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {} + try { $smiAsicOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('static','--asic') } catch {} $_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) if ($_asicGfxTokens.Count -gt 0) { $ROCmGfxArch = if ($_smiVisIdx -lt $_asicGfxTokens.Count) { $_asicGfxTokens[$_smiVisIdx] } else { $_asicGfxTokens[0] } @@ -1346,9 +1532,12 @@ shell.Run cmd, 0, False } catch {} } # ── Arch resolution: env-var override → name inference ────────────── - # Covers users whose amd-smi is too old to report the GFX target and - # who don't have hipinfo (HIP-runtime-only, common on Strix Halo / iGPU). - if ($HasROCm -and -not $ROCmGfxArch) { + # Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime + # ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the + # studio setup forward --rocm-gfx and pull a GPU-accelerated (lemonade) + # llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels + # still require a confirmed HIP SDK -- they stay gated on $HasROCm below. + if (-not $ROCmGfxArch) { # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. if ($env:UNSLOTH_ROCM_GFX_ARCH) { $ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower() @@ -1356,15 +1545,20 @@ shell.Run cmd, 0, False substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan" } # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). + # Targets only arches the lemonade-sdk ROCm prebuilts cover + # (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 - @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 - @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) - @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) - @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop - @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 - @{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix) + @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080) + @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060) + @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) + @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) + @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) + @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) + @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- lemonade gfx103X + @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X + @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X ) foreach ($row in $nameArchTable) { if ($ROCmGpuLabel -match $row.P) { @@ -1405,11 +1599,11 @@ shell.Run cmd, 0, False } } catch {} } - if (-not $ROCmVersion) { + if (-not $ROCmVersion -and $amdSmiAllowed) { $amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue if ($amdSmiVer) { try { - $smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String + $smiVerOut = Invoke-AmdSmiNoElevate $amdSmiVer.Source @('version') if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') { $ROCmVersion = $Matches[1] } @@ -1419,6 +1613,46 @@ shell.Run cmd, 0, False } } + # ── Optional WSL-ROCm driver hint ──────────────────────────────────────── + # An AMD GPU can also be used inside WSL2, but only with Adrenalin >= 26.2.2 + # (first production ROCDXG/WSL release); native Windows GPU works with any + # recent driver. We can't auto-install it (AMD referrer-gates downloads, no + # winget package), so just point at AMD's page. Shown only when the installed + # driver predates 26.2.2 (Feb 2026); suppress with UNSLOTH_SKIP_AMD_DRIVER_HINT=1. + function Show-AmdWslDriverHint { + if ($env:UNSLOTH_SKIP_AMD_DRIVER_HINT) { return } + try { + $amd = Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue | + Where-Object { $_.Name -match 'AMD|Radeon' } | Select-Object -First 1 + if (-not $amd) { return } + $drvDate = $null + try { + if ($amd.DriverDate -is [datetime]) { + # Get-CimInstance returns DriverDate already parsed. + $drvDate = $amd.DriverDate + } elseif ($amd.DriverDate) { + # Get-WmiObject style WMI datetime string. + $drvDate = [Management.ManagementDateTimeConverter]::ToDateTime([string]$amd.DriverDate) + } + } catch {} + # Older than 26.2.2 (Feb 2026) => can't expose the GPU to WSL ROCm. + # Unreadable date => still show the hint (informational, suppressible). + if ($drvDate -and $drvDate -ge (Get-Date '2026-02-01')) { return } + substep "Tip: to use this GPU inside WSL too, install AMD Adrenalin 26.2.2+ (for WSL2)." "Cyan" + substep " Your current driver predates it; native Windows GPU is unaffected. Get it from AMD:" "Cyan" + substep " https://www.amd.com/en/resources/support-articles/release-notes/RN-RAD-WIN-26-2-2.html" "Cyan" + substep " Then reboot and run this installer inside an Ubuntu-24.04 WSL distro." "Cyan" + # If WSL isn't installed yet, point at the command that provisions it + # (best-effort; wsl.exe absent => no WSL). + $hasWsl = $false + try { $hasWsl = [bool](Get-Command wsl.exe -ErrorAction SilentlyContinue) } catch {} + if (-not $hasWsl) { + substep " No WSL yet? Install it first: wsl --install -d Ubuntu-24.04" "Cyan" + } + substep " (suppress: set UNSLOTH_SKIP_AMD_DRIVER_HINT=1)" "Cyan" + } catch {} + } + if ($HasNvidiaSmi) { step "gpu" "NVIDIA GPU detected" } elseif ($HasROCm) { @@ -1435,15 +1669,24 @@ shell.Run cmd, 0, False substep " This is a driver issue, not an SDK issue." "Yellow" substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" + } elseif ($ROCmGfxArch) { + # Known arch: Studio setup installs AMD's bundled-runtime ROCm PyTorch wheels + # (repo.amd.com), which ship their own runtime -- HIP SDK optional. + step "gpu" "AMD ROCm ($ROCmGfxArch)" "Cyan" + substep "Detected: $ROCmGpuLabel" "Cyan" + substep "GPU PyTorch uses AMD's bundled-runtime ROCm wheels -- HIP SDK not required (optional)." "Cyan" } elseif ($ROCmGpuLabel) { - step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow" + step "gpu" "AMD GPU detected -- arch unknown" "Yellow" substep "Detected: $ROCmGpuLabel" "Yellow" - substep "Install the HIP SDK for ROCm GPU inference:" "Yellow" + substep "Could not determine the GPU arch -- install the HIP SDK or set" "Yellow" + substep "UNSLOTH_ROCM_GFX_ARCH to enable GPU ROCm PyTorch:" "Yellow" substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" } else { step "gpu" "none (chat-only / GGUF)" "Yellow" substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow" } + # On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint. + if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint } # ── Choose the correct PyTorch index URL based on driver CUDA version ── # Mirrors Get-PytorchCudaTag in setup.ps1. @@ -1529,16 +1772,24 @@ shell.Run cmd, 0, False # ── Print CPU-only hint when no GPU detected ── if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") { Write-Host "" - if ($HipSdkInstalled -and -not $HasROCm) { - substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow" - } elseif ($ROCmGpuLabel) { - substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow" + if ($ROCmGfxArch) { + # Known AMD arch: install.ps1 lays down CPU PyTorch as a base, then + # setup.ps1 swaps in AMD's bundled-runtime GPU ROCm wheels (no HIP SDK). + substep "Installing CPU PyTorch as a base -- Studio setup installs GPU ROCm" "Cyan" + substep "wheels for $ROCmGfxArch next (bundled runtime; HIP SDK not required)." "Cyan" } else { - substep "No NVIDIA GPU detected." "Yellow" + if ($HipSdkInstalled -and -not $HasROCm) { + substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow" + } elseif ($ROCmGpuLabel) { + substep "Installing CPU-only PyTorch (AMD GPU arch unknown -- install the HIP SDK" "Yellow" + substep "or set UNSLOTH_ROCM_GFX_ARCH to enable GPU ROCm)." "Yellow" + } else { + substep "No NVIDIA GPU detected." "Yellow" + } + substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow" + substep "re-run with --no-torch for a faster, lighter install:" "Yellow" + substep ".\install.ps1 --no-torch" "Yellow" } - substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow" - substep "re-run with --no-torch for a faster, lighter install:" "Yellow" - substep ".\install.ps1 --no-torch" "Yellow" Write-Host "" } diff --git a/install.sh b/install.sh index d7de8a946d..af74410397 100755 --- a/install.sh +++ b/install.sh @@ -1207,6 +1207,15 @@ STUB_EOF # Escape single quotes for PowerShell single-quoted string embedding _css_sc_args_ps=$(printf '%s' "$_css_sc_args" | sed "s/'/''/g") + # DISTINCT shortcut name so the WSL launcher never clobbers a native + # install's "Unsloth Studio.lnk" in the same folder. Per-distro suffix. + if [ -n "$_css_distro" ]; then + _css_lnk_name="Unsloth Studio (WSL - ${_css_distro}).lnk" + else + _css_lnk_name="Unsloth Studio (WSL).lnk" + fi + _css_lnk_name_ps=$(printf '%s' "$_css_lnk_name" | sed "s/'/''/g") + # Create shortcuts via a temp PowerShell script to avoid escaping issues _css_ps1_tmp=$(mktemp /tmp/unsloth-shortcut-XXXXXX.ps1 2>/dev/null) || true if [ -n "$_css_ps1_tmp" ]; then @@ -1214,19 +1223,57 @@ STUB_EOF \$WshShell = New-Object -ComObject WScript.Shell \$targetExe = (Get-Command '$_css_sc_target' -ErrorAction SilentlyContinue).Source if (-not \$targetExe) { exit 1 } +# Best-effort: fetch the Unsloth icon to a stable Windows path (shared with a +# native install if one exists) so the WSL shortcut shows the proper icon. +\$iconDir = Join-Path \$env:LOCALAPPDATA 'Unsloth Studio' +\$iconPath = Join-Path \$iconDir 'unsloth.ico' +if (-not (Test-Path -LiteralPath \$iconPath)) { + try { + New-Item -ItemType Directory -Force -Path \$iconDir | Out-Null + Invoke-WebRequest -Uri 'https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth.ico' -OutFile \$iconPath -UseBasicParsing -ErrorAction Stop + } catch {} +} +\$hasIcon = \$false +if (Test-Path -LiteralPath \$iconPath) { + try { \$b = [System.IO.File]::ReadAllBytes(\$iconPath); if (\$b.Length -ge 4 -and \$b[0] -eq 0 -and \$b[1] -eq 0 -and \$b[2] -eq 1 -and \$b[3] -eq 0) { \$hasIcon = \$true } } catch {} +} \$locations = @( [Environment]::GetFolderPath('Desktop'), (Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs') ) +\$created = @() foreach (\$dir in \$locations) { if (-not \$dir -or -not (Test-Path \$dir)) { continue } - \$linkPath = Join-Path \$dir 'Unsloth Studio.lnk' + \$linkPath = Join-Path \$dir '$_css_lnk_name_ps' \$shortcut = \$WshShell.CreateShortcut(\$linkPath) \$shortcut.TargetPath = \$targetExe \$shortcut.Arguments = '$_css_sc_args_ps' - \$shortcut.Description = 'Launch Unsloth Studio' + \$shortcut.Description = 'Launch Unsloth Studio (WSL)' + if (\$hasIcon) { \$shortcut.IconLocation = "\$iconPath,0" } \$shortcut.Save() + \$created += \$linkPath } +# Force Explorer to re-read EACH new shortcut's icon so it renders immediately +# instead of a stale/blank (generic) icon. The reliable, NON-disruptive fix +# (no explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, +# SHCNF_PATHW, ) -- the global SHCNE_ASSOCCHANGED alone does not recover a +# stale item. Also clear the on-disk icon cache for heavier staleness. +try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {} +try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {} +try { + Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int e, uint f, string a, System.IntPtr b);' -ErrorAction SilentlyContinue + foreach (\$p in \$created) { try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, \$p, [System.IntPtr]::Zero) } catch {} } + [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, \$null, [System.IntPtr]::Zero) +} catch {} +# Win11 Start Menu keeps its own tile-icon cache (preserve start2.bin). +try { + \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState' + if (Test-Path -LiteralPath \$smeh) { + Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } +} catch {} WSLPS1_EOF # Convert WSL path to Windows path for powershell.exe @@ -1236,6 +1283,13 @@ WSLPS1_EOF fi rm -f "$_css_ps1_tmp" fi + # If WSL interop is disabled (powershell.exe "Exec format error"), the + # shortcut wasn't created; tell the user how to launch / re-enable it. + if [ "$_css_created" -ne 1 ]; then + substep "Couldn't create the Windows shortcut (WSL interop may be disabled)." "$C_WARN" + substep " Launch Studio from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN" + substep " (re-enable shortcuts: turn WSL interop back on, e.g. run 'wsl --shutdown' then reopen WSL.)" "$C_WARN" + fi fi if [ "$_css_created" -eq 1 ]; then @@ -1650,9 +1704,21 @@ _find_no_torch_runtime() { } # ── AMD ROCm GPU detection helper ── +# WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg only when +# HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and /opt/rocm/bin can be +# off PATH outside login shells (the profile.d drop-in). Seed both before any +# rocminfo probe or a ROCDXG WSL host is misdetected as CPU-only. +_ensure_rocm_probe_env() { + export HSA_ENABLE_DXG_DETECTION="${HSA_ENABLE_DXG_DETECTION:-1}" + if ! command -v rocminfo >/dev/null 2>&1 && [ -x /opt/rocm/bin/rocminfo ]; then + PATH="$PATH:/opt/rocm/bin" + fi +} + # Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs # KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices). _has_amd_rocm_gpu() { + _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then return 0 @@ -1914,6 +1980,122 @@ _pick_radeon_wheel() { esac } +# ── ROCm-on-WSL bootstrap for AMD Strix Halo (gfx1151) ─────────────────────── +# No-op everywhere except: WSL + GPU wanted + no usable GPU yet + /dev/dxg + +# Strix Halo APU. Every other config (NVIDIA, native-Linux ROCm, macOS, Windows, +# CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts +# the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 + +# librocdxg), then sources the env it persisted so detection finds the GPU. +_maybe_bootstrap_rocm_wsl() { + [ "${OS:-}" = "wsl" ] || return 0 + [ "${SKIP_TORCH:-false}" = "false" ] || return 0 + [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 + # Leave any already-usable GPU completely alone (NVIDIA, or working ROCm). + if _has_usable_nvidia_gpu; then return 0; fi + # "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the + # generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and + # would skip this bootstrap while the real GPU is still unusable. awk consumes + # all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail. + _ensure_rocm_probe_env + if command -v rocminfo >/dev/null 2>&1 && \ + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + return 0 + fi + # WSL GPU passthrough device must exist (present on any WSL2 GPU host). + [ -e /dev/dxg ] || return 0 + # Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match + # the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S"). + grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + command -v bash >/dev/null 2>&1 || return 0 + + # Fast path: already configured (librocdxg present) but launched from a + # non-login shell so the persisted env wasn't loaded -- just load it. + if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then + if [ -r /etc/profile.d/unsloth-rocm-wsl.sh ]; then + # shellcheck disable=SC1091 + . /etc/profile.d/unsloth-rocm-wsl.sh || true + else + # librocdxg present but the env drop-in is gone (e.g. a Studio + # uninstall removed it while keeping shared ROCm). Restore the FULL + # env inline (so rocminfo is on PATH) and recreate the drop-in. + _rw_rocm=/opt/rocm + export HSA_ENABLE_DXG_DETECTION=1 + export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 + export PATH="${_rw_rocm}/bin:${PATH}" + export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" + # Persist the drop-in so later non-login Studio launches get the env + # too. /etc/profile.d is root-owned: a plain redirect fails for a + # non-root reinstall (ROCm would silently disappear after this shell), + # so tee through sudo when not root. Best-effort -- the current shell + # already has the env, so the install proceeds either way. + _rw_dropin="$( + printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' + printf 'export HSA_ENABLE_DXG_DETECTION=1\n' + printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' + printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" + printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" + printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' + )" + if [ "$(id -u)" = "0" ]; then + printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true + elif command -v sudo >/dev/null 2>&1; then + printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true + fi + fi + return 0 + fi + + echo "" + substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN" + substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU." + substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)" + + # Locate the helper: prefer the copy shipped beside install.sh, else fetch it. + _rw_helper="${_REPO_ROOT:-.}/scripts/install_rocm_wsl_strixhalo.sh" + _rw_tmp="" + if [ ! -r "$_rw_helper" ]; then + _rw_tmp="$(mktemp 2>/dev/null || echo /tmp/_unsloth_rocm_wsl.sh)" + if download "https://raw.githubusercontent.com/unslothai/unsloth/main/scripts/install_rocm_wsl_strixhalo.sh" "$_rw_tmp" 2>/dev/null; then + _rw_helper="$_rw_tmp" + else + substep "Could not fetch the ROCm-on-WSL helper; using CPU fallback." "$C_WARN" + [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" + return 0 + fi + fi + + # Consent: the narrow guarded case is exactly the GPU setup the user ran the + # installer for, so it proceeds AUTOMATICALLY by default (works with no TTY, + # e.g. `curl ... | sh`). Opt out via UNSLOTH_SKIP_ROCM_WSL_SETUP=1 (top of + # function). The Tauri app drives its own consent UI, so under TAURI_MODE it + # only runs when the app passes UNSLOTH_ROCM_WSL_AUTO=1; else surface and wait. + _rw_go=1 + if [ "${TAURI_MODE:-false}" = "true" ] && [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" != "1" ]; then + tauri_log "ROCM_WSL_AVAILABLE" "strixhalo" + substep "Enable the GPU from the desktop app (or set UNSLOTH_ROCM_WSL_AUTO=1)." "$C_WARN" + _rw_go=0 + fi + + if [ "$_rw_go" = "1" ]; then + # Helper does its own sudo + is idempotent. SMOKE_TEST=0: install.sh + # installs torch itself right after, into the real venv. + if UNSLOTH_WSL_SMOKE_TEST=0 bash "$_rw_helper"; then + # Pull the helper's persisted env into THIS shell so detection + # (rocminfo) now enumerates the GPU and routes to gfx1151. + if [ -r /etc/profile.d/unsloth-rocm-wsl.sh ]; then + # shellcheck disable=SC1091 + . /etc/profile.d/unsloth-rocm-wsl.sh || true + fi + substep "ROCm-on-WSL ready; continuing with GPU install." "$C_OK" + else + substep "ROCm-on-WSL setup did not complete; falling back to CPU-only." "$C_WARN" + fi + fi + [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" + return 0 +} +_maybe_bootstrap_rocm_wsl || true + TORCH_INDEX_URL=$(get_torch_index_url) # rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. @@ -2018,6 +2200,7 @@ if _has_usable_nvidia_gpu; then step "gpu" "NVIDIA GPU detected" elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then # Probe gfx arch for the display label, honouring HIP_VISIBLE_DEVICES + _ensure_rocm_probe_env _gpu_disp_gfx_all="" _gpu_disp_mkt="" if command -v rocminfo >/dev/null 2>&1; then @@ -2048,14 +2231,20 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx" # Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable) elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then + # Kept in sync with the nameArchTable in install.ps1 / setup.ps1. + # gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on + # gfx1102 (bash case has no negative lookahead like the PS tables). case "$_gpu_disp_mkt" in - *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 - *"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 iGPU - *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 iGPU - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop - *"RX 7600"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 - *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU + *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 + *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 + *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point) + *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _gpu_disp_gfx="gfx1030" ;; # RDNA 2 (Navi 21) + *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _gpu_disp_gfx="gfx1032" ;; # RDNA 2 (Navi 23) + *"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) _gpu_disp_gfx="gfx1034" ;; # RDNA 2 (Navi 24) esac if [ -n "$_gpu_disp_gfx" ]; then substep "gfx arch inferred from GPU name: $_gpu_disp_gfx" @@ -2089,7 +2278,34 @@ case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" - substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" + if [ "$OS" = "wsl" ]; then + # WSL + no GPU detected (detection above found nothing). Common + # cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet -- + # /dev/dxg present (graphics) but no ROCm runtime. + _wsl_ubu_ver="" + [ -r /etc/os-release ] && _wsl_ubu_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}") + if [ -e /dev/dxg ]; then + substep "A GPU is plumbed into WSL (/dev/dxg) but no ROCm runtime is exposed to it." "$C_WARN" + fi + substep "For an AMD GPU, ROCm-on-WSL currently needs ALL of:" + substep " 1. AMD Adrenalin Edition 26.1.1+ on Windows (26.2.2+ for Strix Halo / Ryzen AI Max+)." + substep " Older drivers lack production ROCDXG/WSL support, so ROCm can't see the GPU." + substep " Get it from AMD (open in a browser -- direct downloads are referrer-gated):" + substep " https://www.amd.com/en/resources/support-articles/release-notes/RN-RAD-WIN-26-2-2.html" + substep " 2. ROCm 7.2.1 + librocdxg inside WSL (with HSA_ENABLE_DXG_DETECTION=1)." + substep " 3. A WSL distro AMD supports for ROCm -- Ubuntu 24.04 is the known-good one." + if [ -n "$_wsl_ubu_ver" ] && [ "$_wsl_ubu_ver" != "24.04" ]; then + substep " This distro is Ubuntu $_wsl_ubu_ver, which AMD may not support for ROCm-on-WSL yet." "$C_WARN" + fi + substep "Set up the GPU in WSL with a dedicated Ubuntu 24.04 distro:" + substep " wsl --install Ubuntu-24.04 # run in Windows PowerShell, then reopen WSL" + substep " # then re-run this installer inside Ubuntu-24.04 -- it will detect the GPU." + substep "AMD ROCm-on-WSL docs: https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/" + substep "Strix Halo (gfx1151): this installer auto-offers ROCm-on-WSL setup once the" + substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself." + else + substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" + fi substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):" substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" fi diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh new file mode 100644 index 0000000000..5ef9ee386a --- /dev/null +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -0,0 +1,289 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# ────────────────────────────────────────────────────────────────────────────── +# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151) +# ────────────────────────────────────────────────────────────────────────────── +# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime +# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG +# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04 +# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via +# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies). +# +# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with +# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once +# installed + rebooted, /dev/dxg is exposed to WSL and this script builds the rest. +# +# HOW ROCDXG WORKS (and why older /usr/lib/wsl/lib notes are wrong): librocdxg.so +# is AMD's user-mode bridge between the Linux HSA runtime and the Windows driver +# over /dev/dxg. The STANDARD hsa-rocr runtime (NOT the gone "roc4wsl" package) +# loads it when HSA_ENABLE_DXG_DETECTION=1. No hsa/rocm libs need injecting into +# /usr/lib/wsl/lib (it holds only d3d12/dxcore), yet rocminfo enumerates gfx1151 +# fine -- so we gate on /dev/dxg, not on WSL lib injection. +# +# KNOWN CAVEAT (ROCm/ROCm#6022): librocdxg can cap usable ROCm VRAM at the WSL +# VM's RAM (.wslconfig [wsl2] memory=) on some BIOS UMA layouts, and amd-smi +# doesn't work in WSL. On OOM below capacity, raise memory= (then wsl --shutdown) +# and watch GPU use from Windows. Large-UMA BIOS exposes the full pool regardless. +# +# Verified on Ryzen AI Max+ PRO 395 / Radeon 8060S (gfx1151) with ROCm 7.2.1 + +# Ubuntu 24.04 + WSL2 + Adrenalin. These pins MOVE; bump + re-verify on newer ROCm. +# ────────────────────────────────────────────────────────────────────────────── +set -euo pipefail + +# ── Tunables (override via env) ────────────────────────────────────────────── +ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install +GFX="gfx1151" +LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build +# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test. +TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/" +# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs +# torch itself into the real venv right after, so a duplicate download is wasteful. +SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}" +# REQUIRED constraint -- without it pip prefers PyPI's newer CUDA torch over the +# gfx1151 ROCm wheel. 2.11 carries AMD's real gfx1151 fix (matches install.sh). +TORCH_CONSTRAINT="${UNSLOTH_WSL_TORCH_CONSTRAINT:-torch>=2.11.0,<2.12.0}" +ROCM_DIR="" # resolved after install + +say() { printf '\n\033[1;36m== %s\033[0m\n' "$*"; } +note() { printf ' %s\n' "$*"; } +die() { printf '\n\033[1;31m[BLOCKED] %s\033[0m\n' "$*" >&2; exit 1; } + +# sudo only if not already root (WSL distros often run as root) +SUDO="" +if [ "$(id -u)" -ne 0 ]; then + command -v sudo >/dev/null 2>&1 || die "Need root or sudo to install ROCm." + SUDO="sudo" +fi + +# ── Windows 11 SDK (headers for the librocdxg build) ───────────────────────── +# librocdxg's cmake build needs the Windows SDK 'shared' headers, which live on +# the Windows HOST under C:\Program Files (x86)\Windows Kits\10\Include\\. +_WIN_SDK_INC_BASE="/mnt/c/Program Files (x86)/Windows Kits/10/Include" + +# Print the newest installed SDK include dir with 'shared' headers, or nothing. +# find + read loop (not `for ... in $(ls)`) since the base path has a space. +_find_win_sdk() { + [ -d "$_WIN_SDK_INC_BASE" ] || return 0 + while IFS= read -r _inc; do + [ -n "$_inc" ] || continue + if [ -d "$_inc/shared" ]; then printf '%s' "$_inc"; return 0; fi + done < <(find "$_WIN_SDK_INC_BASE" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort -Vr) + return 0 +} + +# Best-effort: install the Windows 11 SDK on the Windows HOST via winget so the +# build has its headers with no manual step. Elevates -> ONE UAC prompt; headers +# appear under /mnt/c immediately (no reboot). Never fatal -- failure falls +# through to a manual-install message. Opt out: UNSLOTH_SKIP_WIN_SDK_INSTALL=1. +_install_windows_sdk_via_winget() { + [ "${UNSLOTH_SKIP_WIN_SDK_INSTALL:-0}" = "1" ] && { note "Skipping Windows SDK auto-install (UNSLOTH_SKIP_WIN_SDK_INSTALL=1)."; return 0; } + command -v powershell.exe >/dev/null 2>&1 || return 0 + # `command -v` succeeds even with WSL interop OFF (.exe on PATH but fails + # with "Exec format error"); verify it actually executes. + powershell.exe -NoProfile -Command "exit 0" >/dev/null 2>&1 || return 0 + if ! powershell.exe -NoProfile -Command "if (Get-Command winget -ErrorAction SilentlyContinue) { exit 0 } else { exit 1 }" >/dev/null 2>&1; then + note "winget not available on the Windows host -- cannot auto-install the Windows SDK." + return 0 + fi + say "Installing the Windows 11 SDK on the Windows host via winget" + note "librocdxg needs its headers. Approve the UAC prompt on the Windows desktop." + note "One-time (~1-3 GB download); opt out with UNSLOTH_SKIP_WIN_SDK_INSTALL=1." + # Newest SDK first, then a fallback. Header presence is the source of truth + # (re-check each attempt), not winget's exit code. /dev/null || true +if [ "${VERSION_ID:-}" != "24.04" ]; then + die "This targets Ubuntu 24.04 (found '${VERSION_ID:-unknown}'). AMD's ROCm-on-WSL supports 24.04; create a dedicated distro: wsl --install Ubuntu-24.04 (do not run on 26.04 -- ROCm 7.2 does not target it yet)." +fi + +if [ ! -e /dev/dxg ]; then + die "/dev/dxg missing -- WSL GPU paravirtualization not present. Ensure this is WSL2 (not WSL1) on a recent Windows build, and that an AMD GPU + ROCDXG-capable Adrenalin driver is installed on the Windows host (then reboot)." +fi +note "Ubuntu 24.04 + /dev/dxg present." +# Don't block on hsa/rocm libs in /usr/lib/wsl/lib: a working ROCDXG setup +# doesn't need them (only d3d12/dxcore). Real readiness is checked via rocminfo. + +# ── Step 1: build/runtime prerequisites ────────────────────────────────────── +say "Installing build prerequisites" +export DEBIAN_FRONTEND=noninteractive +$SUDO apt-get update -y +# `make` is explicit: cmake shells out to it but Ubuntu only *recommends* it, so +# minimal images lack it and the librocdxg `make -j` build would fail. +$SUDO apt-get install -y cmake make gcc g++ git wget gpg ca-certificates python3-venv python3-pip + +# ── Step 2: ROCm ${ROCM_VER} userspace (no DKMS -- WSL uses the Windows driver) ─ +say "Installing ROCm ${ROCM_VER} userspace" +if ! command -v rocminfo >/dev/null 2>&1 && [ ! -x /opt/rocm/bin/rocminfo ]; then + # Direct apt-repo install (leaner than amdgpu-install; repo is indexed by + # ROCm version, e.g. .../apt/7.2.1). + $SUDO mkdir -p /etc/apt/keyrings + wget -qO- https://repo.radeon.com/rocm/rocm.gpg.key \ + | gpg --dearmor | $SUDO tee /etc/apt/keyrings/rocm.gpg >/dev/null + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VER} noble main" \ + | $SUDO tee /etc/apt/sources.list.d/rocm.list >/dev/null + printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' \ + | $SUDO tee /etc/apt/preferences.d/rocm-pin-600 >/dev/null + $SUDO apt-get update -y + # rocm-libs pulls everything torch links at runtime (rocblas, hipblas, + # miopen-hip, rccl, ...); hsa-rocr + rocminfo come as deps. Large (~5 GB + # download / ~23 GB installed). + $SUDO apt-get install -y rocm-libs rocminfo hip-runtime-amd +else + note "ROCm already present -- skipping apt install." +fi + +# Resolve the real ROCm dir and ensure the canonical /opt/rocm symlink. apt lays +# ROCm under /opt/rocm- and rocm-core symlinks /opt/rocm -> that; repair if +# an earlier partial run left /opt/rocm as a real dir blocking the symlink. +_real="$(ls -d /opt/rocm-* 2>/dev/null | sort -V | tail -1 || true)" +if [ -n "$_real" ] && [ ! -L /opt/rocm ] && [ -d /opt/rocm ]; then + # /opt/rocm is a real dir blocking the symlink. Only treat it as a removable + # stray stub if it's NOT a real ROCm install (a real one has bin/rocminfo / + # bin/hipcc / .info/version) -- this protects a user's pre-existing ROCm. Even + # then we MOVE IT ASIDE, never rm -rf, so a wrong guess can't lose data. + if [ -e /opt/rocm/bin/rocminfo ] || [ -e /opt/rocm/bin/hipcc ] || [ -e /opt/rocm/.info/version ]; then + note "/opt/rocm is a real ROCm install -- leaving it untouched (will install librocdxg into it)." + else + note "Moving stray /opt/rocm stub aside -> $_real (not deleting it)" + $SUDO cp -an /opt/rocm/. "$_real"/ 2>/dev/null || true + $SUDO mv /opt/rocm "/opt/rocm.unsloth-stub-bak.$(date +%s)" 2>/dev/null || true + [ -e /opt/rocm ] || $SUDO ln -s "$_real" /opt/rocm + fi +elif [ -n "$_real" ] && [ ! -e /opt/rocm ]; then + $SUDO ln -s "$_real" /opt/rocm +fi +if [ -L /opt/rocm ] || [ -d /opt/rocm ]; then ROCM_DIR="/opt/rocm"; else ROCM_DIR="$_real"; fi +{ [ -n "$ROCM_DIR" ] && [ -d "$ROCM_DIR" ]; } || die "ROCm not found under /opt after install." +note "ROCm at ${ROCM_DIR}" + +# ── Step 3: build librocdxg (DXG <-> HSA bridge; not yet an apt package) ────── +say "Building librocdxg (${LIBROCDXG_REF})" +if [ -e "${ROCM_DIR}/lib/librocdxg.so" ]; then + note "librocdxg already installed -- skipping build." +else + # Discover the newest installed Win11 SDK (version differs per machine). If + # absent, auto-install via winget (one UAC prompt) and re-discover; only if + # that ALSO fails do we stop with manual instructions. + _win_sdk="$(_find_win_sdk)" + if [ -z "$_win_sdk" ]; then + note "Windows 11 SDK headers not found -- attempting automatic install..." + _install_windows_sdk_via_winget + _win_sdk="$(_find_win_sdk)" + fi + [ -n "$_win_sdk" ] || die "Windows 11 SDK headers not found under 'C:\\Program Files (x86)\\Windows Kits\\10\\Include\\*\\shared', and the automatic winget install did not complete. Install it on the Windows host (e.g. 'winget install Microsoft.WindowsSDK.10.0.26100') and re-run." + note "Windows SDK: ${_win_sdk}" + _src="${HOME}/.unsloth/librocdxg" + rm -rf "$_src" + git clone --depth 1 --branch "$LIBROCDXG_REF" https://github.com/ROCm/librocdxg.git "$_src" \ + || git clone "https://github.com/ROCm/librocdxg.git" "$_src" + ( + cd "$_src" + git checkout "$LIBROCDXG_REF" 2>/dev/null || true + mkdir -p build && cd build + cmake .. -DWIN_SDK="${_win_sdk}/shared" + make -j"$(nproc)" + $SUDO make install + ) +fi +# Ensure soname symlinks resolve to whatever version was built (e.g. 1.2.0). +_dxg_real="$(ls -1 "${ROCM_DIR}"/lib/librocdxg.so.*.* 2>/dev/null | sort -V | tail -1 || true)" +if [ -n "$_dxg_real" ]; then + _dxg_base="$(basename "$_dxg_real")" # librocdxg.so.1.2.0 + _dxg_major="$(printf '%s' "$_dxg_base" | sed -E 's/librocdxg\.so\.([0-9]+).*/\1/')" + $SUDO ln -sf "$_dxg_base" "${ROCM_DIR}/lib/librocdxg.so.${_dxg_major}" + $SUDO ln -sf "librocdxg.so.${_dxg_major}" "${ROCM_DIR}/lib/librocdxg.so" +fi +echo "${ROCM_DIR}/lib" | $SUDO tee /etc/ld.so.conf.d/rocm.conf >/dev/null +$SUDO ldconfig + +# ── Step 4: persist environment (system-wide so Studio's worker inherits it) ── +say "Persisting ROCm-on-WSL environment" +_envfile="/etc/profile.d/unsloth-rocm-wsl.sh" +$SUDO tee "$_envfile" >/dev/null <>> Unsloth ROCm-on-WSL (gfx1151) >>> +export HSA_ENABLE_DXG_DETECTION=1 +export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 +export PATH="${ROCM_DIR}/bin:\${PATH}" +export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}" +# <<< Unsloth ROCm-on-WSL (gfx1151) <<< +EOF +# also drop into ~/.bashrc for interactive shells +if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then + cat "$_envfile" >> "${HOME}/.bashrc" +fi +# export into the current process so verification below works immediately +export HSA_ENABLE_DXG_DETECTION=1 +export PATH="${ROCM_DIR}/bin:${PATH}" +export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}" + +# ── Step 5: verify the runtime enumerates the GPU ──────────────────────────── +say "Verifying rocminfo sees ${GFX}" +# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs +# rocminfo on first match, which under `set -o pipefail` turns a successful match +# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a +# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass. +_rocminfo_out="$(rocminfo 2>/dev/null || true)" +if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then + printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true + die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." +fi +# Display-only summary: best-effort (|| true) so head's early pipe-close under +# `set -o pipefail` can't fail the bootstrap after verification already passed. +printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true +note "ROCm-on-WSL runtime is live for ${GFX}." + +# ── Step 6 (optional): torch smoke test from the gfx1151 index ─────────────── +if [ "$SMOKE_TEST" = "1" ]; then + say "Smoke-testing PyTorch on ${GFX} (throwaway venv)" + _venv="${HOME}/.unsloth/rocm-smoketest" + rm -rf "$_venv"; python3 -m venv "$_venv" + "$_venv/bin/pip" install --quiet --upgrade pip + # gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py + # deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch. + "$_venv/bin/pip" install --index-url "$TORCH_INDEX" \ + --extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \ + die "torch install from ${TORCH_INDEX} failed." + "$_venv/bin/python" - <<'PY' +import torch +ok = torch.cuda.is_available() +print("torch:", torch.__version__, "| cuda(rocm) available:", ok) +if ok: + print("device:", torch.cuda.get_device_name(0)) + free, total = torch.cuda.mem_get_info(0) + print(f"mem: free={free/1e9:.1f} GB total={total/1e9:.1f} GB") + import time + a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16) + b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16) + torch.cuda.synchronize(); t0 = time.time() + for _ in range(10): c = a @ b + torch.cuda.synchronize() + print(f"matmul ok ({(time.time()-t0)/10*1e3:.1f} ms/iter)") +raise SystemExit(0 if ok else 1) +PY + rm -rf "$_venv" +fi + +say "Done." +note "ROCm-on-WSL is ready for ${GFX}. If you ran this standalone, install Unsloth" +note "in THIS distro and it will detect the GPU automatically:" +note " curl -fsSL https://unsloth.ai/install.sh | sh" diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 25dd4af01b..f0f6e4fddc 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -16,15 +16,19 @@ function Uninstall-UnslothStudio { function _Step { param([string]$Msg) Write-Host $Msg } function _Substep { param([string]$Msg, [string]$Color = "Gray") Write-Host " $Msg" -ForegroundColor $Color } - # Remove a file/dir/symlink only if it exists. Idempotent. + # Remove a file/dir/symlink if present. Idempotent; retries since a just-killed + # process can briefly hold a handle (Windows refuses the delete until released). function _RemovePath { param([string]$Path) if ([string]::IsNullOrWhiteSpace($Path)) { return } - if (Test-Path -LiteralPath $Path) { + if (-not (Test-Path -LiteralPath $Path)) { return } + for ($attempt = 1; $attempt -le 3; $attempt++) { try { Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop _Substep "removed: $Path" "Green" + return } catch { + if ($attempt -lt 3) { Start-Sleep -Milliseconds 700; continue } _Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow" } } @@ -236,9 +240,58 @@ function Uninstall-UnslothStudio { } catch { } } + # Stop processes that would block deleting the paths we remove. Unlike + # _StopStudioProcesses (venv exe only), this also catches llama-server/llama-cli, + # the unsloth.exe shim, and orphaned mp workers under SYSTEM python holding a + # venv DLL (an open DLL handle blocks the dir delete) -- found by scanning each + # candidate's loaded modules, not just its image path. + function _StopProcessesLockingRoots { + param([string[]]$Roots) + $clean = @($Roots | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\','/') }) + if ($clean.Count -eq 0) { return } + $underRoot = { + param($p) + if (-not $p) { return $false } + foreach ($r in $clean) { if ($p -ieq $r -or $p -ilike "$r\*") { return $true } } + return $false + } + # 1. Image path under a target root (venv python, shim, llama-server). + try { + foreach ($proc in (Get-CimInstance Win32_Process -ErrorAction SilentlyContinue)) { + if ((& $underRoot $proc.ExecutablePath)) { + try { Stop-Process -Id $proc.ProcessId -Force -ErrorAction SilentlyContinue } catch { } + } + } + } catch { } + # 2. A loaded module under a target root (orphaned mp-fork python holding a + # venv DLL). Scoped to names that load our DLLs to keep the scan fast. + try { + $cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli -ErrorAction SilentlyContinue + foreach ($proc in $cands) { + $hit = $false + try { + foreach ($m in $proc.Modules) { if ((& $underRoot $m.FileName)) { $hit = $true; break } } + } catch { } # access denied enumerating modules -> skip + if ($hit) { try { Stop-Process -Id $proc.Id -Force -ErrorAction SilentlyContinue } catch { } } + } + } catch { } + } + # Default install root + default data dir. $defaultStudioHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth\studio" } else { $null } $defaultDataDir = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth Studio" } else { $null } + # Default-mode ~/.unsloth holds a SHARED llama.cpp build + .cache that are + # siblings of studio (not under it), so deleting misses them -- handle + # explicitly. No-op in env/custom mode (nested under the custom root, removed + # with it). A user-set UNSLOTH_LLAMA_CPP_PATH is left alone. + $defaultUnslothHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth" } else { $null } + $defaultLlamaCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "llama.cpp" } else { $null } + $defaultCache = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".cache" } else { $null } + # llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging, + # sibling of the install dir). Usually pruned after activate, but an interrupted + # build can leave a ".staging-XXXX" tree; removing it lets the empty-dir + # cleanup of ~/.unsloth below succeed. No-op in env/custom mode and when absent. + $defaultStaging = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".staging" } else { $null } # Build known-root list FIRST so the port-file kill can verify ownership. $customRoots = @(_CustomStudioRoots) @@ -255,6 +308,9 @@ function Uninstall-UnslothStudio { _StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots } _StopStudioProcesses -KnownRoots $knownRoots + # Also stop anything holding a handle on the exact paths we delete (llama-server, + # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache)) # ── Remove custom-root install trees ── _Step "Removing data and install directories..." @@ -273,6 +329,16 @@ function Uninstall-UnslothStudio { if ($defaultStudioHome) { _RemovePath $defaultStudioHome } # Default data dir. if ($defaultDataDir) { _RemovePath $defaultDataDir } + # Default-mode shared llama.cpp build + cache (siblings of studio under + # ~/.unsloth). No-op in env/custom mode and when absent. + if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp } + if ($defaultCache) { _RemovePath $defaultCache } + if ($defaultStaging) { _RemovePath $defaultStaging } + # Drop ~/.unsloth itself, but ONLY if now empty -- never nuke unrelated content. + if ($defaultUnslothHome -and (Test-Path -LiteralPath $defaultUnslothHome) -and + -not (Get-ChildItem -LiteralPath $defaultUnslothHome -Force -ErrorAction SilentlyContinue)) { + _RemovePath $defaultUnslothHome + } # ── Remove desktop and Start Menu shortcuts ── _Step "Removing desktop and Start Menu shortcuts..." @@ -283,6 +349,18 @@ function Uninstall-UnslothStudio { if ($env:APPDATA) { _RemovePath (Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk") } + # Invalidate the Win11 Start Menu tile cache so the removed shortcut's tile + # disappears promptly instead of lingering stale (mirrors install.ps1's + # New-StudioShortcuts). Preserves start2.bin (the pin layout). + try { + $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" + if (Test-Path -LiteralPath $smehTemp) { + Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue | + Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } + } catch { } # ── Clean user PATH and registry backup ── _Step "Cleaning user PATH and registry..." diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index bb830151ae..e97b28799b 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -212,6 +212,21 @@ _custom_studio_roots | while IFS= read -r _custom_root; do _remove_path "$_custom_root" done _remove_path "$HOME/.unsloth/studio" +# Default-mode shared llama.cpp build + cache are siblings of studio (not removed +# by deleting it). No-op in env/custom mode (they nest under the custom root) and +# when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept. +_remove_path "$HOME/.unsloth/llama.cpp" +_remove_path "$HOME/.unsloth/.cache" +# llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging). +# Normally pruned after activate, but an interrupted build can leave it behind; +# removing it lets the rmdir below succeed. No-op in env/custom mode and absent. +_remove_path "$HOME/.unsloth/.staging" +# ROCm-on-WSL helper artifacts (librocdxg build clone + smoke-test venv). No-op +# where they don't exist; removing them lets the rmdir below succeed. +_remove_path "$HOME/.unsloth/librocdxg" +_remove_path "$HOME/.unsloth/rocm-smoketest" +# Drop ~/.unsloth only if now empty (rmdir refuses non-empty, so user content is kept). +rmdir "$HOME/.unsloth" 2>/dev/null || true _remove_path "$HOME/.local/share/unsloth" # CLI shim: only the symlink Studio created, never a pip-installed file. _remove_cli_shim @@ -244,23 +259,102 @@ case "$_os" in Linux) if [ "$_is_wsl" = "1" ]; then echo "Removing WSL Windows-side shortcuts..." - # install.sh creates 'Unsloth Studio.lnk' on the Windows Desktop and - # Start Menu Programs folder via powershell.exe; mirror that path. - if command -v powershell.exe >/dev/null 2>&1; then + # install.sh creates per-distro 'Unsloth Studio (WSL - ).lnk' + # on the Windows Desktop + Start Menu via powershell.exe. Scope removal + # to THIS distro (passed as $args[0]) so a multi-distro install keeps the + # other distros' launchers; the TARGET=wsl.exe check still spares a + # native install's "Unsloth Studio.lnk". Prefer powershell.exe; test it + # can EXECUTE (`command -v` succeeds even with interop OFF -- .exe then + # fails "Exec format error", common on systemd-enabled distros). + _wsl_distro="${WSL_DISTRO_NAME:-}" + _ps_ran=0 + if command -v powershell.exe >/dev/null 2>&1 && \ + powershell.exe -NoProfile -Command "exit 0" >/dev/null 2>&1; then + _ps_ran=1 + # Inject the distro into the command: a -Command string does not + # receive trailing tokens as $args. WSL distro names are safe to + # embed (no quotes/$/backtick). # shellcheck disable=SC2016 - # $env:APPDATA is a PowerShell expansion; intentionally literal at shell level. - powershell.exe -NoProfile -Command ' - $names = @("Desktop","StartMenu"); + powershell.exe -NoProfile -Command '$distro = "'"$_wsl_distro"'"; $dirs = @( [Environment]::GetFolderPath("Desktop"), (Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs") ); + $ws = New-Object -ComObject WScript.Shell; foreach ($d in $dirs) { - if (-not $d) { continue } - $p = Join-Path $d "Unsloth Studio.lnk"; - if (Test-Path -LiteralPath $p) { Remove-Item -LiteralPath $p -Force } + if (-not $d -or -not (Test-Path -LiteralPath $d)) { continue } + Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio*.lnk" -ErrorAction SilentlyContinue | ForEach-Object { + try { + $sc = $ws.CreateShortcut($_.FullName); + if ("$($sc.TargetPath) $($sc.Arguments)" -notmatch "wsl\.exe") { return } + # When the distro is known, require the per-distro + # name for this distro or its -d "" argument + # so launchers for other distros are not removed. + if ($distro) { + $nameMatch = ($_.Name -eq "Unsloth Studio (WSL - $distro).lnk"); + $argMatch = ($sc.Arguments -match ("-d\s+`"?" + [regex]::Escape($distro) + "`"?")); + if (-not ($nameMatch -or $argMatch)) { return } + } + Remove-Item -LiteralPath $_.FullName -Force -ErrorAction SilentlyContinue + } catch { } + } }' >/dev/null 2>&1 || true fi + # Fallback when powershell.exe can't run (interop disabled): remove the + # WSL .lnk files via drvfs. The "Unsloth Studio (WSL..." name is + # WSL-specific, so a native install's "Unsloth Studio.lnk" never matches. + if [ "$_ps_ran" = "0" ]; then + for _drive in /mnt/c /mnt/d /mnt/e; do + [ -d "$_drive/Users" ] || continue + for _udir in "$_drive"/Users/*; do + [ -d "$_udir" ] || continue + for _scdir in \ + "$_udir/Desktop" \ + "$_udir/OneDrive/Desktop" \ + "$_udir"/OneDrive*/Desktop \ + "$_udir/AppData/Roaming/Microsoft/Windows/Start Menu/Programs"; do + [ -d "$_scdir" ] || continue + if [ -n "$_wsl_distro" ]; then + # Exact per-distro name (no glob) so other distros survive. + _lnk="$_scdir/Unsloth Studio (WSL - ${_wsl_distro}).lnk" + [ -e "$_lnk" ] && rm -f "$_lnk" 2>/dev/null && echo " removed: $_lnk" || true + else + # Distro unknown: fall back to the broad WSL prefix. + for _lnk in "$_scdir"/"Unsloth Studio (WSL"*.lnk; do + [ -e "$_lnk" ] && rm -f "$_lnk" 2>/dev/null && echo " removed: $_lnk" || true + done + fi + done + done + done + fi + # ── ROCm-on-WSL config (install_rocm_wsl_strixhalo.sh) ── + # Remove Unsloth's own ROCDXG config (the env it persisted). The system + # ROCm userspace is a shared prereq (like CUDA) and is LEFT IN PLACE by + # default; set UNSLOTH_UNINSTALL_ROCM=1 to remove it too. + echo "Removing ROCm-on-WSL config..." + _sudo="" + if [ "$_uid" != "0" ] && command -v sudo >/dev/null 2>&1; then _sudo="sudo"; fi + $_sudo rm -f /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true + if [ -f "$HOME/.bashrc" ] && grep -q "Unsloth ROCm-on-WSL" "$HOME/.bashrc" 2>/dev/null; then + _bk=$(mktemp 2>/dev/null || echo "$HOME/.bashrc.unsloth.tmp") + if sed '/# >>> Unsloth ROCm-on-WSL/,/# <<< Unsloth ROCm-on-WSL/d' "$HOME/.bashrc" > "$_bk" 2>/dev/null; then + cat "$_bk" > "$HOME/.bashrc" 2>/dev/null || true + echo " cleaned ROCm-on-WSL block from ~/.bashrc" + fi + rm -f "$_bk" 2>/dev/null || true + fi + if [ "${UNSLOTH_UNINSTALL_ROCM:-0}" = "1" ]; then + echo " removing system ROCm (UNSLOTH_UNINSTALL_ROCM=1)..." + $_sudo rm -f /etc/apt/sources.list.d/rocm.list /etc/apt/preferences.d/rocm-pin-600 \ + /etc/apt/keyrings/rocm.gpg /etc/ld.so.conf.d/rocm.conf 2>/dev/null || true + $_sudo sh -c 'rm -rf /opt/rocm /opt/rocm-*' 2>/dev/null || true + if command -v ldconfig >/dev/null 2>&1; then $_sudo ldconfig 2>/dev/null || true; fi + elif [ -d /opt/rocm ]; then + echo " Note: ROCm userspace (/opt/rocm*) left in place (shared prereq)." + echo " Remove it by re-running with UNSLOTH_UNINSTALL_ROCM=1, or manually:" + echo " sudo rm -rf /opt/rocm /opt/rocm-* && sudo ldconfig" + fi fi echo "Removing Linux .desktop entry..." _remove_path "$HOME/.local/share/applications/unsloth-studio.desktop" diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 953c1d3141..317c9c7f21 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3174,46 +3174,88 @@ class LlamaCppBackend: except Exception as e: logger.debug("Failed to set ROCm visibility env vars for child: %s", e) - # Defensive kill: drop an orphan Popen a concurrent load may - # have stored before we overwrite the reference (#5161). - self._kill_process() + # One-shot --fit off retry: recent llama.cpp runs a "fitting + # params to device memory" step by default (--fit defaults to + # 'on') even when -ngl is explicit. That step has aborted on + # some ROCm hosts (ggml-cuda.cu ROCm error during worst-case + # estimation, e.g. MTP + mmproj models on gfx1151). When + # Studio's own VRAM math already placed the model + # (use_fit=False), the step is redundant second-guessing -- + # retry once with --fit off before declaring the load failed. + # Never retry when fit was requested (use_fit) or the caller + # passed an explicit fit flag via extra args. + _fit_retry_allowed = self._fit_off_retry_eligible(cmd, use_fit) + for _spawn_attempt in (0, 1): + # Defensive kill: drop an orphan Popen a concurrent load may + # have stored before we overwrite the reference (#5161). + # Also reaps the crashed first attempt on the retry pass. + self._kill_process() - self._stdout_lines = [] - # Tee llama-server output to a dedicated log file so a - # post-mortem has the full trail even when the parent only kept - # the last 50 lines. Path is under the studio home. - self._llama_log_fh = None - try: - log_dir = _swa_cache_path().parent / "logs" / "llama-server" - log_dir.mkdir(parents = True, exist_ok = True) - self._llama_log_path = ( - log_dir / f"llama-{int(time.time())}-port-{self._port}.log" + self._stdout_lines = [] + # Tee llama-server output to a dedicated log file so a + # post-mortem has the full trail even when the parent only + # kept the last 50 lines. Path is under the studio home. + self._llama_log_fh = None + try: + log_dir = _swa_cache_path().parent / "logs" / "llama-server" + log_dir.mkdir(parents = True, exist_ok = True) + # Include the attempt index: the --fit off retry can + # respawn within the same epoch second, and reusing the + # name would truncate the crash log the retry warning + # just pointed the user at. + self._llama_log_path = log_dir / ( + f"llama-{int(time.time())}-port-{self._port}" + f"-try{_spawn_attempt}.log" + ) + self._llama_log_fh = open( + self._llama_log_path, + "w", + encoding = "utf-8", + buffering = 1, + ) + logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}") + except OSError as e: + # Best-effort; never block the load on logging. + logger.debug(f"Could not open llama-server log file: {e}") + self._llama_log_path = None + self._process = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **_windows_hidden_subprocess_kwargs(), ) - self._llama_log_fh = open( - self._llama_log_path, - "w", - encoding = "utf-8", - buffering = 1, - ) - logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}") - except OSError as e: - # Best-effort; never block the load on logging. - logger.debug(f"Could not open llama-server log file: {e}") - self._llama_log_path = None - self._process = subprocess.Popen( - cmd, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - text = True, - env = env, - **_windows_hidden_subprocess_kwargs(), - ) - # Background thread to drain stdout (prevents pipe deadlock) - self._stdout_thread = threading.Thread( - target = self._drain_stdout, daemon = True, name = "llama-stdout" - ) - self._stdout_thread.start() + # Background thread to drain stdout (prevents pipe deadlock) + self._stdout_thread = threading.Thread( + target = self._drain_stdout, daemon = True, name = "llama-stdout" + ) + self._stdout_thread.start() + if self._wait_for_health(timeout = 600.0): + break + _startup_crashed = ( + self._process.poll() is not None and self._process.returncode != 0 + ) + if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed: + logger.warning( + "llama-server crashed during startup (exit code %s) " + "with the default memory-fit step enabled; Studio " + "already verified the model fits, retrying once " + "with --fit off. Crash log: %s", + self._process.returncode, + self._llama_log_path, + ) + cmd = [*cmd, "--fit", "off"] + continue + self._kill_process() + raise RuntimeError( + self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + model_identifier, + ) + ) # Store the resolved on-disk path, not the caller's kwarg: in # HF mode gguf_path is None and ``model_path`` is what @@ -3243,17 +3285,7 @@ class LlamaCppBackend: max_available_ctx if max_available_ctx > 0 else self._effective_context_length ) - # Wait for llama-server to become healthy. - if not self._wait_for_health(timeout = 600.0): - self._kill_process() - raise RuntimeError( - self._classify_llama_start_failure( - "\n".join(self._stdout_lines[-50:]), - gguf_path, - self._model_identifier, - ) - ) - + # Health was confirmed inside the spawn/retry loop above. self._healthy = True # Commit caller intent only after _healthy=True so a failed start @@ -3917,6 +3949,23 @@ class LlamaCppBackend: """atexit handler to ensure llama-server is terminated.""" self._kill_process() + @staticmethod + def _fit_off_retry_eligible(cmd: "list[str]", use_fit: bool) -> bool: + """Whether a llama-server startup crash may be retried with --fit off. + + Only when Studio's own VRAM math placed the model (use_fit=False) + and nothing on the command line set the fit mode explicitly + (-fit / --fit, space- or equals-form). --fit-ctx / --fit-target / + -fitc / -fitt tune the fit step but do not select the mode, so + they do not block the retry. + """ + if use_fit: + return False + for a in cmd: + if a in ("-fit", "--fit") or a.startswith(("-fit=", "--fit=")): + return False + return True + def _wait_for_health( self, timeout: float = 120.0, @@ -3933,9 +3982,17 @@ class LlamaCppBackend: if self._stdout_thread is not None: self._stdout_thread.join(timeout = 2) output = "\n".join(self._stdout_lines[-50:]) + # Keep the TAIL: crash details (abort reason, ROCm/CUDA error + # text) print last, after the long startup banner. Head + # truncation has cut off exactly the diagnostic line before. + _log_hint = ( + f" Full log: {self._llama_log_path}" + if getattr(self, "_llama_log_path", None) + else "" + ) logger.error( f"llama-server exited with code {self._process.returncode}. " - f"Output: {output[:2000]}" + f"Output (tail): {output[-2000:]}{_log_hint}" ) return False diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 50323478b2..f7fbaa0b69 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -26,6 +26,22 @@ import subprocess as _sp from pathlib import Path from typing import Any, Callable +# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ────── +# Mirrors main.py. In WSL the AMD GPU is reached via the ROCDXG bridge +# (librocdxg.so over /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_ +# DETECTION=1 is set before torch touches the GPU. A worker spawned outside a +# login shell misses the installer's persisted env and falls back to CPU. +# Gated to no-op unless BOTH /dev/dxg and librocdxg.so exist, so native Linux +# ROCm, NVIDIA, macOS and Windows are unaffected. +if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ: + try: + if os.path.exists("/dev/dxg") and any( + os.path.exists(_p + "/librocdxg.so") for _p in ("/opt/rocm/lib", "/opt/rocm/lib64") + ): + os.environ["HSA_ENABLE_DXG_DETECTION"] = "1" + except Exception: + pass + logger = get_logger(__name__) from utils.hardware import apply_gpu_ids from utils.wheel_utils import ( @@ -678,8 +694,11 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: ``set_per_process_memory_fraction`` cap to leave OS headroom. Classification priority: - 1. ``gcnArchName`` / variant spellings (stable, naming-independent). - 2. Device-name substring match (last resort when all arch attrs absent; + 1. ``props.is_integrated`` truthy (hipDeviceProp_t.integrated -- the + driver's own unified-memory answer; covers APUs beyond the hardcoded + arch set, e.g. gfx1103 Phoenix iGPUs). Only ever upgrades to unified. + 2. ``gcnArchName`` / variant spellings (stable, naming-independent). + 3. Device-name substring match (last resort when all arch attrs absent; AMD SDK / Radeon wheels may not populate them): - gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M`` - gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395), @@ -692,6 +711,16 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: gcn_arch = _v break + # Driver's own answer first: hipDeviceProp_t.integrated (exposed as + # props.is_integrated; same gate PR #5988's UMA safetensors fast-load + # uses). Strictly additive -- only a truthy value upgrades to unified; + # 0/absent falls through to the arch/name logic below, so a wheel that + # omits or zeroes the field can never downgrade the known APU set. This + # covers unified APUs outside the hardcoded arches (gfx1103 Phoenix + # iGPUs, future parts) with one universal signal. + if getattr(props, "is_integrated", 0): + return gcn_arch, True + if gcn_arch: return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"} @@ -1982,6 +2011,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> os.environ["TORCHDYNAMO_DISABLE"] = "1" logger.info("Windows ROCm: torch.compile (dynamo) disabled") + # bitsandbytes' import-time get_rocm_gpu_arch() probe runs + # `hipinfo.exe` from PATH; the AMD torch wheel ships it in the venv + # Scripts dir, which is on PATH only for activated venvs. Prepend + # it so the probe succeeds instead of logging a scary (harmless) + # "Could not detect ROCm GPU architecture" ERROR on every import. + # Normally inherited from main.py's env, but workers can also be + # spawned standalone (tests, CLI) -- keep the guard here too. + _scripts_dir = os.path.dirname(sys.executable) + if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")): + import shutil as _shutil + if not _shutil.which("hipinfo.exe"): + os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "") + # BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB # wheel may ship a DLL whose suffix doesn't match. Detect the actual # DLL name and override; "72" is a safe fallback. Callers may @@ -2178,10 +2220,25 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if _is_unified and not _gcn_arch: logger.debug( "ROCm OOM guard: gcnArchName absent -- inferred " - "unified memory from device name %r; applying 0.80 cap", + "unified memory from device name %r; applying unified cap", _dev_name, ) - _mem_fraction = 0.80 if _is_unified else 0.90 + # Unified hosts on native Windows: mem_get_info's total is the + # WDDM budget the driver grants HIP (BIOS carve + ~half of the + # remaining RAM) -- the OS share is already outside it, so the + # Linux 0.80 starve-protection double-taxes (48.49 GiB budget → + # 38.79 allowed) and blocks loads that fit in free memory. + # 1.0 removes the double-tax. Current AMD Windows wheels only + # enforce sub-1.0 fractions (measured on gfx1151: 0.5 caps, + # 1.0 still allocates past the budget via WDDM overcommit), so + # 1.0 behaves like torch's uncapped default, with WDDM + # arbitrating residency; on wheels that do enforce it, it caps + # at exactly the driver-granted budget. On Linux the total + # spans nearly all RAM, so keep the 0.80 OS headroom there. + if _is_unified: + _mem_fraction = 1.0 if sys.platform == "win32" else 0.80 + else: + _mem_fraction = 0.90 _torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction) logger.info( "ROCm OOM guard: set_per_process_memory_fraction(%.2f) — " @@ -2191,6 +2248,28 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> _dev_name, _gcn_arch or "unknown arch", ) + # Unified Windows APUs: the WDDM budget is user-raisable, but + # nothing on the box says so -- users see "48 GB VRAM" on a + # 96 GB machine and assume a Studio bug. Say where the limit + # comes from and how to raise it. + if _is_unified and sys.platform == "win32": + try: + import psutil as _psutil + + _phys = _psutil.virtual_memory().total + _granted = _torch_mem.cuda.mem_get_info(0)[1] + if _granted < 0.75 * _phys: + logger.info( + "Windows grants the GPU %.1f GiB of %.1f GiB " + "system RAM (driver/WDDM budget). To raise it: " + "increase the BIOS UMA frame buffer size, or " + "AMD Software > Performance > Tuning > " + "Variable Graphics Memory.", + _granted / 1024**3, + _phys / 1024**3, + ) + except Exception: + pass except Exception as _oom_guard_err: logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err) diff --git a/studio/backend/main.py b/studio/backend/main.py index 35c955405e..38605a5009 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -63,6 +63,23 @@ if sys.platform == "win32": _add_rocm_dll_dirs() del _add_rocm_dll_dirs + # ── Windows AMD ROCm: make hipInfo.exe resolvable for subprocess probes ── + # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via PATH at import + # time; the AMD torch wheel ships it in the venv Scripts dir, which is on + # PATH only when the venv is activated -- Studio launches python directly. + # Without this, every bitsandbytes import logs a scary (but harmless) + # "Could not detect ROCm GPU architecture: [WinError 2]" ERROR + WARNING. + # Gated on the file existing: only AMD ROCm wheels ship hipInfo.exe, so + # NVIDIA/CPU hosts are untouched. os.add_dll_directory above does not help + # here -- subprocess PATH resolution ignores DLL search directories. + _scripts_dir = os.path.dirname(sys.executable) + if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")): + import shutil as _shutil + if not _shutil.which("hipinfo.exe"): + os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "") + del _shutil + del _scripts_dir + # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─ # bitsandbytes derives the rocm.dll name from torch.version.hip, but the # wheel ships rocm72.dll, so the server crashes ("Configured ROCm binary not @@ -108,6 +125,28 @@ if sys.platform == "win32": _bnb_rocm_ver_final, ) +# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ────── +# In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over +# /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE +# torch touches the GPU. A worker launched outside a login shell (e.g. +# `wsl.exe -d Ubuntu-24.04 python ...`) misses the installer's persisted env +# and silently falls back to CPU. Set it here, gated to no-op unless BOTH +# /dev/dxg AND librocdxg.so exist -- native Linux ROCm, NVIDIA, macOS and +# Windows are unaffected. +elif sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ: + try: + if os.path.exists("/dev/dxg") and any( + os.path.exists(os.path.join(_p, "librocdxg.so")) + for _p in ("/opt/rocm/lib", "/opt/rocm/lib64") + ): + os.environ["HSA_ENABLE_DXG_DETECTION"] = "1" + import logging as _logging + _logging.getLogger(__name__).info( + "WSL ROCm: set HSA_ENABLE_DXG_DETECTION=1 (librocdxg bridge present)" + ) + except Exception: + pass + # Put backend dir on sys.path so _platform_compat is importable when main.py # is launched directly (e.g. `uvicorn main:app`). _backend_dir = str(_Path(__file__).parent) diff --git a/studio/backend/run.py b/studio/backend/run.py index c992a4bec2..4619cb849c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -8,6 +8,7 @@ Self-contained; can be moved to any directory. import os import sys +import time from pathlib import Path from typing import Optional @@ -673,6 +674,94 @@ def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Pa return None, attempted +class _TeeStream: + """Mirror writes to the original stream and a session log file. + + Console behavior is unchanged (writes/returns delegate to the original + stream; Tauri's structured-stdout protocol and isatty probes see exactly + what they saw before). The file copy is best-effort: a full disk or a + closed handle must never break the console.""" + + def __init__(self, stream, log_fh): + self._stream = stream + self._log_fh = log_fh + + def write(self, data): + try: + self._log_fh.write(data) + except Exception: + pass + return self._stream.write(data) + + def flush(self): + try: + self._log_fh.flush() + except Exception: + pass + try: + self._stream.flush() + except Exception: + pass + + def __getattr__(self, name): + return getattr(self._stream, name) + + +def _setup_server_disk_logging(): + """Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim + faulthandler at the same file so hard crashes (access violations / + SIGSEGV in the GPU runtime) leave a stack trace on disk. + + Also exports PYTHONFAULTHANDLER=1 so child Python processes (training + workers) dump native-crash stacks to their captured stderr. Keeps the + newest 20 session logs. Opt out with UNSLOTH_STUDIO_NO_FILE_LOG=1. + Returns the log path, or None when disabled/unavailable. + """ + if os.environ.get("UNSLOTH_STUDIO_NO_FILE_LOG") == "1": + return None + try: + from utils.paths import studio_root + log_dir = Path(studio_root()) / "logs" / "server" + except Exception: + home = ( + os.environ.get("UNSLOTH_STUDIO_HOME") + or os.environ.get("STUDIO_HOME") + or os.path.join(os.path.expanduser("~"), ".unsloth", "studio") + ) + log_dir = Path(home) / "logs" / "server" + try: + log_dir.mkdir(parents = True, exist_ok = True) + stamp = time.strftime("%Y%m%d-%H%M%S") + log_path = log_dir / f"server-{stamp}-pid{os.getpid()}.log" + # Line-buffered so the tail survives a hard kill; errors="replace" + # so a console encoding quirk can never take the server down. + log_fh = open(log_path, "w", encoding = "utf-8", errors = "replace", buffering = 1) + except Exception: + return None + + import faulthandler + + try: + faulthandler.enable(file = log_fh, all_threads = True) + except Exception: + pass + # Children (training workers) inherit: their native-crash stacks land on + # the stderr the server already captures. + os.environ.setdefault("PYTHONFAULTHANDLER", "1") + + sys.stdout = _TeeStream(sys.stdout, log_fh) + sys.stderr = _TeeStream(sys.stderr, log_fh) + + # Best-effort retention: keep the newest 20 session logs. + try: + logs = sorted(log_dir.glob("server-*.log"), key = lambda p: p.stat().st_mtime) + for old in logs[:-20]: + old.unlink(missing_ok = True) + except Exception: + pass + return log_path + + def run_server( host: str = "127.0.0.1", port: int = 8888, @@ -705,6 +794,16 @@ def run_server( except Exception: pass + # Persist a session log + native-crash stacks BEFORE importing main, so + # even import-time failures leave evidence on disk. Field report: Studio + # "terminates without a warning" -- a native crash in the GPU runtime + # kills the process with no Python traceback, and a desktop-shortcut + # console closes before anything can be read. Console-only logging made + # that undiagnosable. + _session_log = _setup_server_disk_logging() + if _session_log is not None and not silent: + print(f"Session log: {_session_log}") + # Set env var BEFORE importing main so CORS middleware picks it up. if api_only: os.environ["UNSLOTH_API_ONLY"] = "1" diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 33f9e9d803..1ba6c9f7b5 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -145,3 +145,111 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", should_not_be_called) assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False assert called["n"] == 0 + + +class TestCrashLogTail: + """The "exited with code X" log must keep the TAIL of the output. + + Crash diagnostics (abort reason, ROCm/CUDA error text) print last, + after the long startup banner; head truncation has cut off exactly + the diagnostic line in field reports (gfx1151 fit-step abort).""" + + @staticmethod + def _capture_error_logs(monkeypatch) -> list: + """Capture module-logger .error() messages directly -- immune to + whatever logging/structlog config sibling test modules installed.""" + import core.inference.llama_cpp as _llama_mod + + records: list = [] + fake_logger = mock.Mock() + fake_logger.error = mock.Mock(side_effect = lambda msg, *a, **k: records.append(msg)) + monkeypatch.setattr(_llama_mod, "logger", fake_logger) + return records + + def test_crash_log_keeps_tail_not_head(self, monkeypatch): + records = self._capture_error_logs(monkeypatch) + b = _make_backend() + b._process.poll.return_value = 1 + b._process.returncode = 1 + # >2000 chars of banner, diagnostic on the final line. + banner = [f"load_model: tensor blk.{i} buffer ROCm0" for i in range(80)] + diagnostic = "ggml-cuda.cu:103: ROCm error: out of memory" + b._stdout_lines = banner + [diagnostic] + + assert b._wait_for_health(timeout = 1.0, interval = 0.01) is False + + crash_logs = [m for m in records if "exited with code" in m] + assert crash_logs, "crash must produce an exited-with-code log" + assert diagnostic in crash_logs[-1] + assert "Output (tail)" in crash_logs[-1] + # The head of the banner must be the part sacrificed to truncation. + assert "blk.0 buffer" not in crash_logs[-1] + + def test_crash_log_mentions_log_file_when_present(self, monkeypatch): + records = self._capture_error_logs(monkeypatch) + b = _make_backend() + b._process.poll.return_value = 1 + b._process.returncode = 1 + b._stdout_lines = ["boom"] + b._llama_log_path = Path("C:/logs/llama-123-port-1234.log") + + assert b._wait_for_health(timeout = 1.0, interval = 0.01) is False + + crash_logs = [m for m in records if "exited with code" in m] + assert crash_logs and "llama-123-port-1234.log" in crash_logs[-1] + + +class TestRetryLogFilenameUnique: + """The --fit off retry can respawn within the same epoch second; the log + filename must carry the attempt index or the second open ("w") truncates + the crash log the retry warning just referenced (found by simulation: + frozen time.time -> single file, crash evidence gone).""" + + def test_log_name_includes_attempt_index(self): + src = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + ).read_text(encoding = "utf-8") + assert "-try{_spawn_attempt}.log" in src + + +class TestFitOffRetryEligible: + """Gate for the one-shot --fit off startup-crash retry. + + Retry only when Studio's own VRAM math placed the model and nothing + on the command line chose the fit mode explicitly.""" + + def test_eligible_for_plain_ngl_launch(self): + cmd = ["llama-server", "-m", "x.gguf", "-ngl", "-1", "--jinja"] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is True + + def test_not_eligible_when_use_fit(self): + cmd = ["llama-server", "-m", "x.gguf", "--fit", "on"] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = True) is False + + @pytest.mark.parametrize( + "fit_args", + [ + ["--fit", "on"], + ["--fit", "off"], + ["-fit", "off"], + ["--fit=on"], + ["-fit=off"], + ], + ) + def test_not_eligible_with_explicit_fit_flag(self, fit_args): + cmd = ["llama-server", "-m", "x.gguf", *fit_args] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is False + + @pytest.mark.parametrize( + "tuning_args", + [ + ["--fit-ctx", "8192"], + ["--fit-target", "1024"], + ["-fitc", "4096"], + ["-fitt", "512"], + ["--fit-ctx=8192"], + ], + ) + def test_fit_tuning_flags_do_not_block_retry(self, tuning_args): + cmd = ["llama-server", "-m", "x.gguf", *tuning_args] + assert LlamaCppBackend._fit_off_retry_eligible(cmd, use_fit = False) is True diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index 21edabd142..6e70c7cde4 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -13,6 +13,7 @@ wrong headroom factor on a 128 GiB unified-memory pool. from __future__ import annotations +from pathlib import Path from types import SimpleNamespace import pytest @@ -28,6 +29,46 @@ def _props(**kwargs) -> SimpleNamespace: return SimpleNamespace(**kwargs) +# ── Path 0: props.is_integrated (driver's own unified-memory answer) ───────── + + +class TestIsIntegratedSignal: + """hipDeviceProp_t.integrated wins when truthy; 0/absent never downgrades. + + Same universal gate PR #5988's UMA safetensors fast-load uses -- keeps + Studio's two unified-memory consumers on one signal.""" + + def test_integrated_upgrades_unknown_apu(self) -> None: + # gfx1103 Phoenix iGPU: outside the hardcoded arch set, but the + # driver says integrated -> unified. + props = _props(gcnArchName = "gfx1103", name = "Radeon 780M", is_integrated = 1) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1103" + assert is_unified is True + + def test_integrated_wins_without_any_arch(self) -> None: + props = _props(name = "Some Future APU", is_integrated = 1) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is True + + def test_zero_does_not_downgrade_known_apu(self) -> None: + # A wheel that zeroes the field must not flip Strix Halo to discrete. + props = _props(gcnArchName = "gfx1151", name = "x", is_integrated = 0) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert is_unified is True + + def test_absent_keeps_existing_behavior(self) -> None: + props = _props(gcnArchName = "gfx1201", name = "RX 9070 XT") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert is_unified is False + + def test_discrete_with_zero_stays_discrete(self) -> None: + props = _props(gcnArchName = "gfx1100", name = "RX 7900 XTX", is_integrated = 0) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert is_unified is False + + # ── Path 1: canonical gcnArchName ──────────────────────────────────────────── @@ -168,3 +209,35 @@ class TestDeviceNameFallback: gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "" assert is_unified is False + + +# ── Fraction selection (source-pinned) ─────────────────────────────────────── + + +_WORKER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + + +class TestMemFractionSelection: + """Pin the per-platform fraction policy in worker.py section 1g. + + On native Windows, torch.cuda.mem_get_info's total is the WDDM budget + the driver grants HIP -- the OS share of RAM is already outside it, so + a 0.80 cap double-taxes (field report: 48.49 GiB budget -> '38.79 GiB + allowed' OOM denying a 47.29 GiB load that fit in free memory). 1.0 + removes the double-tax; current AMD Windows wheels enforce only + sub-1.0 fractions, so it behaves like torch's uncapped default with + WDDM arbitrating residency (measured on gfx1151).""" + + def test_unified_win32_uses_budget_exact_fraction(self) -> None: + source = _WORKER_PY.read_text(encoding = "utf-8") + assert '1.0 if sys.platform == "win32" else 0.80' in source + + def test_discrete_keeps_090(self) -> None: + source = _WORKER_PY.read_text(encoding = "utf-8") + assert "_mem_fraction = 0.90" in source + + def test_win32_unified_logs_vgm_hint(self) -> None: + """Users must learn the WDDM budget is raisable (BIOS UMA / AMD + Software Variable Graphics Memory) instead of assuming a bug.""" + source = _WORKER_PY.read_text(encoding = "utf-8") + assert "Variable Graphics Memory" in source diff --git a/studio/backend/tests/test_server_disk_logging.py b/studio/backend/tests/test_server_disk_logging.py new file mode 100644 index 0000000000..05d03d869c --- /dev/null +++ b/studio/backend/tests/test_server_disk_logging.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the server session log + native-crash capture in run.py. + +Field regression: Studio "terminates without a warning" -- a native crash in +the GPU runtime kills the process with no Python traceback, and a desktop- +shortcut console closes before anything can be read. The server must tee its +console output to disk and aim faulthandler at the same file so even hard +crashes leave evidence. +""" + +from __future__ import annotations + +import io +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import run as run_mod # noqa: E402 + + +class TestTeeStream: + def test_writes_reach_both_and_return_original(self): + console, log = io.StringIO(), io.StringIO() + tee = run_mod._TeeStream(console, log) + n = tee.write("hello") + assert console.getvalue() == "hello" == log.getvalue() + assert n == 5 # delegate's return value, console contract unchanged + + def test_log_failure_never_breaks_console(self): + class Broken: + def write(self, data): + raise OSError("disk full") + + def flush(self): + raise OSError("disk full") + + console = io.StringIO() + tee = run_mod._TeeStream(console, Broken()) + assert tee.write("still works") == len("still works") + tee.flush() # must not raise + assert console.getvalue() == "still works" + + def test_attribute_proxy(self): + console, log = io.StringIO(), io.StringIO() + tee = run_mod._TeeStream(console, log) + # isatty / encoding probes must see the original stream's answers. + assert tee.isatty() == console.isatty() + + +class TestSetupServerDiskLogging: + def test_opt_out_env(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_FILE_LOG", "1") + assert run_mod._setup_server_disk_logging() is None + + def test_creates_log_and_enables_faulthandler(self, monkeypatch, tmp_path): + import faulthandler + + monkeypatch.delenv("UNSLOTH_STUDIO_NO_FILE_LOG", raising = False) + monkeypatch.delenv("PYTHONFAULTHANDLER", raising = False) + # Both resolution paths (utils.paths.studio_root and the env + # fallback) honor UNSLOTH_STUDIO_HOME, so this redirects the log dir. + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + orig_out, orig_err = sys.stdout, sys.stderr + was_enabled = faulthandler.is_enabled() + try: + log_path = run_mod._setup_server_disk_logging() + assert log_path is not None + assert Path(log_path).is_file() + assert "logs" in str(log_path) + # faulthandler armed at the file; children inherit the env switch. + assert faulthandler.is_enabled() + import os + + assert os.environ.get("PYTHONFAULTHANDLER") == "1" + print("tee-capture-marker") + sys.stdout.flush() + assert "tee-capture-marker" in Path(log_path).read_text( + encoding = "utf-8", errors = "replace" + ) + finally: + sys.stdout, sys.stderr = orig_out, orig_err + if not was_enabled: + faulthandler.disable() + + def test_run_server_wires_logging_before_main_import(self): + src = (Path(_BACKEND_DIR) / "run.py").read_text(encoding = "utf-8") + call_idx = src.index("_setup_server_disk_logging()", src.index("def run_server")) + main_import_idx = src.index("from main import app", src.index("def run_server")) + assert call_idx < main_import_idx, ( + "disk logging must be armed before importing main so import-time " + "failures leave evidence on disk" + ) diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index 39aafe0489..27dfe187cc 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -12,6 +12,7 @@ import math import os import platform import re +import shutil import subprocess import sys from typing import Any, Optional @@ -33,24 +34,85 @@ _amd_smi_consecutive_failures = 0 _amd_smi_disabled = False +def _hip_sdk_present() -> bool: + """True if a HIP SDK is detectable (hipinfo on PATH or under HIP_PATH/ + ROCM_PATH), meaning amd-smi has a working runtime and runs un-elevated.""" + if shutil.which("hipinfo"): + return True + for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + root = os.environ.get(var) + if root and os.path.exists(os.path.join(root, "bin", "hipinfo.exe")): + return True + return False + + +def _amd_smi_allowed() -> bool: + """Whether it is safe to spawn amd-smi here. + + On Windows without a working HIP runtime, amd-smi elevates a child at + runtime -- popping a UAC/DiskPart prompt that RunAsInvoker can't suppress + (its manifest is asInvoker). So only call it on Windows with a HIP SDK + present or UNSLOTH_ENABLE_AMD_SMI=1. Linux amd-smi never elevates. + """ + if platform.system() != "Windows": + return True + flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower() + if flag in ("1", "true", "yes", "on"): + return True + if flag in ("0", "false", "no", "off"): + return False + return _hip_sdk_present() + + def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]: """Run amd-smi with the given args and return parsed JSON, or None.""" global _amd_smi_consecutive_failures, _amd_smi_disabled if _amd_smi_disabled: return None + if not _amd_smi_allowed(): + # Permanently skip amd-smi on Windows w/o a HIP SDK: every call would + # pop a UAC/DiskPart prompt (see _amd_smi_allowed). VRAM polling is then + # unavailable, but that beats the prompt. Opt back in with + # UNSLOTH_ENABLE_AMD_SMI=1. + if not _amd_smi_disabled: + logger.info( + "amd-smi disabled on Windows (no HIP SDK detected) to avoid a " + "UAC/DiskPart elevation prompt; GPU VRAM polling unavailable. " + "Set UNSLOTH_ENABLE_AMD_SMI=1 to force amd-smi." + ) + _amd_smi_disabled = True + return None + if shutil.which("amd-smi") is None: + # amd-smi does not exist on Windows (neither Adrenalin nor the HIP SDK + # ship a CLI) and can be absent on minimal Linux installs. Disable the + # poller in one step instead of burning the 3-strike circuit breaker + # on guaranteed FileNotFoundError spawns. Studio's VRAM display falls + # back to torch mem_get_info. + if not _amd_smi_disabled: + logger.info( + "amd-smi not found on PATH; GPU utilization polling via " + "amd-smi unavailable (VRAM falls back to torch mem_get_info)." + ) + _amd_smi_disabled = True + return None + _amd_env = child_env_without_native_path_secret() + if platform.system() == "Windows": + # RunAsInvoker belt-and-suspenders for any manifest-elevating helper; + # the real guard is _amd_smi_allowed() above. Mirrors install scripts. + _amd_env = {**_amd_env, "__COMPAT_LAYER": "RunAsInvoker"} try: result = subprocess.run( ["amd-smi", *args, "--json"], capture_output = True, text = True, timeout = timeout, - env = child_env_without_native_path_secret(), + env = _amd_env, **windows_hidden_subprocess_kwargs(), ) except (OSError, subprocess.TimeoutExpired) as e: if isinstance(e, FileNotFoundError): - # amd-smi ships with Adrenalin, not the HIP SDK; absence is expected - # on HIP SDK-only Windows setups. + # Raced a PATH change after the which() check above; absence is + # expected on Windows (no AMD product ships an amd-smi CLI there). logger.debug("amd-smi not found (not in PATH): %s", e) else: logger.warning("amd-smi query failed: %s", e) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 5d3f0ed339..4509ea58ed 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -47,6 +47,38 @@ EXIT_FALLBACK = 2 EXIT_ERROR = 1 EXIT_BUSY = 3 +# DiskPart-prompt suppression. RunAsInvoker does NOT stop amd-smi's runtime +# elevation (its manifest is asInvoker), so this is just harmless belt-and- +# suspenders for manifest-elevating tools. The real guard is _amd_smi_allowed(): +# we don't spawn amd-smi on Windows w/o a HIP SDK (or opt-in). +if platform.system() == "Windows": + os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker") + + +def _amd_smi_allowed() -> bool: + """Whether it is safe to spawn amd-smi here. + + On Windows w/o a working HIP runtime, amd-smi elevates a child and pops a + UAC/DiskPart prompt RunAsInvoker can't suppress. Only call it on Windows + when a HIP SDK is detectable (hipinfo present) or UNSLOTH_ENABLE_AMD_SMI=1; + Linux/macOS always allowed. When skipped, the gfx arch still arrives via the + forwarded --rocm-gfx, so prebuilt selection is unaffected. + """ + if platform.system() != "Windows": + return True + flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower() + if flag in ("1", "true", "yes", "on"): + return True + if flag in ("0", "false", "no", "off"): + return False + if shutil.which("hipinfo"): + return True + for _var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + _root = os.environ.get(_var) + if _root and os.path.isfile(os.path.join(_root, "bin", "hipinfo.exe")): + return True + return False + def windows_hidden_subprocess_kwargs() -> dict[str, object]: """Return Windows-only subprocess kwargs that suppress console windows.""" @@ -2672,6 +2704,15 @@ def run_capture( check: bool = False, env: dict[str, str] | None = None, ) -> subprocess.CompletedProcess[str]: + # amd-smi on Windows auto-elevates and pops a UAC/DiskPart prompt mid-install; + # RunAsInvoker forces it un-elevated. Callers already fall back to WMI/name + # detection. Mirrors install.ps1's Invoke-AmdSmiNoElevate; Windows-only. + if ( + command + and platform.system() == "Windows" + and os.path.basename(command[0]).lower().startswith("amd-smi") + ): + env = {**(os.environ if env is None else env), "__COMPAT_LAYER": "RunAsInvoker"} result = subprocess.run( command, capture_output = True, @@ -2857,6 +2898,13 @@ def detect_host() -> HostInfo: has_rocm = False rocm_gfx_target: str | None = None if is_linux: + # WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg + # only when HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and + # rocminfo can live only under /opt/rocm/bin (the profile.d PATH + # drop-in reaches login shells only). Probe accordingly or a ROCDXG + # WSL host is misdetected as CPU-only. + _dxg_probe_env = {**os.environ} + _dxg_probe_env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") for _cmd, _check in ( # rocminfo: a real gfx GPU id (3-4 chars, nonzero first digit). # gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines @@ -2869,10 +2917,18 @@ def detect_host() -> HostInfo: (["amd-smi", "list"], _amd_smi_has_gpu), ): _exe = shutil.which(_cmd[0]) + if not _exe and _cmd[0] == "rocminfo": + _opt_rocminfo = "/opt/rocm/bin/rocminfo" + if os.access(_opt_rocminfo, os.X_OK): + _exe = _opt_rocminfo if not _exe: continue try: - _result = run_capture([_exe, *_cmd[1:]], timeout = 10) + _result = run_capture( + [_exe, *_cmd[1:]], + timeout = 10, + env = _dxg_probe_env if _cmd[0] == "rocminfo" else None, + ) except Exception: continue if _result.returncode == 0 and _result.stdout.strip(): @@ -2897,12 +2953,20 @@ def detect_host() -> HostInfo: _candidate = os.path.join(_root, "bin", f"{name}.exe") if os.path.isfile(_candidate): return _candidate + # AMD torch wheels ship hipInfo.exe into the venv Scripts dir + # (next to python.exe) -- resolvable on driver-only hosts where no + # SDK dir exists, so a standalone rerun can still detect the GPU. + _venv_candidate = os.path.join(os.path.dirname(sys.executable), f"{name}.exe") + if os.path.isfile(_venv_candidate): + return _venv_candidate return None - for _cmd, _check in ( - (["hipinfo"], lambda out: "gcnarchname" in out.lower()), - (["amd-smi", "list"], _amd_smi_has_gpu), - ): + _win_probes = [(["hipinfo"], lambda out: "gcnarchname" in out.lower())] + if _amd_smi_allowed(): + # Skipped on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); + # gfx arch still arrives via --rocm-gfx, so has_rocm is set by override. + _win_probes.append((["amd-smi", "list"], _amd_smi_has_gpu)) + for _cmd, _check in _win_probes: _exe = _resolve_exe(_cmd[0]) if not _exe: continue @@ -3522,16 +3586,12 @@ def _detect_host_rocm_version() -> tuple[int, int] | None: return int(parts[0]), int(parts[1]) except Exception: pass - amd_smi = shutil.which("amd-smi") + amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None if amd_smi: try: - result = subprocess.run( - [amd_smi, "version"], - stdout = subprocess.PIPE, - stderr = subprocess.DEVNULL, - text = True, - timeout = 5, - ) + # Off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); + # hipconfig below and the version-file reads above cover that case. + result = run_capture([amd_smi, "version"], timeout = 5) if result.returncode == 0: m = re.search(r"ROCm version:\s*(\d+)\.(\d+)", result.stdout) if m: @@ -4809,6 +4869,46 @@ def validated_validation_model_bytes(data: bytes) -> bytes: return data +def _hf_resolve_url_parts(url: str) -> tuple[str, str, str] | None: + """Parse a huggingface.co .../resolve// URL into + (repo_id, revision, filename); None if it is not such a URL.""" + try: + parsed = urllib.parse.urlparse(url) + except Exception: + return None + if (parsed.netloc or "").lower() not in ("huggingface.co", "www.huggingface.co"): + return None + parts = parsed.path.strip("/").split("/") + # //resolve// + if len(parts) >= 5 and parts[2] == "resolve": + return f"{parts[0]}/{parts[1]}", parts[3], "/".join(parts[4:]) + return None + + +def _fetch_validation_model_bytes() -> bytes: + """Fetch the tiny GGUF validation model. Prefer huggingface_hub (completes + TLS chains via AIA fetching that bare urllib can't on some Windows/proxy + setups); fall back to the direct URL when hf_hub is unavailable or fails.""" + parts = _hf_resolve_url_parts(TEST_MODEL_URL) + if parts is not None: + repo_id, revision, filename = parts + try: + from huggingface_hub import hf_hub_download + local = hf_hub_download(repo_id = repo_id, filename = filename, revision = revision) + return validated_validation_model_bytes(Path(local).read_bytes()) + except Exception as exc: + log( + f"huggingface_hub fetch of validation model failed ({exc}); " + "falling back to direct URL" + ) + return validated_validation_model_bytes( + download_bytes( + TEST_MODEL_URL, + progress_label = f"Downloading {download_label_from_url(TEST_MODEL_URL)}", + ) + ) + + def download_validation_model(path: Path, cache_path: Path | None = None) -> None: try: data: bytes | None = None @@ -4821,12 +4921,7 @@ def download_validation_model(path: Path, cache_path: Path | None = None) -> Non data = None if data is None: log("downloading tiny GGUF validation model") - data = validated_validation_model_bytes( - download_bytes( - TEST_MODEL_URL, - progress_label = f"Downloading {download_label_from_url(TEST_MODEL_URL)}", - ) - ) + data = _fetch_validation_model_bytes() if cache_path is not None: atomic_write_bytes(cache_path, data) atomic_write_bytes(path, data) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2c8f6228d4..e742d85eed 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -40,6 +40,14 @@ IS_MACOS = sys.platform == "darwin" IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64" IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64" IS_LINUX = sys.platform.startswith("linux") + +# DiskPart-prompt suppression: amd-smi auto-elevates on Windows, popping a +# UAC/DiskPart prompt mid-install. This installer only spawns probes and pip/uv +# (none need elevation), so set __COMPAT_LAYER=RunAsInvoker process-wide -- every +# amd-smi subprocess then runs un-elevated, no per-call guard needed. setup.ps1 +# keeps per-call guards since it ALSO spawns winget installers that need elevation. +if IS_WINDOWS: + os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker") # torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64, # and win_amd64. On other hosts the audio extras must be filtered out (the # extras-no-deps step would otherwise fail), regardless of NO_TORCH. @@ -137,6 +145,39 @@ def _bnb_rocm_prerelease_url() -> str | None: return _BNB_ROCM_PRERELEASE_URLS.get(arch) +def _amd_smi_env() -> dict[str, str] | None: + """On Windows, env with __COMPAT_LAYER=RunAsInvoker; None elsewhere. + NB: RunAsInvoker doesn't stop amd-smi's runtime elevation (its manifest is + asInvoker -- it elevates a child via ShellExecute). The real guard is + _amd_smi_allowed() below; this is harmless belt-and-suspenders.""" + if platform.system() != "Windows": + return None + return {**os.environ, "__COMPAT_LAYER": "RunAsInvoker"} + + +def _amd_smi_allowed() -> bool: + """Whether it is safe to spawn amd-smi here. + + On Windows w/o a working HIP runtime, amd-smi elevates a child and pops a + UAC/DiskPart prompt RunAsInvoker can't suppress. Only call it on Windows with + a HIP SDK (hipinfo present) or UNSLOTH_ENABLE_AMD_SMI=1; Linux/macOS always. + """ + if platform.system() != "Windows": + return True + flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower() + if flag in ("1", "true", "yes", "on"): + return True + if flag in ("0", "false", "no", "off"): + return False + if shutil.which("hipinfo"): + return True + for _var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + _root = os.environ.get(_var) + if _root and os.path.isfile(os.path.join(_root, "bin", "hipinfo.exe")): + return True + return False + + def _detect_rocm_version() -> tuple[int, int] | None: """Return (major, minor) of the installed ROCm stack, or None.""" rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm" @@ -155,8 +196,10 @@ def _detect_rocm_version() -> tuple[int, int] | None: except Exception: pass - # Try amd-smi version (outputs "... | ROCm version: X.Y.Z") - amd_smi = shutil.which("amd-smi") + # Try amd-smi version (outputs "... | ROCm version: X.Y.Z"). + # Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); + # hipconfig below covers that case. + amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None if amd_smi: try: result = subprocess.run( @@ -165,6 +208,7 @@ def _detect_rocm_version() -> tuple[int, int] | None: stderr = subprocess.DEVNULL, text = True, timeout = 5, + env = _amd_smi_env(), ) if result.returncode == 0: import re @@ -282,6 +326,14 @@ def _detect_windows_gfx_arch() -> str | None: if os.path.isfile(_candidate): hipinfo = _candidate break + if not hipinfo: + # 2b. AMD torch wheels ship hipInfo.exe into the venv Scripts dir + # (next to python.exe); resolvable even on driver-only hosts with no + # SDK install at all. Lets `studio update` re-detect the arch on a + # venv that already has the AMD wheel. + _venv_hipinfo = os.path.join(os.path.dirname(sys.executable), "hipInfo.exe") + if os.path.isfile(_venv_hipinfo): + hipinfo = _venv_hipinfo if hipinfo: try: result = subprocess.run( @@ -304,7 +356,9 @@ def _detect_windows_gfx_arch() -> str | None: pass # 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo. - amd_smi = shutil.which("amd-smi") + # Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); the arch + # arrives via --rocm-gfx / name inference there, so this is only needed when safe. + amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None if amd_smi: for _args in (("static", "--asic"), ("list",)): try: @@ -313,6 +367,7 @@ def _detect_windows_gfx_arch() -> str | None: stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, timeout = 10, + env = _amd_smi_env(), ) if result.returncode != 0: continue @@ -330,6 +385,74 @@ def _detect_windows_gfx_arch() -> str | None: return _pick except Exception: continue + + # 4. Last resort: GPU marketing name via WMI → arch table. Driver-only + # hosts (Adrenalin, no HIP SDK) have neither hipinfo nor amd-smi + # (amd-smi does not exist on Windows at all), but the display driver + # always knows the GPU name. Mirrors setup.ps1's $nameArchTable so a + # standalone `studio update` can repair a CPU-only venv on such hosts. + try: + result = subprocess.run( + [ + "powershell.exe", + "-NoProfile", + "-NonInteractive", + "-Command", + "(Get-CimInstance Win32_VideoController).Name", + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 30, + creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0), + ) + if result.returncode == 0: + _tokens = [] + for _name in result.stdout.decode(errors = "replace").splitlines(): + _arch = _gfx_arch_from_gpu_name(_name.strip()) + if _arch: + _tokens.append(_arch) + _pick = _dedup_pick(_tokens) + if _pick: + print(f" gfx arch inferred from GPU name (WMI): {_pick}") + return _pick + except Exception: + pass + return None + + +# GPU marketing-name → gfx arch table, mirroring setup.ps1's $nameArchTable. +# Most-specific first; first match wins. Covers only arches the lemonade-sdk +# prebuilts / AMD Windows torch indexes support; unknown names return None +# (callers then fall back cleanly to CPU). +_WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [ + (r"9070 XT|9080", "gfx1201"), # RDNA 4 (Radeon RX 9070 XT / 9080) + (r"9070|9060", "gfx1200"), # RDNA 4 (Radeon RX 9070 / 9060) + # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + (r"8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"), + # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) + ( + r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]" + r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", + "gfx1150", + ), + # RDNA 3 desktop / workstation (Navi 31) + (r"RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700", "gfx1100"), + (r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710", "gfx1102"), # Navi 33 + # RDNA 3 iGPU (Phoenix / Hawk Point) + (r"780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme", "gfx1103"), + (r"RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900", "gfx1030"), # Navi 21 + (r"RX 6650|RX 6600|PRO W6600|PRO W6650", "gfx1032"), # Navi 23 + (r"RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500", "gfx1034"), # Navi 24 +] + + +def _gfx_arch_from_gpu_name(name: str) -> "str | None": + """Map a GPU marketing name to its gfx arch via _WIN_GPU_NAME_ARCH_TABLE.""" + if not name: + return None + for _pat, _arch in _WIN_GPU_NAME_ARCH_TABLE: + if re.search(_pat, name, re.IGNORECASE): + return _arch return None @@ -386,6 +509,10 @@ def _has_rocm_gpu() -> bool: exe = shutil.which(cmd[0]) if not exe: continue + # Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); + # rely on rocminfo / the sysfs fallback there. + if cmd[0] == "amd-smi" and not _amd_smi_allowed(): + continue try: result = subprocess.run( [exe, *cmd[1:]], @@ -393,6 +520,7 @@ def _has_rocm_gpu() -> bool: stderr = subprocess.DEVNULL, text = True, timeout = 10, + env = _amd_smi_env() if cmd[0] == "amd-smi" else None, ) except Exception: continue @@ -454,7 +582,8 @@ def _detect_amd_gfx_codes() -> list[str]: probes: list[list[str]] = [] if shutil.which("rocminfo"): probes.append(["rocminfo"]) - if shutil.which("amd-smi"): + # Gate amd-smi off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt). + if shutil.which("amd-smi") and _amd_smi_allowed(): probes.append(["amd-smi", "list"]) probes.append(["amd-smi", "static", "--asic"]) for cmd in probes: @@ -465,6 +594,7 @@ def _detect_amd_gfx_codes() -> list[str]: stderr = subprocess.DEVNULL, text = True, timeout = 15, + env = _amd_smi_env() if cmd[0] == "amd-smi" else None, ) except Exception: continue @@ -515,6 +645,18 @@ def _install_bnb_windows_rocm() -> bool: if "BNB_ROCM_VERSION" not in os.environ: _ver = _detect_bnb_rocm_dll_ver() or "72" os.environ["BNB_ROCM_VERSION"] = _ver + # Make hipInfo.exe (shipped into the venv Scripts dir by the AMD torch + # wheel) resolvable via PATH for this process and every child python the + # installer spawns (import checks, precompile): bitsandbytes runs + # `hipinfo.exe` at import time to detect the GPU arch and logs a scary + # (harmless) ERROR + WARNING on every import when it is missing. The venv + # Scripts dir is on PATH only when the venv is activated, which neither + # Studio nor the installer's child processes ever do. + _scripts_dir = os.path.dirname(sys.executable) + if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which( + "hipinfo.exe" + ): + os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "") return True @@ -1569,6 +1711,12 @@ def install_python_stack() -> int: _wexe = shutil.which(_wcmd[0]) if not _wexe: continue + # Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart + # prompt), as _has_rocm_gpu()/_detect_amd_gfx_codes do. The only loss + # is the best-effort "AMD GPU detected" note; ROCm-torch state below + # comes from the install itself. + if _wcmd[0] == "amd-smi" and not _amd_smi_allowed(): + continue try: _wr = subprocess.run( [_wexe, *_wcmd[1:]], @@ -1576,6 +1724,7 @@ def install_python_stack() -> int: stderr = subprocess.DEVNULL, text = True, timeout = 10, + env = _amd_smi_env() if _wcmd[0] == "amd-smi" else None, ) except Exception: continue diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 9a919e31e3..ce94f651f8 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -700,6 +700,55 @@ if (-not $HasNvidiaSmi) { } } } +# ── Helper: run amd-smi without triggering a UAC elevation prompt ── +# amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing +# DiskPart UAC prompt mid-install (Studio backend amd.py hits the same). RunAsInvoker +# forces it (and helpers it spawns) to run un-elevated; on failure the WMI name -> +# gfx fallback still resolves the arch. +function Invoke-AmdSmiNoElevate { + param( + [Parameter(Mandatory = $true, Position = 0)][string]$Exe, + [Parameter(Position = 1)][string[]]$SmiArgs = @(), + [int]$TimeoutSec = 30 + ) + # RunAsInvoker blocks the auto-elevation/UAC prompt; the timeout bounds a flaky + # amd-smi that can otherwise spin for minutes (30s mirrors the backend amd.py). + $prevCompat = [Environment]::GetEnvironmentVariable('__COMPAT_LAYER', 'Process') + $env:__COMPAT_LAYER = 'RunAsInvoker' + try { + # [Process]::Start, NOT Start-Process -PassThru: the latter leaves .ExitCode + # $null after WaitForExit on PS 5.1, so $LASTEXITCODE (checked by callers) + # reads non-zero and kills detection. Async reads drain the pipes (no + # deadlock); amd-smi args have no spaces so a plain join is safe. + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $Exe + $psi.Arguments = ($SmiArgs -join ' ') + $psi.UseShellExecute = $false + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + if (-not $proc.WaitForExit($TimeoutSec * 1000)) { + try { $proc.Kill() } catch {} + $global:LASTEXITCODE = 124 + return "" + } + $global:LASTEXITCODE = $proc.ExitCode + return ($outTask.Result + "`n" + $errTask.Result) + } catch { + $global:LASTEXITCODE = 1 + return "" + } finally { + if ($null -eq $prevCompat) { + Remove-Item Env:__COMPAT_LAYER -ErrorAction SilentlyContinue + } else { + $env:__COMPAT_LAYER = $prevCompat + } + } +} + # ── AMD ROCm detection (Windows): probe hipinfo/amd-smi for actual GPU ── $HasROCm = $false $HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility) @@ -750,14 +799,24 @@ if (-not $HasNvidiaSmi) { } catch {} } # amd-smi fallback: HIP runtime present but hipinfo unavailable (no full HIP SDK). - # Confirms GPU visibility via 'list', then attempts 'static --asic' to extract - # the gfx arch that hipinfo would have provided. Critical for Strix Halo - # (gfx1151) and other iGPUs where only the HIP runtime is installed. - if (-not $HasROCm) { + # 'list' confirms GPU visibility, 'static --asic' extracts the gfx arch hipinfo + # would give. Critical for Strix Halo (gfx1151) and other HIP-runtime-only iGPUs. + # + # BUT on hosts without a working HIP runtime amd-smi elevates a child at runtime, + # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (its manifest is + # asInvoker; even 'amd-smi version' hangs). So only probe when a HIP SDK is present + # (hipinfo found -> un-elevated) or the user opts in; else fall through to WMI name + # inference (enough to pick ROCm wheels + lemonade llama.cpp). + # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the HIP-SDK + # heuristic: a HIP SDK binary with a broken runtime can still pop the prompt, so + # $HipSdkInstalled must NOT silently re-enable it. + $amdSmiOptOut = $env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(0|false|no|off)$' + $amdSmiAllowed = (-not $amdSmiOptOut) -and ($HipSdkInstalled -or ($env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(1|true|yes|on)$')) + if (-not $HasROCm -and $amdSmiAllowed) { $amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue if ($amdSmiExe) { try { - $smiOut = & $amdSmiExe.Source list 2>&1 | Out-String + $smiOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('list') if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { $HasROCm = $true # Attempt 1: newer amd-smi versions embed the gfx arch in list output. @@ -791,7 +850,7 @@ if (-not $HasNvidiaSmi) { # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, # including the GFX target needed for wheel index selection. $smiAsicOut = "" - try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {} + try { $smiAsicOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('static','--asic') } catch {} if ($smiAsicOut -match "(?i)\b(gfx\d+[a-z]?)\b") { $script:ROCmGfxArch = $Matches[1].ToLower() $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" @@ -819,27 +878,34 @@ if (-not $HasNvidiaSmi) { } catch {} } # ── Arch resolution: env-var override → name inference ────────────────── - # Runs after all probe methods. Covers users whose amd-smi version is too - # old to report the GFX target and who don't have hipinfo (HIP-runtime-only - # installs, common on Strix Halo / iGPU systems). - if ($HasROCm -and -not $script:ROCmGfxArch) { + # Runs after all probes, even when none confirmed a ROCm runtime ($HasROCm false): + # the Adrenalin driver alone runs the lemonade-sdk llama.cpp prebuilt (bundles its + # own runtime), and all it needs is the gfx arch, inferable from the WMI GPU name. + # Resolving it here lets setup.ps1 forward --rocm-gfx so a GPU llama.cpp is pulled + # instead of CPU. (PyTorch ROCm wheels still require a HIP SDK -- gated on $HasROCm + # below -- so this only affects llama.cpp / inference.) + if (-not $script:ROCmGfxArch) { # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. if ($env:UNSLOTH_ROCM_GFX_ARCH) { $script:ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower() $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $script:ROCmGfxArch" "Cyan" } - # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). - # Ordered most-specific first; first match wins. + # 2. Best-effort name → arch lookup (amd-smi / WMI). Most-specific first, + # first match wins. Covers only arches the lemonade-sdk prebuilts support + # (gfx120X/110X/1151/1150/103X); unknown names fall back cleanly to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 - @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 - @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) - @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) - @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop - @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 - @{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix) + @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (Radeon RX 9070 XT / 9080) + @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (Radeon RX 9070 / 9060) + @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) + @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31) + @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) + @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) + @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- lemonade gfx103X + @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X + @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X ) foreach ($row in $nameArchTable) { if ($ROCmGpuLabel -match $row.P) { @@ -881,11 +947,11 @@ if (-not $HasNvidiaSmi) { } } catch {} } - if (-not $script:ROCmVersion) { + if (-not $script:ROCmVersion -and $amdSmiAllowed) { $amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue if ($amdSmiVer) { try { - $smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String + $smiVerOut = Invoke-AmdSmiNoElevate $amdSmiVer.Source @('version') if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') { $script:ROCmVersion = $Matches[1] } } catch {} } @@ -910,11 +976,20 @@ if ($HasNvidiaSmi) { substep " This is a driver issue, not an SDK issue." "Yellow" substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" +} elseif ($script:ROCmGfxArch) { + # Known arch: PyTorch comes from AMD's bundled-runtime ROCm wheels (repo.amd.com), + # which ship their own runtime -- HIP SDK optional (only adds the system toolchain). + Write-Host "" + step "gpu" "AMD ROCm ($script:ROCmGfxArch)" "Cyan" + substep "Detected: $ROCmGpuLabel" "Cyan" + substep "GPU PyTorch uses AMD's bundled-runtime ROCm wheels -- HIP SDK not required (optional)." "Cyan" + Write-Host "" } elseif ($ROCmGpuLabel) { Write-Host "" - step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow" + step "gpu" "AMD GPU detected -- arch unknown" "Yellow" substep "Detected: $ROCmGpuLabel" "Yellow" - substep "Install the HIP SDK for ROCm GPU inference:" "Yellow" + substep "Could not determine the GPU arch (gfx...). Install the HIP SDK or set" "Yellow" + substep "UNSLOTH_ROCM_GFX_ARCH to enable GPU ROCm PyTorch:" "Yellow" substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" Write-Host "" } else { @@ -1192,7 +1267,7 @@ if (-not $NvccPath -and $RequireOrExit) { $drMajor = [int]$DriverMaxCuda.Split('.')[0] $AvailableVersions = @() try { - $rawOutput = winget show Nvidia.CUDA --versions --accept-source-agreements 2>&1 | Out-String + $rawOutput = winget show Nvidia.CUDA --versions --source winget --accept-source-agreements 2>&1 | Out-String # Parse version lines (e.g. "12.6", "12.5", "11.8") foreach ($line in $rawOutput -split "`n") { $line = $line.Trim() @@ -1345,8 +1420,12 @@ $script:CudaToolkitReady = $true if ($HasROCm) { $rocmVerLabel = if ($script:ROCmVersionFull) { "ROCm $script:ROCmVersionFull" } elseif ($script:ROCmVersion) { "ROCm $script:ROCmVersion" } else { "ROCm (version unknown)" } step "rocm" $rocmVerLabel +} elseif ($script:ROCmGfxArch) { + # GPU training/inference works via AMD's bundled-runtime ROCm PyTorch wheels; + # the HIP SDK is optional (only the system ROCm toolchain). + step "rocm" "GPU via bundled ROCm wheels ($script:ROCmGfxArch) -- HIP SDK optional" "Cyan" } elseif ($ROCmGpuLabel) { - step "rocm" "HIP SDK not found -- GPU-accelerated training unavailable" "Yellow" + step "rocm" "AMD GPU detected -- arch unknown; HIP SDK not found" "Yellow" } # ============================================ @@ -1997,6 +2076,21 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true + # ...but not if an AMD GPU is present and installed PyTorch is CPU-only + # (host predates ROCm-wheel support, or GPU added later): the fast "up to + # date" path would leave the user on CPU torch with Train/Export disabled. + # Force the dependency pass so the ROCm wheels get installed. + if ($script:ROCmGfxArch) { + $_torchIsCpu = $true + try { + & python -c "import torch, sys; sys.exit(0 if torch.cuda.is_available() else 1)" 2>$null + if ($LASTEXITCODE -eq 0) { $_torchIsCpu = $false } + } catch {} + if ($_torchIsCpu) { + substep "AMD GPU ($script:ROCmGfxArch) detected but installed PyTorch is CPU-only -- reinstalling ROCm PyTorch" "Cyan" + $SkipPythonDeps = $false + } + } } elseif ($InstalledVer -and $LatestVer) { substep "$_PkgName $InstalledVer -> $LatestVer available, updating..." } elseif (-not $LatestVer) { @@ -2060,7 +2154,13 @@ if ($HasNvidiaSmi) { # Wheels bundle their own ROCm runtime; HIP SDK version is irrelevant. $ROCmGfxArch = $script:ROCmGfxArch $ROCmIndexUrl = $null -if ($HasROCm -and $CuTag -eq "cpu") { +# Install AMD ROCm PyTorch wheels when ROCm is confirmed OR a gfx arch is known +# (name-inferred on Adrenalin-only hosts). The per-arch wheels bundle the runtime +# (rocm-sdk-libraries-), so torch.cuda.is_available() is True without a HIP +# SDK -- which flips Studio out of chat-only (CHAT_ONLY) and enables Train/Export. +# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed +# ROCm install still falls back to CPU below, so this is safe. +if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -2450,7 +2550,6 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { substep "Skipping prebuilt install -- falling back to source build" "Yellow" } else { Write-Host "" - substep "installing prebuilt llama.cpp bundle (preferred path)..." if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" # If the existing install is the wrong kind (e.g. windows-cpu on a ROCm @@ -2461,7 +2560,11 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { try { $existingMeta = Get-Content $existingMetaPath -Raw | ConvertFrom-Json $existingKind = $existingMeta.install_kind - $expectedKind = if ($HasROCm) { "windows-hip" } elseif ($HasNvidiaSmi) { "windows-cuda" } else { "windows-cpu" } + # A name-inferred gfx arch (Adrenalin-only, no confirmed runtime) + # still wants the GPU (windows-hip) build -- the lemonade prebuilt + # bundles its own runtime. Treat a known arch as ROCm-capable here, + # mirroring the --rocm-gfx forward below. + $expectedKind = if ($HasROCm -or $script:ROCmGfxArch) { "windows-hip" } elseif ($HasNvidiaSmi) { "windows-cuda" } else { "windows-cpu" } if ($existingKind -and $existingKind -ne $expectedKind) { substep "Removing mismatched llama.cpp install (found '$existingKind', need '$expectedKind')..." Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue @@ -2471,6 +2574,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } } } + substep "installing prebuilt llama.cpp bundle (preferred path)..." # 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. @@ -2486,12 +2590,14 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { ) if ($HasROCm) { $prebuiltArgs += "--has-rocm" - # Forward the resolved gfx arch so the lemonade HIP prebuilt is picked - # even when the installer's own probe cannot report it (amd-smi-only - # hosts, name-inferred arch). - if ($script:ROCmGfxArch) { - $prebuiltArgs += @("--rocm-gfx", $script:ROCmGfxArch) - } + } + # Forward the resolved gfx arch so the lemonade HIP prebuilt is picked even + # when the installer's probe can't confirm the runtime (amd-smi-only / + # Adrenalin-only, name-inferred arch). --rocm-gfx is authoritative and + # implies ROCm in install_llama_prebuilt.py, so the GPU prebuilt is selected + # even with $HasROCm false. Gating on $HasROCm gave Strix Halo / 8060S CPU. + if ($script:ROCmGfxArch) { + $prebuiltArgs += @("--rocm-gfx", $script:ROCmGfxArch) } if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) @@ -2586,7 +2692,7 @@ if ($NeedLlamaSourceBuild) { substep "installing OpenSSL dev (for HTTPS in llama-server)..." $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue) if ($HasWinget) { - winget install -e --id ShiningLight.OpenSSL.Dev --accept-package-agreements --accept-source-agreements + winget install -e --id ShiningLight.OpenSSL.Dev --source winget --accept-package-agreements --accept-source-agreements # Re-check after install foreach ($root in $OpenSslRoots) { if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) { @@ -2668,6 +2774,18 @@ if (-not $NeedLlamaSourceBuild) { Write-Host "" if ($HasNvidiaSmi) { substep "building llama.cpp with CUDA support..." + } elseif ($HasROCm -or $script:ROCmGfxArch) { + # AMD GPU present but in the CPU-only source-build fallback: a HIP source + # build needs the full HIP SDK + ROCm clang toolchain. AMD GPU acceleration + # comes from the lemonade prebuilt (bundles the runtime, no SDK) -- reaching + # here means it couldn't be installed. Warn loudly, don't ship a slow CPU build. + $_amdArch = if ($script:ROCmGfxArch) { $script:ROCmGfxArch } else { "ROCm" } + substep "[WARN] AMD GPU ($_amdArch) detected, but the GPU-accelerated lemonade" "Yellow" + substep " llama.cpp prebuilt could not be installed -- falling back to a CPU build." "Yellow" + substep " The prebuilt is the AMD GPU path (no HIP SDK required). To restore GPU" "Yellow" + substep " acceleration: re-run the installer (check your network / proxy), or set" "Yellow" + substep " UNSLOTH_LLAMA_RELEASE_TAG to a tag with a gfx prebuilt for your GPU." "Yellow" + substep "building llama.cpp (CPU-only fallback)..." "Yellow" } else { substep "building llama.cpp (CPU-only, no NVIDIA GPU detected)..." } diff --git a/studio/setup.sh b/studio/setup.sh index 15a20a572f..42e9f23626 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -773,6 +773,14 @@ fi fi # ── GPU detection summary (mirrors setup.ps1 step "gpu" block) ── +# WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg only when +# HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and /opt/rocm/bin can be +# off PATH outside login shells (the profile.d drop-in). Seed both before the +# probes or a ROCDXG WSL host is misdetected as CPU-only. +export HSA_ENABLE_DXG_DETECTION="${HSA_ENABLE_DXG_DETECTION:-1}" +if ! command -v rocminfo >/dev/null 2>&1 && [ -x /opt/rocm/bin/rocminfo ]; then + PATH="$PATH:/opt/rocm/bin" +fi _setup_amd_detected=false _setup_gfx_all="" _setup_mkt="" @@ -810,14 +818,20 @@ elif [ "$_setup_amd_detected" = true ]; then substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_setup_gfx" # Name-based arch inference when tools don't report gfx (mirrors setup.ps1 nameArchTable) elif [ -z "$_setup_gfx" ] && [ -n "$_setup_mkt" ]; then + # Kept in sync with the table in install.sh (and the PS nameArchTable). + # gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on + # gfx1102 (bash case has no negative lookahead like the PS tables). case "$_setup_mkt" in - *"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 - *"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 iGPU - *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 iGPU - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop - *"RX 7600"*) _setup_gfx="gfx1102" ;; # RDNA 3 - *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU + *"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 + *9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 + *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) + *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _setup_gfx="gfx1102" ;; # RDNA 3 (Navi 33) + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point) + *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _setup_gfx="gfx1030" ;; # RDNA 2 (Navi 21) + *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _setup_gfx="gfx1032" ;; # RDNA 2 (Navi 23) + *"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) _setup_gfx="gfx1034" ;; # RDNA 2 (Navi 24) esac if [ -n "$_setup_gfx" ]; then substep "gfx arch inferred from GPU name: $_setup_gfx" diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py new file mode 100644 index 0000000000..583532eb74 --- /dev/null +++ b/tests/studio/install/test_pr5940_followups.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Tests for the AMD-Windows installer follow-ups (PR #5940): + + * the huggingface_hub validation-model fetch + its urllib fallback, + * run_capture's Windows-only amd-smi __COMPAT_LAYER=RunAsInvoker injection, + * parity of the name->arch table between install.ps1 and setup.ps1. + +Mock-only; no AMD hardware or network required. +""" + +import importlib.util +import re +import subprocess +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +# install_llama_prebuilt.py is self-contained (stdlib + optional filelock), so it +# loads without the studio backend on sys.path. +_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +_SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_pr5940", _PREBUILT_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +prebuilt = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = prebuilt +_SPEC.loader.exec_module(prebuilt) + +_INSTALL_PS1 = PACKAGE_ROOT / "install.ps1" +_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" +_INSTALL_SH = PACKAGE_ROOT / "install.sh" + + +# ── _hf_resolve_url_parts ──────────────────────────────────────────────────── + + +def test_hf_resolve_url_parts_valid(): + assert prebuilt._hf_resolve_url_parts( + "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" + ) == ("ggml-org/models", "main", "tinyllamas/stories260K.gguf") + + +@pytest.mark.parametrize( + "url", + [ + "https://github.com/owner/repo/releases/download/x.gguf", # not huggingface + "https://huggingface.co/owner/repo", # no /resolve// + "https://huggingface.co/owner/repo/blob/main/x.gguf", # /blob/ not /resolve/ + "not even a url", + ], +) +def test_hf_resolve_url_parts_non_hf_returns_none(url): + assert prebuilt._hf_resolve_url_parts(url) is None + + +# ── _fetch_validation_model_bytes ──────────────────────────────────────────── + + +def test_fetch_validation_model_prefers_huggingface_hub(tmp_path): + model = tmp_path / "stories260K.gguf" + model.write_bytes(b"GGUF-via-hf") + fake_hf = MagicMock(return_value = str(model)) + with ( + patch.object(prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b), + patch.dict(sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}), + ): + assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-hf" + assert fake_hf.called # hf path was taken, urllib not needed + + +def test_fetch_validation_model_falls_back_to_urllib_on_hf_failure(): + fake_hf = MagicMock(side_effect = RuntimeError("hf unreachable")) + with ( + patch.object(prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b), + patch.dict(sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}), + patch.object(prebuilt, "download_bytes", return_value = b"GGUF-via-urllib") as dl, + ): + assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-urllib" + assert dl.called # fell back to the direct URL download + + +# ── run_capture amd-smi RunAsInvoker injection ─────────────────────────────── + + +def _capture_env(command, system): + captured = {"env": "sentinel"} + + def fake_run(cmd, **kwargs): + captured["env"] = kwargs.get("env") + return subprocess.CompletedProcess(cmd, 0, "", "") + + with ( + patch.object(prebuilt.subprocess, "run", side_effect = fake_run), + patch.object(prebuilt.platform, "system", return_value = system), + ): + prebuilt.run_capture(command) + return captured["env"] + + +def test_run_capture_injects_runasinvoker_for_amd_smi_on_windows(): + env = _capture_env(["amd-smi", "list"], "Windows") + assert env is not None and env.get("__COMPAT_LAYER") == "RunAsInvoker" + + +def test_run_capture_injects_for_full_path_amd_smi_exe_on_windows(): + env = _capture_env(["amd-smi.exe", "version"], "Windows") + assert env is not None and env.get("__COMPAT_LAYER") == "RunAsInvoker" + + +def test_run_capture_no_injection_for_non_amd_smi_on_windows(): + assert _capture_env(["rocminfo"], "Windows") is None + + +def test_run_capture_no_injection_on_linux(): + # amd-smi does not auto-elevate on Linux, so no env override is applied. + assert _capture_env(["amd-smi", "list"], "Linux") is None + + +# ── name->arch table parity (install.ps1 vs setup.ps1) ─────────────────────── + + +def _ps_name_arch_rows(text): + return re.findall(r'@\{\s*P\s*=\s*"([^"]*)"\s*;\s*A\s*=\s*"(gfx[0-9a-z]+)"', text) + + +def test_ps_name_arch_tables_in_sync(): + t1 = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8")) + t2 = _ps_name_arch_rows(_SETUP_PS1.read_text(encoding = "utf-8")) + assert t1, "no nameArchTable found in install.ps1" + assert t1 == t2, f"name->arch tables drifted:\ninstall.ps1={t1}\nsetup.ps1={t2}" + + +def test_rx_7700s_resolves_to_gfx1102_not_gfx1100(): + rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8")) + name = "AMD Radeon RX 7700S" + matched = next((arch for pattern, arch in rows if re.search(pattern, name)), None) + assert matched == "gfx1102", f"RX 7700S matched {matched!r}, expected gfx1102" + + +def test_radeon_8060s_resolves_to_gfx1151(): + rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8")) + name = "AMD Radeon(TM) 8060S Graphics" + matched = next((arch for pattern, arch in rows if re.search(pattern, name)), None) + assert matched == "gfx1151" + + +def _sh_name_arch_rows(text, var = "_gpu_disp_gfx"): + """Parse a bash `case "$..._mkt" in ... ) ="gfxNNNN"` name->arch + table into [(substr_tokens, arch), ...] preserving order.""" + rows = [] + for line in text.splitlines(): + m = re.search(var + r'="(gfx[0-9a-z]+)"', line) + if not m or '*"' not in line: + continue + tokens = re.findall(r'\*"([^"]+)"\*', line) + if tokens: + rows.append((tokens, m.group(1))) + return rows + + +def _sh_resolve(rows, name): + for tokens, arch in rows: + if any(tok in name for tok in tokens): # bash *"X"* == substring + return arch + return None + + +def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd(): + """The bash install.sh name->arch table must agree with the PowerShell + source-of-truth for the Strix Halo (gfx1151) vs Strix Point (gfx1150) + split, and must never misclassify NVIDIA/Intel as an AMD gfx.""" + sh_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8")) + ps_rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8")) + assert sh_rows, "no name->arch case table found in install.sh" + cases = { + "AMD Radeon(TM) 8060S Graphics": "gfx1151", # Strix Halo + "AMD Ryzen AI Max+ PRO 395 w/ Radeon 8060S": "gfx1151", + "AMD Radeon 890M Graphics": "gfx1150", # Strix Point (NOT gfx1151) + "AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150", + "AMD Radeon RX 7700S": "gfx1102", + "NVIDIA GeForce RTX 4090": None, + "Intel(R) Arc A770 Graphics": None, + } + for name, expect in cases.items(): + sh = _sh_resolve(sh_rows, name) + assert sh == expect, f"install.sh: {name!r} -> {sh!r}, expected {expect!r}" + if expect is not None: # cross-check bash agrees with the PowerShell table + ps = next((a for p, a in ps_rows if re.search(p, name)), None) + assert sh == ps, f"install.sh/install.ps1 drift for {name!r}: {sh!r} vs {ps!r}" + + +def test_setup_sh_name_arch_table_in_sync_with_install_sh(): + """studio/setup.sh keeps its own copy of the bash name->arch table (over + `_setup_gfx`); it must stay row-for-row identical to install.sh's, both in + tokens and in match order (order carries the RX 7700S -> gfx1102 rule).""" + install_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8")) + setup_rows = _sh_name_arch_rows( + (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8"), + var = "_setup_gfx", + ) + assert setup_rows, "no name->arch case table found in studio/setup.sh" + assert install_rows == setup_rows, ( + "bash name->arch tables drifted:\n" + f"install.sh={install_rows}\nstudio/setup.sh={setup_rows}" + ) + # The historical drift this guards against: Strix Point SKUs must be + # gfx1150, and the spaceless RX 7700S must match gfx1102 before gfx1100. + for name, expect in { + "AMD Radeon 890M Graphics": "gfx1150", + "AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150", + "AMD Radeon(TM) 8060S Graphics": "gfx1151", + "AMD Radeon RX 7700S": "gfx1102", + }.items(): + got = _sh_resolve(setup_rows, name) + assert got == expect, f"setup.sh: {name!r} -> {got!r}, expected {expect!r}" + + +# ── amd-smi gating (DiskPart UAC-prompt avoidance) ─────────────────────────── +# On Windows w/o a HIP SDK, amd-smi elevates and pops a UAC/DiskPart prompt +# RunAsInvoker can't suppress, so _amd_smi_allowed() skips it by default; +# HIP-SDK hosts and an explicit opt-in keep it. + + +def _amd_smi_allowed_under(system, hipinfo_present, env): + which = ( + (lambda name: r"C:\hip\bin\hipinfo.exe" if name == "hipinfo" else None) + if hipinfo_present + else (lambda name: None) + ) + with ( + patch.object(prebuilt.platform, "system", return_value = system), + patch.object(prebuilt.shutil, "which", side_effect = which), + patch.dict(prebuilt.os.environ, env, clear = True), + ): + return prebuilt._amd_smi_allowed() + + +def test_amd_smi_allowed_on_linux_regardless(): + # Linux amd-smi does not elevate -> always allowed (no regression on Linux). + assert _amd_smi_allowed_under("Linux", hipinfo_present = False, env = {}) is True + + +def test_amd_smi_skipped_on_windows_without_hip_sdk(): + # The DiskPart fix: no HIP SDK + no opt-in -> do not spawn amd-smi. + assert _amd_smi_allowed_under("Windows", hipinfo_present = False, env = {}) is False + + +def test_amd_smi_allowed_on_windows_with_hip_sdk(): + # hipinfo present => amd-smi runs un-elevated, so it is allowed (no regression + # for HIP-SDK Windows users, who never saw the prompt). + assert _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {}) is True + + +def test_amd_smi_opt_in_forces_on_windows_no_sdk(): + assert ( + _amd_smi_allowed_under( + "Windows", hipinfo_present = False, env = {"UNSLOTH_ENABLE_AMD_SMI": "1"} + ) + is True + ) + + +def test_amd_smi_opt_out_overrides_hip_sdk(): + assert ( + _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {"UNSLOTH_ENABLE_AMD_SMI": "0"}) + is False + ) + + +def test_ps_installers_gate_amd_smi_on_windows(): + # Both PowerShell installers must gate amd-smi behind HIP-SDK presence + the + # UNSLOTH_ENABLE_AMD_SMI opt-in, mirroring _amd_smi_allowed(). + for ps in (_INSTALL_PS1, _SETUP_PS1): + text = ps.read_text(encoding = "utf-8") + assert "UNSLOTH_ENABLE_AMD_SMI" in text, f"{ps.name} missing amd-smi opt-in gate" + assert "amdSmiAllowed" in text, f"{ps.name} missing amd-smi gate variable" + + +def test_install_python_stack_gates_every_amd_smi_spawn(): + # Regression for the DiskPart UAC prompt: every function that both names the + # `amd-smi` command AND spawns a subprocess must gate it behind + # _amd_smi_allowed(). The "ROCm torch missing" probe once spawned `amd-smi + # list` ungated on Adrenalin-only hosts; not-spawning is the only fix. + import ast + + src = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text(encoding = "utf-8") + tree = ast.parse(src) + + def _names_amd_smi_command(node): + # Exact "amd-smi"/"amd-smi.exe" constant, not a substring in a log message. + return any( + isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.lower() in ("amd-smi", "amd-smi.exe") + for n in ast.walk(node) + ) + + def _spawns_subprocess(node): + for n in ast.walk(node): + if ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "subprocess" + ): + return True + return False + + def _references_gate(node): + return any(isinstance(n, ast.Name) and n.id == "_amd_smi_allowed" for n in ast.walk(node)) + + offenders = [ + node.name + for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and _names_amd_smi_command(node) + and _spawns_subprocess(node) + and not _references_gate(node) + ] + assert not offenders, ( + "install_python_stack.py spawns amd-smi without an _amd_smi_allowed() " + f"gate in: {offenders} -- this pops the Windows UAC/DiskPart prompt on " + "Adrenalin-only (no HIP SDK) hosts." + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index e0309dbd35..c761cd090f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -550,9 +550,13 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) def test_no_rocm_skips(self, mock_nvidia, mock_pip): """No ROCm toolchain should skip entirely.""" - with patch("os.path.isdir", return_value = False): - with patch("shutil.which", return_value = None): - _ensure_rocm_torch() + # _detect_windows_gfx_arch pinned to None: on a real AMD test host its + # WMI name fallback would otherwise answer and defeat the "no ROCm + # anywhere" premise of this test. + with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None): + with patch("os.path.isdir", return_value = False): + with patch("shutil.which", return_value = None): + _ensure_rocm_torch() mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") @@ -692,8 +696,12 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_rocm_gpu", return_value = False) def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip): """ROCm tools present but no actual AMD GPU should skip entirely.""" - with patch("os.path.isdir", return_value = True): - _ensure_rocm_torch() + # Pin the Windows arch probe to None: on a real AMD test host the WMI + # name fallback would otherwise answer and defeat the "no actual GPU" + # premise (the Linux path under test uses _has_rocm_gpu, mocked False). + with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None): + with patch("os.path.isdir", return_value = True): + _ensure_rocm_torch() mock_pip.assert_not_called() @@ -862,7 +870,14 @@ class TestInstallShStructure: """Verify install.sh structural properties without running it.""" def test_no_here_strings(self): - """install.sh must not use <<< (not POSIX).""" + """install.sh must not use the bash-only `<<<` here-string operator. + + `<<<` inside a quoted literal (e.g. a marker in a printf) is just data, + not a here-string, so strip quoted spans first: this still catches a + real `cmd <<< word` (outside quotes) without false-positiving on data. + """ + import re + sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") # <<< is bash-only; breaks dash @@ -870,7 +885,11 @@ class TestInstallShStructure: stripped = line.lstrip() if stripped.startswith("#"): continue - assert "<<<" not in line, f"install.sh:{i} uses non-POSIX <<< here-string" + # Remove quoted string literals so `<<<` inside them is ignored; + # a genuine here-string operator lives outside any quotes. + unquoted = re.sub(r"'[^']*'", "", line) + unquoted = re.sub(r'"[^"]*"', "", unquoted) + assert "<<<" not in unquoted, f"install.sh:{i} uses non-POSIX <<< here-string" def test_rocm_detection_present(self): """install.sh should have ROCm detection in get_torch_index_url.""" @@ -1202,6 +1221,10 @@ class TestAmdGpuMonitoring: ): monkeypatch.delenv(var, raising = False) + # amd-smi is gated off on Windows w/o a HIP SDK; this test mocks it as + # available, so opt in so the gate allows it on every platform. + monkeypatch.setenv("UNSLOTH_ENABLE_AMD_SMI", "1") + mock_json = json.dumps( [ { @@ -1216,8 +1239,12 @@ class TestAmdGpuMonitoring: mock_result.returncode = 0 mock_result.stdout = mock_json - with patch.object(subprocess, "run", return_value = mock_result): - result = amd_mod.get_primary_gpu_utilization() + # The premise is "amd-smi exists and answers": the absence guard + # which()-checks before spawning, so hosts without a real amd-smi + # (Linux CI, driver-only Windows) need which mocked too. + with patch.object(amd_mod.shutil, "which", return_value = "/usr/bin/amd-smi"): + with patch.object(subprocess, "run", return_value = mock_result): + result = amd_mod.get_primary_gpu_utilization() assert result["available"] is True assert result["gpu_utilization_pct"] == 50.0 assert result["temperature_c"] == 65.0 @@ -1238,7 +1265,12 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module") - with patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")): + # Opt in so the call reaches subprocess.run (gated off on Windows w/o a + # HIP SDK); testing the OSError handling here. + with ( + patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}), + patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")), + ): result = amd_mod.get_primary_gpu_utilization() assert result["available"] is False @@ -1258,10 +1290,15 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module") - with patch.object( - subprocess, - "run", - side_effect = subprocess.TimeoutExpired("amd-smi", 5), + # Opt in so the call reaches subprocess.run (gated off on Windows w/o a + # HIP SDK); testing the timeout handling here. + with ( + patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}), + patch.object( + subprocess, + "run", + side_effect = subprocess.TimeoutExpired("amd-smi", 5), + ), ): result = amd_mod.get_primary_gpu_utilization() assert result["available"] is False @@ -1508,8 +1545,13 @@ class TestDetectWindowsGfxArch: """Verify hipinfo parsing for GPU arch detection on Windows.""" def test_returns_none_when_hipinfo_not_on_path(self): + # Also neutralise the venv-hipInfo and WMI-name fallbacks: this test + # pins "no probe source available -> None", and the suite may run on a + # real AMD host where WMI would legitimately answer. with patch("shutil.which", return_value = None): - result = stack_mod._detect_windows_gfx_arch() + with patch("os.path.isfile", return_value = False): + with patch("subprocess.run", side_effect = FileNotFoundError): + result = stack_mod._detect_windows_gfx_arch() assert result is None def test_parses_gcnarchname_from_hipinfo_output(self): @@ -1531,11 +1573,21 @@ class TestDetectWindowsGfxArch: assert result is None def test_returns_none_when_no_gcnarchname_in_output(self): + # hipinfo answers but without a gcnArchName line. Route only the + # hipinfo/amd-smi probes to the mock; the WMI fallback must get + # nothing (FileNotFoundError) -- otherwise the mocked device name + # would legitimately resolve via the name table. mock_result = MagicMock() mock_result.returncode = 0 - mock_result.stdout = b"deviceName : Radeon RX 9060 XT\n" + mock_result.stdout = b"deviceName : SomeUnknownDevice\n" + + def _run(cmd, **kwargs): + if cmd and "powershell" in str(cmd[0]).lower(): + raise FileNotFoundError(cmd[0]) + return mock_result + with patch("shutil.which", return_value = "/usr/bin/hipinfo"): - with patch("subprocess.run", return_value = mock_result): + with patch("subprocess.run", side_effect = _run): result = stack_mod._detect_windows_gfx_arch() assert result is None @@ -1558,6 +1610,99 @@ class TestDetectWindowsGfxArch: assert result == "gfx1201" +# TEST: install_python_stack.py -- GPU-name / WMI fallback (no amd-smi, no hipinfo) + + +class TestGfxArchNameFallback: + """amd-smi does not exist on Windows (neither Adrenalin consistently nor + the HIP SDK ship a CLI) and driver-only hosts lack hipinfo too. The + detection chain must still resolve the arch from the GPU marketing name + (WMI), mirroring setup.ps1's $nameArchTable.""" + + @pytest.mark.parametrize( + "name, expected", + [ + ("AMD Radeon(TM) 8060S Graphics", "gfx1151"), + ("AMD Ryzen AI MAX+ 395 w/ Radeon 8060S", "gfx1151"), + ("AMD Radeon(TM) 890M", "gfx1150"), + ("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"), + ("AMD Radeon RX 9070 XT", "gfx1201"), + ("AMD Radeon RX 9070", "gfx1200"), + ("AMD Radeon RX 7700S", "gfx1102"), # (?!S) lookahead must not hit gfx1100 + ("AMD Radeon RX 7700 XT", "gfx1100"), + ("AMD Radeon(TM) 780M", "gfx1103"), + ("NVIDIA GeForce RTX 4090", None), + ("Microsoft Basic Display Adapter", None), + ("", None), + ], + ) + def test_name_to_arch_mapping(self, name, expected): + assert stack_mod._gfx_arch_from_gpu_name(name) == expected + + def test_wmi_fallback_resolves_arch_without_any_tools(self): + """hipinfo absent everywhere + amd-smi absent -> WMI name fallback.""" + ps_result = MagicMock() + ps_result.returncode = 0 + ps_result.stdout = b"AMD Radeon(TM) 8060S Graphics\r\nMicrosoft Basic Display Adapter\r\n" + + def _run(cmd, **kwargs): + if cmd and "powershell.exe" in str(cmd[0]).lower(): + return ps_result + raise FileNotFoundError(cmd[0]) + + with patch.dict(os.environ, {}, clear = False): + for _v in ( + "HIP_PATH", + "ROCM_PATH", + "UNSLOTH_ROCM_GFX_ARCH", + "UNSLOTH_ENABLE_AMD_SMI", + ): + os.environ.pop(_v, None) + with patch("shutil.which", return_value = None): + with patch("os.path.isfile", return_value = False): + with patch("subprocess.run", side_effect = _run): + result = stack_mod._detect_windows_gfx_arch() + assert result == "gfx1151" + + def test_wmi_fallback_returns_none_for_non_amd_hosts(self): + ps_result = MagicMock() + ps_result.returncode = 0 + ps_result.stdout = b"NVIDIA GeForce RTX 4090\r\n" + + def _run(cmd, **kwargs): + if cmd and "powershell.exe" in str(cmd[0]).lower(): + return ps_result + raise FileNotFoundError(cmd[0]) + + with patch.dict(os.environ, {}, clear = False): + for _v in ("HIP_PATH", "ROCM_PATH", "UNSLOTH_ROCM_GFX_ARCH"): + os.environ.pop(_v, None) + with patch("shutil.which", return_value = None): + with patch("os.path.isfile", return_value = False): + with patch("subprocess.run", side_effect = _run): + result = stack_mod._detect_windows_gfx_arch() + assert result is None + + def test_stack_probes_venv_hipinfo(self): + """The venv Scripts dir hipInfo.exe (shipped by AMD torch wheels) must + be a probe candidate so `studio update` works on driver-only hosts.""" + source = _STACK_PATH.read_text(encoding = "utf-8") + assert 'os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")' in source + + def test_prebuilt_resolve_exe_probes_venv_dir(self): + """install_llama_prebuilt's _resolve_exe must include the venv Scripts + candidate for the same driver-only standalone-rerun scenario.""" + source = _PREBUILT_PATH.read_text(encoding = "utf-8") + assert "_venv_candidate" in source + + def test_runtime_monitor_guards_amd_smi_absence(self): + """amd.py must which()-check amd-smi before spawning so absence + disables the poller in one step (no FileNotFoundError strikes).""" + amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + source = amd_path.read_text(encoding = "utf-8") + assert 'shutil.which("amd-smi") is None' in source + + # TEST: install_python_stack.py -- _install_bnb_windows_rocm @@ -2350,6 +2495,38 @@ class TestServerStartupRocmFixes: source = _MAIN_PY_PATH.read_text(encoding = "utf-8") assert '"BNB_ROCM_VERSION" not in os.environ' in source + # ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ──────────────── + # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via subprocess PATH + # at import time. The AMD torch wheel ships hipInfo.exe in the venv + # Scripts dir, which is on PATH only for activated venvs -- Studio and the + # installer launch python directly, so without the prepend every bnb + # import logs "Could not detect ROCm GPU architecture: [WinError 2]". + + def test_main_py_prepends_hipinfo_dir_to_path(self): + """main.py must make hipInfo.exe resolvable before bnb imports.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert "hipInfo.exe" in source + # The prepend must come before the BNB_ROCM_VERSION block (both run + # pre-import; order documents that bnb sees the fixed PATH). + assert source.find("hipInfo.exe") < source.find("BNB_ROCM_VERSION") + + def test_main_py_hipinfo_prepend_gated_on_file_presence(self): + """Only AMD ROCm wheels ship hipInfo.exe; NVIDIA/CPU hosts must be + untouched, so the prepend must check the file exists first.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert 'os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe"))' in source + + def test_worker_py_prepends_hipinfo_dir_to_path(self): + """worker.py must mirror the prepend for standalone-spawned workers.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "hipInfo.exe" in source + + def test_install_stack_prepends_hipinfo_dir_to_path(self): + """install_python_stack.py must prepend so the installer's child + import checks inherit a PATH where bnb's probe succeeds.""" + source = _STACK_PATH.read_text(encoding = "utf-8") + assert "hipInfo.exe" in source + # ── torch._C._distributed_c10d stubs in hardware.py ────────────────────── def test_hardware_py_injects_distributed_c10d_stub(self): diff --git a/tests/studio/test_cli_studio_stop_windows.py b/tests/studio/test_cli_studio_stop_windows.py new file mode 100644 index 0000000000..7ba3d12214 --- /dev/null +++ b/tests/studio/test_cli_studio_stop_windows.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for `unsloth studio stop` on Windows (PR #5940). + +`stop` once used the POSIX `os.kill(pid, 0)` probe, which raises OSError +(WinError 87) for every pid on Windows -- crashing before reaching taskkill. +The fix adds a cross-platform `_pid_alive(pid)` (tasklist on Windows, signal-0 +elsewhere). + +AST + mock-only; no real process management, no Studio deps imported. +""" + +import ast +import os +import subprocess +import sys +import types +from pathlib import Path + +import pytest + +_STUDIO_CMD_PY = Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py" +_SOURCE = _STUDIO_CMD_PY.read_text(encoding = "utf-8") + + +def _func_source(name: str) -> str: + """Return the source of a top-level function `name` in studio.py.""" + tree = ast.parse(_SOURCE) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name: + return ast.get_source_segment(_SOURCE, node) + raise AssertionError(f"function {name!r} not found in studio.py") + + +def _load_pid_alive(platform: str, fake_run = None): + """Exec just `_pid_alive` with injectable sys/subprocess, so we can drive + the win32 branch on any host without importing the full unsloth_cli.""" + src = _func_source("_pid_alive") + fake_sys = types.SimpleNamespace(platform = platform) + fake_sub = types.SimpleNamespace(run = fake_run) if fake_run is not None else subprocess + ns = {"os": os, "sys": fake_sys, "subprocess": fake_sub} + exec(src, ns) + return ns["_pid_alive"] + + +# ── AST: stop() must not use the broken bare liveness probe ────────────────── + + +def test_stop_does_not_use_bare_oskill_liveness_probe(): + """stop() must not call os.kill(pid, 0) -- it crashes on Windows.""" + stop_src = _func_source("stop") + tree = ast.parse(stop_src) + for call in ast.walk(tree): + if not isinstance(call, ast.Call): + continue + f = call.func + is_os_kill = ( + isinstance(f, ast.Attribute) + and f.attr == "kill" + and isinstance(f.value, ast.Name) + and f.value.id == "os" + ) + if is_os_kill and len(call.args) == 2: + sig = call.args[1] + if isinstance(sig, ast.Constant) and sig.value == 0: + raise AssertionError( + "stop() still uses os.kill(pid, 0); it raises WinError 87 on " + "Windows. Use the cross-platform _pid_alive() helper instead." + ) + + +def test_pid_alive_helper_is_defined_and_used_by_stop(): + assert "def _pid_alive(" in _SOURCE, "_pid_alive helper missing" + assert "_pid_alive(pid)" in _func_source("stop"), "stop() must use _pid_alive" + # The helper must special-case Windows via tasklist (os.kill(pid,0) is invalid there). + helper = _func_source("_pid_alive") + assert 'sys.platform == "win32"' in helper + assert "tasklist" in helper + + +# ── Behavioral: the win32 tasklist branch ──────────────────────────────────── + + +def _fake_tasklist(returns_pid: int | None, *, raises: bool = False): + def _run( + cmd, + capture_output = False, + text = False, + timeout = None, + ): + assert cmd[0] == "tasklist" + assert "/FI" in cmd # filtered by PID + if raises: + raise OSError("boom") + if returns_pid is None: + stdout = "INFO: No tasks are running which match the specified criteria.\n" + else: + stdout = f'"python.exe","{returns_pid}","Console","1","12,345 K"\n' + return types.SimpleNamespace(stdout = stdout, returncode = 0) + + return _run + + +def test_pid_alive_windows_true_when_tasklist_lists_pid(): + pid_alive = _load_pid_alive("win32", fake_run = _fake_tasklist(4242)) + assert pid_alive(4242) is True + + +def test_pid_alive_windows_false_when_tasklist_empty(): + pid_alive = _load_pid_alive("win32", fake_run = _fake_tasklist(None)) + assert pid_alive(4242) is False + + +def test_pid_alive_windows_assumes_alive_when_tasklist_errors(): + # Can't determine -> assume alive; taskkill is the source of truth. + pid_alive = _load_pid_alive("win32", fake_run = _fake_tasklist(None, raises = True)) + assert pid_alive(4242) is True + + +# ── Behavioral: the POSIX signal-0 branch (skip on Windows runners) ─────────── + + +@pytest.mark.skipif(sys.platform == "win32", reason = "POSIX os.kill(pid,0) branch") +def test_pid_alive_posix_true_for_self_false_for_dead(): + pid_alive = _load_pid_alive("linux") + assert pid_alive(os.getpid()) is True + assert pid_alive(2_000_000_000) is False diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index a458f195b4..b506f764d7 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1130,6 +1130,33 @@ def run( _PID_FILE = STUDIO_HOME / "studio.pid" +def _pid_alive(pid: int) -> bool: + """Return True if a process with ``pid`` exists. + + ``os.kill(pid, 0)`` raises OSError (WinError 87) for every pid on Windows, + so use ``tasklist`` there and the signal-0 probe elsewhere. + """ + if sys.platform == "win32": + try: + out = subprocess.run( + ["tasklist", "/FI", f"PID eq {int(pid)}", "/NH", "/FO", "CSV"], + capture_output = True, + text = True, + timeout = 10, + ).stdout + except Exception: + # Can't determine -- assume alive; taskkill no-ops if already gone. + return True + return f'"{int(pid)}"' in out + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + @studio_app.command() def stop(): """Stop a running Unsloth Studio server. @@ -1151,15 +1178,11 @@ def stop(): pid = int(pid_text) - # Check if the process is still alive - try: - os.kill(pid, 0) - except ProcessLookupError: + # Check if still alive (os.kill(pid, 0) is invalid on Windows -- see _pid_alive). + if not _pid_alive(pid): typer.echo(f"Studio server (PID {pid}) is not running. Cleaning up stale PID file.") _PID_FILE.unlink(missing_ok = True) raise typer.Exit(0) - except PermissionError: - pass # process exists but we may not own it; try to signal anyway # Send SIGTERM (graceful shutdown) or TerminateProcess on Windows try: @@ -1176,17 +1199,13 @@ def stop(): typer.echo(f"Failed to stop Studio server (PID {pid}): {e}", err = True) raise typer.Exit(1) - # Wait briefly for the process to exit and clean up + # Wait briefly for the process to exit and clean up. for _ in range(10): time.sleep(0.5) - try: - os.kill(pid, 0) - except ProcessLookupError: + if not _pid_alive(pid): _PID_FILE.unlink(missing_ok = True) typer.echo("Studio server stopped.") raise typer.Exit(0) - except PermissionError: - break typer.echo("Studio server is shutting down (may take a few seconds).")