diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 25eeaedd3c..f3b5987a9c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.7 + rev: v0.15.8 hooks: - id: ruff args: diff --git a/README.md b/README.md index 8f783bf661..26a578656c 100644 --- a/README.md +++ b/README.md @@ -57,19 +57,20 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**. ```bash curl -fsSL https://unsloth.ai/install.sh | sh ``` -If you don't have `curl`, use `wget`. Launch after setup via: -```bash -source unsloth_studio/bin/activate -unsloth studio -H 0.0.0.0 -p 8888 -``` - #### Windows: ```powershell irm https://unsloth.ai/install.ps1 | iex ``` -Launch after setup via: -```powershell -& .\unsloth_studio\Scripts\unsloth.exe studio -H 0.0.0.0 -p 8888 + +#### Launch +```bash +unsloth studio -H 0.0.0.0 -p 8888 +``` + +#### Update +To update, use the same install commands as above. Or run (does not work on Windows): +```bash +unsloth studio update ``` #### Docker @@ -82,64 +83,8 @@ docker run -d -e JUPYTER_PASSWORD="mypassword" \ unsloth/unsloth ``` -#### macOS, Linux, WSL developer installs: -```bash -curl -LsSf https://astral.sh/uv/install.sh | sh -uv venv unsloth_studio --python 3.13 -source unsloth_studio/bin/activate -uv pip install unsloth --torch-backend=auto -unsloth studio setup -unsloth studio -H 0.0.0.0 -p 8888 -``` - -#### Windows PowerShell developer installs: -```powershell -winget install -e --id Python.Python.3.13 -winget install --id=astral-sh.uv -e -uv venv unsloth_studio --python 3.13 -.\unsloth_studio\Scripts\activate -uv pip install unsloth --torch-backend=auto -unsloth studio setup -unsloth studio -H 0.0.0.0 -p 8888 -``` - -#### Nightly - MacOS, Linux, WSL: -```bash -curl -LsSf https://astral.sh/uv/install.sh | sh -git clone --filter=blob:none https://github.com/unslothai/unsloth.git unsloth_studio -cd unsloth_studio -uv venv --python 3.13 -source .venv/bin/activate -uv pip install -e . --torch-backend=auto -unsloth studio setup -unsloth studio -H 0.0.0.0 -p 8888 -``` -Then to launch every time: -```bash -cd unsloth_studio -source .venv/bin/activate -unsloth studio -H 0.0.0.0 -p 8888 -``` - -#### Nightly - Windows: -Run in Windows Powershell: -```bash -winget install -e --id Python.Python.3.13 -winget install --id=astral-sh.uv -e -git clone --filter=blob:none https://github.com/unslothai/unsloth.git unsloth_studio -cd unsloth_studio -uv venv --python 3.13 -.\.venv\Scripts\activate -uv pip install -e . --torch-backend=auto -unsloth studio setup -unsloth studio -H 0.0.0.0 -p 8888 -``` -Then to launch every time: -```bash -cd unsloth_studio -.\.venv\Scripts\activate -unsloth studio -H 0.0.0.0 -p 8888 -``` +#### Developer, Nightly, Uninstall +To see developer, nightly and uninstallation etc. instructions, see [advanced installation](#-advanced-installation). ### Unsloth Core (code-based) #### Linux, WSL: @@ -197,6 +142,76 @@ Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get- - **FP8 & Vision RL**: You can now do FP8 & VLM GRPO on consumer GPUs. [FP8 Blog](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) • [Vision RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/vision-reinforcement-learning-vlm-rl) - **gpt-oss** by OpenAI: Read our [RL blog](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/gpt-oss-reinforcement-learning), [Flex Attention](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/long-context-gpt-oss-training) blog and [Guide](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune). +## 📥 Advanced Installation +The below advanced instructions are for Unsloth Studio. For Unsloth Core advanced installation, [view our docs](https://unsloth.ai/docs/get-started/install/pip-install#advanced-pip-installation). +#### Developer installs: macOS, Linux, WSL: +```bash +git clone https://github.com/unslothai/unsloth +cd unsloth +./install.sh --local +unsloth studio -H 0.0.0.0 -p 8888 +``` +Then to update : +```bash +unsloth studio update +``` + +#### Developer installs: Windows PowerShell: +```powershell +git clone https://github.com/unslothai/unsloth.git +cd unsloth +Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass +.\install.ps1 --local +unsloth studio -H 0.0.0.0 -p 8888 +``` +Then to update : +```bash +unsloth studio update +``` + +#### Nightly: MacOS, Linux, WSL: +```bash +git clone https://github.com/unslothai/unsloth +cd unsloth +git checkout nightly +./install.sh --local +unsloth studio -H 0.0.0.0 -p 8888 +``` +Then to launch every time: +```bash +unsloth studio -H 0.0.0.0 -p 8888 +``` + +#### Nightly: Windows: +Run in Windows Powershell: +```bash +git clone https://github.com/unslothai/unsloth.git +cd unsloth +git checkout nightly +Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass +.\install.ps1 --local +unsloth studio -H 0.0.0.0 -p 8888 +``` +Then to launch every time: +```bash +unsloth studio -H 0.0.0.0 -p 8888 +``` + +#### Uninstall +You can uninstall Unsloth Studio by deleting its install folder usually located under `$HOME/.unsloth/studio` on Mac/Linux/WSL and `%USERPROFILE%\.unsloth\studio` on Windows. Using the `rm -rf` commands will **delete everything**, including your history, cache: + +* ​ **MacOS, WSL, Linux:** `rm -rf ~/.unsloth/studio` +* ​ **Windows (PowerShell):** `Remove-Item -Recurse -Force "$HOME\.unsloth\studio"` + +For more info, [see our docs](https://unsloth.ai/docs/new/studio/install#uninstall). + +#### Deleting model files + +You can delete old model files either from the bin icon in model search or by removing the relevant cached model folder from the default Hugging Face cache directory. By default, HF uses: + +* ​ **MacOS, Linux, WSL:** `~/.cache/huggingface/hub/` +* ​ **Windows:** `%USERPROFILE%\.cache\huggingface\hub\` + ## 💚 Community and Links | Type | Links | | ----------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ | diff --git a/build.sh b/build.sh index 3118e8810a..cf8aa02910 100644 --- a/build.sh +++ b/build.sh @@ -29,7 +29,22 @@ _restore_gitignores() { } trap _restore_gitignores EXIT -npm install +# Use bun for install if available (faster), fall back to npm. +_install_ok=false +if command -v bun &>/dev/null; then + if bun install; then + _install_ok=true + else + echo "⚠ bun install failed, falling back to npm" + rm -rf node_modules + fi +fi +if [ "$_install_ok" != "true" ]; then + if ! npm install; then + echo "❌ ERROR: package install failed" >&2 + exit 1 + fi +fi npm run build # outputs to studio/frontend/dist/ _restore_gitignores diff --git a/install.ps1 b/install.ps1 index 1613ec6258..0c36046195 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1,17 +1,102 @@ # Unsloth Studio Installer for Windows PowerShell # Usage: irm https://raw.githubusercontent.com/unslothai/unsloth/main/install.ps1 | iex -# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 +# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local +# NoTorch: .\install.ps1 --no-torch (skip PyTorch, GGUF-only mode) +# Test: .\install.ps1 --package roland-sloth function Install-UnslothStudio { $ErrorActionPreference = "Stop" + $script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1") + + # ── Parse flags ── + $StudioLocalInstall = $false + $PackageName = "unsloth" + $RepoRoot = "" + $SkipTorch = $false + $argList = $args + for ($i = 0; $i -lt $argList.Count; $i++) { + switch ($argList[$i]) { + "--local" { $StudioLocalInstall = $true } + "--no-torch" { $SkipTorch = $true } + "--verbose" { $script:UnslothVerbose = $true } + "-v" { $script:UnslothVerbose = $true } + "--package" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red + return + } + $PackageName = $argList[$i] + } + } + } + # Propagate to child processes so they also respect verbose mode. + # Process-scoped -- does not persist. + if ($script:UnslothVerbose) { + $env:UNSLOTH_VERBOSE = '1' + } + + if ($StudioLocalInstall) { + $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path + if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) { + Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red + return + } + } - $VenvName = "unsloth_studio" $PythonVersion = "3.13" + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + $VenvDir = Join-Path $StudioHome "unsloth_studio" + + $Rule = [string]::new([char]0x2500, 52) + $Sloth = [char]::ConvertFromUtf32(0x1F9A5) + + function Enable-StudioVirtualTerminal { + if ($env:NO_COLOR) { return $false } + try { + if (-not ("StudioVT.Native" -as [type])) { + Add-Type -Namespace StudioVT -Name Native -MemberDefinition @' +[DllImport("kernel32.dll")] public static extern IntPtr GetStdHandle(int nStdHandle); +[DllImport("kernel32.dll")] public static extern bool GetConsoleMode(IntPtr h, out uint m); +[DllImport("kernel32.dll")] public static extern bool SetConsoleMode(IntPtr h, uint m); +'@ -ErrorAction Stop + } + $h = [StudioVT.Native]::GetStdHandle(-11) + [uint32]$mode = 0 + if (-not [StudioVT.Native]::GetConsoleMode($h, [ref]$mode)) { return $false } + $mode = $mode -bor 0x0004 + return [StudioVT.Native]::SetConsoleMode($h, $mode) + } catch { + return $false + } + } + $script:StudioVtOk = Enable-StudioVirtualTerminal + + function Get-StudioAnsi { + param( + [Parameter(Mandatory = $true)] + [ValidateSet('Title', 'Dim', 'Ok', 'Warn', 'Err', 'Reset')] + [string]$Kind + ) + $e = [char]27 + switch ($Kind) { + 'Title' { return "${e}[38;5;150m" } + 'Dim' { return "${e}[38;5;245m" } + 'Ok' { return "${e}[38;5;108m" } + 'Warn' { return "${e}[38;5;136m" } + 'Err' { return "${e}[91m" } + 'Reset' { return "${e}[0m" } + } + } Write-Host "" - Write-Host "=========================================" - Write-Host " Unsloth Studio Installer (Windows)" - Write-Host "=========================================" + if ($script:StudioVtOk -and -not $env:NO_COLOR) { + Write-Host (" " + (Get-StudioAnsi Title) + $Sloth + " Unsloth Studio Installer (Windows)" + (Get-StudioAnsi Reset)) + Write-Host (" {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset)) + } else { + Write-Host (" {0} Unsloth Studio Installer (Windows)" -f $Sloth) -ForegroundColor DarkGreen + Write-Host " $Rule" -ForegroundColor DarkGray + } Write-Host "" # ── Helper: refresh PATH from registry (deduplicating entries) ── @@ -31,13 +116,96 @@ function Install-UnslothStudio { $env:Path = $unique -join ";" } + function step { + param( + [Parameter(Mandatory = $true)][string]$Label, + [Parameter(Mandatory = $true)][string]$Value, + [string]$Color = "Green" + ) + if ($script:StudioVtOk -and -not $env:NO_COLOR) { + $dim = Get-StudioAnsi Dim + $rst = Get-StudioAnsi Reset + $val = switch ($Color) { + 'Green' { Get-StudioAnsi Ok } + 'Yellow' { Get-StudioAnsi Warn } + 'Red' { Get-StudioAnsi Err } + 'DarkGray' { Get-StudioAnsi Dim } + default { Get-StudioAnsi Ok } + } + $padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) } + Write-Host (" {0}{1}{2}{3}{4}{2}" -f $dim, $padded, $rst, $val, $Value) + } else { + $padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) } + Write-Host (" {0}" -f $padded) -NoNewline -ForegroundColor DarkGray + $fc = switch ($Color) { + 'Green' { 'DarkGreen' } + 'Yellow' { 'Yellow' } + 'Red' { 'Red' } + 'DarkGray' { 'DarkGray' } + default { 'DarkGreen' } + } + Write-Host $Value -ForegroundColor $fc + } + } + + function substep { + param( + [Parameter(Mandatory = $true)][string]$Message, + [string]$Color = "DarkGray" + ) + if ($script:StudioVtOk -and -not $env:NO_COLOR) { + $msgCol = switch ($Color) { + 'Yellow' { (Get-StudioAnsi Warn) } + 'Red' { (Get-StudioAnsi Err) } + default { (Get-StudioAnsi Dim) } + } + $pad = "".PadRight(15) + Write-Host (" {0}{1}{2}{3}" -f $msgCol, $pad, $Message, (Get-StudioAnsi Reset)) + } else { + $fc = switch ($Color) { + 'Yellow' { 'Yellow' } + 'Red' { 'Red' } + default { 'DarkGray' } + } + Write-Host (" {0,-15}{1}" -f "", $Message) -ForegroundColor $fc + } + } + + # Run native commands quietly by default to match install.sh behavior. + # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. + function Invoke-InstallCommand { + param( + [Parameter(Mandatory = $true)][ScriptBlock]$Command + ) + $prevEap = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { + # Reset to avoid stale values from prior native commands. + $global:LASTEXITCODE = 0 + if ($script:UnslothVerbose) { + # Merge stderr into stdout so progress/warning output stays visible + # without flipping $? on successful native commands (PS 5.1 treats + # stderr records as errors that set $? = $false even on exit code 0). + & $Command 2>&1 | Out-Host + } else { + $output = & $Command 2>&1 | Out-String + if ($LASTEXITCODE -ne 0) { + Write-Host $output -ForegroundColor Red + } + } + return [int]$LASTEXITCODE + } finally { + $ErrorActionPreference = $prevEap + } + } + function New-StudioShortcuts { param( [Parameter(Mandatory = $true)][string]$UnslothExePath ) if (-not (Test-Path $UnslothExePath)) { - Write-Host "[WARN] Cannot create shortcuts: unsloth.exe not found at $UnslothExePath" -ForegroundColor Yellow + substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow" return } try { @@ -50,7 +218,7 @@ function Install-UnslothStudio { $localAppDataDir = $env:LOCALAPPDATA if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) { - Write-Host "[WARN] LOCALAPPDATA path unavailable; skipped shortcut creation" -ForegroundColor Yellow + substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow" return } $appDir = Join-Path $localAppDataDir "Unsloth Studio" @@ -73,10 +241,10 @@ function Install-UnslothStudio { $null } if (-not $desktopLink) { - Write-Host "[WARN] Desktop path unavailable; skipped desktop shortcut creation" -ForegroundColor Yellow + substep "Desktop path unavailable; skipped desktop shortcut creation" "Yellow" } if (-not $startMenuLink) { - Write-Host "[WARN] APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" -ForegroundColor Yellow + substep "APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" "Yellow" } $iconPath = Join-Path $appDir "unsloth.ico" $bundledIcon = $null @@ -135,6 +303,44 @@ function Find-HealthyStudioPort { return `$null } +function Test-PortBusy { + param([Parameter(Mandatory = `$true)][int]`$Port) + `$listener = `$null + try { + `$listener = [System.Net.Sockets.TcpListener]::new([System.Net.IPAddress]::Any, `$Port) + `$listener.Start() + return `$false + } catch { + return `$true + } finally { + if (`$listener) { try { `$listener.Stop() } catch {} } + } +} + +function Find-FreeLaunchPort { + `$maxPort = `$basePort + `$maxPortOffset + try { + `$listening = Get-NetTCPConnection -State Listen -ErrorAction Stop | + Where-Object { `$_.LocalPort -ge `$basePort -and `$_.LocalPort -le `$maxPort } | + Select-Object -ExpandProperty LocalPort + for (`$offset = 0; `$offset -le `$maxPortOffset; `$offset++) { + `$candidate = `$basePort + `$offset + if (`$candidate -notin `$listening) { + return `$candidate + } + } + } catch { + # Get-NetTCPConnection unavailable or restricted; probe ports directly + for (`$offset = 0; `$offset -le `$maxPortOffset; `$offset++) { + `$candidate = `$basePort + `$offset + if (-not (Test-PortBusy -Port `$candidate)) { + return `$candidate + } + } + } + return `$null +} + # If Studio is already healthy on any expected port, just open it and exit. `$existingPort = Find-HealthyStudioPort if (`$existingPort) { @@ -163,7 +369,16 @@ try { `$powershellExe = Join-Path `$env:SystemRoot 'System32\WindowsPowerShell\v1.0\powershell.exe' `$studioExe = '$SingleQuotedExePath' - `$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$basePort + `$launchPort = Find-FreeLaunchPort + if (-not `$launchPort) { + `$msg = "No free port found in range `$basePort-`$(`$basePort + `$maxPortOffset)" + try { + Add-Type -AssemblyName System.Windows.Forms -ErrorAction Stop + [System.Windows.Forms.MessageBox]::Show(`$msg, 'Unsloth Studio') | Out-Null + } catch {} + exit 1 + } + `$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$launchPort `$launchArgs = @( '-NoExit', '-NoProfile', @@ -284,27 +499,27 @@ shell.Run cmd, 0, False $shortcut.Save() $createdShortcutCount++ } catch { - Write-Host "[WARN] Could not create shortcut at ${linkPath}: $($_.Exception.Message)" -ForegroundColor Yellow + substep "could not create shortcut at ${linkPath}: $($_.Exception.Message)" "Yellow" } } if ($createdShortcutCount -gt 0) { - Write-Host "[OK] Created Unsloth Studio shortcut(s): $createdShortcutCount" -ForegroundColor Green + substep "Created Unsloth Studio shortcut" } else { - Write-Host "[WARN] No Unsloth Studio shortcuts were created" -ForegroundColor Yellow + substep "no Unsloth Studio shortcuts were created" "Yellow" } } catch { - Write-Host "[WARN] Shortcut creation unavailable: $($_.Exception.Message)" -ForegroundColor Yellow + substep "shortcut creation unavailable: $($_.Exception.Message)" "Yellow" } } catch { - Write-Host "[WARN] Shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" -ForegroundColor Yellow + substep "shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" "Yellow" } } # ── Check winget ── if (-not (Get-Command winget -ErrorAction SilentlyContinue)) { - Write-Host "Error: winget is not available." -ForegroundColor Red - Write-Host " Install it from https://aka.ms/getwinget" -ForegroundColor Yellow - Write-Host " or install Python $PythonVersion and uv manually, then re-run." -ForegroundColor Yellow + step "winget" "not available" "Red" + substep "Install it from https://aka.ms/getwinget" "Yellow" + substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow" return } @@ -382,10 +597,10 @@ shell.Run cmd, 0, False # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. $DetectedPython = Find-CompatiblePython if ($DetectedPython) { - Write-Host "==> Python already installed: Python $($DetectedPython.Version)" + step "python" "Python $($DetectedPython.Version) already installed" } if (-not $DetectedPython) { - Write-Host "==> Installing Python ${PythonVersion}..." + substep "installing Python ${PythonVersion}..." $pythonPackageId = "Python.Python.$PythonVersion" # Temporarily lower ErrorActionPreference so that winget stderr # (progress bars, warnings) does not become a terminating error @@ -407,7 +622,7 @@ shell.Run cmd, 0, False # 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). - Write-Host " Python not found on PATH after winget. Retrying with --force..." + substep "Python not found on PATH after winget. Retrying with --force..." "Yellow" $ErrorActionPreference = "Continue" try { winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements --force @@ -429,7 +644,7 @@ shell.Run cmd, 0, False # ── Install uv if not present ── if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - Write-Host "==> Installing uv package manager..." + substep "installing uv package manager..." $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} @@ -437,32 +652,73 @@ shell.Run cmd, 0, False Refresh-SessionPath # Fallback: if winget didn't put uv on PATH, try the PowerShell installer if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - Write-Host " Trying alternative uv installer..." + substep "trying alternative uv installer..." "Yellow" powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" Refresh-SessionPath } } if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - Write-Host "Error: uv could not be installed." -ForegroundColor Red - Write-Host " Install it from https://docs.astral.sh/uv/" -ForegroundColor Yellow + step "uv" "could not be installed" "Red" + substep "Install it from https://docs.astral.sh/uv/" "Yellow" return } - # ── Create venv (skip if it already exists and has a valid interpreter) ── + # ── Create venv (migrate old layout if possible, otherwise fresh) ── # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. - $VenvPython = Join-Path $VenvName "Scripts\python.exe" + if (-not (Test-Path $StudioHome)) { + New-Item -ItemType Directory -Path $StudioHome -Force | Out-Null + } + + $VenvPython = Join-Path $VenvDir "Scripts\python.exe" + $_Migrated = $false + + if (Test-Path $VenvPython) { + # New layout already exists -- nuke for fresh install + substep "removing existing environment for fresh install..." + Remove-Item -Recurse -Force $VenvDir + } elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) { + # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating + $OldVenv = Join-Path $StudioHome ".venv" + $OldPy = Join-Path $OldVenv "Scripts\python.exe" + substep "found legacy Studio environment, validating..." + $prevEAP2 = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { + & $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null + $torchOk = ($LASTEXITCODE -eq 0) + } catch { $torchOk = $false } + $ErrorActionPreference = $prevEAP2 + if ($torchOk) { + substep "legacy environment is healthy -- migrating..." + Move-Item -Path $OldVenv -Destination $VenvDir -Force + substep "moved .venv -> unsloth_studio" + $_Migrated = $true + } else { + substep "legacy environment failed validation -- creating fresh environment" "Yellow" + Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue + } + } elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) { + # CWD-relative venv from old install.ps1 -- migrate to absolute path + $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio" + substep "found CWD-relative Studio environment, migrating to $VenvDir..." + Move-Item -Path $CwdVenv -Destination $VenvDir -Force + substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio" + $_Migrated = $true + } + if (-not (Test-Path $VenvPython)) { - if (Test-Path $VenvName) { Remove-Item -Recurse -Force $VenvName } - Write-Host "==> Creating Python $($DetectedPython.Version) virtual environment (${VenvName})..." - uv venv $VenvName --python "$($DetectedPython.Path)" - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] Failed to create virtual environment (exit code $LASTEXITCODE)" -ForegroundColor Red + step "venv" "creating Python $($DetectedPython.Version) virtual environment" + substep "$VenvDir" + $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } + if ($venvExit -ne 0) { + Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red return } } else { - Write-Host "==> Virtual environment ${VenvName} already exists, skipping creation." + step "venv" "using migrated environment" + substep "$VenvDir" } # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── @@ -471,7 +727,7 @@ shell.Run cmd, 0, False try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue if ($nvSmiCmd) { - & $nvSmiCmd.Source 2>&1 | Out-Null + & $nvSmiCmd.Source *> $null if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } } } catch {} @@ -482,18 +738,18 @@ shell.Run cmd, 0, False )) { if (Test-Path $p) { try { - & $p 2>&1 | Out-Null + & $p *> $null if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } } catch {} } } } if ($HasNvidiaSmi) { - Write-Host "[OK] NVIDIA GPU detected" -ForegroundColor Green + step "gpu" "NVIDIA GPU detected" } else { - Write-Host "[WARN] No NVIDIA GPU detected. Studio will run in chat-only (GGUF) mode." -ForegroundColor Yellow - Write-Host " Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow - Write-Host " https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow + step "gpu" "none (chat-only / GGUF)" "Yellow" + substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" + substep "https://www.nvidia.com/Download/index.aspx" "Yellow" } # ── Choose the correct PyTorch index URL based on driver CUDA version ── @@ -513,11 +769,21 @@ shell.Run cmd, 0, False return "$baseUrl/cpu" } } catch {} - Write-Host "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" -ForegroundColor Yellow + substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow" return "$baseUrl/cu126" } $TorchIndexUrl = Get-TorchIndexUrl + # ── Print CPU-only hint when no GPU detected ── + if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") { + Write-Host "" + Write-Host " NOTE: No NVIDIA GPU detected." -ForegroundColor Yellow + Write-Host " Installing CPU-only PyTorch. If you only need GGUF chat/inference," + Write-Host " re-run with --no-torch for a faster, lighter install:" + Write-Host " .\install.ps1 --no-torch" + Write-Host "" + } + # ── Install PyTorch first, then unsloth separately ── # # Why two steps? @@ -536,25 +802,116 @@ shell.Run cmd, 0, False # CUDA wheels. Missing dependencies (transformers, trl, peft, etc.) # are still pulled in because they are new, not upgrades. # - Write-Host "==> Installing PyTorch ($TorchIndexUrl)..." - uv pip install --python $VenvPython torch torchvision torchaudio --index-url $TorchIndexUrl - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] Failed to install PyTorch (exit code $LASTEXITCODE)" -ForegroundColor Red - return + # ── Helper: find no-torch-runtime.txt ── + function Find-NoTorchRuntimeFile { + if ($StudioLocalInstall -and (Test-Path (Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"))) { + return Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt" + } + $installed = Get-ChildItem -Path $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue | + Where-Object { $_.FullName -like "*studio*backend*requirements*no-torch-runtime.txt" } | + Select-Object -ExpandProperty FullName -First 1 + return $installed } - Write-Host "==> Installing unsloth (this may take a few minutes)..." - uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11" - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red - return + if ($_Migrated) { + # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state + # in the new venv location, while preserving existing torch/CUDA + substep "upgrading unsloth in migrated environment..." + if ($SkipTorch) { + # No-torch: install unsloth + unsloth-zoo with --no-deps, then + # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + if ($baseInstallExit -eq 0) { + $NoTorchReq = Find-NoTorchRuntimeFile + if ($NoTorchReq) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + } + } + } else { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } + if ($StudioLocalInstall) { + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } + } + } elseif ($TorchIndexUrl) { + if ($SkipTorch) { + substep "skipping PyTorch (--no-torch flag set)." "Yellow" + } else { + substep "installing PyTorch ($TorchIndexUrl)..." + $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + if ($torchInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red + return + } + } + + substep "installing unsloth (this may take a few minutes)..." + if ($SkipTorch) { + # No-torch: install unsloth + unsloth-zoo with --no-deps, then + # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + if ($baseInstallExit -eq 0) { + $NoTorchReq = Find-NoTorchRuntimeFile + if ($NoTorchReq) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + } + } + } elseif ($StudioLocalInstall) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo } + } else { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" } + } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } + + if ($StudioLocalInstall) { + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } + } + } else { + # Fallback: GPU detection failed to produce a URL -- let uv resolve torch + substep "installing unsloth (this may take a few minutes)..." + if ($StudioLocalInstall) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } + } else { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "$PackageName" --torch-backend=auto } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } + } } # ── Run studio setup ── # setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools, # CUDA Toolkit, Node.js, and other dependencies automatically via winget. - Write-Host "==> Running unsloth studio setup..." - $UnslothExe = Join-Path $VenvName "Scripts\unsloth.exe" + step "setup" "running unsloth studio setup..." + $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" if (-not (Test-Path $UnslothExe)) { Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow @@ -562,35 +919,57 @@ shell.Run cmd, 0, False Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow return } - & $UnslothExe studio setup - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red + # Tell setup.ps1 to skip base package installation (install.ps1 already did it) + $env:SKIP_STUDIO_BASE = "1" + $env:STUDIO_PACKAGE_NAME = $PackageName + $env:UNSLOTH_NO_TORCH = if ($SkipTorch) { "true" } else { "false" } + # Always set STUDIO_LOCAL_INSTALL explicitly to avoid stale values from + # a previous --local run in the same PowerShell session. + if ($StudioLocalInstall) { + $env:STUDIO_LOCAL_INSTALL = "1" + $env:STUDIO_LOCAL_REPO = $RepoRoot + } else { + $env:STUDIO_LOCAL_INSTALL = "0" + Remove-Item Env:STUDIO_LOCAL_REPO -ErrorAction SilentlyContinue + } + # Use 'studio setup' (not 'studio update') because 'update' pops + # SKIP_STUDIO_BASE, which would cause redundant package reinstallation + # and bypass the fast-path version check from PR #4667. + $studioArgs = @('studio', 'setup') + if ($script:UnslothVerbose) { $studioArgs += '--verbose' } + & $UnslothExe @studioArgs + $setupExit = $LASTEXITCODE + if ($setupExit -ne 0) { + Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red return } New-StudioShortcuts -UnslothExePath $UnslothExe - Write-Host "" - Write-Host "=========================================" - Write-Host " Unsloth Studio installed!" - Write-Host "=========================================" - Write-Host "" + # ── Add venv Scripts dir to User PATH so `unsloth studio` works from any terminal ── + $ScriptsDir = Join-Path $VenvDir "Scripts" + $UserPath = [System.Environment]::GetEnvironmentVariable("Path", "User") + if (-not $UserPath -or $UserPath -notlike "*$ScriptsDir*") { + if ($UserPath) { + [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir;$UserPath", "User") + } else { + [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User") + } + Refresh-SessionPath + step "path" "added unsloth to PATH" + } # Launch studio automatically in interactive terminals; # in non-interactive environments (CI, Docker) just print instructions. $IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) if ($IsInteractive) { - Write-Host "==> Launching Unsloth Studio..." - Write-Host "" - $UnslothExe = Join-Path $VenvName "Scripts\unsloth.exe" & $UnslothExe studio -H 0.0.0.0 -p 8888 } else { - Write-Host " To launch, run:" - Write-Host "" - Write-Host " .\${VenvName}\Scripts\activate" - Write-Host " unsloth studio -H 0.0.0.0 -p 8888" + step "launch" "manual commands:" + substep "& `"$VenvDir\Scripts\Activate.ps1`"" + substep "unsloth studio -H 0.0.0.0 -p 8888" Write-Host "" } } -Install-UnslothStudio +Install-UnslothStudio @args diff --git a/install.sh b/install.sh index 0893955939..9ea80bc161 100755 --- a/install.sh +++ b/install.sh @@ -1,11 +1,111 @@ #!/bin/sh # Unsloth Studio Installer -# Usage (curl): curl -fsSL https://raw.githubusercontent.com/unslothai/unsloth/main/install.sh | sh -# Usage (wget): wget -qO- https://raw.githubusercontent.com/unslothai/unsloth/main/install.sh | sh +# Usage (curl): curl -fsSL https://unsloth.ai/install.sh | sh +# Usage (wget): wget -qO- https://unsloth.ai/install.sh | sh +# Usage (local): ./install.sh --local (install from local repo instead of PyPI) +# Usage (no-torch): ./install.sh --no-torch (skip PyTorch, GGUF-only mode) +# Usage (test): ./install.sh --package roland-sloth (install a different package name) +# Usage (py): ./install.sh --python 3.12 (override auto-detected Python version) set -e -VENV_NAME="unsloth_studio" -PYTHON_VERSION="3.13" +# ── Output style (aligned with studio/setup.sh) ── +RULE="" +_rule_i=0 +while [ "$_rule_i" -lt 52 ]; do + RULE="${RULE}─" + _rule_i=$((_rule_i + 1)) +done +if [ -n "${NO_COLOR:-}" ]; then + C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= +elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then + _ESC="$(printf '\033')" + C_TITLE="${_ESC}[38;5;150m" + C_DIM="${_ESC}[38;5;245m" + C_OK="${_ESC}[38;5;108m" + C_WARN="${_ESC}[38;5;136m" + C_ERR="${_ESC}[91m" + C_RST="${_ESC}[0m" +else + C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= +fi + +step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; } +substep() { printf " ${C_DIM}%-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; } + +# ── Parse flags ── +STUDIO_LOCAL_INSTALL=false +PACKAGE_NAME="unsloth" +_USER_PYTHON="" +_NO_TORCH_FLAG=false +_VERBOSE=false +_next_is_package=false +_next_is_python=false +for arg in "$@"; do + if [ "$_next_is_package" = true ]; then + PACKAGE_NAME="$arg" + _next_is_package=false + continue + fi + if [ "$_next_is_python" = true ]; then + _USER_PYTHON="$arg" + _next_is_python=false + continue + fi + case "$arg" in + --local) STUDIO_LOCAL_INSTALL=true ;; + --package) _next_is_package=true ;; + --python) _next_is_python=true ;; + --no-torch) _NO_TORCH_FLAG=true ;; + --verbose|-v) _VERBOSE=true ;; + esac +done + +if [ "$_VERBOSE" = true ]; then + export UNSLOTH_VERBOSE=1 +fi + +_is_verbose() { + [ "${UNSLOTH_VERBOSE:-0}" = "1" ] +} + +run_maybe_quiet() { + if _is_verbose; then + "$@" + else + "$@" > /dev/null 2>&1 + fi +} + +run_install_cmd() { + _label="$1" + shift + if _is_verbose; then + "$@" && return 0 + _rc=$? + step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 + return "$_rc" + fi + _log=$(mktemp) + "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } + _rc=$? + step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 + cat "$_log" >&2 + rm -f "$_log" + return $_rc +} + +if [ "$_next_is_package" = true ]; then + echo "❌ ERROR: --package requires an argument." >&2 + exit 1 +fi +if [ "$_next_is_python" = true ]; then + echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2 + exit 1 +fi + +PYTHON_VERSION="" # resolved after platform detection +STUDIO_HOME="$HOME/.unsloth/studio" +VENV_DIR="$STUDIO_HOME/unsloth_studio" # ── Helper: download a URL to a file (supports curl and wget) ── download() { @@ -92,17 +192,12 @@ _smart_apt_install() { # ── Helper: create desktop shortcuts and launcher script ── # Usage: create_studio_shortcuts # Creates ~/.local/share/unsloth/launch-studio.sh (shared launcher), -# plus platform-specific shortcuts (Linux .desktop / macOS .app bundle). -# Skipped on WSL (no native desktop). +# plus platform-specific shortcuts (Linux .desktop / macOS .app bundle / +# WSL Windows Desktop+Start Menu .lnk). create_studio_shortcuts() { _css_exe="$1" _css_os="$2" - # Skip on WSL -- no native desktop environment - if [ "$_css_os" = "wsl" ]; then - return 0 - fi - # Validate exe if [ ! -x "$_css_exe" ]; then echo "[WARN] Cannot create shortcuts: unsloth not found at $_css_exe" @@ -231,6 +326,17 @@ _open_browser() { _url="$1" if [ "$(uname)" = "Darwin" ] && command -v open >/dev/null 2>&1; then open "$_url" + elif grep -qi microsoft /proc/version 2>/dev/null; then + # WSL: xdg-open is unreliable; use Windows browser via PowerShell or cmd + if command -v powershell.exe >/dev/null 2>&1; then + powershell.exe -NoProfile -Command "Start-Process '$_url'" >/dev/null 2>&1 & + elif command -v cmd.exe >/dev/null 2>&1; then + cmd.exe /c start "" "$_url" >/dev/null 2>&1 & + elif command -v xdg-open >/dev/null 2>&1; then + xdg-open "$_url" >/dev/null 2>&1 & + else + echo "Open in your browser: $_url" >&2 + fi elif command -v xdg-open >/dev/null 2>&1; then xdg-open "$_url" >/dev/null 2>&1 & else @@ -294,6 +400,8 @@ _acquire_lock() { } _release_lock() { + [ -d "$LOCK_DIR" ] || return 0 + [ "$(cat "$LOCK_DIR/pid" 2>/dev/null)" = "$$" ] || return 0 rm -rf "$LOCK_DIR" } @@ -319,24 +427,48 @@ _launch_port=$(_find_launch_port) || { exit 1 } -# Launch studio in a terminal -_launch_cmd=$(printf '%q ' "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port") -_launch_cmd=${_launch_cmd% } -_spawn_terminal "$_launch_cmd" +if [ -t 1 ]; then + # ── Foreground mode (TTY available) ── + # Background subshell: wait for studio to become healthy, release the + # single-instance lock, then open the browser. The lock stays held until + # health is confirmed so a second launcher cannot race during startup. + ( + _obwr_deadline=$(($(date +%s) + TIMEOUT_SEC)) + while [ "$(date +%s)" -lt "$_obwr_deadline" ]; do + if _check_health "$_launch_port"; then + _release_lock + _open_browser "http://localhost:$_launch_port" + exit 0 + fi + sleep "$POLL_INTERVAL_SEC" + done + # Timed out -- release the lock anyway so future launches are not blocked + _release_lock + ) & + # Clear traps so exec does not trigger _release_lock (the subshell owns it) + trap - EXIT INT TERM + exec "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port" +else + # ── Background mode (no TTY) ── + # Used by macOS .app and headless invocations. + _launch_cmd=$(printf '%q ' "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port") + _launch_cmd=${_launch_cmd% } + _spawn_terminal "$_launch_cmd" -# Poll for health -_deadline=$(($(date +%s) + TIMEOUT_SEC)) -while [ "$(date +%s)" -lt "$_deadline" ]; do - _port=$(_find_healthy_port) && { - _open_browser "http://localhost:$_port" - exit 0 - } - sleep "$POLL_INTERVAL_SEC" -done + # Poll for health on the specific port we launched on + _deadline=$(($(date +%s) + TIMEOUT_SEC)) + while [ "$(date +%s)" -lt "$_deadline" ]; do + if _check_health "$_launch_port"; then + _open_browser "http://localhost:$_launch_port" + exit 0 + fi + sleep "$POLL_INTERVAL_SEC" + done -echo "Unsloth Studio did not become healthy within ${TIMEOUT_SEC}s." >&2 -echo "Check logs at: $LOG_FILE" >&2 -exit 1 + echo "Unsloth Studio did not become healthy within ${TIMEOUT_SEC}s." >&2 + echo "Check logs at: $LOG_FILE" >&2 + exit 1 +fi LAUNCHER_EOF chmod +x "$_css_launcher" @@ -348,44 +480,33 @@ LAUNCHER_EOF printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" > "$_css_data_dir/studio.conf" # ── Icon: try bundled, then download ── - # favicon.png (small, for Linux) and unsloth-gem.png (large, for macOS icns) + # rounded-512.png used for both Linux and macOS icons _css_script_dir="" if [ -n "${0:-}" ] && [ -f "$0" ]; then _css_script_dir=$(cd "$(dirname "$0")" 2>/dev/null && pwd) || true fi - # Try to find favicon.png from installed package (site-packages) or local repo - _css_found_favicon="" - _css_found_gem="" + # Try to find rounded-512.png from installed package (site-packages) or local repo + _css_found_icon="" _css_venv_dir=$(dirname "$(dirname "$_css_exe")") # Check site-packages for _sp in "$_css_venv_dir"/lib/python*/site-packages/unsloth/studio/frontend/public; do - if [ -f "$_sp/favicon.png" ]; then - _css_found_favicon="$_sp/favicon.png" - fi - if [ -f "$_sp/unsloth-gem.png" ]; then - _css_found_gem="$_sp/unsloth-gem.png" + if [ -f "$_sp/rounded-512.png" ]; then + _css_found_icon="$_sp/rounded-512.png" fi done # Check local repo (when running from clone) - if [ -z "$_css_found_favicon" ] && [ -n "$_css_script_dir" ] && [ -f "$_css_script_dir/studio/frontend/public/favicon.png" ]; then - _css_found_favicon="$_css_script_dir/studio/frontend/public/favicon.png" - fi - if [ -z "$_css_found_gem" ] && [ -n "$_css_script_dir" ] && [ -f "$_css_script_dir/studio/frontend/public/unsloth-gem.png" ]; then - _css_found_gem="$_css_script_dir/studio/frontend/public/unsloth-gem.png" + if [ -z "$_css_found_icon" ] && [ -n "$_css_script_dir" ] && [ -f "$_css_script_dir/studio/frontend/public/rounded-512.png" ]; then + _css_found_icon="$_css_script_dir/studio/frontend/public/rounded-512.png" fi - # Copy or download favicon.png - if [ -n "$_css_found_favicon" ]; then - cp "$_css_found_favicon" "$_css_icon_png" 2>/dev/null || true - elif [ ! -f "$_css_icon_png" ]; then - download "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/favicon.png" "$_css_icon_png" 2>/dev/null || true - fi - # Copy or download unsloth-gem.png (for macOS icns) - if [ -n "$_css_found_gem" ]; then - cp "$_css_found_gem" "$_css_gem_png" 2>/dev/null || true - elif [ ! -f "$_css_gem_png" ]; then - download "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth-gem.png" "$_css_gem_png" 2>/dev/null || true + # Copy or download rounded-512.png (used for both Linux icon and macOS icns) + if [ -n "$_css_found_icon" ]; then + cp "$_css_found_icon" "$_css_icon_png" 2>/dev/null || true + cp "$_css_found_icon" "$_css_gem_png" 2>/dev/null || true + else + download "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/rounded-512.png" "$_css_icon_png" 2>/dev/null || true + cp "$_css_icon_png" "$_css_gem_png" 2>/dev/null || true fi # Validate PNG header (first 4 bytes: \x89PNG) @@ -421,7 +542,7 @@ Name=Unsloth Studio Comment=Launch Unsloth Studio Exec="$_css_exec_escaped" Icon=$_css_icon_escaped -Terminal=false +Terminal=true StartupNotify=true Categories=Development;Science; DESKTOP_EOF @@ -517,17 +638,78 @@ STUB_EOF ln -sf "$_css_app" "$HOME/Desktop/Unsloth Studio" 2>/dev/null || true fi _css_created=1 + + elif [ "$_css_os" = "wsl" ]; then + # ── WSL: create Windows Desktop and Start Menu shortcuts ── + # Detect current WSL distro for targeted shortcut + _css_distro="${WSL_DISTRO_NAME:-}" + + # Build the wsl.exe arguments. + # Double-quote distro name and launcher path for Windows command line + # parsing so values with spaces (e.g. "Ubuntu Preview") are kept as + # single arguments. + _css_wsl_args="" + if [ -n "$_css_distro" ]; then + _css_wsl_args="-d \"$_css_distro\" " + fi + _css_wsl_args="${_css_wsl_args}-- bash -l -c \"exec \\\"$_css_launcher\\\"\"" + + # Detect whether Windows Terminal (wt.exe) is available (better UX) + _css_use_wt=false + if command -v wt.exe >/dev/null 2>&1; then + _css_use_wt=true + fi + + if [ "$_css_use_wt" = true ]; then + _css_sc_target='wt.exe' + _css_sc_args="wsl.exe $_css_wsl_args" + else + _css_sc_target='wsl.exe' + _css_sc_args="$_css_wsl_args" + fi + + # Escape single quotes for PowerShell single-quoted string embedding + _css_sc_args_ps=$(printf '%s' "$_css_sc_args" | 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 + cat > "$_css_ps1_tmp" << WSLPS1_EOF +\$WshShell = New-Object -ComObject WScript.Shell +\$targetExe = (Get-Command '$_css_sc_target' -ErrorAction SilentlyContinue).Source +if (-not \$targetExe) { exit 1 } +\$locations = @( + [Environment]::GetFolderPath('Desktop'), + (Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs') +) +foreach (\$dir in \$locations) { + if (-not \$dir -or -not (Test-Path \$dir)) { continue } + \$linkPath = Join-Path \$dir 'Unsloth Studio.lnk' + \$shortcut = \$WshShell.CreateShortcut(\$linkPath) + \$shortcut.TargetPath = \$targetExe + \$shortcut.Arguments = '$_css_sc_args_ps' + \$shortcut.Description = 'Launch Unsloth Studio' + \$shortcut.Save() +} +WSLPS1_EOF + + # Convert WSL path to Windows path for powershell.exe + _css_ps1_win=$(wslpath -w "$_css_ps1_tmp" 2>/dev/null) + if [ -n "$_css_ps1_win" ]; then + powershell.exe -NoProfile -ExecutionPolicy Bypass -File "$_css_ps1_win" >/dev/null 2>&1 && _css_created=1 + fi + rm -f "$_css_ps1_tmp" + fi fi if [ "$_css_created" -eq 1 ]; then - echo "[OK] Created Unsloth Studio shortcut(s)" + substep "Created Unsloth Studio shortcut" fi } echo "" -echo "=========================================" -echo " Unsloth Studio Installer" -echo "=========================================" +printf " ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Installer" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" # ── Detect platform ── @@ -537,7 +719,48 @@ if [ "$(uname)" = "Darwin" ]; then elif grep -qi microsoft /proc/version 2>/dev/null; then OS="wsl" fi -echo "==> Platform: $OS" +step "platform" "$OS" + +# ── Architecture detection & Python version ── +_ARCH=$(uname -m) +MAC_INTEL=false +if [ "$OS" = "macos" ] && [ "$_ARCH" = "x86_64" ]; then + # Guard against Apple Silicon running under Rosetta (reports x86_64). + # sysctl hw.optional.arm64 returns "1" on Apple Silicon even in Rosetta. + if [ "$(sysctl -in hw.optional.arm64 2>/dev/null || echo 0)" = "1" ]; then + echo "" + echo " WARNING: Apple Silicon detected, but this shell is running under Rosetta (x86_64)." + echo " Re-run install.sh from a native arm64 terminal for full PyTorch support." + echo " Continuing in GGUF-only mode for now." + echo "" + fi + MAC_INTEL=true +fi + +if [ -n "$_USER_PYTHON" ]; then + PYTHON_VERSION="$_USER_PYTHON" + echo " Using user-specified Python $PYTHON_VERSION (--python override)" +elif [ "$MAC_INTEL" = true ]; then + PYTHON_VERSION="3.12" +else + PYTHON_VERSION="3.13" +fi + +if [ "$MAC_INTEL" = true ]; then + echo "" + echo " NOTE: Intel Mac (x86_64) detected." + echo " PyTorch is unavailable for this platform (dropped Jan 2024)." + echo " Studio will install in GGUF-only mode." + echo " Chat, inference via GGUF, and data recipes will work." + echo " Training requires Apple Silicon or Linux with GPU." + echo "" +fi + +# ── Unified SKIP_TORCH: --no-torch flag OR Intel Mac auto-detection ── +SKIP_TORCH=false +if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then + SKIP_TORCH=true +fi # ── Check system dependencies ── # cmake and git are needed by unsloth studio setup to build the GGUF inference @@ -576,8 +799,8 @@ MISSING=$(echo "$MISSING" | sed 's/^ *//') if [ -n "$MISSING" ]; then echo "" - echo "==> Unsloth Studio needs these packages: $MISSING" - echo " These are needed to build the GGUF inference engine." + step "deps" "missing: $MISSING" "$C_WARN" + substep "These are needed to build the GGUF inference engine." case "$OS" in macos) @@ -602,7 +825,7 @@ if [ -n "$MISSING" ]; then esac echo "" else - echo "==> All system dependencies found." + step "deps" "all system dependencies found" fi # ── Install uv ── @@ -648,10 +871,10 @@ _uv_version_ok() { } if ! command -v uv >/dev/null 2>&1 || ! _uv_version_ok uv; then - echo "==> Installing uv package manager..." + substep "installing uv package manager..." _uv_tmp=$(mktemp) download "https://astral.sh/uv/install.sh" "$_uv_tmp" - sh "$_uv_tmp" /dev/null 2>&1 || ! _uv_version_ok uv; then export PATH="$HOME/.local/bin:$PATH" fi -# ── Create venv (skip if it already exists and has a valid interpreter) ── -if [ ! -x "$VENV_NAME/bin/python" ]; then - [ -e "$VENV_NAME" ] && rm -rf "$VENV_NAME" - echo "==> Creating Python ${PYTHON_VERSION} virtual environment (${VENV_NAME})..." - uv venv "$VENV_NAME" --python "$PYTHON_VERSION" -else - echo "==> Virtual environment ${VENV_NAME} already exists, skipping creation." +# ── Create venv (migrate old layout if possible, otherwise fresh) ── +mkdir -p "$STUDIO_HOME" + +_MIGRATED=false + +if [ -x "$VENV_DIR/bin/python" ]; then + # New layout already exists — nuke for fresh install + rm -rf "$VENV_DIR" +elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then + # Old layout exists — validate before migrating + substep "found legacy Studio environment, validating..." + if "$STUDIO_HOME/.venv/bin/python" -c " +import torch +device = 'cuda' if torch.cuda.is_available() else 'cpu' +A = torch.ones((10, 10), device=device) +B = torch.ones((10, 10), device=device) +C = torch.ones((10, 10), device=device) +D = A + B +E = D @ C +torch.testing.assert_close(torch.unique(E), torch.tensor((20,), device=E.device, dtype=E.dtype)) +" >/dev/null 2>&1; then + echo "✅ Legacy environment is healthy — migrating..." + mv "$STUDIO_HOME/.venv" "$VENV_DIR" + echo " Moved ~/.unsloth/studio/.venv → $VENV_DIR" + _MIGRATED=true + else + echo "⚠️ Legacy environment failed validation — creating fresh environment" + rm -rf "$STUDIO_HOME/.venv" + fi fi +# If an Intel Mac has a stale 3.13 venv from a previous failed install, recreate +# (skip when the user explicitly chose a version via --python) +if [ "$SKIP_TORCH" = true ] && [ "$MAC_INTEL" = true ] && [ -z "$_USER_PYTHON" ] && [ -x "$VENV_DIR/bin/python" ]; then + _PY_MM=$("$VENV_DIR/bin/python" -c \ + "import sys; print('{}.{}'.format(*sys.version_info[:2]))" 2>/dev/null || echo "") + if [ "$_PY_MM" != "3.12" ]; then + echo " Recreating Intel Mac environment with Python 3.12 (was $_PY_MM)..." + rm -rf "$VENV_DIR" + fi +fi + +if [ ! -x "$VENV_DIR/bin/python" ]; then + step "venv" "creating Python ${PYTHON_VERSION} virtual environment" + substep "$VENV_DIR" + run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" +fi + +# Guard against Python 3.13.8 torch import bug on Apple Silicon +# (skip when the user explicitly chose a version via --python) +if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _PY_VER=$("$VENV_DIR/bin/python" -c \ + "import sys; print('{}.{}.{}'.format(*sys.version_info[:3]))" 2>/dev/null || echo "") + if [ "$_PY_VER" = "3.13.8" ]; then + echo " WARNING: Python 3.13.8 has a known torch import bug." + echo " Recreating venv with Python 3.12..." + rm -rf "$VENV_DIR" + PYTHON_VERSION="3.12" + run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + fi +fi + +if [ -x "$VENV_DIR/bin/python" ]; then + step "venv" "using environment" + substep "${VENV_DIR}" +fi + +# ── Resolve repo root (for --local installs) ── +_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" + +# ── Helper: find no-torch-runtime.txt (local repo or site-packages) ── +_find_no_torch_runtime() { + # Check local repo first (for --local installs) + if [ -f "$_REPO_ROOT/studio/backend/requirements/no-torch-runtime.txt" ]; then + echo "$_REPO_ROOT/studio/backend/requirements/no-torch-runtime.txt" + return + fi + # Check inside installed package + _rt=$(find "$VENV_DIR" -path "*/studio/backend/requirements/no-torch-runtime.txt" -print -quit 2>/dev/null || echo "") + if [ -n "$_rt" ]; then + echo "$_rt" + return + fi +} + +# ── Detect GPU and choose PyTorch index URL ── +# Mirrors Get-TorchIndexUrl in install.ps1. +# On CPU-only machines this returns the cpu index, avoiding the solver +# dead-end where --torch-backend=auto resolves to unsloth==2024.8. +get_torch_index_url() { + _base="https://download.pytorch.org/whl" + # macOS: always CPU (no CUDA support) + case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac + # Try nvidia-smi + _smi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _smi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _smi="/usr/bin/nvidia-smi" + fi + if [ -z "$_smi" ]; then echo "$_base/cpu"; return; fi + # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P) + _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ + | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ + | head -1) + if [ -z "$_cuda_ver" ]; then + echo "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" >&2 + echo "$_base/cu126"; return + fi + _major=${_cuda_ver%%.*} + _minor=${_cuda_ver#*.} + if [ "$_major" -ge 13 ]; then echo "$_base/cu130" + elif [ "$_major" -eq 12 ] && [ "$_minor" -ge 8 ]; then echo "$_base/cu128" + elif [ "$_major" -eq 12 ] && [ "$_minor" -ge 6 ]; then echo "$_base/cu126" + elif [ "$_major" -ge 12 ]; then echo "$_base/cu124" + elif [ "$_major" -ge 11 ]; then echo "$_base/cu118" + else echo "$_base/cpu"; fi +} +TORCH_INDEX_URL=$(get_torch_index_url) + +# ── Print CPU-only hint when no GPU detected ── +case "$TORCH_INDEX_URL" in + */cpu) + if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then + echo "" + echo " NOTE: No NVIDIA GPU detected (nvidia-smi not found)." + echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference," + echo " re-run with --no-torch for a faster, lighter install:" + echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" + echo "" + fi + ;; +esac + # ── Install unsloth directly into the venv (no activation needed) ── -echo "==> Installing unsloth (this may take a few minutes)..." -uv pip install --python "$VENV_NAME/bin/python" "unsloth>=2026.3.11" --torch-backend=auto +_VENV_PY="$VENV_DIR/bin/python" +if [ "$_MIGRATED" = true ]; then + # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state + # in the new venv location, while preserving existing torch/CUDA + substep "upgrading unsloth in migrated environment..." + if [ "$SKIP_TORCH" = true ]; then + # No-torch: install unsloth + unsloth-zoo with --no-deps (current + # PyPI metadata still declares torch as a hard dep), then install + # runtime deps (typer, safetensors, transformers, etc.) with --no-deps + # to prevent transitive torch resolution. + run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ + --reinstall-package unsloth --reinstall-package unsloth-zoo \ + "unsloth>=2026.3.16" unsloth-zoo + _NO_TORCH_RT="$(_find_no_torch_runtime)" + if [ -n "$_NO_TORCH_RT" ]; then + run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" + fi + else + run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + --reinstall-package unsloth --reinstall-package unsloth-zoo \ + "unsloth>=2026.3.16" unsloth-zoo + fi + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + fi +elif [ -n "$TORCH_INDEX_URL" ]; then + # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) + if [ "$SKIP_TORCH" = true ]; then + substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN" + else + substep "installing PyTorch ($TORCH_INDEX_URL)..." + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" + fi + # Fresh: Step 2 - install unsloth, preserving pre-installed torch + substep "installing unsloth (this may take a few minutes)..." + if [ "$SKIP_TORCH" = true ]; then + # No-torch: install unsloth + unsloth-zoo with --no-deps, then + # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. + run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ + --upgrade-package unsloth --upgrade-package unsloth-zoo \ + "unsloth>=2026.3.16" unsloth-zoo + _NO_TORCH_RT="$(_find_no_torch_runtime)" + if [ -n "$_NO_TORCH_RT" ]; then + run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" + fi + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + fi + elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then + run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ + --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + else + run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \ + --upgrade-package unsloth "$PACKAGE_NAME" + fi +else + # Fallback: GPU detection failed to produce a URL -- let uv resolve torch + substep "installing unsloth (this may take a few minutes)..." + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + else + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "$PACKAGE_NAME" --torch-backend=auto + fi +fi # ── Run studio setup ── -# Ensure the venv's Python is on PATH for setup.sh's Python discovery. -# On macOS the system Python may be outside the 3.11-3.13 range that -# setup.sh requires, but uv already installed a compatible interpreter -# inside the venv. -VENV_ABS_BIN="$(cd "$VENV_NAME/bin" && pwd)" +# When --local, use the repo's own setup.sh directly. +# Otherwise, find it inside the installed package. +SETUP_SH="" +if [ "$STUDIO_LOCAL_INSTALL" = true ] && [ -f "$_REPO_ROOT/studio/setup.sh" ]; then + SETUP_SH="$_REPO_ROOT/studio/setup.sh" +fi + +if [ -z "$SETUP_SH" ] || [ ! -f "$SETUP_SH" ]; then + SETUP_SH=$("$VENV_DIR/bin/python" -c " +import importlib.resources +print(importlib.resources.files('studio') / 'setup.sh') +" 2>/dev/null || echo "") +fi + +# Fallback: search site-packages +if [ -z "$SETUP_SH" ] || [ ! -f "$SETUP_SH" ]; then + SETUP_SH=$(find "$VENV_DIR" -path "*/studio/setup.sh" -print -quit 2>/dev/null || echo "") +fi + +if [ -z "$SETUP_SH" ] || [ ! -f "$SETUP_SH" ]; then + echo "❌ ERROR: Could not find studio/setup.sh in the installed package." + exit 1 +fi + +# Ensure the venv's Python is on PATH so setup.sh can find it. +VENV_ABS_BIN="$(cd "$VENV_DIR/bin" && pwd)" if [ -n "$VENV_ABS_BIN" ]; then export PATH="$VENV_ABS_BIN:$PATH" fi -echo "==> Running unsloth studio setup..." -REQUESTED_PYTHON_VERSION="$(cd "$VENV_NAME/bin" && pwd)/python" \ -"$VENV_NAME/bin/unsloth" studio setup /dev/null 2>&1; then + step "setup" "bash is required to run studio setup" "$C_ERR" + substep "Please install bash and re-run install.sh" + exit 1 +fi + +step "setup" "running unsloth studio update..." +# install.sh already installs base packages (unsloth + unsloth-zoo) and +# no-torch-runtime.txt above, so tell install_python_stack.py to skip +# the base step to avoid redundant reinstallation. +_SKIP_BASE=1 +# Run setup.sh outside set -e so that a llama.cpp build failure (exit 1) +# does not skip PATH setup, shortcuts, and launch below. We capture the +# exit code and propagate it after post-install steps finish. +_SETUP_EXIT=0 +if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + SKIP_STUDIO_BASE="$_SKIP_BASE" \ + STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \ + STUDIO_LOCAL_INSTALL=1 \ + STUDIO_LOCAL_REPO="$_REPO_ROOT" \ + UNSLOTH_NO_TORCH="$SKIP_TORCH" \ + bash "$SETUP_SH" /dev/null; then + echo '' >> "$_SHELL_PROFILE" + echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE" + echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE" + step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE" + fi + fi + export PATH="$_LOCAL_BIN:$PATH" + ;; +esac create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" +# If setup.sh failed, report and exit now. +# PATH and shortcuts are already set up so the user can fix and retry. +if [ "$_SETUP_EXIT" -ne 0 ]; then + echo "" + step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR" + substep "Check the output above for details, then re-run:" + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + substep " unsloth studio update --local" + else + substep " unsloth studio update" + fi + echo "" + exit "$_SETUP_EXIT" +fi + echo "" -echo "=========================================" -echo " Unsloth Studio installed!" -echo "=========================================" +printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" -echo " To launch, run:" -echo "" -echo " source ${VENV_NAME}/bin/activate" -echo " unsloth studio -H 0.0.0.0 -p 8888" -echo "" +# Launch studio automatically in interactive terminals; +# in non-interactive environments (Docker, CI, cloud-init) just print instructions. +if [ -t 1 ]; then + step "launch" "starting Unsloth Studio..." + "$VENV_DIR/bin/unsloth" studio -H 0.0.0.0 -p 8888 + _LAUNCH_EXIT=$? + if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then + echo "" + echo "⚠️ Unsloth Studio failed to start after migration." + echo " Your migrated environment may be incompatible." + echo " To fix, remove the environment and reinstall:" + echo "" + echo " rm -rf $VENV_DIR" + echo " curl -fsSL https://unsloth.ai/install.sh | sh" + echo "" + fi + exit "$_LAUNCH_EXIT" +else + step "launch" "manual commands:" + substep "unsloth studio -H 0.0.0.0 -p 8888" + substep "or activate env first:" + substep "source ${VENV_DIR}/bin/activate" + substep "unsloth studio -H 0.0.0.0 -p 8888" + echo "" +fi diff --git a/pyproject.toml b/pyproject.toml index 0c7dd0b962..b06131021a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,8 @@ studio = [ ] [tool.setuptools.packages.find] -exclude = ["images*", "tests*", "kernels/moe*"] +include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"] +exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"] [project.optional-dependencies] triton = [ @@ -87,7 +88,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.3.5", + "unsloth_zoo>=2026.3.6", "torchvision", "unsloth[triton]", ] @@ -577,7 +578,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.3.5", + "unsloth_zoo>=2026.3.6", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.3.0", diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 7191dfde41..46e2067ba7 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -1,157 +1,153 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "view-in-github", - "colab_type": "text" - }, - "source": [ - "\"Open" - ] - }, - { - "cell_type": "markdown", - "id": "6b87de59", - "metadata": { - "id": "6b87de59" - }, - "source": [ - "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", - "
\n", - "\n", - "\n", - " Join Discord if you need help + ⭐ Star us on Github ⭐\n", - "
\n", - "\n", - "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", - "\n", - "### Unsloth Studio\n", - "\n", - "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). Currently, installation may take 30+ mins so use a newer GPU.\n", - "\n", - "\n", - "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", - "\n", - "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" - ] - }, - { - "cell_type": "markdown", - "id": "e4206349", - "metadata": { - "id": "e4206349" - }, - "source": [ - "

