ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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
commit 19e9c60a8e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 877 additions and 301 deletions
Download patch file Download diff file Expand all files Collapse all files

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")