Elgato_dark/unsloth
1
0
Fork 0
mirror of https://github.com/unslothai/unsloth synced 2026-04-21 13:37:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Consolidate dual venvs and separate install from update (#4530)

Browse source 
* refactor: consolidate dual venvs into single ~/.unsloth/studio/unsloth_studio

* refactor: separate install.sh (first-time) from setup.sh (smart update with PyPI version check)

* fix: install.sh calls setup.sh directly, keep both setup and update CLI commands

* fix: use importlib.resources.files() directly without _path attribute

* fix: bootstrap uv before pip upgrade to handle uv venvs without pip

* fix: frontend 404 when launched via CLI, add global symlink to ~/.local/bin

* feat: add --local flag to install.sh and unsloth studio update for branch testing

* fix: resolve repo root from script location for --local installs

* feat: add --package flag to install.sh for testing with custom package names

* feat: add --package flag to unsloth studio update

* fix: always nuke venv in install.sh for clean installs

* revert: remove Windows changes, will handle in separate PR

* fix: error when --package is passed without an argument

* revert: restore Windows scripts to current main

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: always explicitly set STUDIO_LOCAL_INSTALL and STUDIO_PACKAGE_NAME env vars

* fix: pass explicit STUDIO_LOCAL_REPO env var for --local installs

* fix: align banner box for Setup vs Update labels

* deprecate: hide 'unsloth studio setup' command, point users to update/install.sh

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: check stdout not stdin for auto-launch detection (curl pipe fix)

* fix: update install URL to unsloth.ai/install.sh

* fix: update install.sh usage comments to unsloth.ai/install.sh

* fix: use --upgrade-package for base deps to preserve existing torch/CUDA installs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: --local install now also installs unsloth-zoo via base.txt before editable overlay

* fix: don't skip base packages for --local installs (editable needs unsloth-zoo)

* refactor: move --local full dep install to install.sh, keep SKIP_STUDIO_BASE for all paths

* feat: add migration support for old .venv and CWD-based installs in setup.sh

* Revert "feat: add migration support for old .venv and CWD-based installs in setup.sh"

This reverts commit 301291d002.

* feat: migrate old .venv layout in install.sh instead of always nuking

* feat: validate old .venv with torch CUDA test before migration, recovery message on launch failure

* fix: try CUDA then fall back to CPU for migration validation

* fix: upgrade unsloth/unsloth-zoo with --reinstall-package on migration to preserve torch

* remove: delete unused unsloth ui command (use unsloth studio instead)

* Fix Windows venv path mismatch between install.ps1, setup.ps1, and studio.py

install.ps1 was creating the venv CWD-relative ($VenvName = "unsloth_studio"),
setup.ps1 was using an absolute path to ".unsloth\studio\.venv", and studio.py
looks for ".unsloth\studio\unsloth_studio". All three paths were different, so
the Windows installer would never produce a working Studio setup.

install.ps1:
- Use absolute $StudioHome + $VenvDir matching the Linux install.sh layout
- Add 3-way migration: old .venv at STUDIO_HOME, CWD-relative ~/unsloth_studio
  from the previous install.ps1, or fresh creation with torch validation
- For migrated envs, upgrade unsloth while preserving existing torch/CUDA wheels
- Set SKIP_STUDIO_BASE=1 before calling setup.ps1 (matches install.sh behavior)
- Fix launch instructions to use the absolute venv path

setup.ps1:
- Change $VenvDir from ".unsloth\studio\.venv" to ".unsloth\studio\unsloth_studio"
- Add SKIP_STUDIO_BASE guard: error out if venv is missing when called from
  install.ps1 (which should have already created it)
- Differentiate "Setup" vs "Update" in banners based on SKIP_STUDIO_BASE

* setup.ps1: unconditionally error if venv missing, matching setup.sh

setup.sh always errors out if the venv does not exist (line 224-228),
telling the user to run install.sh first. setup.ps1 was conditionally
creating a bare venv with python -m venv when SKIP_STUDIO_BASE was not
set, which would produce an empty venv with no torch or unsloth. Now
setup.ps1 matches setup.sh: always error, always point to install.ps1.

* Fix --torch-backend=auto CPU solver dead-end on Linux, macOS, and Windows

On CPU-only machines, `uv pip install unsloth --torch-backend=auto`
falls back to unsloth==2024.8 because the CPU solver cannot satisfy
newer unsloth's dependencies. install.ps1 already solved this with a
two-step approach; this applies the same fix to install.sh and
install_python_stack.py.

install.sh: add get_torch_index_url() that detects GPU via nvidia-smi
and maps CUDA versions to PyTorch index URLs (matching install.ps1's
Get-TorchIndexUrl). Fresh installs now install torch first via explicit
--index-url, then install unsloth with --upgrade-package to preserve
the pre-installed torch. All 5 --torch-backend=auto removed from
primary paths.

install.ps1: add fallback else-branch when TorchIndexUrl is empty,
using --torch-backend=auto as last resort (matching install.sh).

install_python_stack.py: remove unconditional --torch-backend=auto
from _build_uv_cmd. Torch is pre-installed by install.sh/setup.ps1
by the time this runs. Callers that need it can set UV_TORCH_BACKEND.

Both install.sh and install.ps1 now share the same three-branch logic:
migrated env (upgrade-package only), normal (torch-first + index-url),
and fallback (--torch-backend=auto if URL detection fails).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use --reinstall-package for migrated envs on both Linux and Windows

For migrated environments (moved from legacy venv location),
--reinstall-package is better than --upgrade-package because it forces
a clean reinstall even if the same version is already installed. This
ensures proper .dist-info and .pyc state in the new venv location.

--upgrade-package remains correct for the fresh install path where
torch is already installed and we just want to add unsloth without
re-resolving torch.

* Address review findings: portability, parity, and stale comments

- Replace grep -oP (GNU Perl regex) with POSIX sed in
  get_torch_index_url() so the script works on BSD grep (macOS is
  already guarded by the Darwin early-return, but Alpine/BusyBox
  would silently get the wrong CUDA tag)
- Add LC_ALL=C before nvidia-smi invocation to prevent locale-dependent
  output parsing issues
