ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/install_python_stack.py
Daniel Han 7d722faedc Fix Windows ROCm install: idempotency guard, nvidia-smi fallback, Python version gates 
- Add torch.version.hip probe to _ensure_rocm_torch_windows() so steps 2b
  and 13 skip the 2.1-3.9 GB reinstall when ROCm torch is already healthy.
  Mirrors the existing Linux idempotency guard.

- Add Windows fallback paths for nvidia-smi.exe in _has_usable_nvidia_gpu()
  (NVSMI dir + System32) matching the PowerShell scripts, so NVIDIA wins on
  mixed GPU systems even when nvidia-smi is not on PATH.

- install.ps1: only narrow Python to 3.12 when AMD + torch (not --no-torch);
  set $PythonVersion = "3.12" so the existing winget auto-install flow works
  on clean AMD machines instead of hard-stopping.

- install.ps1: handle empty $venvPyVer as an error in the AMD version check
  instead of silently continuing into a guaranteed pip failure.

- setup.ps1: upgrade Python 3.12 version check from warning to hard error
  (exit 1), preventing multi-GB downloads that pip will reject on non-3.12.

- install_python_stack.py: use os.environ["ProgramFiles"] for ROCm root
  fallback instead of hardcoded C:\Program Files.

- Remove unused _HIP_SDK_DOWNLOAD_URL constant.
2026-04-12 21:22:45 +00:00

1461 lines
54 KiB
Python
Raw Blame History