" - ] - }, - { - "cell_type": "markdown", - "id": "27da2957", - "metadata": { - "id": "27da2957" - }, - "source": [ - "### Setup: Clone repo and run setup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "27e68f91", - "metadata": { - "id": "27e68f91" - }, - "outputs": [], - "source": [ - "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n", - "%cd /content/unsloth\n", - "!chmod +x studio/setup.sh && ./studio/setup.sh" - ] - }, - { - "cell_type": "markdown", - "id": "3e1771a9", - "metadata": { - "id": "3e1771a9" - }, - "source": [ - "### Start Unsloth Studio" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "277e431e", - "metadata": { - "id": "277e431e" - }, - "outputs": [], - "source": [ - "import sys, time\n", - "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", - "from colab import start\n", - "start()" - ] - }, - { - "cell_type": "code", - "source": [ - "from google.colab import output\n", - "output.serve_kernel_port_as_iframe(8888, height = 1200, width = \"100%\")\n", - "for _ in range(10000): time.sleep(300), print(\"=\", end = \"\")" - ], - "metadata": { - "id": "wb9UELh--XzX" - }, - "id": "wb9UELh--XzX", - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "f2b0c6a1", - "metadata": { - "id": "f2b0c6a1" - }, - "source": [ - "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", - "\n", - "Some other resources:\n", - "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", - "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", - "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", - "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", - "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", - "\n", - "
\n", - " \n", - " \n", - " \n", - "\n", - " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", - "\n", - " This notebook is licensed AGPL-3.0\n", - "
" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "gpuType": "T4", - "provenance": [], - "include_colab_link": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "name": "python" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file + { + "cell_type": "markdown", + "id": "6b87de59", + "metadata": { + "id": "6b87de59" + }, + "source": [ + "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", + "
\n", + "\n", + "\n", + " Join Discord if you need help + ⭐ Star us on Github ⭐\n", + "
\n", + "\n", + "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", + "\n", + "### Unsloth Studio\n", + "\n", + "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", + "\n", + "\n", + "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", + "\n", + "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" + ] + }, + { + "cell_type": "markdown", + "id": "e4206349", + "metadata": { + "id": "e4206349" + }, + "source": [ + "