- Add warning on stderr when nvidia-smi output is unparseable, matching
  install.ps1's [WARN] message
- Add explicit unsloth-zoo positional arg to install.ps1 migrated path,
  matching install.sh (--reinstall-package alone won't install it if it
  was never present in the migrated env)
- Fix stale comment in install_python_stack.py line 392 that still
  claimed --torch-backend=auto is added by _build_uv_cmd
- Add sed to test tools directory (function now uses sed instead of grep)

* Add --index-url to migrated env path to prevent CPU torch resolution

The migrated path runs uv pip install with --reinstall-package for
unsloth/unsloth-zoo. While uv should keep existing torch as satisfied,
the resolver could still re-resolve torch as a transitive dependency.
Without --index-url pointing at the correct CUDA wheel index, the
resolver would fall back to plain PyPI and potentially pull CPU-only
torch. Adding --index-url $TORCH_INDEX_URL ensures CUDA wheels are
available if the resolver needs them.

Applied to both install.sh and install.ps1.

* Revert --index-url on migrated env path

The original install.ps1 on main already handles the migrated path
without --index-url and it works correctly. --reinstall-package only
forces reinstall of the named packages while uv keeps existing torch
as satisfied. No need for the extra flag.

* Fix unsloth studio update --local not installing local checkout

studio.py sets STUDIO_LOCAL_REPO when --local is passed, but
install_python_stack.py never read it. The update path always
installed from PyPI regardless of the --local flag.

Add a local_repo branch that first updates deps from base.txt
(with --upgrade-package to preserve torch), then overlays the
local checkout as an editable install with --no-deps.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Roland Tannous 2026-03-25 16:24:21 +04:00 committed by GitHub
parent 3446e0c489
commit 19e9c60a8e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 877 additions and 301 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

88
install.ps1
View file

