Extract shared install_python_stack.py for cross-platform setup

install_python_stack.py

@@ -0,0 +1,182 @@
+#!/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 urllib.request
+from pathlib import Path
+
+# ── 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"
+
+# ── Helpers ────────────────────────────────────────────────────────────
+
+def _green(msg: str) -> str:
+    return f"\033[92m{msg}\033[0m"
+
+def _cyan(msg: str) -> str:
+    return f"\033[96m{msg}\033[0m"
+
+def _red(msg: str) -> str:
+    return f"\033[91m{msg}\033[0m"
+
+
+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)
+
+
+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)])
+    if req is not None:
+        cmd.extend(["-r", str(req)])
+    run(label, cmd)
+
+
+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())