2dce261f4a
The previous Windows AMD install was incomplete in two material ways, both surfaced while verifying whether PyTorch has upstream Windows ROCm wheels (it does not -- pytorch.org's get-started page states "ROCm is not available on Windows" and every wheel under download.pytorch.org /whl/rocm7.x is manylinux_2_28_x86_64 only; upstream work is tracked in pytorch/pytorch#159520 and targeted for a future release). repo.radeon.com therefore remains the only source for Windows ROCm torch until that RFC lands, and AMD's install docs at rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html document a two-step pip procedure that we were only implementing half of. Fixing both bugs here so the PR actually produces a working torch.version.hip import on a fresh Windows host. Bug 1: missing ROCm SDK wheels install step AMD's procedure first installs rocm_sdk_core, rocm_sdk_devel, rocm_sdk_libraries_custom, and rocm-<ver>.tar.gz (about 1.4 GB total). These wheels ship the ROCm runtime libraries that torch links against at import time. Without them `import torch` fails with missing DLL errors even when the torch wheels themselves are installed. Both 7.1.1 and 7.2.1 need this step; 7.1.1 stamps its SDK wheels with the `0.1.dev0` version string while 7.2.1 uses `7.2.1`. This adds the 4 SDK URLs per release to the wheel map and installs them as Step 1 ahead of the existing torch install (Step 2). Both steps are passed in a single pip call each so pip's dep resolver does not reset torch between wheels (matches AMD's troubleshooting guidance for the same failure mode). Bug 2: HIP SDK was a hard prerequisite, should be optional install.ps1 / setup.ps1 / install_python_stack.py all errored out when $env:HIP_PATH was absent, pointing users at the HIP SDK download page. But the HIP SDK developer toolkit at C:\Program Files\AMD\ROCm\ is for people compiling HIP kernels, not for running PyTorch. AMD's docs list only (a) the AMD graphics driver 26.2.2+ and (b) Python 3.12 as prerequisites. Gating on HIP_PATH was blocking the exact audience that #4280 is about -- regular Radeon users running Unsloth. HIP_PATH is now an optional version hint. When present and valid we use it to select the matching ROCm release; when absent or unsupported (e.g. HIP 6.4) we fall back to the newest stable release (_DEFAULT_WINDOWS_ROCM_VERSION = (7, 2)) and print a visible note pointing at the graphics driver download page. The HIP SDK install prompts have been removed from all three files. Use pip (not uv) for the Radeon wheels Both the SDK and torch install steps call `python -m pip install` directly via the new force_pip=True path added in the previous commit. AMD's documented procedure uses pip, uv has known wheel-corruption issues on similar large ROCm wheels (unslothai/unsloth#4966 for bitsandbytes), and pip is the combination AMD validates. This matches the fix applied to bitsandbytes on Linux ROCm.
1407 lines
51 KiB
Python
1407 lines
51 KiB
Python
#!/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. These wheels ship the ROCm runtime libraries
|
|
# that torch links against at import time. They total about 1.4
|
|
# GB. The HIP SDK developer toolkit (HIP_PATH) is NOT a substitute
|
|
# -- torch imports the Python-packaged runtime from rocm_sdk_*.
|
|
# Step 2: torch + torchvision + torchaudio. About 820 MB.
|
|
#
|
|
# 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"
|
|
)
|
|
_HIP_SDK_DOWNLOAD_URL = (
|
|
"https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.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 = 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:
|
|
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
|
|
|
|
# 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}"
|
|
)
|
|
)
|
|
|
|
# Step 1: ROCm SDK wheels (runtime libraries torch imports). ~1.4 GB
|
|
# download on a clean venv.
|
|
pip_install(
|
|
f"ROCm SDK (Windows, {ver[0]}.{ver[1]})",
|
|
"--force-reinstall",
|
|
"--no-cache-dir",
|
|
wheels["sdk_core"],
|
|
wheels["sdk_devel"],
|
|
wheels["sdk_libraries"],
|
|
wheels["sdk_tarball"],
|
|
constrain = False,
|
|
force_pip = True,
|
|
)
|
|
|
|
# Step 2: torch wheels. ~820 MB download.
|
|
pip_install(
|
|
f"ROCm torch (Windows, {ver[0]}.{ver[1]})",
|
|
"--force-reinstall",
|
|
"--no-cache-dir",
|
|
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())
|