" + ] + }, + { + "cell_type": "markdown", + "id": "27da2957", + "metadata": { + "id": "27da2957" + }, + "source": [ + "### Setup: Clone repo and run setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27e68f91", + "metadata": { + "id": "27e68f91" + }, + "outputs": [], + "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh" + }, + { + "cell_type": "markdown", + "id": "3e1771a9", + "metadata": { + "id": "3e1771a9" + }, + "source": [ + "### Start Unsloth Studio" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "277e431e", + "metadata": { + "id": "277e431e" + }, + "outputs": [], + "source": [ + "import sys, time\n", + "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", + "from colab import start\n", + "start()" + ] + }, + { + "cell_type": "code", + "source": [ + "from google.colab import output\n", + "output.serve_kernel_port_as_iframe(8888, height = 1200, width = \"100%\")\n", + "for _ in range(10000): time.sleep(300), print(\"=\", end = \"\")" + ], + "metadata": { + "id": "wb9UELh--XzX" + }, + "id": "wb9UELh--XzX", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "f2b0c6a1", + "metadata": { + "id": "f2b0c6a1" + }, + "source": [ + "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", + "\n", + "Some other resources:\n", + "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", + "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", + "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", + "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", + "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", + "\n", + "
\n", + " \n", + " \n", + " \n", + "\n", + " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", + "\n", + " This notebook is licensed AGPL-3.0\n", + "
" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/studio/backend/assets/configs/full_finetune.yaml b/studio/backend/assets/configs/full_finetune.yaml index 7536ed1f11..e398515f61 100644 --- a/studio/backend/assets/configs/full_finetune.yaml +++ b/studio/backend/assets/configs/full_finetune.yaml @@ -10,13 +10,13 @@ training: load_in_4bit: false output_dir: outputs num_epochs: 1 - learning_rate: 0.0002 + learning_rate: 2e-5 batch_size: 1 gradient_accumulation_steps: 4 warmup_steps: 5 max_steps: 0 save_steps: 0 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/lora_text.yaml b/studio/backend/assets/configs/lora_text.yaml index 7101a00e85..9cb6b8c700 100644 --- a/studio/backend/assets/configs/lora_text.yaml +++ b/studio/backend/assets/configs/lora_text.yaml @@ -16,7 +16,7 @@ training: warmup_steps: 5 max_steps: 0 save_steps: 0 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml index d96e5077b2..12566019b8 100644 --- a/studio/backend/assets/configs/model_defaults/default.yaml +++ b/studio/backend/assets/configs/model_defaults/default.yaml @@ -6,13 +6,13 @@ training: max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 - learning_rate: 5e-5 + learning_rate: 2e-4 batch_size: 2 gradient_accumulation_steps: 4 warmup_ratio: 0.1 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: true diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml index 1219648a8a..f7b49c75b7 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml @@ -12,7 +12,7 @@ training: warmup_ratio: 0.03 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml index db742e11b5..be7da0f624 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml @@ -11,7 +11,7 @@ training: warmup_ratio: 0.03 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml index 499c112929..d9e49bc0d5 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml @@ -11,7 +11,7 @@ training: warmup_ratio: 0.03 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml index 016d284a16..c3422d399f 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml @@ -11,7 +11,7 @@ training: warmup_ratio: 0.03 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml index eefaecc8cd..529a56a527 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml @@ -11,7 +11,7 @@ training: warmup_ratio: 0.03 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml index c45b71b4ae..fa7bd8c1ea 100644 --- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml @@ -13,7 +13,7 @@ training: warmup_steps: 5 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml index f8f78f5edc..a4acbe9262 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml @@ -13,7 +13,7 @@ training: warmup_steps: 5 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: true diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml index 3938f10627..2bc3f6f871 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml @@ -13,7 +13,7 @@ training: warmup_steps: 0 max_steps: 30 save_steps: 30 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: true diff --git a/studio/backend/assets/configs/vision_lora.yaml b/studio/backend/assets/configs/vision_lora.yaml index 60641b16e8..063a970316 100644 --- a/studio/backend/assets/configs/vision_lora.yaml +++ b/studio/backend/assets/configs/vision_lora.yaml @@ -16,7 +16,7 @@ training: warmup_steps: 5 max_steps: 0 save_steps: 0 - weight_decay: 0.01 + weight_decay: 0.001 random_seed: 3407 packing: false train_on_completions: false diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 25a9408ccb..efd0e10bdb 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -18,31 +18,6 @@ if _backend_dir not in sys.path: import _platform_compat # noqa: F401 -def _bootstrap_studio_venv() -> None: - """Expose the Studio venv's site-packages to the current interpreter. - - On Colab, notebook cells run outside the venv subshell. Instead of - installing the full stack into system Python, we prepend the venv's - site-packages so that packages like structlog, fastapi, etc. are - importable from notebook cells and take priority over system copies. - """ - venv_lib = Path.home() / ".unsloth" / "studio" / ".venv" / "lib" - if not venv_lib.exists(): - import warnings - - warnings.warn( - f"Studio venv not found at {venv_lib.parent} -- run 'unsloth studio setup' first", - stacklevel = 2, - ) - return - for sp in venv_lib.glob("python*/site-packages"): - sp_str = str(sp) - if sp_str not in sys.path: - sys.path.insert(0, sp_str) - - -_bootstrap_studio_venv() - from loggers import get_logger logger = get_logger(__name__) diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index a9fbe659b3..500bc9e706 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -217,6 +217,7 @@ class ExportOrchestrator: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -227,6 +228,7 @@ class ExportOrchestrator: "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, + "hf_token": hf_token, } # Always kill existing subprocess and spawn fresh. diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 1a265690ff..867bdefc62 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -18,7 +18,14 @@ from typing import Optional, Union, Generator, Tuple from utils.models import ModelConfig, get_base_model_from_lora from utils.paths import is_model_cached from utils.utils import format_error_message -from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory +from utils.hardware import ( + get_device, + clear_gpu_cache, + log_gpu_memory, + get_device_map, + raise_if_offloaded, + get_visible_gpu_count, +) from core.inference.audio_codecs import AudioCodecManager from io import StringIO import structlog @@ -241,6 +248,7 @@ class InferenceBackend: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """ Load any model: base, LoRA adapter, text, or vision. @@ -260,6 +268,10 @@ class InferenceBackend: return False self.loading_models.add(model_name) + device_map = get_device_map(gpu_ids) + logger.info( + f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" + ) self.models[model_name] = { "is_vision": config.is_vision, @@ -290,6 +302,7 @@ class InferenceBackend: config.path, auto_model = CsmForConditionalGeneration, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -325,6 +338,7 @@ class InferenceBackend: config.path, dtype = torch.float32, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -345,6 +359,7 @@ class InferenceBackend: llm_path, dtype = torch.float32, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -361,6 +376,7 @@ class InferenceBackend: config.path, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -378,6 +394,7 @@ class InferenceBackend: whisper_language = "English", whisper_task = "transcribe", load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -405,6 +422,7 @@ class InferenceBackend: model_name = config.path, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -420,6 +438,11 @@ class InferenceBackend: audio_type, self.device, model_repo_path = model_repo_path ) + # Reject CPU/disk offload for audio models too + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) + self.active_model_name = model_name self.loading_models.discard(model_name) logger.info(f"Successfully loaded audio model: {model_name}") @@ -441,6 +464,7 @@ class InferenceBackend: max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -497,6 +521,7 @@ class InferenceBackend: max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -507,6 +532,10 @@ class InferenceBackend: self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) + # Load chat template info self._load_chat_template_info(model_name) @@ -615,6 +644,7 @@ class InferenceBackend: dtype = None, load_in_4bit: bool = True, hf_token: Optional[str] = None, + gpu_ids: Optional[list[int]] = None, ) -> Tuple[bool, Optional[str], Optional[str]]: """ Final Corrected Version: @@ -639,7 +669,12 @@ class InferenceBackend: base_model_name, None, is_lora = False ) if not self.load_model( - base_config, max_seq_length, dtype, load_in_4bit, hf_token + base_config, + max_seq_length, + dtype, + load_in_4bit, + hf_token, + gpu_ids = gpu_ids, ): return False, None, None @@ -1037,12 +1072,12 @@ class InferenceBackend: input_text, add_special_tokens = False, return_tensors = "pt", - ).to(self.device) + ).to(model.device) else: # Text-only for vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to( - self.device + model.device ) # Stream with TextIteratorStreamer + background thread @@ -1182,7 +1217,7 @@ class InferenceBackend: return_dict = True, return_tensors = "pt", truncation = False, - ).to(self.device) + ).to(model.device) try: from transformers import TextIteratorStreamer diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7b1db8fd04..c1f87ff936 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10,12 +10,13 @@ through its OpenAI-compatible /v1/chat/completions endpoint. import atexit import contextlib +import hashlib import json +import re import struct import structlog from loggers import get_logger import shutil -import signal import socket import subprocess import threading @@ -48,15 +49,24 @@ class LlamaCppBackend: self._is_vision: bool = False self._healthy = False self._context_length: Optional[int] = None + self._effective_context_length: Optional[int] = None + self._max_context_length: Optional[int] = None self._chat_template: Optional[str] = None self._supports_reasoning: bool = False + self._reasoning_always_on: bool = False self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None self._reasoning_default: bool = True + # KV-cache estimation fields (populated by _read_gguf_metadata) + self._n_layers: Optional[int] = None + self._n_kv_heads: Optional[int] = None + self._n_heads: Optional[int] = None + self._embedding_length: Optional[int] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None self._cancel_event = threading.Event() + self._api_key: Optional[str] = None self._kill_orphaned_servers() atexit.register(self._cleanup) @@ -90,7 +100,13 @@ class LlamaCppBackend: @property def context_length(self) -> Optional[int]: - return self._context_length + """Return the effective context length the server is running at.""" + return self._effective_context_length or self._context_length + + @property + def max_context_length(self) -> Optional[int]: + """Return the maximum context currently available on this hardware.""" + return self._max_context_length or self._context_length @property def chat_template(self) -> Optional[str]: @@ -100,6 +116,10 @@ class LlamaCppBackend: def supports_reasoning(self) -> bool: return self._supports_reasoning + @property + def reasoning_always_on(self) -> bool: + return self._reasoning_always_on + @property def reasoning_default(self) -> bool: return self._reasoning_default @@ -275,7 +295,8 @@ class LlamaCppBackend: continue gpus.append((idx, free_mib)) return gpus - except Exception: + except Exception as e: + logger.debug(f"Failed to query GPU free memory via nvidia-smi: {e}") return [] @staticmethod @@ -283,11 +304,11 @@ class LlamaCppBackend: model_size_bytes: int, gpus: list[tuple[int, int]], ) -> tuple[Optional[list[int]], bool]: - """Pick GPU(s) for a model based on file size and free memory. + """Pick GPU(s) for a model based on estimated VRAM and free memory. - Uses GGUF file size as a rough proxy for VRAM usage (actual usage - is higher due to KV cache and compute buffers, but 70% threshold - accounts for that). + ``model_size_bytes`` should include both model weights and estimated + KV cache. The 70% threshold provides headroom for compute buffers, + CUDA context, and other runtime overhead. Returns (gpu_indices, use_fit): - ([1], False) model fits on 1 GPU at 70% of free @@ -316,8 +337,115 @@ class LlamaCppBackend: return sorted(selected), False # Model is too large even for all GPUs, let --fit handle it + logger.debug( + "Model does not fit in available GPU memory, falling back to --fit", + model_size_mib = round(model_size_mib, 2), + ranked_gpus = ranked, + ) return None, True + # ── KV cache VRAM estimation ───────────────────────────────────── + + def _can_estimate_kv(self) -> bool: + """True if we have enough GGUF metadata to estimate KV cache size.""" + return ( + self._n_layers is not None + and self._embedding_length is not None + and (self._n_kv_heads is not None or self._n_heads is not None) + ) + + def _estimate_kv_cache_bytes( + self, n_ctx: int, cache_type_kv: Optional[str] = None + ) -> int: + """Estimate KV cache VRAM for a given context length. + + Returns 0 if metadata is insufficient for estimation. + """ + if not self._can_estimate_kv() or n_ctx <= 0: + return 0 + + n_layers = self._n_layers # type: ignore[assignment] + n_kv_heads = self._n_kv_heads or self._n_heads # type: ignore[assignment] + head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + + # Bytes per element depends on KV cache quantization + bpe = { + "f32": 4.0, + "f16": 2.0, + "bf16": 2.0, + "q8_0": 34 / 32, + "q5_1": 0.75, + "q5_0": 0.6875, + "q4_1": 0.625, + "q4_0": 0.5625, + "iq4_nl": 0.5625, + }.get(cache_type_kv or "f16", 2.0) + + # K + V caches: 2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe + return int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe) + + def _fit_context_to_vram( + self, + requested_ctx: int, + available_mib: int, + model_size_bytes: int, + cache_type_kv: Optional[str] = None, + min_ctx: int = 4096, + ) -> int: + """Return the largest context length that fits in GPU VRAM. + + Uses 70% of available VRAM as the budget (matching _select_gpus + threshold -- 30% reserved for compute buffers, CUDA context, + scratch space, flash-attn workspace, etc.). + If the model weights alone don't fit, returns min_ctx unchanged. + """ + if not self._can_estimate_kv(): + logger.debug( + "Skipping context fit because KV cache metadata is unavailable", + requested_ctx = requested_ctx, + available_mib = available_mib, + ) + return requested_ctx + + budget_bytes = available_mib * 1024 * 1024 * 0.70 + model_footprint = model_size_bytes + + # Check if requested context already fits + kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv) + if model_footprint + kv <= budget_bytes: + return requested_ctx + + # Model weights alone exceed budget -- can't help by reducing ctx. + # Return requested_ctx unchanged; --fit will handle VRAM management. + if model_footprint >= budget_bytes: + logger.debug( + "Model footprint exceeds GPU budget before KV cache", + requested_ctx = requested_ctx, + available_mib = available_mib, + model_size_gb = round(model_footprint / (1024**3), 2), + ) + return requested_ctx + + # Binary search for max context that fits + remaining = budget_bytes - model_footprint + effective_min = min(min_ctx, requested_ctx) + lo, hi = effective_min, requested_ctx + best = effective_min + while lo <= hi: + mid = (lo + hi) // 2 + kv = self._estimate_kv_cache_bytes(mid, cache_type_kv) + if kv <= remaining: + best = mid + lo = mid + 1 + else: + hi = mid - 1 + + # Round down to nearest 256 for alignment, but never exceed requested_ctx + best = (best // 256) * 256 + best = max(effective_min, best) + best = min(best, requested_ctx) + return best + # ── Variant fallback ──────────────────────────────────────────── @staticmethod @@ -442,7 +570,7 @@ class LlamaCppBackend: LlamaCppBackend._gguf_skip_value(f, atype) def _read_gguf_metadata(self, gguf_path: str) -> None: - """Read context_length and chat_template from a GGUF file's KV header. + """Read context_length, architecture params, and chat_template from a GGUF header. Parses only the KV pairs we need (~30ms even for multi-GB files). For split GGUFs, metadata is always in shard 1. @@ -452,12 +580,19 @@ class LlamaCppBackend: self._context_length = None self._chat_template = None self._supports_reasoning = False + self._reasoning_always_on = False self._supports_tools = False + self._n_layers = None + self._n_kv_heads = None + self._n_heads = None + self._embedding_length = None try: WANTED = {"general.architecture", "tokenizer.chat_template"} + # Additional arch-specific keys are added dynamically once + # we know the architecture name. + arch_keys: dict[str, str] = {} # gguf_key -> attribute name arch = None - ctx_key = None with open(gguf_path, "rb") as f: magic = struct.unpack(" tags or reasoning_content + # in their chat template always produce thinking output + # (no toggle to disable it). + if not self._supports_reasoning: + if ( + "" in tpl + and "" in tpl + or "reasoning_content" in tpl + ): + self._supports_reasoning = True + self._reasoning_always_on = True + logger.info( + "GGUF metadata: model always reasons ( tags in template)" + ) # Detect tool calling support from chat template tool_markers = [ "{%- if tools %}", @@ -828,18 +985,149 @@ class LlamaCppBackend: self._port = self._find_free_port() - # Select GPU(s) based on model size and free memory + # Select GPU(s) based on model size + estimated KV cache. + # Seed safe defaults before GPU probing so the except path + # still has valid state to publish. + effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) + max_available_ctx = self._context_length or effective_ctx try: model_size = self._get_gguf_size_bytes(model_path) gpus = self._get_gpu_free_memory() - gpu_indices, use_fit = self._select_gpus(model_size, gpus) + + # Resolve effective context: 0 means let llama-server use the + # model's native length. Only expand to a known native length + # if metadata is available; otherwise preserve 0 as a sentinel. + if n_ctx > 0: + effective_ctx = n_ctx + elif self._context_length is not None: + effective_ctx = self._context_length + else: + effective_ctx = 0 + original_ctx = effective_ctx + # Default UI ceiling to the model's native context length. + # GPU/VRAM-fit logic below may shrink this if hardware is limited. + max_available_ctx = self._context_length or effective_ctx + + # Auto-cap context to fit in GPU VRAM and select GPUs. + # + # Two policies depending on whether the user set n_ctx: + # + # Explicit n_ctx (user chose a context length): + # Honor it. Try the full requested context with _select_gpus + # (which uses as many GPUs as needed). Only cap if it doesn't + # fit on any GPU combination. + # + # Auto n_ctx=0 (model's native context): + # Prefer fewer GPUs with reduced context over more GPUs, + # since multi-GPU is slower and the user didn't ask for a + # specific context length. + gpu_indices, use_fit = None, True + explicit_ctx = n_ctx > 0 + + if gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. + native_ctx_for_cap = self._context_length or effective_ctx + if native_ctx_for_cap > 0: + ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + best_cap = 0 + for n_gpus in range(1, len(ranked_for_cap) + 1): + subset = ranked_for_cap[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + native_ctx_for_cap, + pool_mib, + model_size, + cache_type_kv, + ) + kv = self._estimate_kv_cache_bytes(capped, cache_type_kv) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.70: + best_cap = max(best_cap, capped) + if best_cap > 0: + max_available_ctx = best_cap + + if explicit_ctx: + # Try to honor the user's requested context exactly. + requested_total = model_size + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv + ) + gpu_indices, use_fit = self._select_gpus(requested_total, gpus) + + # Full context doesn't fit anywhere -- cap it on the + # best GPU subset we can find (fewest GPUs first). + if use_fit: + ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + effective_ctx, + pool_mib, + model_size, + cache_type_kv, + ) + kv = self._estimate_kv_cache_bytes( + capped, cache_type_kv + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.70: + effective_ctx = capped + gpu_indices = sorted(idx for idx, _ in subset) + use_fit = False + break + else: + # Auto context: prefer fewer GPUs, cap context to fit. + ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + effective_ctx, + pool_mib, + model_size, + cache_type_kv, + ) + kv = self._estimate_kv_cache_bytes(capped, cache_type_kv) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.70: + effective_ctx = capped + gpu_indices = sorted(idx for idx, _ in subset) + use_fit = False + break + + elif gpus: + # Can't estimate KV -- fall back to file-size-only check. + # Without KV estimation we cannot prove a hardware cap, so + # keep the ceiling at the native context (already the default). + logger.debug( + "Falling back to file-size-only GPU selection", + model_size_gb = round(model_size / (1024**3), 2), + ) + gpu_indices, use_fit = self._select_gpus(model_size, gpus) + + if effective_ctx < original_ctx: + kv_est = self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv) + logger.info( + f"Context auto-reduced: {original_ctx} -> {effective_ctx} " + f"(model: {model_size / (1024**3):.1f} GB, " + f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + ) + + kv_cache_bytes = self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv + ) logger.info( f"GGUF size: {model_size / (1024**3):.1f} GB, " + f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"context: {effective_ctx}, " f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) except Exception as e: logger.warning(f"GPU selection failed ({e}), using --fit on") gpu_indices, use_fit = None, True + effective_ctx = n_ctx # fall back to original cmd = [ binary, @@ -848,7 +1136,7 @@ class LlamaCppBackend: "--port", str(self._port), "-c", - "0", # 0 = use model's native context size + str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", "1", # Single-user studio, saves VRAM "--flash-attn", @@ -857,6 +1145,9 @@ class LlamaCppBackend: if use_fit: cmd.extend(["--fit", "on"]) + elif gpu_indices is not None: + # Model fits on selected GPU(s) -- offload all layers + cmd.extend(["-ngl", "-1"]) if n_threads is not None: cmd.extend(["--threads", str(n_threads)]) @@ -935,7 +1226,23 @@ class LlamaCppBackend: cmd.extend(["--mmproj", mmproj_path]) logger.info(f"Using mmproj for vision: {mmproj_path}") - logger.info(f"Starting llama-server: {' '.join(cmd)}") + # Option C: add --api-key for direct client access when enabled + import os as _os + import secrets as _secrets + + if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": + self._api_key = _secrets.token_urlsafe(32) + cmd.extend(["--api-key", self._api_key]) + logger.info("llama-server started with --api-key for direct streaming") + else: + self._api_key = None + + _log_cmd = list(cmd) + if "--api-key" in _log_cmd: + _ki = _log_cmd.index("--api-key") + 1 + if _ki < len(_log_cmd): + _log_cmd[_ki] = "" + logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") # Set library paths so llama-server can find its shared libs and CUDA DLLs import os @@ -966,6 +1273,46 @@ class LlamaCppBackend: lib_dirs = [binary_dir] _arch = platform.machine() # x86_64, aarch64, etc. + + # Pip-installed nvidia CUDA runtime libs (e.g. torch's + # bundled cuda-bindings). The prebuilt llama.cpp binary + # links against libcudart.so.13 / libcublas.so.13 which + # live here, not in /usr/local/cuda. + import glob as _glob + + for _nv_pattern in [ + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cu*", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cudnn", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "nvjitlink", + "lib", + ), + ]: + for _nv_dir in _glob.glob(_nv_pattern): + if os.path.isdir(_nv_dir): + lib_dirs.append(_nv_dir) + for cuda_lib in [ "/usr/local/cuda/lib64", f"/usr/local/cuda/targets/{_arch}-linux/lib", @@ -1006,12 +1353,35 @@ class LlamaCppBackend: self._gguf_path = gguf_path self._hf_repo = hf_repo - self._hf_variant = hf_variant + # For local GGUF files, extract variant from filename if not provided + if hf_variant: + self._hf_variant = hf_variant + elif gguf_path: + try: + from utils.models.model_config import _extract_quant_label + + self._hf_variant = _extract_quant_label(gguf_path) + except Exception: + self._hf_variant = None + else: + self._hf_variant = None self._is_vision = is_vision self._model_identifier = model_identifier + # Store the effective (possibly capped) context separately. + # Do NOT overwrite _context_length -- it holds the model's native + # context length from GGUF metadata and is used for display/info. + self._effective_context_length = ( + effective_ctx if effective_ctx > 0 else self._context_length + ) + self._max_context_length = ( + 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 = 120.0): + if not self._wait_for_health(timeout = 600.0): self._kill_process() raise RuntimeError( "llama-server failed to start. " @@ -1042,10 +1412,17 @@ class LlamaCppBackend: self._port = None self._healthy = False self._context_length = None + self._effective_context_length = None + self._max_context_length = None self._chat_template = None self._supports_reasoning = False + self._reasoning_always_on = False self._supports_tools = False self._cache_type_kv = None + self._n_layers = None + self._n_kv_heads = None + self._n_heads = None + self._embedding_length = None # Clean up temp chat template file if hasattr(self, "_chat_template_file") and self._chat_template_file: try: @@ -1088,42 +1465,161 @@ class LlamaCppBackend: def _kill_orphaned_servers(): """Kill orphaned llama-server processes started by studio. - Only kills processes whose binary lives under ~/.unsloth/llama.cpp/ - to avoid terminating unrelated llama-server instances on the machine. + Only kills processes whose resolved binary lives under a known + Studio install directory (or matches an exact env-var override) + to avoid terminating unrelated llama-server instances. + + Mirrors every location that _find_llama_server_binary() can + return from so that orphans from any supported install path + are still cleaned up. + + Uses psutil for cross-platform support (Linux, macOS, Windows). + Falls back to pgrep + /proc//exe on Linux when psutil is + not installed. """ import os import signal + import sys try: - # Use pgrep with full command match to identify studio-managed servers - result = subprocess.run( - ["pgrep", "-a", "-f", "llama-server"], - capture_output = True, - text = True, - timeout = 5, - ) - if result.returncode != 0: - return - for line in result.stdout.strip().splitlines(): - parts = line.strip().split(None, 1) - if len(parts) < 2: - continue - pid = int(parts[0]) - cmdline = parts[1] - if pid == os.getpid(): - continue - # Only kill if it's a studio-managed server (lives under .unsloth/) - if ".unsloth/" not in cmdline and "unsloth" not in cmdline.lower(): - continue + # -- Build the ownership allowlist -------------------------------- + # Two kinds of matches: + # exact_binaries -- env var overrides (exact path match only) + # install_roots -- directory trees that are Studio-owned + # (binary must be *under* one of these) + install_roots: list[Path] = [] + + # Primary install dir (setup.sh / prebuilt installer) + install_roots.append(Path.home() / ".unsloth" / "llama.cpp") + + # Legacy in-tree build dirs (older setup.sh versions) + project_root = Path(__file__).resolve().parents[4] + install_roots.append(project_root / "llama.cpp") + + # Legacy: extracted binary + install_roots.append(project_root / "bin") + + # UNSLOTH_LLAMA_CPP_PATH env var (custom install dir) + custom_dir = os.environ.get("UNSLOTH_LLAMA_CPP_PATH") + if custom_dir: + install_roots.append(Path(custom_dir)) + + # LLAMA_SERVER_PATH env var (exact binary path) + exact_binaries: list[Path] = [] + env_binary = os.environ.get("LLAMA_SERVER_PATH") + if env_binary: try: - os.kill(pid, signal.SIGKILL) - logger.info(f"Killed orphaned llama-server process (pid={pid})") - except ProcessLookupError: + exact_binaries.append(Path(env_binary).resolve()) + except OSError: pass - except PermissionError: + + # Resolve all roots so is_relative_to works reliably + resolved_roots: list[Path] = [] + for root in install_roots: + try: + resolved_roots.append(root.resolve()) + except OSError: pass + + my_pid = os.getpid() + + # -- Enumerate processes ------------------------------------------- + # Prefer psutil (cross-platform). Fall back to pgrep + /proc on + # Linux when psutil is not installed. + try: + import psutil + + has_psutil = True + except ImportError: + has_psutil = False + + if has_psutil: + for proc in psutil.process_iter(["pid", "name", "exe"]): + try: + if proc.info["pid"] == my_pid: + continue + + name = proc.info.get("name") or "" + if not name.lower().startswith("llama-server"): + continue + + exe = proc.info.get("exe") + if not exe: + continue + + exe_path = Path(exe).resolve() + + # Check ownership: exact binary match OR binary is + # under a known install root (proper ancestry, not + # substring). + is_ours = exe_path in exact_binaries or any( + exe_path.is_relative_to(root) for root in resolved_roots + ) + if not is_ours: + continue + + proc.kill() + logger.info( + f"Killed orphaned llama-server process " + f"(pid={proc.info['pid']})" + ) + except ( + psutil.NoSuchProcess, + psutil.AccessDenied, + psutil.ZombieProcess, + ): + pass + else: + # -- Fallback: pgrep + /proc//exe (Linux only) ----------- + if sys.platform != "linux": + return + result = subprocess.run( + ["pgrep", "-a", "-f", "llama-server"], + capture_output = True, + text = True, + timeout = 5, + ) + if result.returncode != 0: + return + + for line in result.stdout.strip().splitlines(): + parts = line.strip().split(None, 1) + if len(parts) < 2: + continue + pid = int(parts[0]) + if pid == my_pid: + continue + + # Resolve the actual executable. /proc//exe is a + # symlink to the real binary and avoids all cmdline- + # parsing ambiguities (spaces in paths, argv rewriting). + # Fall back to the first cmdline token when /proc is + # unavailable. + proc_exe = Path(f"/proc/{pid}/exe") + try: + binary = proc_exe.resolve(strict = True) + except (OSError, ValueError): + cmdline = parts[1] + token = cmdline.split()[0] if cmdline.strip() else "" + if not token: + continue + binary = Path(token).resolve(strict = False) + + owned = binary in exact_binaries or any( + binary.is_relative_to(root) for root in resolved_roots + ) + if not owned: + continue + + try: + os.kill(pid, signal.SIGKILL) + logger.info(f"Killed orphaned llama-server process (pid={pid})") + except ProcessLookupError: + pass + except PermissionError: + pass except Exception: - pass + logger.warning("Error during orphan server cleanup", exc_info = True) def _cleanup(self): """atexit handler to ensure llama-server is terminated.""" @@ -1364,6 +1860,7 @@ class LlamaCppBackend: url: str, payload: dict, cancel_event: Optional[threading.Event] = None, + headers: Optional[dict] = None, ): """Open an httpx streaming POST with cancel support. @@ -1430,7 +1927,11 @@ class LlamaCppBackend: pool = 10, ) with client.stream( - "POST", url, json = payload, timeout = prefill_timeout + "POST", + url, + json = payload, + timeout = prefill_timeout, + headers = headers, ) as response: _response_ref[0] = response if cancel_event is not None and cancel_event.is_set(): @@ -1504,9 +2005,16 @@ class LlamaCppBackend: # can finish. Cancel during streaming is handled by the # watcher thread (closes the response on cancel_event). stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, payload, cancel_event + client, + url, + payload, + cancel_event, + headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -1614,7 +2122,7 @@ class LlamaCppBackend: stop: Optional[list[str]] = None, cancel_event: Optional[threading.Event] = None, enable_thinking: Optional[bool] = None, - max_tool_iterations: int = 10, + max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, @@ -1638,14 +2146,67 @@ class LlamaCppBackend: _accumulated_predicted_ms = 0.0 _accumulated_predicted_n = 0 + # ── Shared patterns for stripping tool XML from streamed content ── + import re as _re_tool + + _TOOL_CLOSED_PATTERNS = [ + _re_tool.compile(r".*?", _re_tool.DOTALL), + _re_tool.compile(r".*?", _re_tool.DOTALL), + ] + _TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [ + _re_tool.compile(r".*$", _re_tool.DOTALL), + _re_tool.compile(r".*$", _re_tool.DOTALL), + ] + + def _strip_tool_markup(text: str, *, final: bool = False) -> str: + if not auto_heal_tool_calls: + return text + patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS + for pat in patterns: + text = pat.sub("", text) + return text.strip() if final else text + + # XML prefixes that signal a tool call in content. + # Empty when auto_heal is disabled so the buffer never + # speculatively holds content for XML detection. + _TOOL_XML_SIGNALS = ( + ("", " str: + raw = json.dumps({"t": name, "a": args}, sort_keys = True) + return hashlib.md5(raw.encode()).hexdigest() + + def _is_duplicate_call(name: str, args: dict) -> bool: + """Block if the immediately previous call was identical and succeeded.""" + if not _tool_call_history: + return False + key = _tool_call_key(name, args) + last_key, last_failed = _tool_call_history[-1] + return last_key == key and not last_failed + + def _record_tool_call(name: str, args: dict, failed: bool) -> None: + key = _tool_call_key(name, args) + _tool_call_history.append((key, failed)) + for iteration in range(max_tool_iterations): if cancel_event is not None and cancel_event.is_set(): return - # Build payload for non-streaming tool detection pass + # Build payload -- stream: True so we detect tool signals + # in the first 1-2 chunks without a non-streaming penalty. payload = { "messages": conversation, - "stream": False, + "stream": True, + "stream_options": {"include_usage": True}, "temperature": temperature, "top_p": top_p, "top_k": top_k if top_k >= 0 else 0, @@ -1663,96 +2224,433 @@ class LlamaCppBackend: payload["stop"] = stop try: - with httpx.Client(timeout = None) as client: - resp = client.post(url, json = payload) - if resp.status_code != 200: - raise RuntimeError( - f"llama-server returned {resp.status_code}: {resp.text}" + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} + if self._api_key + else None + ) + + # ── Speculative buffer state machine ────────────────── + # BUFFERING: accumulating content, checking for tool signals + # STREAMING: no tool detected, yielding tokens to caller + # DRAINING: tool signal found, silently consuming rest + _S_BUFFERING = 0 + _S_STREAMING = 1 + _S_DRAINING = 2 + + detect_state = _S_BUFFERING + content_buffer = "" # Raw content held during BUFFERING + content_accum = "" # All content tokens (for tool parsing) + reasoning_accum = "" + cumulative_display = "" # Cumulative text yielded (with ) + in_thinking = False + has_content_tokens = False + tool_calls_acc = {} # Structured delta.tool_calls fragments + has_structured_tc = False + _iter_usage = None + _iter_timings = None + _stream_done = False + _last_emitted = "" + + stream_timeout = httpx.Timeout( + connect = 10, + read = 0.5, + write = 10, + pool = 10, + ) + with httpx.Client(timeout = stream_timeout) as client: + with self._stream_with_retry( + client, + url, + payload, + cancel_event, + headers = _auth_headers, + ) as response: + if response.status_code != 200: + error_body = response.read().decode() + raise RuntimeError( + f"llama-server returned {response.status_code}: " + f"{error_body}" + ) + + raw_buf = "" + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + ): + raw_buf += raw_chunk + while "\n" in raw_buf: + line, raw_buf = raw_buf.split("\n", 1) + line = line.strip() + + if not line: + continue + if line == "data: [DONE]": + # Flush thinking state for STREAMING + if detect_state == _S_STREAMING and in_thinking: + if has_content_tokens: + cumulative_display += "" + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative_display, + final = True, + ), + } + else: + cumulative_display = reasoning_accum + yield { + "type": "content", + "text": cumulative_display, + } + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue + + try: + chunk_data = json.loads(line[6:]) + _ct = chunk_data.get("timings") + if _ct: + _iter_timings = _ct + _cu = chunk_data.get("usage") + if _cu: + _iter_usage = _cu + + choices = chunk_data.get("choices", []) + if not choices: + continue + + delta = choices[0].get("delta", {}) + + # ── Structured tool_calls ── + tc_deltas = delta.get("tool_calls") + if tc_deltas: + # Once visible content has been + # emitted, do not reclassify this + # turn as a tool call. + if _last_emitted: + continue + has_structured_tc = True + detect_state = _S_DRAINING + for tc_d in tc_deltas: + idx = tc_d.get("index", 0) + if idx not in tool_calls_acc: + tool_calls_acc[idx] = { + "id": tc_d.get("id", f"call_{idx}"), + "type": "function", + "function": { + "name": "", + "arguments": "", + }, + } + elif tc_d.get("id"): + # Update ID if real one + # arrives on a later delta + tool_calls_acc[idx]["id"] = tc_d["id"] + func = tc_d.get("function", {}) + if func.get("name"): + tool_calls_acc[idx]["function"][ + "name" + ] += func["name"] + if func.get("arguments"): + tool_calls_acc[idx]["function"][ + "arguments" + ] += func["arguments"] + continue + + # ── Reasoning tokens ── + # Only yield in STREAMING state. In BUFFERING + # and DRAINING, accumulate silently so we don't + # corrupt the consumer's prev_text tracker + # (routes/inference.py never resets prev_text + # between tool iterations). + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_accum += reasoning + if detect_state == _S_STREAMING: + if not in_thinking: + cumulative_display += "" + in_thinking = True + cumulative_display += reasoning + yield { + "type": "content", + "text": cumulative_display, + } + + # ── Content tokens ── + token = delta.get("content", "") + if token: + has_content_tokens = True + content_accum += token + + if detect_state == _S_DRAINING: + pass # accumulate silently + + elif detect_state == _S_STREAMING: + if in_thinking: + cumulative_display += "" + in_thinking = False + cumulative_display += token + cleaned = _strip_tool_markup( + cumulative_display, + ) + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + yield { + "type": "content", + "text": cleaned, + } + + elif detect_state == _S_BUFFERING: + content_buffer += token + stripped_buf = content_buffer.lstrip() + if not stripped_buf: + continue + + # Check tool signal prefixes + is_prefix = False + is_match = False + for sig in _TOOL_XML_SIGNALS: + if stripped_buf.startswith(sig): + is_match = True + break + if sig.startswith(stripped_buf): + is_prefix = True + break + + if is_match: + detect_state = _S_DRAINING + elif ( + is_prefix + and len(stripped_buf) + < _MAX_BUFFER_CHARS + ): + pass # keep buffering + else: + # Not a tool -- flush buffer + detect_state = _S_STREAMING + # Flush any reasoning accumulated + # during BUFFERING phase + if reasoning_accum: + cumulative_display += "" + cumulative_display += ( + reasoning_accum + ) + cumulative_display += "" + cumulative_display += content_buffer + cleaned = _strip_tool_markup( + cumulative_display, + ) + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + yield { + "type": "content", + "text": cleaned, + } + + except json.JSONDecodeError: + logger.debug( + f"Skipping malformed SSE line: " f"{line[:100]}" + ) + if _stream_done: + break # exit outer for + + # ── Resolve BUFFERING at stream end ── + if detect_state == _S_BUFFERING: + stripped_buf = content_buffer.lstrip() + if ( + stripped_buf + and auto_heal_tool_calls + and any(s in stripped_buf for s in _TOOL_XML_SIGNALS) + ): + detect_state = _S_DRAINING + elif content_accum or reasoning_accum: + detect_state = _S_STREAMING + if content_buffer: + # Flush any reasoning accumulated first + if reasoning_accum: + cumulative_display += "" + cumulative_display += reasoning_accum + cumulative_display += "" + cumulative_display += content_buffer + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative_display, + final = True, + ), + } + elif reasoning_accum and not has_content_tokens: + # Reasoning-only response (no content tokens): + # show reasoning as plain text, matching + # the final streaming pass behavior for + # models that put everything in reasoning. + cumulative_display = reasoning_accum + yield { + "type": "content", + "text": cumulative_display, + } + else: + return + + # ── STREAMING path: no tool call ── + if detect_state == _S_STREAMING: + # Safety net: check for XML tool signals in content. + # The route layer resets prev_text on tool_start, so + # post-tool synthesis streams correctly even if + # content was already emitted before the tool XML. + _safety_tc = None + if auto_heal_tool_calls and any( + s in content_accum for s in _TOOL_XML_SIGNALS + ): + _safety_tc = self._parse_tool_calls_from_text( + content_accum, ) - data = resp.json() - except httpx.ConnectError: - raise RuntimeError("Lost connection to llama-server") + if not _safety_tc: + # Content was already streamed. Yield metadata. + yield {"type": "status", "text": ""} + _fu = _iter_usage or {} + _fc = _fu.get("completion_tokens", 0) + _fp = _fu.get("prompt_tokens", 0) + _tc = _fc + _accumulated_completion_tokens + if ( + _iter_usage + or _iter_timings + or _accumulated_completion_tokens + ): + _mt = dict(_iter_timings) if _iter_timings else {} + if _accumulated_predicted_ms or _accumulated_predicted_n: + _mt["predicted_ms"] = ( + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + _accumulated_predicted_n + ) + _mt["predicted_n"] = _tn + _tms = _mt["predicted_ms"] + if _tms > 0: + _mt["predicted_per_second"] = _tn / (_tms / 1000.0) + yield { + "type": "metadata", + "usage": { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + }, + "timings": _mt, + } + return - choices = data.get("choices", []) - if not choices: - return - - choice = choices[0] - finish_reason = choice.get("finish_reason", "") - message = choice.get("message", {}) - - # If model wants to call tools - tool_calls = message.get("tool_calls") - - # Fallback: detect tool calls embedded as XML/text in content - # Some models output XML instead of structured tool_calls, - # or bare tags without wrapper. - content_text = message.get("content", "") or "" - if ( - auto_heal_tool_calls - and not tool_calls - and ("" in content_text or " blocks since they - # can contain arbitrary content including code. - import re - - # Strip ... blocks (greedy inside) - content_text = re.sub( - r".*?", - "", - content_text, - flags = re.DOTALL, + # Safety net caught tool XML -- treat as tool call + tool_calls = _safety_tc + content_text = _strip_tool_markup( + content_accum, + final = True, ) - # Strip unterminated ... to end - content_text = re.sub( - r".*$", - "", - content_text, - flags = re.DOTALL, - ) - # Strip bare ... blocks - content_text = re.sub( - r".*?", - "", - content_text, - flags = re.DOTALL, - ) - # Strip unterminated bare to end - content_text = re.sub( - r".*$", - "", - content_text, - flags = re.DOTALL, - ).strip() logger.info( - f"Parsed {len(tool_calls)} tool call(s) from content text" + f"Safety net: parsed {len(tool_calls)} tool call(s) " + f"from streamed content" ) + else: + # ── DRAINING path: assemble tool_calls ── + tool_calls = None + content_text = content_accum + if has_structured_tc: + # Filter out incomplete fragments (e.g. from + # truncation by max_tokens or disconnect). + tool_calls = [ + tool_calls_acc[i] + for i in sorted(tool_calls_acc) + if ( + tool_calls_acc[i] + .get("function", {}) + .get("name", "") + .strip() + ) + ] or None + if ( + not tool_calls + and auto_heal_tool_calls + and any(s in content_accum for s in _TOOL_XML_SIGNALS) + ): + tool_calls = self._parse_tool_calls_from_text( + content_accum, + ) + if tool_calls and not has_structured_tc: + content_text = _strip_tool_markup( + content_text, + final = True, + ) + if tool_calls: + logger.info( + f"Parsed {len(tool_calls)} tool call(s) from " + f"{'structured delta' if has_structured_tc else 'content text'}" + ) + if not tool_calls: + # DRAINING but no tool calls (false positive). + # Merge accumulated metrics from prior tool + # iterations so they are not silently dropped. + yield {"type": "status", "text": ""} + if content_accum: + # Strip leaked tool-call XML before yielding + content_accum = _strip_tool_markup( + content_accum, final = True + ) + if content_accum: + yield {"type": "content", "text": content_accum} + _fu = _iter_usage or {} + _fc = _fu.get("completion_tokens", 0) + _fp = _fu.get("prompt_tokens", 0) + _tc = _fc + _accumulated_completion_tokens + if ( + _iter_usage + or _iter_timings + or _accumulated_completion_tokens + ): + _mt = dict(_iter_timings) if _iter_timings else {} + if _accumulated_predicted_ms or _accumulated_predicted_n: + _mt["predicted_ms"] = ( + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + _accumulated_predicted_n + ) + _mt["predicted_n"] = _tn + _tms = _mt["predicted_ms"] + if _tms > 0: + _mt["predicted_per_second"] = _tn / (_tms / 1000.0) + yield { + "type": "metadata", + "usage": { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + }, + "timings": _mt, + } + return - if finish_reason == "tool_calls" or (tool_calls and len(tool_calls) > 0): - # Only accumulate metrics for responses that are actually used - _accumulated_completion_tokens += data.get("usage", {}).get( + # ── Execute tool calls ── + _accumulated_completion_tokens += (_iter_usage or {}).get( "completion_tokens", 0 ) - _iter_timings = data.get("timings", {}) - _accumulated_predicted_ms += _iter_timings.get("predicted_ms", 0) - _accumulated_predicted_n += _iter_timings.get("predicted_n", 0) - # Append the assistant message with tool_calls to conversation + _it = _iter_timings or {} + _accumulated_predicted_ms += _it.get("predicted_ms", 0) + _accumulated_predicted_n += _it.get("predicted_n", 0) + assistant_msg = {"role": "assistant", "content": content_text} if tool_calls: assistant_msg["tool_calls"] = tool_calls conversation.append(assistant_msg) - # Execute each tool call for tc in tool_calls or []: func = tc.get("function", {}) tool_name = func.get("name", "") raw_args = func.get("arguments", {}) - # Handle arguments as either string or dict if isinstance(raw_args, str): try: arguments = json.loads(raw_args) @@ -1764,7 +2662,6 @@ class LlamaCppBackend: else: arguments = raw_args - # Yield status update if tool_name == "web_search": status_text = f"Searching: {arguments.get('query', '')}" elif tool_name == "python": @@ -1787,7 +2684,6 @@ class LlamaCppBackend: status_text = f"Calling: {tool_name}" yield {"type": "status", "text": status_text} - # Emit tool_start so the frontend can record inputs yield { "type": "tool_start", "tool_name": tool_name, @@ -1795,19 +2691,28 @@ class LlamaCppBackend: "arguments": arguments, } - # Execute the tool - _effective_timeout = ( - None if tool_call_timeout >= 9999 else tool_call_timeout - ) - result = execute_tool( - tool_name, - arguments, - cancel_event = cancel_event, - timeout = _effective_timeout, - session_id = session_id, - ) + # ── Duplicate call detection ────────────── + if _is_duplicate_call(tool_name, arguments): + result = ( + "You already made this exact call. " + "Do not repeat the same tool call. " + "Try a different approach: fetch a URL " + "from previous results, use Python to " + "process data you already have, or " + "provide your final answer now." + ) + else: + _effective_timeout = ( + None if tool_call_timeout >= 9999 else tool_call_timeout + ) + result = execute_tool( + tool_name, + arguments, + cancel_event = cancel_event, + timeout = _effective_timeout, + session_id = session_id, + ) - # Emit tool_end so the frontend can record outputs yield { "type": "tool_end", "tool_name": tool_name, @@ -1815,11 +2720,32 @@ class LlamaCppBackend: "result": result, } - # Append tool result to conversation + # Nudge model to try a different approach on errors + _error_prefixes = ( + "Error", + "Search failed", + "Execution error", + "Blocked:", + "Exit code", + "Failed to fetch", + "Failed to resolve", + "No query provided", + ) + _is_error = isinstance(result, str) and result.lstrip().startswith( + _error_prefixes + ) + _record_tool_call(tool_name, arguments, failed = _is_error) + _result_content = result + if _is_error: + _result_content = ( + result + "\n\nThe tool call encountered an issue. " + "Please try a different approach or rephrase your request." + ) + tool_msg = { "role": "tool", "name": tool_name, - "content": result, + "content": _result_content, } tool_call_id = tc.get("id") if tool_call_id: @@ -1829,26 +2755,28 @@ class LlamaCppBackend: # Continue the loop to let model respond with context continue - # No tool calls -- model answered directly. - # If no tools were executed at all, just yield the content - # from this response instead of making a redundant second request. - if iteration == 0 and content_text: - yield {"type": "status", "text": ""} - yield {"type": "content", "text": content_text} - _direct_usage = data.get("usage") - _direct_timings = data.get("timings") - if _direct_usage or _direct_timings: - yield { - "type": "metadata", - "usage": _direct_usage, - "timings": _direct_timings, - } - return + except httpx.ConnectError: + raise RuntimeError("Lost connection to llama-server") + except Exception as e: + if cancel_event is not None and cancel_event.is_set(): + return + raise - # Tools were called in previous iterations; do a final - # streaming pass so the model can synthesize a response - # incorporating the tool results. - break + # ── Tool iteration cap reached -- synthesize final answer ── + # The model used all iterations without producing a final text + # response. Inject a nudge so the final streaming pass produces + # a useful answer instead of continuing to request tools. + if max_tool_iterations > 0: + conversation.append( + { + "role": "user", + "content": ( + "You have used all available tool calls. Based on " + "everything you have found so far, provide your final " + "answer now. Do not call any more tools." + ), + } + ) # Clear status yield {"type": "status", "text": ""} @@ -1874,28 +2802,6 @@ class LlamaCppBackend: stream_payload["stop"] = stop stream_payload["stream_options"] = {"include_usage": True} - import re as _re_final - - # Closed blocks only -- safe to strip mid-stream without shrinking later. - _TOOL_CLOSED_PATTERNS = [ - _re_final.compile(r".*?", _re_final.DOTALL), - _re_final.compile(r".*?", _re_final.DOTALL), - ] - # Open-ended patterns strip from an opening tag to end-of-string. - # Only applied on the final flush to avoid non-monotonic shrinking. - _TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [ - _re_final.compile(r".*$", _re_final.DOTALL), - _re_final.compile(r".*$", _re_final.DOTALL), - ] - - def _strip_tool_markup(text: str, *, final: bool = False) -> str: - if not auto_heal_tool_calls: - return text - patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text - cumulative = "" _last_emitted = "" in_thinking = False @@ -1907,9 +2813,16 @@ class LlamaCppBackend: try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, stream_payload, cancel_event + client, + url, + stream_payload, + cancel_event, + headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -2035,7 +2948,10 @@ class LlamaCppBackend: if not self.is_loaded: return None try: - with httpx.Client(timeout = 10) as client: + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) + with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _detok(tid: int) -> str: r = client.post( @@ -2153,7 +3069,12 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - with httpx.Client(timeout = httpx.Timeout(300, connect = 10)) as client: + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) + with httpx.Client( + timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers + ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError( diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 6ff7fd2cbf..f293a0dcd8 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -17,6 +17,7 @@ Pattern follows core/training/training.py. import atexit import base64 +import os import structlog from loggers import get_logger import multiprocessing as mp @@ -27,11 +28,17 @@ import uuid from io import BytesIO from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union +from utils.hardware import prepare_gpu_selection logger = get_logger(__name__) _CTX = mp.get_context("spawn") + +class DownloadStallError(RuntimeError): + """Raised when the worker reports no download progress for too long.""" + + # Dispatcher timeout constants (seconds) _DISPATCH_READ_TIMEOUT = 30.0 _DISPATCH_POLL_INTERVAL = 0.5 @@ -262,12 +269,17 @@ class InferenceOrchestrator: except (EOFError, OSError, ValueError): return None - def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict: + def _wait_response(self, expected_type: str, timeout: float = 300.0) -> dict: """Block until a response of the expected type arrives. Also handles 'status' and 'error' events during the wait. Returns the matching response dict. Raises RuntimeError on timeout or subprocess crash. + + The *timeout* is an **inactivity** timeout: it resets whenever the + subprocess sends a status message, so long-running operations (large + downloads, slow model loads) won't be killed as long as the subprocess + keeps reporting progress. """ deadline = time.monotonic() + timeout @@ -292,8 +304,15 @@ class InferenceOrchestrator: if rtype == "status": logger.info("Subprocess status: %s", resp.get("message", "")) + # Reset deadline — subprocess is still alive and working + deadline = time.monotonic() + timeout continue + if rtype == "stall": + msg = resp.get("message", "Download stalled") + logger.warning("Subprocess reported stall: %s", msg) + raise DownloadStallError(msg) + # Other response types during wait — skip logger.debug( "Skipping response type '%s' while waiting for '%s'", @@ -302,7 +321,8 @@ class InferenceOrchestrator: ) raise RuntimeError( - f"Timeout waiting for '{expected_type}' response after {timeout}s" + f"Timeout waiting for '{expected_type}' response " + f"(no activity for {timeout}s)" ) def _drain_queue(self) -> list: @@ -571,6 +591,7 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """Load a model for inference. @@ -594,7 +615,16 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "gpu_ids": gpu_ids, } + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + gpu_ids, + model_name = model_name, + hf_token = hf_token, + load_in_4bit = load_in_4bit, + ) + sub_config["resolved_gpu_ids"] = resolved_gpu_ids + sub_config["gpu_selection"] = gpu_selection # Always kill existing subprocess and spawn fresh. # Reusing a subprocess after unsloth patches torch internals @@ -608,36 +638,66 @@ class InferenceOrchestrator: # Dead subprocess — clean up self._shutdown_subprocess(timeout = 2) - logger.info( - "Spawning fresh inference subprocess for '%s' (transformers %s.x)", - model_name, - needed_major, + disable_xet = sub_config.get("disable_xet", False) or ( + os.environ.get("HF_HUB_DISABLE_XET") == "1" ) - self._spawn_subprocess(sub_config) - resp = self._wait_response("loaded", timeout = 180) - # Update local state from response - if resp.get("success"): - self._current_transformers_major = needed_major - model_info = resp.get("model_info", {}) - self.active_model_name = model_info.get("identifier", model_name) - self.models[self.active_model_name] = { - "is_vision": model_info.get("is_vision", False), - "is_lora": model_info.get("is_lora", False), - "display_name": model_info.get("display_name", model_name), - "is_audio": model_info.get("is_audio", False), - "audio_type": model_info.get("audio_type"), - "has_audio_input": model_info.get("has_audio_input", False), - } - self.loading_models.discard(model_name) - logger.info("Model '%s' loaded successfully in subprocess", model_name) - return True - else: - error = resp.get("error", "Failed to load model") - self.loading_models.discard(model_name) - self.active_model_name = None - self.models.clear() - raise Exception(error) + for attempt in range(2): + logger.info( + "Spawning fresh inference subprocess for '%s' " + "(transformers %s.x, attempt %d/2%s)", + model_name, + needed_major, + attempt + 1, + ", xet disabled" if disable_xet else "", + ) + sub_config["disable_xet"] = disable_xet + self._spawn_subprocess(sub_config) + + try: + resp = self._wait_response("loaded") + except DownloadStallError: + # First stall and Xet was enabled -> retry with Xet disabled + if attempt == 0 and not disable_xet: + logger.warning( + "Download stalled for '%s' -- retrying with " + "HF_HUB_DISABLE_XET=1", + model_name, + ) + self._shutdown_subprocess(timeout = 5) + disable_xet = True + continue + # Second stall (or already had xet disabled) -> give up + self._shutdown_subprocess(timeout = 5) + raise RuntimeError( + f"Download stalled for '{model_name}' even with " + f"HF_HUB_DISABLE_XET=1 -- check your network connection" + ) + + # Got a response — check success + if resp.get("success"): + self._current_transformers_major = needed_major + model_info = resp.get("model_info", {}) + self.active_model_name = model_info.get("identifier", model_name) + self.models[self.active_model_name] = { + "is_vision": model_info.get("is_vision", False), + "is_lora": model_info.get("is_lora", False), + "display_name": model_info.get("display_name", model_name), + "is_audio": model_info.get("is_audio", False), + "audio_type": model_info.get("audio_type"), + "has_audio_input": model_info.get("has_audio_input", False), + } + self.loading_models.discard(model_name) + logger.info( + "Model '%s' loaded successfully in subprocess", model_name + ) + return True + else: + error = resp.get("error", "Failed to load model") + self.loading_models.discard(model_name) + self.active_model_name = None + self.models.clear() + raise Exception(error) except Exception: self.loading_models.discard(model_name) @@ -661,7 +721,7 @@ class InferenceOrchestrator: "model_name": model_name, } ) - resp = self._wait_response("unloaded", timeout = 30) + resp = self._wait_response("unloaded") # Update local state self.models.pop(model_name, None) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 55bfa095f9..65302fe2f3 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -57,16 +57,23 @@ WEB_SEARCH_TOOL = { "type": "function", "function": { "name": "web_search", - "description": "Search the web for current information, recent events, or facts you are uncertain about.", + "description": ( + "Search the web and fetch page content. Returns snippets for all results. " + "Use the url parameter to fetch full page text from a specific URL." + ), "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "The search query", - } + }, + "url": { + "type": "string", + "description": "A URL to fetch full page content from (instead of searching). Use this to read a page found in search results.", + }, }, - "required": ["query"], + "required": [], }, }, } @@ -131,7 +138,11 @@ def execute_tool( ) effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout if name == "web_search": - return _web_search(arguments.get("query", ""), timeout = effective_timeout) + return _web_search( + arguments.get("query", ""), + url = arguments.get("url"), + timeout = effective_timeout, + ) if name == "python": return _python_exec( arguments.get("code", ""), cancel_event, effective_timeout, session_id @@ -143,9 +154,180 @@ def execute_tool( return f"Unknown tool: {name}" -def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) -> str: - """Search the web using DuckDuckGo and return formatted results.""" - if not query.strip(): +_MAX_PAGE_CHARS = 16000 # limit fetched page text +_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1 # cap raw download size + + +def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]: + """Resolve *hostname*, reject non-public IPs, return a pinned IP string. + + Returns ``(ok, reason_or_empty, resolved_ip)``. The caller should + connect to *resolved_ip* (with a ``Host`` header) to prevent DNS + rebinding between validation and the actual fetch. + """ + import ipaddress + import socket + + try: + infos = socket.getaddrinfo(hostname, port, type = socket.SOCK_STREAM) + except OSError as e: + return False, f"Failed to resolve host: {e}", "" + + if not infos: + return False, f"Failed to resolve host: no addresses for {hostname!r}", "" + + for *_, sockaddr in infos: + ip = ipaddress.ip_address(sockaddr[0]) + if ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + ): + return False, f"Blocked: refusing to fetch non-public address {ip}.", "" + + # Return the first resolved address for pinning + first_ip = infos[0][4][0] + return True, "", first_ip + + +def _fetch_page_text( + url: str, max_chars: int = _MAX_PAGE_CHARS, timeout: int = 30 +) -> str: + """Fetch a URL and return plain text content (HTML tags stripped). + + Blocks private/loopback/link-local targets (SSRF protection) and caps + the download size to avoid unbounded memory usage. + """ + import re as _re + from urllib.parse import urlparse + + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r})." + if not parsed.hostname: + return "Blocked: URL is missing a hostname." + + port = parsed.port or (443 if parsed.scheme == "https" else 80) + ok, reason, pinned_ip = _validate_and_resolve_host(parsed.hostname, port) + if not ok: + return reason + + try: + import urllib.request + from urllib.error import HTTPError as _HTTPError + from urllib.parse import urljoin, urlunparse + + # Disable auto-redirect so we can validate each hop for SSRF. + # urllib raises HTTPError for 3xx when the handler returns None, + # so we catch that and extract the Location header manually. + class _NoRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None + + opener = urllib.request.build_opener(_NoRedirect) + max_bytes = max_chars * 4 + 1 + current_url = url + current_host = parsed.hostname + + for _hop in range(5): + # Pin to the validated IP to prevent DNS rebinding. + # Rewrite the URL to use the IP and set the Host header. + cp = urlparse(current_url) + ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip + pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + + req = urllib.request.Request( + pinned_url, + headers = { + "User-Agent": "UnslothStudio/1.0", + "Host": current_host, + }, + ) + try: + resp = opener.open(req, timeout = timeout) + except _HTTPError as e: + if e.code not in (301, 302, 303, 307, 308): + return ( + f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}" + ) + location = e.headers.get("Location") + if not location: + return "Failed to fetch URL: redirect missing Location header." + current_url = urljoin(current_url, location) + rp = urlparse(current_url) + if rp.scheme not in ("http", "https") or not rp.hostname: + return "Blocked: redirect target is not a valid http/https URL." + rp_port = rp.port or (443 if rp.scheme == "https" else 80) + ok2, reason2, pinned_ip = _validate_and_resolve_host( + rp.hostname, + rp_port, + ) + if not ok2: + return reason2 + current_host = rp.hostname + continue + # Success -- read capped body + raw_bytes = resp.read(max_bytes) + break + else: + return "Failed to fetch URL: too many redirects." + + charset = resp.headers.get_content_charset() or "utf-8" + raw_html = raw_bytes.decode(charset, errors = "replace") + except _HTTPError as e: + return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}" + except Exception as e: + return f"Failed to fetch URL: {e}" + + # Convert HTML to text -- prefer html2text for clean markdown output + try: + import html2text as _h2t + + converter = _h2t.HTML2Text() + converter.ignore_links = False + converter.ignore_images = True + converter.body_width = 0 # no wrapping + text = converter.handle(raw_html).strip() + except ImportError: + # Fallback: regex-based stripping + text = _re.sub( + r"]*>.*?]*>", + "", + raw_html, + flags = _re.DOTALL | _re.IGNORECASE, + ) + text = _re.sub( + r"]*>.*?]*>", "", text, flags = _re.DOTALL | _re.IGNORECASE + ) + text = _re.sub(r"<[^>]+>", " ", text) + text = _re.sub(r"\s+", " ", text).strip() + + if not text: + return "(page returned no readable text)" + if len(text) > max_chars: + text = text[:max_chars] + f"\n\n... (truncated, {len(text)} chars total)" + return text + + +def _web_search( + query: str, + max_results: int = 5, + timeout: int = _EXEC_TIMEOUT, + url: str | None = None, +) -> str: + """Search the web using DuckDuckGo and return formatted results. + + If ``url`` is provided, fetches that page directly instead of searching. + """ + # Direct URL fetch mode + if url and url.strip(): + fetch_timeout = 60 if timeout is None else min(timeout, 60) + return _fetch_page_text(url.strip(), timeout = fetch_timeout) + + if not query or not query.strip(): return "No query provided." try: from ddgs import DDGS @@ -160,7 +342,13 @@ def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) f"URL: {r.get('href', '')}\n" f"Snippet: {r.get('body', '')}" ) - return "\n\n---\n\n".join(parts) + text = "\n\n---\n\n".join(parts) + text += ( + "\n\n---\n\nIMPORTANT: These are only short snippets. " + "To get the full page content, call web_search with " + 'the url parameter (e.g. {"url": ""}).' + ) + return text except Exception as e: return f"Search failed: {e}" diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 2eb46f3217..e2513f43de 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -22,6 +22,7 @@ from loggers import get_logger import os import queue as _queue import sys +import threading import time import traceback from io import BytesIO @@ -29,6 +30,7 @@ from pathlib import Path from typing import Any logger = get_logger(__name__) +from utils.hardware import apply_gpu_ids def _activate_transformers_version(model_name: str) -> None: @@ -113,6 +115,154 @@ def _build_model_config(config: dict): return mc +def _get_hf_download_state( + model_names: list[str] | None = None, +) -> tuple[int, bool] | None: + """Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error. + + When *model_names* is provided, only those models' ``blobs/`` + directories are checked instead of scanning every cached model -- + much faster on systems with many models. Accepts multiple names so + that LoRA loads can watch both the adapter repo and the base model + repo simultaneously. + + *has_incomplete* is True when any ``*.incomplete`` files exist in the + watched blobs directories, indicating that ``huggingface_hub`` is + actively downloading. + + Returns None if the state cannot be determined (import error, + permission error, etc.) so callers can skip stall logic. + """ + try: + from huggingface_hub.constants import HF_HUB_CACHE + + cache = Path(HF_HUB_CACHE) + if not cache.exists(): + return (0, False) + + total = 0 + has_incomplete = False + blobs_dirs: list[Path] = [] + + if model_names: + for name in model_names: + if not name: + continue + # Skip local filesystem paths -- HF model IDs use forward + # slashes (org/model) but never start with / . ~ or contain + # backslashes. This distinguishes them from absolute paths, + # relative paths, and Windows paths. + if name.startswith(("/", ".", "~")) or "\\" in name: + continue + # HF cache dir format: models--org--name (slashes -> --) + cache_dir_name = "models--" + name.replace("/", "--") + blobs_dir = cache / cache_dir_name / "blobs" + if blobs_dir.exists(): + blobs_dirs.append(blobs_dir) + else: + blobs_dirs = list(cache.glob("models--*/blobs")) + + for bdir in blobs_dirs: + for f in bdir.iterdir(): + try: + if f.is_file(): + total += f.stat().st_size + if f.name.endswith(".incomplete"): + has_incomplete = True + except OSError: + pass + + return (total, has_incomplete) + except Exception as e: + logger.debug("Failed to determine HF download state: %s", e) + return None + + +def _start_heartbeat( + resp_queue: Any, + interval: float = 30.0, + stall_timeout: float = 180.0, + xet_disabled: bool = False, + model_names: list[str] | None = None, +) -> threading.Event: + """Start a daemon thread that sends periodic status heartbeats. + + Monitors the HF Hub cache directory for download activity. A stall + is only reported when ``*.incomplete`` files are present (indicating + ``huggingface_hub`` is actively downloading) **and** the total cache + size has not changed for *stall_timeout* seconds. + + Once the download finishes (no more ``.incomplete`` files), the stall + timer resets, so post-download initialization (quantization, GPU + weight loading) is never misclassified as a stalled download. + + Returns a stop event -- set it to terminate the heartbeat thread. + """ + stop = threading.Event() + transport = "https" if xet_disabled else "xet" + + def _beat(): + state = _get_hf_download_state(model_names) + last_size = state[0] if state is not None else 0 + last_change = time.monotonic() + + while not stop.wait(interval): + state = _get_hf_download_state(model_names) + now = time.monotonic() + + # Skip stall logic if we cannot measure the cache + if state is None: + _send_response( + resp_queue, + { + "type": "status", + "message": f"Loading model ({transport} transport)...", + "ts": time.time(), + }, + ) + continue + + current_size, has_incomplete = state + + if current_size != last_size: + last_size = current_size + last_change = now + + # Only fire stall when .incomplete files are present, + # confirming a download is actively in progress. + # Once downloads finish (no .incomplete), reset the timer + # so model init time is not counted as a stall. + if not has_incomplete: + last_change = now + elif now - last_change >= stall_timeout: + _send_response( + resp_queue, + { + "type": "stall", + "message": ( + f"Download appears stalled ({transport} transport) " + f"-- no progress for {int(now - last_change)}s" + ), + "ts": time.time(), + }, + ) + # Only fire once -- the orchestrator will kill us + return + + _send_response( + resp_queue, + { + "type": "status", + "message": f"Loading model ({transport} transport)...", + "ts": time.time(), + }, + ) + + t = threading.Thread(target = _beat, daemon = True) + t.start() + return stop + + def _handle_load(backend, config: dict, resp_queue: Any) -> None: """Handle a load command: load a model into the backend.""" try: @@ -156,13 +306,50 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: except Exception as e: logger.warning("Could not read adapter_config.json: %s", e) - success = backend.load_model( - config = mc, - max_seq_length = config.get("max_seq_length", 2048), - load_in_4bit = load_in_4bit, - hf_token = hf_token, - trust_remote_code = config.get("trust_remote_code", False), + # Auto-enable trust_remote_code for unsloth/* transformers 5.x models + # (matches the training worker logic in core/training/worker.py) + trust_remote_code = config.get("trust_remote_code", False) + if not trust_remote_code: + from utils.transformers_version import needs_transformers_5 + + model_name = config["model_name"] + if needs_transformers_5(model_name) and model_name.lower().startswith( + "unsloth/" + ): + trust_remote_code = True + logger.info( + "Auto-enabled trust_remote_code for unsloth/* transformers 5.x model: %s", + model_name, + ) + + # Send heartbeats every 30s so the orchestrator knows we're still alive + # (download / weight loading can take a long time on slow connections) + xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1" + + # Watch both the model repo and base model repo (for LoRA loads + # where the base model download is the actual bottleneck) + watch_repos = [mc.identifier] + base = getattr(mc, "base_model", None) + if base and str(base) != mc.identifier: + watch_repos.append(str(base)) + + heartbeat_stop = _start_heartbeat( + resp_queue, + interval = 30.0, + xet_disabled = xet_disabled, + model_names = watch_repos, ) + try: + success = backend.load_model( + config = mc, + max_seq_length = config.get("max_seq_length", 2048), + load_in_4bit = load_in_4bit, + hf_token = hf_token, + trust_remote_code = trust_remote_code, + gpu_ids = config.get("resolved_gpu_ids"), + ) + finally: + heartbeat_stop.set() if success: # Build model_info for the parent to mirror @@ -474,6 +661,10 @@ def run_inference_process( "ignore" # Suppress warnings at C-level before imports ) + if config.get("disable_xet"): + os.environ["HF_HUB_DISABLE_XET"] = "1" + logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)") + import warnings from loggers.config import LogConfig @@ -485,6 +676,8 @@ def run_inference_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) + apply_gpu_ids(config.get("resolved_gpu_ids")) + model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 2324916236..ab1825d94a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -33,7 +33,14 @@ if sys.platform in ("win32", "darwin"): sys.path.insert(0, _compile_cache) import torch -from utils.hardware import clear_gpu_cache, safe_num_proc, dataset_map_num_proc +from utils.hardware import ( + clear_gpu_cache, + safe_num_proc, + dataset_map_num_proc, + get_device_map, + raise_if_offloaded, + get_visible_gpu_count, +) torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported @@ -487,6 +494,7 @@ class UnslothTrainer: is_dataset_audio: bool = False, trust_remote_code: bool = False, full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """Load model for training (supports both text and vision models)""" self.load_in_4bit = load_in_4bit # Store for training_meta.json @@ -624,6 +632,11 @@ class UnslothTrainer: self._update_progress(error = friendly, is_training = False) return False + device_map = get_device_map(gpu_ids) + logger.info( + f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" + ) + # Branch based on model type if self._audio_type == "csm": # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False @@ -636,6 +649,7 @@ class UnslothTrainer: dtype = None, auto_model = CsmForConditionalGeneration, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -651,6 +665,7 @@ class UnslothTrainer: model_name = model_name, dtype = None, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, auto_model = WhisperForConditionalGeneration, whisper_language = "English", @@ -672,6 +687,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -711,6 +727,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = torch.float32, # Spark-TTS requires float32 load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -725,6 +742,7 @@ class UnslothTrainer: model_name, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -741,6 +759,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -754,6 +773,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, # Auto-detect load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -786,12 +806,15 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, # Auto-detect load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, ) logger.info("Loaded text model") + raise_if_offloaded(self.model, device_map, "Studio training") + if self.should_stop: return False @@ -824,6 +847,7 @@ class UnslothTrainer: is_dataset_audio = is_dataset_audio, trust_remote_code = trust_remote_code, full_finetuning = full_finetuning, + gpu_ids = gpu_ids, ) error_msg = str(e) error_lower = error_msg.lower() @@ -2634,14 +2658,14 @@ class UnslothTrainer: eval_steps: float = 0.00, output_dir: str | None = None, num_epochs: int = 3, - learning_rate: float = 5e-5, + learning_rate: float = 2e-4, batch_size: int = 2, gradient_accumulation_steps: int = 4, warmup_steps: int = None, warmup_ratio: float = None, max_steps: int = 0, save_steps: int = 0, - weight_decay: float = 0.01, + weight_decay: float = 0.001, random_seed: int = 3407, packing: bool = False, train_on_completions: bool = False, @@ -3010,7 +3034,7 @@ class UnslothTrainer: "fp16": not is_bfloat16_supported(), "bf16": is_bfloat16_supported(), "logging_steps": 1, - "weight_decay": training_args.get("weight_decay", 0.01), + "weight_decay": training_args.get("weight_decay", 0.001), "seed": training_args.get("random_seed", 3407), "output_dir": output_dir, "report_to": _build_report_targets(training_args), diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 4439e4e173..f35c7e8ad3 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -28,6 +28,7 @@ from pathlib import Path from typing import Optional, Tuple, Any import matplotlib.pyplot as plt +from utils.hardware import prepare_gpu_selection logger = get_logger(__name__) @@ -159,7 +160,7 @@ class TrainingBackend: "warmup_ratio": kwargs.get("warmup_ratio"), "max_steps": kwargs.get("max_steps", 0), "save_steps": kwargs.get("save_steps", 0), - "weight_decay": kwargs.get("weight_decay", 0.01), + "weight_decay": kwargs.get("weight_decay", 0.001), "random_seed": kwargs.get("random_seed", 3407), "packing": kwargs.get("packing", False), "optim": kwargs.get("optim", "adamw_8bit"), @@ -185,6 +186,7 @@ class TrainingBackend: "enable_tensorboard": kwargs.get("enable_tensorboard", False), "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), "trust_remote_code": kwargs.get("trust_remote_code", False), + "gpu_ids": kwargs.get("gpu_ids"), } # Derive load_in_4bit from training_type @@ -192,6 +194,22 @@ class TrainingBackend: config["load_in_4bit"] = False # Spawn subprocess — use locals so state is untouched on failure + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + kwargs.get("gpu_ids"), + model_name = config["model_name"], + hf_token = config["hf_token"] or None, + training_type = config["training_type"], + load_in_4bit = config["load_in_4bit"], + batch_size = config.get("batch_size", 4), + max_seq_length = config.get("max_seq_length", 2048), + lora_rank = config.get("lora_r", 16), + target_modules = config.get("target_modules"), + gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), + optimizer = config.get("optim", "adamw_8bit"), + ) + config["resolved_gpu_ids"] = resolved_gpu_ids + config["gpu_selection"] = gpu_selection + from .worker import run_training_process event_queue = _CTX.Queue() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 891dfca8f7..699cfe74f7 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -29,6 +29,7 @@ import urllib.error import urllib.request logger = get_logger(__name__) +from utils.hardware import apply_gpu_ids _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" @@ -367,6 +368,8 @@ def run_training_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) + apply_gpu_ids(config.get("resolved_gpu_ids")) + model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -682,6 +685,7 @@ def run_training_process( is_dataset_image = config.get("is_dataset_image", False), is_dataset_audio = config.get("is_dataset_audio", False), trust_remote_code = config.get("trust_remote_code", False), + gpu_ids = config.get("resolved_gpu_ids"), ) if not success or trainer.should_stop: if trainer.should_stop: @@ -791,7 +795,7 @@ def run_training_process( warmup_ratio = config.get("warmup_ratio"), max_steps = max_steps if max_steps and max_steps > 0 else 0, save_steps = save_steps if save_steps and save_steps > 0 else 0, - weight_decay = config.get("weight_decay", 0.01), + weight_decay = config.get("weight_decay", 0.001), random_seed = config.get("random_seed", 3407), packing = config.get("packing", False), train_on_completions = config.get("train_on_completions", False), @@ -1137,7 +1141,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> "lr_scheduler_type": config.get("lr_scheduler_type", "linear"), "batch_sampler": BatchSamplers.NO_DUPLICATES, "optim": config.get("optim", "adamw_8bit"), - "weight_decay": config.get("weight_decay", 0.01), + "weight_decay": config.get("weight_decay", 0.001), "seed": config.get("random_seed", 3407), } diff --git a/studio/backend/main.py b/studio/backend/main.py index 5e647f6312..c18f18a743 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -23,10 +23,23 @@ if _backend_dir not in sys.path: # See: https://github.com/python/cpython/issues/102396 import _platform_compat # noqa: F401 +import mimetypes import shutil import warnings from contextlib import asynccontextmanager +# Fix broken Windows registry MIME types. Some Windows installs map .js to +# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes +# module reads from the registry, and FastAPI/Starlette's StaticFiles uses +# mimetypes.guess_type() to set Content-Type headers. Browsers enforce strict +# MIME checking for ES module scripts (