@ -5,8 +5,9 @@
function Install-UnslothStudio {
$ErrorActionPreference = "Stop"
$VenvName = "unsloth_studio"
$PythonVersion = "3.13"
$StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
$VenvDir = Join-Path $StudioHome "unsloth_studio"
Write-Host ""
Write-Host "========================================="
@ -449,20 +450,59 @@ shell.Run cmd, 0, False
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
Write-Host "==> 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"
Write-Host "==> 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) {
Write-Host " Legacy environment is healthy -- migrating..."
Move-Item -Path $OldVenv -Destination $VenvDir -Force
Write-Host " Moved .venv -> unsloth_studio"
$_Migrated = $true
} else {
Write-Host " Legacy environment failed validation -- creating fresh environment"
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"
Write-Host "==> Found CWD-relative Studio environment, migrating to $VenvDir..."
Move-Item -Path $CwdVenv -Destination $VenvDir -Force
Write-Host " 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)"
Write-Host "==> Creating Python $($DetectedPython.Version) virtual environment ($VenvDir)..."
uv venv $VenvDir --python "$($DetectedPython.Path)"
if ($LASTEXITCODE -ne 0) {
Write-Host "[ERROR] Failed to create virtual environment (exit code $LASTEXITCODE)" -ForegroundColor Red
return
}
} else {
Write-Host "==> Virtual environment ${VenvName} already exists, skipping creation."
Write-Host "==> Using migrated environment at $VenvDir"
}
# ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ──
@ -536,15 +576,26 @@ 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
}
if ($_Migrated) {
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
Write-Host "==> Upgrading unsloth in migrated environment..."
uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.11" unsloth-zoo
} elseif ($TorchIndexUrl) {
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
}
Write-Host "==> Installing unsloth (this may take a few minutes)..."
uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11"
Write-Host "==> Installing unsloth (this may take a few minutes)..."
uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11"
} else {
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
Write-Host "==> Installing unsloth (this may take a few minutes)..."
uv pip install --python $VenvPython "unsloth>=2026.3.11" --torch-backend=auto
}
if ($LASTEXITCODE -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red
return
@ -554,7 +605,7 @@ shell.Run cmd, 0, False
# 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"
$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,6 +613,8 @@ 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
}
# Tell setup.ps1 to skip base package installation (install.ps1 already did it)
$env:SKIP_STUDIO_BASE = "1"
& $UnslothExe studio setup
if ($LASTEXITCODE -ne 0) {
Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red
@ -582,12 +635,11 @@ shell.Run cmd, 0, False
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 " & `"$VenvDir\Scripts\Activate.ps1`""
Write-Host " unsloth studio -H 0.0.0.0 -p 8888"
Write-Host ""
}

261
install.sh
View file

@ -1,11 +1,35 @@
#!/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 (test): ./install.sh --package roland-sloth (install a different package name)
set -e
VENV_NAME="unsloth_studio"
# ── Parse flags ──
STUDIO_LOCAL_INSTALL=false
PACKAGE_NAME="unsloth"
_next_is_package=false
for arg in "$@"; do
if [ "$_next_is_package" = true ]; then
PACKAGE_NAME="$arg"
_next_is_package=false
continue
fi
case "$arg" in
--local) STUDIO_LOCAL_INSTALL=true ;;
--package) _next_is_package=true ;;
esac
done
if [ "$_next_is_package" = true ]; then
echo "❌ ERROR: --package requires an argument." >&2
exit 1
fi
PYTHON_VERSION="3.13"
STUDIO_HOME="$HOME/.unsloth/studio"
VENV_DIR="$STUDIO_HOME/unsloth_studio"
# ── Helper: download a URL to a file (supports curl and wget) ──
download() {
@ -659,32 +683,195 @@ if ! command -v uv >/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
echo "==> 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 [ ! -x "$VENV_DIR/bin/python" ]; then
echo "==> Creating Python ${PYTHON_VERSION} virtual environment (${VENV_DIR})..."
uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
else
echo "==> Using migrated environment at ${VENV_DIR}"
fi
# ── Resolve repo root (for --local installs) ──
_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
# ── 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)
# ── 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
echo "==> Upgrading unsloth in migrated environment..."
uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.3.11" unsloth-zoo
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
echo "==> Overlaying local repo (editable)..."
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
echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
uv pip install --python "$_VENV_PY" torch torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
# Fresh: Step 2 - install unsloth, preserving pre-installed torch
echo "==> Installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.3.11" unsloth-zoo
echo "==> Overlaying local repo (editable)..."
uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
else
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
echo "==> Installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.11" --torch-backend=auto
echo "==> Overlaying local repo (editable)..."
uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
else
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
echo "==> Running unsloth setup..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
SKIP_STUDIO_BASE=1 \
STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
STUDIO_LOCAL_INSTALL=1 \
STUDIO_LOCAL_REPO="$_REPO_ROOT" \
bash "$SETUP_SH" </dev/null
else
SKIP_STUDIO_BASE=1 \
STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
bash "$SETUP_SH" </dev/null
fi
# ── Make 'unsloth' available globally via ~/.local/bin ──
mkdir -p "$HOME/.local/bin"
ln -sf "$VENV_DIR/bin/unsloth" "$HOME/.local/bin/unsloth"
_LOCAL_BIN="$HOME/.local/bin"
case ":$PATH:" in
*":$_LOCAL_BIN:"*) ;; # already on PATH
*)
_SHELL_PROFILE=""
if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then
_SHELL_PROFILE="$HOME/.zshrc"
elif [ -f "$HOME/.bashrc" ]; then
_SHELL_PROFILE="$HOME/.bashrc"
elif [ -f "$HOME/.profile" ]; then
_SHELL_PROFILE="$HOME/.profile"
fi
if [ -n "$_SHELL_PROFILE" ]; then
if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then
echo '' >> "$_SHELL_PROFILE"
echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE"
echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE"
echo "==> Added ~/.local/bin to PATH in $_SHELL_PROFILE"
fi
fi
export PATH="$_LOCAL_BIN:$PATH"
;;
esac
create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS"
@ -694,8 +881,32 @@ echo " Unsloth Studio installed!"
echo "========================================="
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
echo "==> Launching Unsloth Studio..."
echo ""
"$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
echo " To launch, run:"
echo ""
echo " unsloth studio -H 0.0.0.0 -p 8888"
echo ""
echo " Or activate the environment first:"
echo ""
echo " source ${VENV_DIR}/bin/activate"
echo " unsloth studio -H 0.0.0.0 -p 8888"
echo ""
fi

2
studio/backend/colab.py
View file

@ -26,7 +26,7 @@ def _bootstrap_studio_venv() -> None:
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"
venv_lib = Path.home() / ".unsloth" / "studio" / "unsloth_studio" / "lib"
if not venv_lib.exists():
import warnings

114
studio/install_python_stack.py
View file

@ -225,9 +225,20 @@ def _translate_pip_args_for_uv(args: tuple[str, ...]) -> list[str]:
def _build_pip_cmd(args: tuple[str, ...]) -> list[str]:
"""Build a standard pip install command."""
"""Build a standard pip install command.
Strips uv-only flags like --upgrade-package that pip doesn't understand.
"""
cmd = [sys.executable, "-m", "pip", "install"]
cmd.extend(args)
skip_next = False
for arg in args:
if skip_next:
skip_next = False
continue
if arg == "--upgrade-package":
skip_next = True # skip the flag and its value
continue
cmd.append(arg)
return cmd
@ -241,7 +252,12 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
# the system Python (observed on Colab and similar environments).
cmd.extend(["--python", sys.executable])
cmd.extend(_translate_pip_args_for_uv(args))
cmd.append("--torch-backend=auto")
# Torch is pre-installed by install.sh/setup.ps1. Do not add
# --torch-backend by default -- it can cause solver dead-ends on
# CPU-only machines. Callers that need it can set UV_TORCH_BACKEND.
_tb = os.environ.get("UV_TORCH_BACKEND", "")
if _tb:
cmd.append(f"--torch-backend={_tb}")
return cmd
@ -325,22 +341,90 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
def install_python_stack() -> int:
global USE_UV, _STEP, _TOTAL
_STEP = 0
_TOTAL = 10 if IS_WINDOWS else 11
# 1. Upgrade pip (needed even with uv as fallback and for bootstrapping)
_progress("pip upgrade")
run("Upgrading pip", [sys.executable, "-m", "pip", "install", "--upgrade", "pip"])
# When called from install.sh (which already installed unsloth into the venv),
# SKIP_STUDIO_BASE=1 is set to avoid redundant reinstallation of base packages.
# When called from "unsloth studio update", it is NOT set so base packages
# (unsloth + unsloth-zoo) are always reinstalled to pick up new versions.
skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1"
# When --package is used, install a different package name (e.g. roland-sloth for testing)
package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth")
# When --local is used, overlay a local repo checkout after updating deps
local_repo = os.environ.get("STUDIO_LOCAL_REPO", "")
base_total = 10 if IS_WINDOWS else 11
_TOTAL = (base_total - 1) if skip_base else base_total
# Try to use uv for faster installs
# 1. Try to use uv for faster installs (must happen before pip upgrade
# because uv venvs don't include pip by default)
USE_UV = _bootstrap_uv()
# 2. Core packages: unsloth-zoo + unsloth
_progress("base packages")
pip_install(
"Installing base packages",
"--no-cache-dir",
req = REQ_ROOT / "base.txt",
)
# 2. Ensure pip is available (uv venvs created by install.sh don't include pip)
_progress("pip bootstrap")
if USE_UV:
run(
"Bootstrapping pip via uv",
[
"uv",
"pip",
"install",
"--python",
sys.executable,
"pip",
],
)
else:
run(
"Upgrading pip",
[sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
)
# 3. Core packages: unsloth-zoo + unsloth (or custom package name)
if skip_base:
print(_green(f"✅ {package_name} already installed — skipping base packages"))
elif local_repo:
# Local dev install: update deps from base.txt, then overlay the
# local checkout as an editable install (--no-deps so torch is
# never re-resolved).
_progress("base packages")
pip_install(
"Updating base packages",
"--no-cache-dir",
"--upgrade-package",
"unsloth",
"--upgrade-package",
"unsloth-zoo",
req = REQ_ROOT / "base.txt",
)
pip_install(
"Overlaying local repo (editable)",
"--no-cache-dir",
"--no-deps",
"-e",
local_repo,
constrain = False,
)
elif package_name != "unsloth":
# Custom package name (e.g. roland-sloth for testing) — install directly
_progress("base packages")
pip_install(
f"Installing {package_name}",
"--no-cache-dir",
package_name,
)
else:
# Update path: upgrade only unsloth + unsloth-zoo while preserving
# existing torch/CUDA installations. Torch is pre-installed by
# install.sh / setup.ps1; --upgrade-package targets only base pkgs.
_progress("base packages")
pip_install(
"Updating base packages",
"--no-cache-dir",
"--upgrade-package",
"unsloth",
"--upgrade-package",
"unsloth-zoo",
req = REQ_ROOT / "base.txt",
)
# 3. Extra dependencies
_progress("unsloth extras")

27
studio/setup.ps1
View file

@ -250,9 +250,15 @@ function Find-VsBuildTools {
# ─────────────────────────────────────────────
# Banner
# ─────────────────────────────────────────────
Write-Host "+==============================================+" -ForegroundColor Green
Write-Host "| Unsloth Studio Setup (Windows) |" -ForegroundColor Green
Write-Host "+==============================================+" -ForegroundColor Green
if ($env:SKIP_STUDIO_BASE -eq "1") {
Write-Host "+==============================================+" -ForegroundColor Green
Write-Host "| Unsloth Studio Setup (Windows) |" -ForegroundColor Green
Write-Host "+==============================================+" -ForegroundColor Green
} else {
Write-Host "+==============================================+" -ForegroundColor Green
Write-Host "| Unsloth Studio Update (Windows) |" -ForegroundColor Green
Write-Host "+==============================================+" -ForegroundColor Green
}
# ==========================================================================
# PHASE 1: System-level prerequisites (winget installs, env vars)
@ -1075,9 +1081,9 @@ if (-not $PythonCmd) {
Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundColor Green
# Always create a .venv for isolation -- even for pip installs.
# Created in the repo root (parent of studio/).
$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv"
# The venv must already exist (created by install.ps1).
# This script (setup.ps1 / "unsloth studio update") only updates packages.
$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio"
# Stale-venv detection: if the venv exists but its torch flavor no longer
# matches the current machine, wipe it so we get a clean install.
@ -1140,8 +1146,10 @@ if (Test-Path $VenvDir -PathType Container) {
}
if (-not (Test-Path $VenvDir)) {
Write-Host " Creating virtual environment at $VenvDir..." -ForegroundColor Cyan
& $PythonCmd -m venv $VenvDir
Write-Host "[ERROR] Virtual environment not found at $VenvDir" -ForegroundColor Red
Write-Host " Run install.ps1 first to create the environment:" -ForegroundColor Yellow
Write-Host " irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow
exit 1
} else {
Write-Host " Reusing existing virtual environment at $VenvDir" -ForegroundColor Green
}
@ -1582,8 +1590,9 @@ if ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
# Done
# ============================================
Write-Host ""
$doneLine = if ($env:SKIP_STUDIO_BASE -eq "1") { "Setup Complete!" } else { "Update Complete!" }
Write-Host "+===============================================+" -ForegroundColor Green
Write-Host "| Setup Complete! |" -ForegroundColor Green
Write-Host "| $doneLine |" -ForegroundColor Green
Write-Host "| |" -ForegroundColor Green
Write-Host "| Launch with: |" -ForegroundColor Green
Write-Host "| unsloth studio -H 0.0.0.0 -p 8888 |" -ForegroundColor Green

196
studio/setup.sh
View file

@ -44,9 +44,15 @@ run_quiet_no_exit() {
_run_quiet return "$@"
}
echo "╔══════════════════════════════════════╗"
echo "║ Unsloth Studio Setup Script ║"
echo "╚══════════════════════════════════════╝"
if [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then
echo "╔══════════════════════════════════════╗"
echo "║ Unsloth Studio Setup Script ║"
echo "╚══════════════════════════════════════╝"
else
echo "╔══════════════════════════════════════╗"
echo "║ Unsloth Studio Update Script ║"
echo "╚══════════════════════════════════════╝"
fi
# ── Clean up stale Unsloth compiled caches ──
rm -rf "$REPO_ROOT/unsloth_compiled_cache"
@ -244,114 +250,31 @@ fi
# ── 6. Python venv + deps ──
# ── 6a. Discover best Python >= 3.11 and < 3.14 (i.e. 3.11.x, 3.12.x, or 3.13.x) ──
MIN_PY_MINOR=11 # minimum minor version (>= 3.11)
MAX_PY_MINOR=13 # maximum minor version (< 3.14)
BEST_PY=""
BEST_MINOR=0
# If the caller (e.g. install.sh) already chose a Python, use it directly.
if [ -n "${REQUESTED_PYTHON_VERSION:-}" ] && [ -x "$REQUESTED_PYTHON_VERSION" ]; then
_req_ver=$("$REQUESTED_PYTHON_VERSION" --version 2>&1 | awk '{print $2}')
_req_major=$(echo "$_req_ver" | cut -d. -f1)
_req_minor=$(echo "$_req_ver" | cut -d. -f2)
if [ "$_req_major" -eq 3 ] 2>/dev/null && \
[ "$_req_minor" -ge "$MIN_PY_MINOR" ] 2>/dev/null && \
[ "$_req_minor" -le "$MAX_PY_MINOR" ] 2>/dev/null; then
BEST_PY="$REQUESTED_PYTHON_VERSION"
echo "Using requested Python version: $BEST_PY"
else
echo "Ignoring requested Python $REQUESTED_PYTHON_VERSION ($_req_ver) -- outside supported range"
fi
fi
if [ -z "$BEST_PY" ]; then
# Collect candidate python3 binaries (python3, python3.9, python3.10, …)
for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
if ! command -v "$candidate" &>/dev/null; then
continue
fi
# Get version string, e.g. "Python 3.12.5"
ver_str=$("$candidate" --version 2>&1) || continue
ver_str=$(echo "$ver_str" | awk '{print $2}')
py_major=$(echo "$ver_str" | cut -d. -f1)
py_minor=$(echo "$ver_str" | cut -d. -f2)
# Skip anything that isn't Python 3
if [ "$py_major" -ne 3 ] 2>/dev/null; then
continue
fi
# Skip versions below 3.11
if [ "$py_minor" -lt "$MIN_PY_MINOR" ] 2>/dev/null; then
continue
fi
# Skip versions above 3.13 (require < 3.14)
if [ "$py_minor" -gt "$MAX_PY_MINOR" ] 2>/dev/null; then
continue
fi
# Keep the highest qualifying version
if [ "$py_minor" -gt "$BEST_MINOR" ]; then
BEST_PY="$candidate"
BEST_MINOR="$py_minor"
fi
done
fi
if [ -z "$BEST_PY" ]; then
echo "❌ ERROR: No Python version between 3.${MIN_PY_MINOR} and 3.${MAX_PY_MINOR} found on this system."
echo " Detected Python 3 installations:"
for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
if command -v "$candidate" &>/dev/null; then
echo " - $candidate ($($candidate --version 2>&1))"
fi
done
echo ""
echo " Please install Python 3.${MIN_PY_MINOR} or 3.${MAX_PY_MINOR}."
echo " For example: sudo apt install python3.12 python3.12-venv"
exit 1
fi
BEST_VER=$("$BEST_PY" --version 2>&1 | awk '{print $2}')
echo "✅ Using $BEST_PY ($BEST_VER) — compatible (3.${MIN_PY_MINOR}.x – 3.${MAX_PY_MINOR}.x)"
REQ_ROOT="$SCRIPT_DIR/backend/requirements"
SINGLE_ENV_CONSTRAINTS="$REQ_ROOT/single-env/constraints.txt"
SINGLE_ENV_DATA_DESIGNER="$REQ_ROOT/single-env/data-designer.txt"
SINGLE_ENV_DATA_DESIGNER_DEPS="$REQ_ROOT/single-env/data-designer-deps.txt"
SINGLE_ENV_PATCH="$REQ_ROOT/single-env/patch_metadata.py"
install_python_stack() {
python "$SCRIPT_DIR/install_python_stack.py"
}
# Create venv under ~/.unsloth/studio/ (shared location, not in repo).
# All platforms (including Colab) use the same isolated venv so that
# studio dependencies are never installed into the system Python.
# The venv must already exist (created by install.sh).
# This script (setup.sh / "unsloth studio update") only updates packages.
STUDIO_HOME="$HOME/.unsloth/studio"
VENV_DIR="$STUDIO_HOME/.venv"
VENV_DIR="$STUDIO_HOME/unsloth_studio"
VENV_T5_DIR="$STUDIO_HOME/.venv_t5"
mkdir -p "$STUDIO_HOME"
# Clean up legacy in-repo venvs if they exist
[ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv"
[ -d "$REPO_ROOT/.venv_overlay" ] && rm -rf "$REPO_ROOT/.venv_overlay"
[ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5"
# Note: do NOT delete $STUDIO_HOME/.venv here — install.sh handles migration
rm -rf "$VENV_DIR"
rm -rf "$VENV_T5_DIR"
# Try creating venv with pip; fall back to --without-pip + bootstrap
# (some environments like Colab have broken ensurepip)
if ! "$BEST_PY" -m venv "$VENV_DIR" 2>/dev/null; then
"$BEST_PY" -m venv --without-pip "$VENV_DIR"
source "$VENV_DIR/bin/activate"
curl -sS https://bootstrap.pypa.io/get-pip.py | python > /dev/null
else
source "$VENV_DIR/bin/activate"
if [ ! -x "$VENV_DIR/bin/python" ]; then
echo "❌ ERROR: Virtual environment not found at $VENV_DIR"
echo " Run install.sh first to create the environment:"
echo " curl -fsSL https://unsloth.ai/install.sh | sh"
exit 1
fi
source "$VENV_DIR/bin/activate"
install_python_stack() {
python "$SCRIPT_DIR/install_python_stack.py"
}
# ── Ensure uv is available (much faster than pip) ──
USE_UV=false
if command -v uv &>/dev/null; then
@ -370,22 +293,53 @@ fast_install() {
}
cd "$SCRIPT_DIR"
install_python_stack
# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
# The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
echo ""
echo " Pre-installing transformers 5.x for newer model support..."
mkdir -p "$VENV_T5_DIR"
run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
# tiktoken is needed by Qwen-family tokenizers. Install with deps since
# regex/requests may be missing on Windows.
run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
# ── Check if Python deps need updating ──
# Compare installed package version against PyPI latest.
# Skip all Python dependency work if versions match (fast update path).
_PKG_NAME="${STUDIO_PACKAGE_NAME:-unsloth}"
_SKIP_PYTHON_DEPS=false
if [ "${SKIP_STUDIO_BASE:-0}" != "1" ] && [ "${STUDIO_LOCAL_INSTALL:-0}" != "1" ]; then
# Only check when NOT called from install.sh (which just installed the package)
INSTALLED_VER=$("$VENV_DIR/bin/python" -c "
from importlib.metadata import version
print(version('$_PKG_NAME'))
" 2>/dev/null || echo "")
LATEST_VER=$(curl -fsSL --max-time 5 "https://pypi.org/pypi/$_PKG_NAME/json" 2>/dev/null \
| "$VENV_DIR/bin/python" -c "import sys,json; print(json.load(sys.stdin)['info']['version'])" 2>/dev/null \
|| echo "")
if [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ] && [ "$INSTALLED_VER" = "$LATEST_VER" ]; then
echo "✅ $_PKG_NAME $INSTALLED_VER is up to date (matches PyPI latest)"
_SKIP_PYTHON_DEPS=true
elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then
echo "⬆️ $_PKG_NAME $INSTALLED_VER → $LATEST_VER available, updating dependencies..."
elif [ -z "$LATEST_VER" ]; then
echo "⚠️ Could not reach PyPI, updating dependencies to be safe..."
fi
fi
if [ "$_SKIP_PYTHON_DEPS" = false ]; then
install_python_stack
# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
# The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
echo ""
echo " Pre-installing transformers 5.x for newer model support..."
mkdir -p "$VENV_T5_DIR"
run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
# tiktoken is needed by Qwen-family tokenizers. Install with deps since
# regex/requests may be missing on Windows.
run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
else
echo "✅ Python dependencies up to date — skipping"
fi
# ── 7. WSL: pre-install GGUF build dependencies ──
# On WSL, sudo requires a password and can't be entered during GGUF export
@ -651,9 +605,15 @@ rm -rf "$LLAMA_CPP_DIR"
fi # end _SKIP_GGUF_BUILD check
echo ""
if [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then
_DONE_LINE="║ Setup Complete! ║"
else
_DONE_LINE="║ Update Complete! ║"
fi
if [ "$IS_COLAB" = true ]; then
echo "╔══════════════════════════════════════╗"
echo "║ Setup Complete! ║"
echo "$_DONE_LINE"
echo "╠══════════════════════════════════════╣"
echo "║ Unsloth Studio is ready to start ║"
echo "║ in your Colab notebook! ║"
@ -663,7 +623,7 @@ if [ "$IS_COLAB" = true ]; then
echo "╚══════════════════════════════════════╝"
else
echo "╔══════════════════════════════════════╗"
echo "║ Setup Complete! ║"
echo "$_DONE_LINE"
echo "╠══════════════════════════════════════╣"
echo "║ Launch with: ║"
echo "║ ║"

0
tests/python/__init__.py Normal file
View file

137
tests/python/test_cross_platform_parity.py Normal file
View file

@ -0,0 +1,137 @@
"""Cross-platform parity tests between install.sh and install.ps1."""
from __future__ import annotations
import re
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
INSTALL_SH = REPO_ROOT / "install.sh"
INSTALL_PS1 = REPO_ROOT / "install.ps1"
class TestNoTorchBackendAutoInInstallSh:
"""install.sh primary install paths must not use --torch-backend=auto.
The fallback else-branch (when TORCH_INDEX_URL is empty) is allowed to
use --torch-backend=auto since that is the last-resort recovery path.
"""
def test_no_torch_backend_auto_outside_fallback(self):
lines = INSTALL_SH.read_text().splitlines()
# Find the fallback block: starts with the "else" after the
# TORCH_INDEX_URL check and ends at the next "fi".
fallback_start = None
fallback_end = None
for i, line in enumerate(lines):
if fallback_start is None and "GPU detection failed" in line:
fallback_start = i
elif (
fallback_start is not None
and fallback_end is None
and line.strip() == "fi"
):
fallback_end = i
break
fallback_range = (
range(fallback_start or 0, (fallback_end or 0) + 1)
if fallback_start
else range(0)
)
matches = [
(i + 1, line)
for i, line in enumerate(lines)
if "--torch-backend=auto" in line
and not line.lstrip().startswith("#")
and i not in fallback_range
]
assert matches == [], (
f"install.sh contains --torch-backend=auto outside the fallback block at lines: "
f"{[m[0] for m in matches]}"
)
def test_fallback_uses_torch_backend_auto(self):
"""The fallback branch should use --torch-backend=auto as recovery."""
text = INSTALL_SH.read_text()
assert (
"GPU detection failed" in text
), "install.sh should have a fallback branch for when GPU detection fails"
class TestInstallShHasGpuDetection:
"""install.sh must contain the get_torch_index_url function."""
def test_function_exists(self):
text = INSTALL_SH.read_text()
assert (
"get_torch_index_url()" in text
), "install.sh is missing the get_torch_index_url() function"
def test_torch_index_url_assigned(self):
text = INSTALL_SH.read_text()
assert (
"TORCH_INDEX_URL=$(get_torch_index_url)" in text
), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
class TestCudaMappingParity:
"""CUDA version thresholds must match between install.sh and install.ps1."""
@staticmethod
def _extract_cuda_thresholds_sh(text: str) -> list[str]:
"""Extract cu* suffixes from the major/minor comparison chain in install.sh."""
# Only match lines in the if/elif chain that compare _major/_minor
in_func = False
results = []
for line in text.splitlines():
if "get_torch_index_url()" in line:
in_func = True
continue
if in_func and line.startswith("}"):
break
if in_func and ("_major" in line or "_minor" in line):
m = re.search(r"/(cu\d+|cpu)", line)
if m:
results.append(m.group(1))
return results
@staticmethod
def _extract_cuda_thresholds_ps1(text: str) -> list[str]:
"""Extract cu* suffixes from the major/minor comparison chain in install.ps1."""
in_func = False
depth = 0
results = []
for line in text.splitlines():
if "function Get-TorchIndexUrl" in line:
in_func = True
depth = 1
continue
if in_func:
depth += line.count("{") - line.count("}")
if depth <= 0:
break
# Only match the if-chain lines that compare $major/$minor
if "$major" in line or "$minor" in line:
m = re.search(r"/(cu\d+|cpu)", line)
if m:
results.append(m.group(1))
return results
def test_same_cuda_suffixes(self):
"""Both scripts should produce the same ordered list of CUDA index suffixes."""
sh_text = INSTALL_SH.read_text()
ps1_text = INSTALL_PS1.read_text()
sh_thresholds = self._extract_cuda_thresholds_sh(sh_text)
ps1_thresholds = self._extract_cuda_thresholds_ps1(ps1_text)
assert len(sh_thresholds) > 0, "Could not extract thresholds from install.sh"
assert len(ps1_thresholds) > 0, "Could not extract thresholds from install.ps1"
assert sh_thresholds == ps1_thresholds, (
f"CUDA mapping mismatch:\n"
f" install.sh: {sh_thresholds}\n"
f" install.ps1: {ps1_thresholds}"
)

56
tests/python/test_install_python_stack.py Normal file
View file

@ -0,0 +1,56 @@
"""Tests for install_python_stack._build_uv_cmd torch-backend handling."""
from __future__ import annotations
import importlib
import os
import sys
from pathlib import Path
from unittest import mock
import pytest
# Add the studio directory so we can import install_python_stack
STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
sys.path.insert(0, str(STUDIO_DIR))
# _build_uv_cmd lives at module level; import after path setup.
# We need to mock parts of the module that do work at import time.
import install_python_stack as ips
class TestBuildUvCmdTorchBackend:
"""Verify _build_uv_cmd only adds --torch-backend when UV_TORCH_BACKEND is set."""
def _call(self, args: tuple[str, ...] = ()) -> list[str]:
return ips._build_uv_cmd(args)
def test_default_no_torch_backend(self):
"""Without UV_TORCH_BACKEND env var, no --torch-backend flag."""
env = os.environ.copy()
env.pop("UV_TORCH_BACKEND", None)
with mock.patch.dict(os.environ, env, clear = True):
cmd = self._call(("somepackage",))
assert not any(
a.startswith("--torch-backend") for a in cmd
), f"--torch-backend should not appear by default, got: {cmd}"
def test_uv_torch_backend_auto(self):
"""UV_TORCH_BACKEND=auto adds --torch-backend=auto."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "auto"}):
cmd = self._call(("somepackage",))
assert "--torch-backend=auto" in cmd
def test_uv_torch_backend_cpu(self):
"""UV_TORCH_BACKEND=cpu adds --torch-backend=cpu."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
cmd = self._call(("somepackage",))
assert "--torch-backend=cpu" in cmd
def test_uv_torch_backend_empty(self):
"""UV_TORCH_BACKEND="" (empty string) should NOT add --torch-backend."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": ""}):
cmd = self._call(("somepackage",))
assert not any(
a.startswith("--torch-backend") for a in cmd
), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}"

16
tests/run_all.sh Executable file
View file

@ -0,0 +1,16 @@
#!/bin/sh
# Run all installer tests.
set -e
TESTS_DIR="$(cd "$(dirname "$0")" && pwd)"
echo "=== Bash tests ==="
sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
echo ""
echo "=== Python tests ==="
python -m pytest "$TESTS_DIR/python/test_install_python_stack.py" -v
python -m pytest "$TESTS_DIR/python/test_cross_platform_parity.py" -v
echo ""
echo "All tests passed."

128
tests/sh/test_get_torch_index_url.sh Executable file
View file

@ -0,0 +1,128 @@
#!/bin/bash
# Unit tests for get_torch_index_url() from install.sh
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
# Extract only the get_torch_index_url function from install.sh
# Also replace the hardcoded /usr/bin/nvidia-smi fallback with a
# controllable path so we can test the "no GPU" scenario on GPU machines.
_FUNC_FILE=$(mktemp)
_FAKE_SMI_DIR=$(mktemp -d)
sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH" \
| sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
> "$_FUNC_FILE"
# Save system PATH so we always have basic tools (uname, grep, head, etc.)
_SYS_PATH="/usr/local/bin:/usr/bin:/bin"
assert_eq() {
_label="$1"; _expected="$2"; _actual="$3"
if [ "$_actual" = "$_expected" ]; then
echo " PASS: $_label"
PASS=$((PASS + 1))
else
echo " FAIL: $_label (expected '$_expected', got '$_actual')"
FAIL=$((FAIL + 1))
fi
}
# Helper: create a mock nvidia-smi that prints a given CUDA version string
make_mock_smi() {
_dir=$(mktemp -d)
cat > "$_dir/nvidia-smi" <<MOCK
#!/bin/sh
cat <<'SMI_OUT'
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: $1 |
+-----------------------------------------------------------------------------------------+
SMI_OUT
MOCK
chmod +x "$_dir/nvidia-smi"
echo "$_dir"
}
# Build a minimal tools directory with symlinks to essential commands
# (uname, grep, head, etc.) but WITHOUT nvidia-smi.
_TOOLS_DIR=$(mktemp -d)
for _cmd in uname grep sed head sh bash cat; do
_real=$(command -v "$_cmd" 2>/dev/null || true)
[ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd"
done
# Helper: run get_torch_index_url with a custom PATH
# $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test
run_func() {
_mock_dir="$1"
if [ "$_mock_dir" = "none" ]; then
# Minimal PATH with only basic tools, no nvidia-smi anywhere
PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
else
# Put mock nvidia-smi dir first, then basic tools
PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
fi
}
echo "=== test_get_torch_index_url ==="
# 1) No nvidia-smi available -> cpu
_result=$(run_func "none")
assert_eq "no nvidia-smi -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
# 2) CUDA 12.6 -> cu126
_dir=$(make_mock_smi "12.6")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.6 -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
rm -rf "$_dir"
# 3) CUDA 12.8 -> cu128
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.8 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 4) CUDA 13.0 -> cu130
_dir=$(make_mock_smi "13.0")
_result=$(run_func "$_dir")
assert_eq "CUDA 13.0 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 5) CUDA 12.4 -> cu124
_dir=$(make_mock_smi "12.4")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.4 -> cu124" "https://download.pytorch.org/whl/cu124" "$_result"
rm -rf "$_dir"
# 6) CUDA 11.8 -> cu118
_dir=$(make_mock_smi "11.8")
_result=$(run_func "$_dir")
assert_eq "CUDA 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result"
rm -rf "$_dir"
# 7) CUDA 10.2 (too old) -> cpu
_dir=$(make_mock_smi "10.2")
_result=$(run_func "$_dir")
assert_eq "CUDA 10.2 -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 8) Unparseable nvidia-smi output -> cu126 default
_dir=$(mktemp -d)
cat > "$_dir/nvidia-smi" <<'MOCK'
#!/bin/sh
echo "something completely unexpected"
MOCK
chmod +x "$_dir/nvidia-smi"
_result=$(run_func "$_dir")
assert_eq "unparseable -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
rm -rf "$_dir"
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ] || exit 1

2
unsloth_cli/__init__.py
View file

@ -6,7 +6,6 @@ import typer
from unsloth_cli.commands.train import train
from unsloth_cli.commands.inference import inference
from unsloth_cli.commands.export import export, list_checkpoints
from unsloth_cli.commands.ui import ui
from unsloth_cli.commands.studio import studio_app
app = typer.Typer(
@ -18,5 +17,4 @@ app.command()(train)
app.command()(inference)
app.command()(export)
app.command("list-checkpoints")(list_checkpoints)
app.command()(ui)
app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")

48
unsloth_cli/commands/studio.py
View file

@ -22,9 +22,9 @@ _PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent
def _studio_venv_python() -> Optional[Path]:
"""Return the studio venv Python binary, or None if not set up."""
if platform.system() == "Windows":
p = STUDIO_HOME / ".venv" / "Scripts" / "python.exe"
p = STUDIO_HOME / "unsloth_studio" / "Scripts" / "python.exe"
else:
p = STUDIO_HOME / ".venv" / "bin" / "python"
p = STUDIO_HOME / "unsloth_studio" / "bin" / "python"
return p if p.is_file() else None
@ -44,7 +44,7 @@ def _find_run_py() -> Optional[Path]:
"lib/python*/site-packages/studio/backend/run.py",
"Lib/site-packages/studio/backend/run.py",
):
for match in (STUDIO_HOME / ".venv").glob(pattern):
for match in (STUDIO_HOME / "unsloth_studio").glob(pattern):
return match
return None
@ -64,7 +64,7 @@ def _find_setup_script() -> Optional[Path]:
f"lib/python*/site-packages/studio/{name}",
f"Lib/site-packages/studio/{name}",
):
for match in (STUDIO_HOME / ".venv").glob(pattern):
for match in (STUDIO_HOME / "unsloth_studio").glob(pattern):
return match
return None
@ -85,7 +85,7 @@ def studio_default(
return
# Always use the studio venv if it exists and we're not already in it
studio_venv_dir = STUDIO_HOME / ".venv"
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
if not in_studio_venv:
@ -132,7 +132,7 @@ def studio_default(
else:
os.execvp(str(studio_python), args)
else:
typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
typer.echo("Studio not set up. Run install.sh first.")
raise typer.Exit(1)
from studio.backend.run import run_server
@ -166,12 +166,11 @@ def studio_default(
typer.echo("\nShutting down...")
# ── unsloth studio setup ─────────────────────────────────────────────
# ── unsloth studio setup / update ─────────────────────────────────────
@studio_app.command()
def setup():
"""Run one-time Studio environment setup."""
def _run_setup_script() -> None:
"""Find and run the studio setup/update script."""
script = _find_setup_script()
if not script:
typer.echo("Error: Could not find setup script (setup.sh / setup.ps1).")
@ -188,6 +187,35 @@ def setup():
raise typer.Exit(result.returncode)
@studio_app.command(hidden = True)
def setup():
"""Deprecated: use 'unsloth studio update' or re-run install.sh."""
typer.echo(
"Note: 'unsloth studio setup' is deprecated. Use 'unsloth studio update' or re-run install.sh."
)
_run_setup_script()
@studio_app.command()
def update(
local: bool = typer.Option(
False, "--local", help = "Install from local repo instead of PyPI"
),
package: str = typer.Option(
"unsloth", "--package", help = "Package name to install/update (for testing)"
),
):
"""Update Unsloth Studio dependencies and rebuild."""
os.environ["STUDIO_LOCAL_INSTALL"] = "1" if local else "0"
os.environ["STUDIO_PACKAGE_NAME"] = package
if local:
# Pass the repo root explicitly so install_python_stack.py doesn't
# have to guess from SCRIPT_DIR (which may be inside site-packages).
repo_root = Path(__file__).resolve().parents[2]
os.environ["STUDIO_LOCAL_REPO"] = str(repo_root)
_run_setup_script()
# ── unsloth studio reset-password ────────────────────────────────────

103
unsloth_cli/commands/ui.py
View file

@ -1,103 +0,0 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import os
import sys
import time
from pathlib import Path
from typing import Optional
import typer
def ui(
port: int = typer.Option(
8888, "--port", "-p", help = "Port to run the UI server on."
),
host: str = typer.Option(
"0.0.0.0", "--host", "-H", help = "Host address to bind to."
),
frontend: Optional[Path] = typer.Option(
None, "--frontend", "-f", help = "Path to frontend build directory."
),
silent: bool = typer.Option(
False, "--silent", "-q", help = "Suppress startup messages."
),
):
"""Launch the Unsloth web UI backend server (alias for 'unsloth studio')."""
from unsloth_cli.commands.studio import (
_studio_venv_python,
_find_run_py,
STUDIO_HOME,
)
# Re-execute in studio venv if available and not already inside it
studio_venv_dir = STUDIO_HOME / ".venv"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
if not in_studio_venv:
studio_python = _studio_venv_python()
run_py = _find_run_py()
if studio_python and run_py:
if not silent:
typer.echo("Launching Unsloth Studio... Please wait...")
args = [
str(studio_python),
str(run_py),
"--host",
host,
"--port",
str(port),
]
if frontend:
args.extend(["--frontend", str(frontend)])
if silent:
args.append("--silent")
# On Windows, os.execvp() spawns a child but the parent lingers,
# so Ctrl+C only kills the parent leaving the child orphaned.
# Use subprocess.run() on Windows so the parent waits for the child.
if sys.platform == "win32":
import subprocess as _sp
proc = _sp.Popen(args)
try:
rc = proc.wait()
except KeyboardInterrupt:
# Child has its own signal handler — let it finish
rc = proc.wait()
raise typer.Exit(rc)
else:
os.execvp(str(studio_python), args)
else:
typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
raise typer.Exit(1)
from studio.backend.run import run_server
if not silent:
from studio.backend.run import _resolve_external_ip
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
run_kwargs = dict(host = host, port = port, silent = silent)
if frontend is not None:
run_kwargs["frontend_path"] = frontend
run_server(**run_kwargs)
from studio.backend.run import _shutdown_event
try:
if _shutdown_event is not None:
# NOTE: Event.wait() without a timeout blocks at the C level
# on Linux, preventing Python from delivering SIGINT (Ctrl+C).
while not _shutdown_event.is_set():
_shutdown_event.wait(timeout = 1)
else:
while True:
time.sleep(1)
except KeyboardInterrupt:
from studio.backend.run import _graceful_shutdown, _server
_graceful_shutdown(_server)
typer.echo("\nShutting down...")