From a537ece7eb3dbfce794d7279432ef8c8d8f06a97 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 14 Mar 2026 05:59:10 +0000 Subject: [PATCH] studio: use uv for Python package installs (8x faster) Replace pip with uv in install_python_stack.py to speed up the Python dependency installation phase of `unsloth studio setup`. - Add _bootstrap_uv() that checks for uv on PATH, and if not found, installs it via pip. Falls back to pip if uv is unavailable. - Translate pip flags to uv equivalents (--no-cache-dir dropped since uv caching is fast, --force-reinstall becomes --reinstall). - Add --torch-backend=auto so uv auto-detects CUDA version for PyTorch ecosystem packages. - Per-install fallback: if any uv install step fails, it retries that step with pip before exiting. Measured on clean venv setup: Python packages (pip): 2m 28s Python packages (uv): 18s Speedup: ~8x Total setup time goes from ~4m 35s to ~2m 30s (llama.cpp build is now the bottleneck at 1m 40s). --- studio/install_python_stack.py | 86 ++++++++++++++++++++++++++++++---- 1 file changed, 76 insertions(+), 10 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index f7c62ade9b..d2d071a41a 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -13,6 +13,7 @@ PATH to point at the venv. from __future__ import annotations import os +import shutil import subprocess import sys import tempfile @@ -73,7 +74,7 @@ def _red(msg: str) -> str: return f"\033[91m{msg}\033[0m" if _HAS_COLOR else msg -def run(label: str, cmd: list[str], *, quiet: bool = True) -> None: +def run(label: str, cmd: list[str], *, quiet: bool = True) -> subprocess.CompletedProcess[bytes]: """Run a command; on failure print output and exit.""" print(_cyan(f" {label}...")) result = subprocess.run( @@ -86,11 +87,28 @@ def run(label: str, cmd: list[str], *, quiet: bool = True) -> None: if result.stdout: print(result.stdout.decode(errors = "replace")) sys.exit(result.returncode) + return result # Packages to skip on Windows (require special build steps) WINDOWS_SKIP_PACKAGES = {"open_spiel", "triton_kernels"} +# ── uv bootstrap ────────────────────────────────────────────────────── + +USE_UV = False # Set by _bootstrap_uv() at the start of install_python_stack() + + +def _bootstrap_uv() -> bool: + """Try to use uv for faster installs. Returns True if uv is available.""" + if shutil.which("uv"): + return True + result = subprocess.run( + [sys.executable, "-m", "pip", "install", "--quiet", "uv"], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + return result.returncode == 0 and shutil.which("uv") is not None + def _filter_requirements(req: Path, skip: set[str]) -> Path: """Return a temp copy of a requirements file with certain packages removed.""" @@ -111,26 +129,69 @@ def _filter_requirements(req: Path, skip: set[str]) -> Path: return Path(tmp.name) +def _translate_pip_args_for_uv(args: tuple[str, ...]) -> list[str]: + """Translate pip flags to their uv equivalents.""" + translated: list[str] = [] + for arg in args: + if arg == "--no-cache-dir": + continue # uv cache is fast; drop this flag + elif arg == "--force-reinstall": + translated.append("--reinstall") + else: + translated.append(arg) + return translated + + +def _build_pip_cmd(args: tuple[str, ...]) -> list[str]: + """Build a standard pip install command.""" + cmd = [sys.executable, "-m", "pip", "install"] + cmd.extend(args) + return cmd + + +def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: + """Build a uv pip install command with translated flags.""" + cmd = ["uv", "pip", "install"] + cmd.extend(_translate_pip_args_for_uv(args)) + cmd.append("--torch-backend=auto") + return cmd + + def pip_install( label: str, *args: str, req: Path | None = None, constrain: bool = True, ) -> None: - """Build and run a pip install command.""" - cmd = [sys.executable, "-m", "pip", "install"] - cmd.extend(args) + """Build and run a pip install command (uses uv when available, falls back to pip).""" + constraint_args: list[str] = [] if constrain and CONSTRAINTS.is_file(): - cmd.extend(["-c", str(CONSTRAINTS)]) + constraint_args = ["-c", str(CONSTRAINTS)] + actual_req = req if req is not None and IS_WINDOWS and WINDOWS_SKIP_PACKAGES: actual_req = _filter_requirements(req, WINDOWS_SKIP_PACKAGES) + req_args: list[str] = [] if actual_req is not None: - cmd.extend(["-r", str(actual_req)]) + req_args = ["-r", str(actual_req)] + try: - run(label, cmd) + if USE_UV: + uv_cmd = _build_uv_cmd(args) + constraint_args + req_args + print(_cyan(f" {label}...")) + result = subprocess.run( + uv_cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + if result.returncode != 0: + print(_red(f" uv failed, falling back to pip...")) + pip_cmd = _build_pip_cmd(args) + constraint_args + req_args + run(label, pip_cmd) + else: + pip_cmd = _build_pip_cmd(args) + constraint_args + req_args + run(label, pip_cmd) finally: - # Clean up temp file if we created one if actual_req is not None and actual_req != req: actual_req.unlink(missing_ok = True) @@ -170,11 +231,16 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None: def install_python_stack() -> int: - print(_cyan("── Installing Python stack ──")) + global USE_UV - # 1. Upgrade pip + # 1. Upgrade pip (needed even with uv as fallback and for bootstrapping) run("Upgrading pip", [sys.executable, "-m", "pip", "install", "--upgrade", "pip"]) + # Try to use uv for faster installs + USE_UV = _bootstrap_uv() + installer = "uv" if USE_UV else "pip" + print(_cyan(f"── Installing Python stack (via {installer}) ──")) + # 2. Core packages: unsloth-zoo + unsloth pip_install( "Installing base packages",