unsloth/install_python_stack.py

#!/usr/bin/env python3
"""Cross-platform Python dependency installer for Unsloth Studio.

Called by both setup.sh (Linux / WSL) and setup.ps1 (Windows) after the
virtual environment is already activated.  Expects `pip` and `python` on
PATH to point at the venv.
"""

from __future__ import annotations

import os
import subprocess
import sys
import tempfile
import urllib.request
from pathlib import Path

IS_WINDOWS = sys.platform == "win32"

# ── Paths ──────────────────────────────────────────────────────────────
SCRIPT_DIR = Path(__file__).resolve().parent
REQ_ROOT = SCRIPT_DIR / "studio" / "backend" / "requirements"
SINGLE_ENV = REQ_ROOT / "single-env"
CONSTRAINTS = SINGLE_ENV / "constraints.txt"

# ── Color support ──────────────────────────────────────────────────────

def _enable_colors() -> bool:
    """Try to enable ANSI color support. Returns True if available."""
    if not hasattr(sys.stdout, "fileno"):
        return False
    try:
        if not os.isatty(sys.stdout.fileno()):
            return False
    except Exception:
        return False
    if IS_WINDOWS:
        try:
            import ctypes
            kernel32 = ctypes.windll.kernel32
            # Enable ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x0004) on stdout
            handle = kernel32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE
            mode = ctypes.c_ulong()
            kernel32.GetConsoleMode(handle, ctypes.byref(mode))
            kernel32.SetConsoleMode(handle, mode.value | 0x0004)
            return True
        except Exception:
            return False
    return True  # Unix terminals support ANSI by default

_HAS_COLOR = _enable_colors()

def _green(msg: str) -> str:
    return f"\033[92m{msg}\033[0m" if _HAS_COLOR else msg

def _cyan(msg: str) -> str:
    return f"\033[96m{msg}\033[0m" if _HAS_COLOR else msg

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:
    """Run a command; on failure print output and exit."""
    print(_cyan(f"   {label}..."))
    result = subprocess.run(
        cmd,
        stdout=subprocess.PIPE if quiet else None,
        stderr=subprocess.STDOUT if quiet else None,
    )
    if result.returncode != 0:
        print(_red(f"❌ {label} failed (exit code {result.returncode}):"))
        if result.stdout:
            print(result.stdout.decode(errors="replace"))
        sys.exit(result.returncode)


# Packages to skip on Windows (require special build steps)
WINDOWS_SKIP_PACKAGES = {"open_spiel"}


def _filter_requirements(req: Path, skip: set[str]) -> Path:
    """Return a temp copy of a requirements file with certain packages removed."""
    lines = req.read_text(encoding="utf-8").splitlines(keepends=True)
    filtered = [
        line for line in lines
        if not any(line.strip().lower().startswith(pkg) for pkg in skip)
    ]
    tmp = tempfile.NamedTemporaryFile(
        mode="w", suffix=".txt", delete=False, encoding="utf-8",
    )
    tmp.writelines(filtered)
    tmp.close()
    return Path(tmp.name)


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)
    if constrain and CONSTRAINTS.is_file():
        cmd.extend(["-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)
    if actual_req is not None:
        cmd.extend(["-r", str(actual_req)])
    try:
        run(label, 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)



def download_file(url: str, dest: Path) -> None:
    """Download a file using urllib (no curl dependency)."""
    urllib.request.urlretrieve(url, dest)


def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
    """Download a file from url and overwrite a file inside an installed package."""
    result = subprocess.run(
        [sys.executable, "-m", "pip", "show", package_name],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(_red(f"   ⚠️  Could not find package {package_name}, skipping patch"))
        return

    location = None
    for line in result.stdout.splitlines():
        if line.lower().startswith("location:"):
            location = line.split(":", 1)[1].strip()
            break

    if not location:
        print(_red(f"   ⚠️  Could not determine location of {package_name}"))
        return

    dest = Path(location) / relative_path
    print(_cyan(f"   Patching {dest.name} in {package_name}..."))
    download_file(url, dest)


# ── Main install sequence ─────────────────────────────────────────────

def install_python_stack() -> int:
    print(_cyan("── Installing Python stack ──"))

    # 1. Upgrade pip
    run("Upgrading pip", [sys.executable, "-m", "pip", "install", "--upgrade", "pip"])

    # 2. Core packages: unsloth-zoo + unsloth
    pip_install(
        "Installing unsloth-zoo + unsloth",
        "--no-cache-dir",
        req=REQ_ROOT / "base.txt",
    )

    # 3. Extra dependencies
    pip_install(
        "Installing additional unsloth dependencies",
        "--no-cache-dir",
        req=REQ_ROOT / "extras.txt",
    )

    # 4. Overrides (torchao, transformers) — force-reinstall
    pip_install(
        "Installing torchao + transformers overrides",
        "--force-reinstall", "--no-cache-dir",
        req=REQ_ROOT / "overrides.txt",
    )

    # 5. Triton kernels (no-deps, from source)
    pip_install(
        "Installing triton kernels",
        "--no-deps", "--no-cache-dir",
        req=REQ_ROOT / "triton-kernels.txt",
        constrain=False,
    )

    # 6. Patch: override llama_cpp.py with fix from unsloth-zoo main branch
    patch_package_file(
        "unsloth-zoo",
        os.path.join("unsloth_zoo", "llama_cpp.py"),
        "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py",
    )

    # 7. Patch: override vision.py with fix from unsloth PR #4091
    patch_package_file(
        "unsloth",
        os.path.join("unsloth", "models", "vision.py"),
        "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py",
    )

    # 8. Studio dependencies
    pip_install(
        "Installing studio dependencies",
        "--no-cache-dir",
        req=REQ_ROOT / "studio.txt",
    )

    # 9. Data-designer dependencies
    pip_install(
        "Installing data-designer dependencies",
        "--no-cache-dir",
        req=SINGLE_ENV / "data-designer-deps.txt",
    )

    # 10. Data-designer packages (no-deps to avoid conflicts)
    pip_install(
        "Installing data-designer",
        "--no-cache-dir", "--no-deps",
        req=SINGLE_ENV / "data-designer.txt",
    )

    # 11. Patch metadata for single-env compatibility
    run(
        "Patching single-env metadata",
        [sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
    )

    # 12. Final check
    run("Running pip check", [sys.executable, "-m", "pip", "check"], quiet=False)

    print(_green("✅ Python dependencies installed"))
    return 0


if __name__ == "__main__":
    sys.exit(install_python_stack())