#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""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 platform
import shutil
import subprocess
import sys
import tempfile
import urllib.request
from pathlib import Path
IS_WINDOWS = sys.platform == "win32"
IS_MACOS = sys.platform == "darwin"
IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64"
# ── ROCm / AMD GPU support ─────────────────────────────────────────────────────
# Mapping from detected ROCm (major, minor) to the best PyTorch wheel tag on
# download.pytorch.org. Entries are checked newest-first (>=).
# ROCm 7.2 only has torch 2.11.0 on download.pytorch.org, which exceeds the
# current torch upper bound (<2.11.0). Fall back to rocm7.1 (torch 2.10.0).
# TODO: uncomment rocm7.2 when torch upper bound is bumped to >=2.11.0
_ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
# (7, 2): "rocm7.2", # torch 2.11.0 -- requires torch>=2.11
(7, 1): "rocm7.1",
(7, 0): "rocm7.0",
(6, 4): "rocm6.4",
(6, 3): "rocm6.3",
(6, 2): "rocm6.2",
(6, 1): "rocm6.1",
(6, 0): "rocm6.0",
}
_PYTORCH_WHL_BASE = "https://download.pytorch.org/whl"
# Windows AMD ROCm torch wheels live at repo.radeon.com because PyTorch
# does NOT publish Windows ROCm wheels on download.pytorch.org (and says so
# on pytorch.org/get-started/locally: "ROCm is not available on Windows").
# Every wheel under download.pytorch.org/whl/rocm{6.4,7.1,7.2}/ is
# manylinux_2_28_x86_64 only. Upstream work to ship Windows ROCm wheels is
# tracked at pytorch/pytorch#159520, targeted for torch 2.10/2.11 but not
# yet delivered. Until then repo.radeon.com is the only source.
#
# AMD's official install procedure
# (rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html) is a
# TWO-STEP pip install:
# Step 1: rocm_sdk_core + rocm_sdk_devel + rocm_sdk_libraries_custom +
# rocm-<ver>.tar.gz. The three wheels ship the ROCm runtime
# libraries torch links against at import time; the tar.gz is a
# tiny (~15 KB) Python sdist for a meta-package named "rocm"
# that requires the other three. The HIP SDK developer toolkit
# (HIP_PATH) is NOT a substitute -- torch imports the
# Python-packaged runtime from rocm_sdk_*. Step 1 download size
# varies dramatically by release: 7.2.1 is about 1.3 GB total
# (sdk_core ~615 MB, sdk_devel ~222 MB, sdk_libraries ~467 MB)
# while 7.1.1 is about 3.2 GB total -- its sdk_devel wheel is a
# massive 2.4 GB (debug symbols / unstripped libraries).
# Step 2: torch + torchvision + torchaudio. About 780 MB for 7.2.1
# (torch itself is 783 MB; vision/audio are small), 692 MB for
# 7.1.1.
#
# All URLs in _ROCM_WINDOWS_TORCH_WHEELS below were verified live on
# 2026-04-11 with HEAD requests; see temp/windows_amd_tests for the
# verification script.
#
# As of 2026-04, repo.radeon.com/rocm/windows/ contains four release dirs:
# rocm-rel-6.4.4/ -- PEP 503 simple index (torch/, torchvision/, torchaudio/
# sub-indexes; wheels themselves live at the top of the
# release dir). Wheels carry alpha + git-hash build tags
# like torch-2.8.0a0+gitfc14c65-cp312-cp312-win_amd64.whl,
# so the filename changes whenever AMD rebuilds and we
# cannot hardcode a URL for it. Supporting 6.4.4 would
# require parsing the PEP 503 index at install time --
# out of scope here.
# rocm-rel-7.1.1/ -- flat layout, torch stable `+rocmsdk20251116` tag,
# SDK wheels stamped `0.1.dev0` (pre-release marker).
# rocm-rel-7.2/ -- flat layout, torch stable `+rocmsdk20260116` tag.
# rocm-rel-7.2.1/ -- flat layout, torch stable `+rocm7.2.1` tag, SDK
# wheels stamped `7.2.1`. Newest 7.2.x release as of
# writing; superset of rocm-rel-7.2.
#
# HIP SDK detection via HIP_PATH is an OPTIONAL version hint, not a
# prerequisite. Users only need an AMD graphics driver (26.2.2+ for 7.2.1)
# and Python 3.12. When HIP_PATH is absent or points at an unsupported
# version, we default to the newest stable release (7.2.1).
_ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = {
(7, 2): {
# Step 1: ROCm SDK wheels (ship the runtime torch imports)
"sdk_core": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/"
"rocm_sdk_core-7.2.1-py3-none-win_amd64.whl"
),
"sdk_devel": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/"
"rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl"
),
"sdk_libraries": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/"
"rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl"
),
"sdk_tarball": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" "rocm-7.2.1.tar.gz"
),
# Step 2: torch wheels
"torch": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/"
"torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl"
),
"torchvision": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/"
"torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl"
),
"torchaudio": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/"
"torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl"
),
},
(7, 1): {
# Step 1: ROCm SDK wheels -- note 7.1.1 stamps the SDK wheels with
# `0.1.dev0` while keeping the torch wheels at 2.9.0.
"sdk_core": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl"
),
"sdk_devel": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"rocm_sdk_devel-0.1.dev0-py3-none-win_amd64.whl"
),
"sdk_libraries": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl"
),
"sdk_tarball": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"rocm-0.1.dev0.tar.gz"
),
# Step 2: torch wheels
"torch": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl"
),
"torchvision": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl"
),
"torchaudio": (
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/"
"torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl"
),
},
}
# Default Windows ROCm release when HIP_PATH is absent or unreadable. Users
# only need a recent AMD graphics driver and Python 3.12 -- the HIP SDK
# developer toolkit is NOT required for running torch.
_DEFAULT_WINDOWS_ROCM_VERSION: tuple[int, int] = (7, 2)
_AMD_RADEON_DRIVER_URL = "https://www.amd.com/en/support/download/drivers.html"
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every
# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI.
_BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
"x86_64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl"
),
"aarch64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl"
),
}
_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1"
def _bnb_rocm_prerelease_url() -> str | None:
"""Return the continuous-release_main bnb wheel URL for the current
architecture, or None when no pre-release wheel is available.
"""
arch = platform.machine().lower()
arch = {"amd64": "x86_64", "arm64": "aarch64"}.get(arch, arch)
return _BNB_ROCM_PRERELEASE_URLS.get(arch)
def _detect_rocm_version() -> tuple[int, int] | None:
"""Return (major, minor) of the installed ROCm stack, or None."""
# Check /opt/rocm/.info/version or ROCM_PATH equivalent
rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm"
for path in (
os.path.join(rocm_root, ".info", "version"),
os.path.join(rocm_root, "lib", "rocm_version"),
):
try:
with open(path) as fh:
parts = fh.read().strip().split("-")[0].split(".")
# Explicit length guard avoids relying on the broad except
# below to swallow IndexError when the version file contains
# a single component (e.g. "6\n" on a partial install).
if len(parts) >= 2:
return int(parts[0]), int(parts[1])
except Exception:
pass
# Try amd-smi version (outputs "... | ROCm version: X.Y.Z")
amd_smi = shutil.which("amd-smi")
if amd_smi:
try:
result = subprocess.run(
[amd_smi, "version"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 5,
)
if result.returncode == 0:
import re
m = re.search(r"ROCm version:\s*(\d+)\.(\d+)", result.stdout)
if m:
return int(m.group(1)), int(m.group(2))
except Exception:
pass
# Try hipconfig --version (outputs bare version like "6.3.21234.2")
hipconfig = shutil.which("hipconfig")
if hipconfig:
try:
result = subprocess.run(
[hipconfig, "--version"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 5,
)
if result.returncode == 0:
raw = result.stdout.decode().strip().split("\n")[0]
parts = raw.split(".")
if (
len(parts) >= 2
and parts[0].isdigit()
and parts[1].split("-")[0].isdigit()
):
return int(parts[0]), int(parts[1].split("-")[0])
except Exception:
pass
# Distro package-manager fallbacks. Package-managed ROCm installs can
# expose GPUs via rocminfo / amd-smi but still lack /opt/rocm/.info/version
# and hipconfig, so probe dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE)
# for the rocm-core package version. Matches the chain in
# install.sh::get_torch_index_url so `unsloth studio update` behaves
# the same as a fresh `curl | sh` install.
import re as _re_pkg
for cmd in (
["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"],
["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"],
):
exe = shutil.which(cmd[0])
if not exe:
continue
try:
result = subprocess.run(
[exe, *cmd[1:]],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 5,
)
except Exception:
continue
if result.returncode != 0 or not result.stdout.strip():
continue
raw = result.stdout.strip()
# dpkg can prepend an epoch ("1:6.3.0-1"); strip it before parsing.
raw = _re_pkg.sub(r"^\d+:", "", raw)
m = _re_pkg.match(r"(\d+)[.-](\d+)", raw)
if m:
return int(m.group(1)), int(m.group(2))
return None
def _has_rocm_gpu() -> bool:
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
import re
for cmd, check_fn in (
# rocminfo: look for "Name: gfxNNNN" with nonzero first digit (gfx000 is the CPU agent)
(["rocminfo"], lambda out: bool(re.search(r"gfx[1-9]", out.lower()))),
# amd-smi list: require "GPU: <number>" data rows, not just a header
(
["amd-smi", "list"],
lambda out: bool(re.search(r"(?im)^gpu\s*[:\[]\s*\d", out)),
),
):
exe = shutil.which(cmd[0])
if not exe:
continue
try:
result = subprocess.run(
[exe, *cmd[1:]],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
except Exception:
continue
if result.returncode == 0 and result.stdout.strip():
if check_fn(result.stdout):
return True
return False
def _detect_rocm_version_windows() -> tuple[int, int] | None:
"""Return (major, minor) of the installed HIP SDK on Windows, or None.
Primary signal: HIP_PATH env var set by the HIP SDK installer, e.g.
``C:\\Program Files\\AMD\\ROCm\\7.2\\``. The trailing component is the
version. We also sanity-check that ``<HIP_PATH>\\bin`` exists so we do
not trust a stale env var pointing at a half-uninstalled SDK.
Fallback: scan ``C:\\Program Files\\AMD\\ROCm\\`` and return the highest
numeric subdirectory.
"""
if not IS_WINDOWS:
return None
# Use ntpath explicitly rather than os.path so the logic is also
# testable on Linux runners where os.path is posixpath (basename of
# "C:\\foo\\7.2" returns the whole string on posix, breaking version
# extraction even though real Windows handles it correctly).
import ntpath as _ntpath
import re as _re
def _parse_component(name: str) -> tuple[int, int] | None:
# Accept "7.2" and "7.2.1", ignore anything else.
m = _re.match(r"^(\d+)\.(\d+)", name.strip())
if not m:
return None
return int(m.group(1)), int(m.group(2))
hip_path = os.environ.get("HIP_PATH", "").strip()
if hip_path:
trimmed = hip_path.rstrip("\\/")
bin_dir = _ntpath.join(trimmed, "bin")
# Only trust HIP_PATH when the bin folder is actually present. A
# broken uninstall can leave the env var pointing at a ghost dir.
if os.path.isdir(bin_dir):
ver = _parse_component(_ntpath.basename(trimmed))
if ver is not None:
return ver
rocm_root = os.path.join(
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
)
if os.path.isdir(rocm_root):
best: tuple[int, int] | None = None
try:
entries = os.listdir(rocm_root)
except OSError:
entries = []
for entry in entries:
sub = _ntpath.join(rocm_root, entry)
if not os.path.isdir(sub):
continue
ver = _parse_component(entry)
if ver is None:
continue
if os.path.isdir(_ntpath.join(sub, "bin")):
if best is None or ver > best:
best = ver
if best is not None:
return best
return None
def _has_rocm_gpu_windows() -> bool:
"""Return True when a Radeon/AMD GPU is visible in WMI Win32_VideoController.
We deliberately avoid ``hipinfo.exe`` here because it lives inside the
HIP SDK -- if we used it to decide whether to prompt the user to install
the HIP SDK we would never trigger the prompt on the hosts that need it
most. WMI is always available on Windows and needs no elevation.
"""
if not IS_WINDOWS:
return False
ps_cmd = (
"Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue "
"| Where-Object { $_.Name -match 'AMD|Radeon' } "
"| Measure-Object | Select-Object -ExpandProperty Count"
)
for exe_name in ("pwsh", "powershell"):
exe = shutil.which(exe_name)
if not exe:
continue
try:
result = subprocess.run(
[exe, "-NoProfile", "-NonInteractive", "-Command", ps_cmd],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 15,
)
except Exception:
continue
if result.returncode != 0:
continue
raw = (result.stdout or "").strip()
if raw.isdigit() and int(raw) > 0:
return True
return False
def _has_usable_nvidia_gpu() -> bool:
"""Return True only when nvidia-smi exists AND reports at least one GPU."""
exe = shutil.which("nvidia-smi")
if not exe and IS_WINDOWS:
# nvidia-smi.exe is often absent from PATH on Windows even with a
# valid driver. Match the fallback paths used in install.ps1 /
# setup.ps1 so the NVIDIA-wins-on-mixed-systems rule is consistent
# between the PowerShell and Python install paths.
_candidates = [
os.path.join(
os.environ.get("ProgramFiles", r"C:\Program Files"),
r"NVIDIA Corporation\NVSMI\nvidia-smi.exe",
),
os.path.join(
os.environ.get("SystemRoot", r"C:\Windows"),
r"System32\nvidia-smi.exe",
),
]
for _c in _candidates:
if os.path.isfile(_c):
exe = _c
break
if not exe:
return False
try:
result = subprocess.run(
[exe, "-L"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
except Exception:
return False
return result.returncode == 0 and "GPU " in result.stdout
def _ensure_rocm_torch_windows() -> None:
"""Install Radeon's Windows ROCm SDK + torch wheels when an AMD GPU is
present. Called from _ensure_rocm_torch().
AMD's documented install procedure
(rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html) is
a TWO-STEP sequence: first the rocm_sdk_* wheels (which ship the ROCm
runtime libraries torch links against at import time), then the torch
wheels themselves. Both steps are mandatory -- torch import fails
with missing-DLL errors without the SDK wheels, even on a host that
has the HIP SDK developer toolkit installed, because torch imports
the Python-packaged runtime.
HIP_PATH is treated as an OPTIONAL version hint, not a prerequisite.
Regular users only need the AMD graphics driver (26.2.2+ for 7.2.1)
and Python 3.12, as documented by AMD. When HIP_PATH is missing or
points at an unsupported version, we default to the newest stable
release (_DEFAULT_WINDOWS_ROCM_VERSION) rather than erroring out.
NVIDIA takes precedence on mixed AMD+NVIDIA hosts. Silently returns
when no AMD GPU is visible so NVIDIA and CPU-only Windows hosts are
never touched. Both pip installs pass force_pip=True because uv's
installer has known problems with these wheels -- matches the fix in
unslothai/unsloth#4966 for bitsandbytes on Linux ROCm, and AMD's own
troubleshooting notes flag pip dep-resolver overwrite scenarios on
this procedure.
"""
# NVIDIA wins on mixed hosts -- matches the Linux branch and avoids
# overwriting a freshly installed CUDA torch with ROCm wheels.
if _has_usable_nvidia_gpu():
return
if not _has_rocm_gpu_windows():
return
# Skip when torch already links against ROCm -- mirrors the Linux
# has_hip_torch probe (line ~622) and makes this function idempotent.
# Without this guard, steps 2b and 13 in install_python_stack() would
# each re-download the full 2.1-3.9 GB wheel set even when the first
# call (or a prior setup.ps1 / install.ps1 run) already succeeded.
try:
_probe = subprocess.run(
[sys.executable, "-c",
"import torch; print(getattr(torch.version,'hip','') or '')"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 30,
)
if _probe.returncode == 0 and _probe.stdout.decode().strip():
return # already has ROCm torch
except Exception:
pass
# Radeon wheels are cp312 only. Warn (do not crash) when the venv's
# Python is not 3.12 -- pip will fail anyway with a clearer message.
if (sys.version_info.major, sys.version_info.minor) != (3, 12):
_safe_print(
_red(
f" Radeon Windows ROCm wheels require Python 3.12. "
f"Found Python {sys.version_info.major}.{sys.version_info.minor}. "
f"Install Python 3.12 from https://python.org and re-run."
)
)
return
# Prefer HIP_PATH as a version hint when available, but fall back to
# the newest stable release so users without the developer SDK still
# get a working install.
detected = _detect_rocm_version_windows()
ver: tuple[int, int]
using_default = False
if detected is not None and detected in _ROCM_WINDOWS_TORCH_WHEELS:
ver = detected
elif detected is not None:
# Detected but unsupported (e.g. 6.4): fall back to newest with a
# visible notice so the user knows what happened.
_safe_print(
_dim(
f" HIP SDK {detected[0]}.{detected[1]} is too old; "
f"falling back to newest stable wheels "
f"({_DEFAULT_WINDOWS_ROCM_VERSION[0]}.{_DEFAULT_WINDOWS_ROCM_VERSION[1]})"
)
)
ver = _DEFAULT_WINDOWS_ROCM_VERSION
using_default = True
else:
ver = _DEFAULT_WINDOWS_ROCM_VERSION
using_default = True
wheels = _ROCM_WINDOWS_TORCH_WHEELS.get(ver)
if wheels is None:
_safe_print(
_red(
f" No Windows ROCm wheel map for version {ver[0]}.{ver[1]}. "
f"Please file an issue at github.com/unslothai/unsloth/issues."
)
)
return
source = "default" if using_default else "HIP_PATH"
_safe_print(
_dim(
f" Installing Radeon ROCm wheels for Windows "
f"(rocm-rel-{ver[0]}.{ver[1]}.x, {source}) from repo.radeon.com"
)
)
if using_default:
_safe_print(
_dim(
f" Ensure your AMD graphics driver is recent; get the "
f"latest from {_AMD_RADEON_DRIVER_URL}"
)
)
# Install all 7 Radeon artefacts (4 SDK + 3 torch) in a SINGLE pip call.
#
# Why one call instead of AMD's documented two steps? torch's metadata
# declares `Requires-Dist: rocm[libraries]==<ver>` which cascades to
# `rocm-sdk-libraries-custom==<ver>`. That package does NOT exist on
# PyPI; it is only reachable via the direct URL at repo.radeon.com.
# Similarly, `rocm-sdk-core==<ver>` on PyPI is stuck at 0.1.0, wrong
# version. If we split the install and use --force-reinstall on the
# torch step, pip's resolver cascades to re-resolve all transitive
# deps, searches PyPI for rocm-sdk-libraries-custom, fails to find
# it, and aborts the whole install.
#
# Passing every URL in one command gives pip's resolver the full dep
# graph upfront. pip picks the right sources, builds the tarball via
# default build isolation, and --force-reinstall works correctly.
# Total download: ~2.1 GB for 7.2.1, ~3.9 GB for 7.1.1 (7.1.1's
# sdk_devel wheel is a massive 2.4 GB -- appears to ship debug
# symbols or unstripped libraries).
pip_install(
f"ROCm SDK + PyTorch (Windows, {ver[0]}.{ver[1]})",
"--force-reinstall",
"--no-cache-dir",
# SDK artefacts (AMD docs call this Step 1)
wheels["sdk_core"],
wheels["sdk_devel"],
wheels["sdk_libraries"],
wheels["sdk_tarball"],
# torch artefacts (AMD docs call this Step 2)
wheels["torch"],
wheels["torchvision"],
wheels["torchaudio"],
constrain = False,
force_pip = True,
)
def _ensure_rocm_torch() -> None:
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
Linux x86_64: downloads ROCm wheels from download.pytorch.org for the
detected /opt/rocm version. Windows x86_64: downloads ROCm wheels from
repo.radeon.com for the detected HIP SDK version (requires Python 3.12
because Radeon only publishes cp312 wheels). No-op on macOS, on
non-x86_64 hosts (PyTorch does not publish ROCm wheels for aarch64 /
arm64), or on mixed AMD+NVIDIA hosts (NVIDIA takes precedence). No-op
when torch already links against HIP on Linux.
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
# Explicit OS / architecture guards so the helper is safe to call
# from any context -- PyTorch only publishes ROCm wheels for x86_64,
# so aarch64 / arm64 hosts must skip this repair path instead of
# failing the update with a missing-wheel error.
if IS_MACOS:
return
if platform.machine().lower() not in {"x86_64", "amd64"}:
return
if IS_WINDOWS:
_ensure_rocm_torch_windows()
return
# NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable
if _has_usable_nvidia_gpu():
return
# Rely on _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the
# authoritative "is this actually an AMD ROCm host?" signal. The old
# gate required /opt/rocm or hipcc to exist, which breaks on
# runtime-only ROCm installs (package-managed minimal installs,
# Radeon software) that ship amd-smi/rocminfo without /opt/rocm or
# hipcc, and leaves `unsloth studio update` unable to repair a
# CPU-only venv on those systems.
if not _has_rocm_gpu():
return # no AMD GPU visible
ver = _detect_rocm_version()
if ver is None:
print(" ROCm detected but version unreadable -- skipping torch reinstall")
return
# Probe whether torch already links against HIP (ROCm is already working).
# Do NOT skip for CUDA-only builds since they are unusable on AMD-only
# hosts (the NVIDIA check above already handled mixed AMD+NVIDIA setups).
try:
probe = subprocess.run(
[
sys.executable,
"-c",
"import torch; print(getattr(torch.version,'hip','') or '')",
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 30,
)
except (OSError, subprocess.TimeoutExpired):
probe = None
has_hip_torch = (
probe is not None
and probe.returncode == 0
and probe.stdout.decode().strip() != ""
)
rocm_torch_ready = has_hip_torch
if not has_hip_torch:
# Select best matching wheel tag (newest ROCm version <= installed)
tag = next(
(
t
for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
if ver >= (maj, mn)
),
None,
)
if tag is None:
print(
f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- "
f"skipping torch reinstall"
)
else:
index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
pip_install(
f"ROCm torch ({tag})",
"--force-reinstall",
"--no-cache-dir",
"torch>=2.4,<2.11.0",
"torchvision<0.26.0",
"torchaudio<2.11.0",
"--index-url",
index_url,
constrain = False,
)
rocm_torch_ready = True
# Install bitsandbytes only when torch links against ROCm. Prefers the
# continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls
# back to PyPI when the pre-release URL is unreachable.
#
# Both installs pass force_pip=True so uv is bypassed entirely. uv's
# installer corrupts the bitsandbytes wheel on ROCm even when the
# command reports success, leaving the venv with a broken bnb import
# at runtime. install.sh has the same fix in unslothai/unsloth#4966.
# We apply it to both the pre-release URL and the PyPI fallback so the
# fix stays consistent regardless of which branch runs.
if rocm_torch_ready:
# bitsandbytes has no official Windows ROCm wheel
# (bitsandbytes-foundation/bitsandbytes#1844), so skip it entirely
# on Windows AMD. 4-bit quantization is not available on that path
# yet -- callers should install 16-bit or use GGUF inference.
if IS_WINDOWS:
return
_bnb_url = _bnb_rocm_prerelease_url()
_bnb_installed = False
if _bnb_url is not None:
_bnb_installed = pip_install_try(
"bitsandbytes (AMD, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_url,
constrain = False,
force_pip = True,
)
if not _bnb_installed:
print(
_red(
" bnb pre-release unreachable; falling back to PyPI "
"(4-bit decode will be broken on ROCm)"
)
)
if not _bnb_installed:
pip_install(
"bitsandbytes (AMD)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_BNB_ROCM_PYPI_FALLBACK,
constrain = False,
force_pip = True,
)
def _infer_no_torch() -> bool:
"""Determine whether to run in no-torch (GGUF-only) mode.
Checks UNSLOTH_NO_TORCH env var first. When unset, falls back to
platform detection so that Intel Macs automatically use GGUF-only
mode even when invoked from ``unsloth studio update`` (which does
not inject the env var).
"""
env = os.environ.get("UNSLOTH_NO_TORCH")
if env is not None:
return env.strip().lower() in ("1", "true")
return IS_MAC_INTEL
NO_TORCH = _infer_no_torch()
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal progress bar (one line, in-place).
# Set UNSLOTH_VERBOSE=1 in the environment to restore full per-step output:
# CLI: unsloth studio setup --verbose
# Linux/Mac: UNSLOTH_VERBOSE=1 ./studio/setup.sh
# Windows: $env:UNSLOTH_VERBOSE="1" ; .\studio\setup.ps1
VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1"
# Progress bar state -- updated by _progress() as each install step runs.
# _TOTAL counts: pip-upgrade + 7 shared steps + triton (non-Windows) + local-plugin + finalize
# Update _TOTAL here if you add or remove install steps in install_python_stack().
_STEP: int = 0
_TOTAL: int = 0 # set at runtime in install_python_stack() based on platform
# -- Paths --------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
REQ_ROOT = SCRIPT_DIR / "backend" / "requirements"
SINGLE_ENV = REQ_ROOT / "single-env"
CONSTRAINTS = SINGLE_ENV / "constraints.txt"
LOCAL_DD_UNSTRUCTURED_PLUGIN = (
SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed"
)
# -- Unicode-safe printing ---------------------------------------------
# On Windows the default console encoding can be a legacy code page
# (e.g. CP1252) that cannot represent Unicode glyphs such as ✅ or ❌.
# _safe_print() gracefully degrades to ASCII equivalents so the
# installer never crashes just because of a status glyph.
_UNICODE_TO_ASCII: dict[str, str] = {
"\u2705": "[OK]", # ✅
"\u274c": "[FAIL]", # ❌
"\u26a0\ufe0f": "[!]", # ⚠️ (warning + variation selector)
"\u26a0": "[!]", # ⚠ (warning without variation selector)
}
def _safe_print(*args: object, **kwargs: object) -> None:
"""Drop-in print() replacement that survives non-UTF-8 consoles."""
try:
print(*args, **kwargs)
except UnicodeEncodeError:
# Stringify, then swap emoji for ASCII equivalents
text = " ".join(str(a) for a in args)
for uni, ascii_alt in _UNICODE_TO_ASCII.items():
text = text.replace(uni, ascii_alt)
# Final fallback: replace any remaining unencodable chars
print(
text.encode(sys.stdout.encoding or "ascii", errors = "replace").decode(
sys.stdout.encoding or "ascii", errors = "replace"
),
**kwargs,
)
# ── Color support ──────────────────────────────────────────────────────
# Same logic as startup_banner: NO_COLOR disables, FORCE_COLOR or TTY enables.
def _stdout_supports_color() -> bool:
"""True if we should emit ANSI colors (matches startup_banner)."""
if os.environ.get("NO_COLOR", "").strip():
return False
if os.environ.get("FORCE_COLOR", "").strip():
return True
try:
if not sys.stdout.isatty():
return False
except (AttributeError, OSError, ValueError):
return False
if IS_WINDOWS:
try:
import ctypes
kernel32 = ctypes.windll.kernel32
handle = kernel32.GetStdHandle(-11)
mode = ctypes.c_ulong()
kernel32.GetConsoleMode(handle, ctypes.byref(mode))
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
except (ImportError, AttributeError, OSError):
return False
return True
_HAS_COLOR = _stdout_supports_color()
# Column layout — matches setup.sh step() helper:
# 2-space indent, 15-char label (dim), then value.
_LABEL = "deps"
_COL = 15
def _green(msg: str) -> str:
return f"\033[38;5;108m{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 _dim(msg: str) -> str:
return f"\033[38;5;245m{msg}\033[0m" if _HAS_COLOR else msg
def _title(msg: str) -> str:
return f"\033[38;5;150m{msg}\033[0m" if _HAS_COLOR else msg
_RULE = "\u2500" * 52
def _step(label: str, value: str, color_fn = None) -> None:
"""Print a single step line in the column format."""
if color_fn is None:
color_fn = _green
padded = label[:_COL]
print(f" {_dim(padded)}{' ' * (_COL - len(padded))}{color_fn(value)}")
def _progress(label: str) -> None:
"""Print an in-place progress bar aligned to the step column layout."""
global _STEP
_STEP += 1
if VERBOSE:
return
width = 20
filled = int(width * _STEP / _TOTAL)
bar = "=" * filled + "-" * (width - filled)
pad = " " * (_COL - len(_LABEL))
end = "\n" if _STEP >= _TOTAL else ""
sys.stdout.write(
f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}"
)
sys.stdout.flush()
def run(
label: str, cmd: list[str], *, quiet: bool = True
) -> subprocess.CompletedProcess[bytes]:
"""Run a command; on failure print output and exit."""
if VERBOSE:
_step(_LABEL, f"{label}...", _dim)
result = subprocess.run(
cmd,
stdout = subprocess.PIPE if quiet else None,
stderr = subprocess.STDOUT if quiet else None,
)
if result.returncode != 0:
_step("error", f"{label} failed (exit code {result.returncode})", _red)
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"}
# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode).
# These packages either *are* torch extensions or have unconditional
# ``Requires-Dist: torch`` in their published metadata, so installing
# them would pull torch back into the environment.
NO_TORCH_SKIP_PACKAGES = {
"torch-stoi",
"timm",
"torchcodec",
"torch-c-dlpack-ext",
"openai-whisper",
"transformers-cfg",
}
# -- uv bootstrap ------------------------------------------------------
USE_UV = False # Set by _bootstrap_uv() at the start of install_python_stack()
UV_NEEDS_SYSTEM = False # Set by _bootstrap_uv() via probe
def _bootstrap_uv() -> bool:
"""Check if uv is available and probe whether --system is needed."""
global UV_NEEDS_SYSTEM
if not shutil.which("uv"):
return False
# Probe: try a dry-run install targeting the current Python explicitly.
# Without --python, uv can ignore the activated venv on some platforms.
probe = subprocess.run(
["uv", "pip", "install", "--dry-run", "--python", sys.executable, "pip"],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if probe.returncode != 0:
# Retry with --system (some envs need it when uv can't find a venv)
probe_sys = subprocess.run(
["uv", "pip", "install", "--dry-run", "--system", "pip"],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if probe_sys.returncode != 0:
return False # uv is broken, fall back to pip
UV_NEEDS_SYSTEM = True
return True
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 _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.
Strips uv-only flags like --upgrade-package that pip doesn't understand.
"""
cmd = [sys.executable, "-m", "pip", "install"]
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
def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
"""Build a uv pip install command with translated flags."""
cmd = ["uv", "pip", "install"]
if UV_NEEDS_SYSTEM:
cmd.append("--system")
# Always pass --python so uv targets the correct environment.
# Without this, uv can ignore an activated venv and install into
# the system Python (observed on Colab and similar environments).
cmd.extend(["--python", sys.executable])
cmd.extend(_translate_pip_args_for_uv(args))
# 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
def pip_install_try(
label: str,
*args: str,
constrain: bool = True,
force_pip: bool = False,
) -> bool:
"""Like pip_install but returns False on failure instead of exiting.
For optional installs with a follow-up fallback.
``force_pip`` skips uv entirely and goes straight to ``python -m pip
install``. Used for wheels that uv installs incorrectly -- notably
the bitsandbytes continuous-release_main wheel on ROCm, see
unslothai/unsloth#4966.
"""
constraint_args: list[str] = []
if constrain and CONSTRAINTS.is_file():
constraint_args = ["-c", str(CONSTRAINTS)]
if USE_UV and not force_pip:
cmd = _build_uv_cmd(args) + constraint_args
else:
cmd = _build_pip_cmd(args) + constraint_args
if VERBOSE:
_step(_LABEL, f"{label}...", _dim)
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode == 0:
return True
if VERBOSE and result.stdout:
print(result.stdout.decode(errors = "replace"))
return False
def pip_install(
label: str,
*args: str,
req: Path | None = None,
constrain: bool = True,
force_pip: bool = False,
) -> None:
"""Build and run a pip install command (uses uv when available, falls back to pip).
``force_pip`` skips the uv attempt entirely and goes straight to
``python -m pip install``. Use this for wheels that uv installs
incorrectly -- bitsandbytes pre-release ROCm wheels are the known
culprit, see unslothai/unsloth#4966.
"""
constraint_args: list[str] = []
if constrain and CONSTRAINTS.is_file():
constraint_args = ["-c", str(CONSTRAINTS)]
actual_req = req
temp_reqs: list[Path] = []
if req is not None and IS_WINDOWS and WINDOWS_SKIP_PACKAGES:
actual_req = _filter_requirements(req, WINDOWS_SKIP_PACKAGES)
temp_reqs.append(actual_req)
if actual_req is not None and NO_TORCH and NO_TORCH_SKIP_PACKAGES:
actual_req = _filter_requirements(actual_req, NO_TORCH_SKIP_PACKAGES)
temp_reqs.append(actual_req)
req_args: list[str] = []
if actual_req is not None:
req_args = ["-r", str(actual_req)]
try:
if USE_UV and not force_pip:
uv_cmd = _build_uv_cmd(args) + constraint_args + req_args
if VERBOSE:
print(f" {label}...")
result = subprocess.run(
uv_cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode == 0:
return
print(_red(f" uv failed, falling back to pip..."))
if result.stdout:
print(result.stdout.decode(errors = "replace"))
pip_cmd = _build_pip_cmd(args) + constraint_args + req_args
run(f"{label} (pip)" if USE_UV and not force_pip else label, pip_cmd)
finally:
for temp_req in temp_reqs:
temp_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:
_step(_LABEL, f"package {package_name} not found, skipping patch", _red)
return
location = None
for line in result.stdout.splitlines():
if line.lower().startswith("location:"):
location = line.split(":", 1)[1].strip()
break
if not location:
_step(_LABEL, f"could not locate {package_name}", _red)
return
dest = Path(location) / relative_path
_step(_LABEL, f"patching {dest.name} in {package_name}...", _dim)
download_file(url, dest)
# -- Main install sequence ---------------------------------------------
def install_python_stack() -> int:
global USE_UV, _STEP, _TOTAL
_STEP = 0
# 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
if IS_MACOS:
base_total -= 1 # triton step is skipped on macOS
# ROCm torch check steps (Linux + Windows, non-macOS, non-no-torch):
# one early check (step 2b) and one final repair (step 13).
if not IS_MACOS and not NO_TORCH:
base_total += 2
_TOTAL = (base_total - 1) if skip_base else base_total
# 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. 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:
# pip may not exist yet (uv-created venvs omit it). Try ensurepip
# first, then upgrade. Only fall back to a direct upgrade when pip
# is already present.
_has_pip = (
subprocess.run(
[sys.executable, "-m", "pip", "--version"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
).returncode
== 0
)
if not _has_pip:
run(
"Bootstrapping pip via ensurepip",
[sys.executable, "-m", "ensurepip", "--upgrade"],
)
else:
run(
"Upgrading pip",
[sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
)
# 3. Core packages: unsloth-zoo + unsloth (or custom package name)
if skip_base:
pass
elif NO_TORCH:
# No-torch update path: install unsloth + unsloth-zoo with --no-deps
# (current PyPI metadata still declares torch as a hard dep), then
# runtime deps with --no-deps (avoids transitive torch).
_progress("base packages (no torch)")
pip_install(
f"Updating {package_name} + unsloth-zoo (no-torch mode)",
"--no-cache-dir",
"--no-deps",
"--upgrade-package",
package_name,
"--upgrade-package",
"unsloth-zoo",
package_name,
"unsloth-zoo",
)
pip_install(
"Installing no-torch runtime deps",
"--no-cache-dir",
"--no-deps",
req = REQ_ROOT / "no-torch-runtime.txt",
)
if local_repo:
pip_install(
"Overlaying local repo (editable)",
"--no-cache-dir",
"--no-deps",
"-e",
local_repo,
constrain = False,
)
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",
)
# 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the
# venv received CPU-only torch (common when pip resolves torch from PyPI).
# Must come immediately after base packages so torch is present for inspection.
# On Linux this pulls from download.pytorch.org; on Windows from
# repo.radeon.com. _ensure_rocm_torch() dispatches internally.
if not IS_MACOS and not NO_TORCH:
_progress("ROCm torch check")
_ensure_rocm_torch()
# 3. Extra dependencies
_progress("unsloth extras")
pip_install(
"Installing additional unsloth dependencies",
"--no-cache-dir",
req = REQ_ROOT / "extras.txt",
)
# 3b. Extra dependencies (no-deps) -- audio model support etc.
_progress("extra codecs")
pip_install(
"Installing extras (no-deps)",
"--no-deps",
"--no-cache-dir",
req = REQ_ROOT / "extras-no-deps.txt",
)
# 4. Overrides (torchao, transformers) -- force-reinstall
# Skip entirely when torch is unavailable (e.g. Intel Mac GGUF-only mode)
# because overrides.txt contains torchao which requires torch.
if NO_TORCH:
_progress("dependency overrides (skipped, no torch)")
else:
_progress("dependency overrides")
pip_install(
"Installing dependency overrides",
"--force-reinstall",
"--no-cache-dir",
req = REQ_ROOT / "overrides.txt",
)
# 5. Triton kernels (no-deps, from source)
# Skip on Windows (no support) and macOS (no support).
if not IS_WINDOWS and not IS_MACOS:
_progress("triton kernels")
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 feature/llama-cpp-windows-support 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",
# )
# # 7a. 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",
# )
# # 7b. Patch : override save.py with fix from feature/llama-cpp-windows-support
# patch_package_file(
# "unsloth",
# os.path.join("unsloth", "save.py"),
# "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py",
# )
# 8. Studio dependencies
_progress("studio deps")
pip_install(
"Installing studio dependencies",
"--no-cache-dir",
req = REQ_ROOT / "studio.txt",
)
# 9. Data-designer dependencies
_progress("data designer deps")
pip_install(
"Installing data-designer base dependencies",
"--no-cache-dir",
req = SINGLE_ENV / "data-designer-deps.txt",
)
# 10. Data-designer packages (no-deps to avoid conflicts)
_progress("data designer")
pip_install(
"Installing data-designer",
"--no-cache-dir",
"--no-deps",
req = SINGLE_ENV / "data-designer.txt",
)
# 11. Local Data Designer seed plugin
if not LOCAL_DD_UNSTRUCTURED_PLUGIN.is_dir():
_safe_print(
_red(
f"❌ Missing local plugin directory: {LOCAL_DD_UNSTRUCTURED_PLUGIN}",
),
)
return 1
_progress("local plugin")
pip_install(
"Installing local data-designer unstructured plugin",
"--no-cache-dir",
"--no-deps",
str(LOCAL_DD_UNSTRUCTURED_PLUGIN),
constrain = False,
)
# 12. Patch metadata for single-env compatibility
_progress("finalizing")
run(
"Patching single-env metadata",
[sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
)
# 13. AMD ROCm: final torch repair. Multiple install steps above can
# pull in CUDA / CPU torch from PyPI (base packages, extras,
# overrides, studio deps, etc.). Running the repair as the very
# last step ensures ROCm torch is in place at runtime, regardless
# of which intermediate step clobbered it. Same behavior on Linux
# (download.pytorch.org wheels) and Windows (repo.radeon.com).
if not IS_MACOS and not NO_TORCH:
_progress("ROCm torch (final)")
_ensure_rocm_torch()
# 14. Final check (silent; third-party conflicts are expected)
subprocess.run(
[sys.executable, "-m", "pip", "check"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
)
_step(_LABEL, "installed")
return 0
if __name__ == "__main__":
sys.exit(install_python_stack())
Reference in a new issue View git blame Copy permalink