install: converge torch-index pin detection via a per-venv marker
Introduce a torch-index MARKER that records the exact wheel --index-url used
after each successful torch install, so `unsloth studio update` / repair makes
the "did the pinned index change?" decision by an EXACT string compare rather
than inferring it from the wheel +rocm/+cu version tag. The tag cannot encode
the AMD per-arch gfx family (two 2.11 gfx indexes both install +rocm7.13.0), so
the tag heuristic missed a gfx1151 -> gfx120X-all switch and a custom-URL swap.
Marker path is per-venv (.unsloth-torch-index), one line = the resolved index
URL, written atomically (temp + rename). Path, format and normalization are
shared across all four installers (install.sh, install_python_stack.py,
setup.ps1, install.ps1).
- Reapply gfx pins on a per-arch target change: the marker's exact compare
reinstalls when the pinned index differs, even when both wheels share a tag.
- Honor custom ROCm URL pins during repair: an explicit index whose leaf is not
rocm/gfx/cu/cpu (e.g. simple, current) now reinstalls torch VERBATIM from the
pin when it differs from the marker ("URL wins verbatim").
- Align the KNOWN-2.11 rocm/gfx set to exactly rocm7.2 plus the gfx allowlist
gfx120x-all/gfx1151/gfx1150 in every language; stop treating an unknown newer
rocm (rocm7.3, which does not exist) as the 2.11 line speculatively.
Backward compatible: with no marker (old venvs, torch installed out-of-band) the
existing +rocm/version-tag heuristics still decide, and a matching marker never
reinstall-loops. A cu128 CUDA pin stays a CUDA pin; custom and current leaves are
not CUDA. Adds marker tests (py/sh/ps) plus cross-installer parity checks.
This commit is contained in:
Daniel Han
parent
1213135cd1
commit
753ef2f5cf
10 changed files with 1149 additions and 125 deletions
41
install.ps1
41
install.ps1
|
|
@ -1983,6 +1983,30 @@ exit 0
|
|||
return "$baseUrl/cu126"
|
||||
}
|
||||
|
||||
# ── Torch-index marker ───────────────────────────────────────────────────
|
||||
# After a successful torch install, record the exact wheel --index-url used at
|
||||
# a stable per-venv path so a later `unsloth studio update` (setup.ps1 /
|
||||
# install_python_stack.py) can detect a pin change by an EXACT string compare
|
||||
# rather than the wheel version-tag heuristic. Path/format MUST match
|
||||
# install.sh, studio/setup.ps1 and install_python_stack.py:
|
||||
# <VenvDir>\.unsloth-torch-index (single line = the resolved index URL)
|
||||
# install.ps1 only WRITES the marker; setup.ps1 reads it during stale detection.
|
||||
function Write-TorchIndexMarker {
|
||||
param([string]$VenvDir, [string]$IndexUrl)
|
||||
if ([string]::IsNullOrWhiteSpace($VenvDir)) { return }
|
||||
if ([string]::IsNullOrWhiteSpace($IndexUrl)) { return }
|
||||
if (-not (Test-Path -LiteralPath $VenvDir -PathType Container)) { return }
|
||||
$marker = Join-Path $VenvDir ".unsloth-torch-index"
|
||||
$tmp = "$marker.$PID.tmp"
|
||||
try {
|
||||
# Single LF-terminated line, no BOM (parity with the sh/py writers).
|
||||
[System.IO.File]::WriteAllText($tmp, ($IndexUrl.Trim() + "`n"), (New-Object System.Text.UTF8Encoding($false)))
|
||||
Move-Item -LiteralPath $tmp -Destination $marker -Force -ErrorAction Stop
|
||||
} catch {
|
||||
try { if (Test-Path -LiteralPath $tmp) { Remove-Item -LiteralPath $tmp -Force -ErrorAction SilentlyContinue } } catch {}
|
||||
}
|
||||
}
|
||||
|
||||
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
|
||||
# torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
|
||||
# matching setup.ps1's stale-venv parse.
|
||||
|
|
@ -2117,7 +2141,11 @@ exit 0
|
|||
$_pinLeaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLower()
|
||||
$_pinRocm211 = $false
|
||||
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') {
|
||||
$_pinRocm211 = ([int]$Matches[1] -gt 7) -or ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -ge 2)
|
||||
# Only KNOWN-2.11 rocm indexes (rocm7.2) get the 2.11 floor; do not floor
|
||||
# an unknown newer rocm speculatively (rocm7.3 does not exist). Matches
|
||||
# install.sh's rocm7.2 KNOWN-2.11 leaf, setup.ps1's Test-RocmKnown211Version
|
||||
# and _ROCM_KNOWN_TORCH211_VERSIONS in install_python_stack.py.
|
||||
$_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
|
||||
}
|
||||
# Only the gfx families the AMD arch map above pins to torch 2.11 need the
|
||||
# floor here (gfx120X-all, gfx1151, gfx1150 -- the _grouped_mm bug arches).
|
||||
|
|
@ -2407,6 +2435,17 @@ exit 0
|
|||
}
|
||||
}
|
||||
|
||||
# ── Record the resolved torch wheel index (marker) ──
|
||||
# Torch is now resolved; write the exact --index-url used so setup.ps1 /
|
||||
# install_python_stack.py can detect a later pin change by an EXACT string
|
||||
# compare instead of the version-tag heuristic. Reflects the installed family:
|
||||
# $ROCmIndexUrl when the ROCm path ran, else the CUDA/CPU/pinned $TorchIndexUrl.
|
||||
# Skipped for --no-torch (nothing installed). Matches install.sh / setup.ps1.
|
||||
if (-not $SkipTorch) {
|
||||
$MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchIndexUrl }
|
||||
Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl
|
||||
}
|
||||
|
||||
# Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped
|
||||
# for --local: the editable install above already makes _PACKAGE_ROOT in
|
||||
# unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__).
|
||||
|
|
|
|||
64
install.sh
64
install.sh
|
|
@ -2215,6 +2215,61 @@ _torch_index_repairable() {
|
|||
esac
|
||||
}
|
||||
|
||||
# ── Torch-index marker ───────────────────────────────────────────────────────
|
||||
# After a successful torch install this records the exact wheel --index-url used
|
||||
# at a stable per-venv path so `unsloth studio update` / setup.ps1 /
|
||||
# install_python_stack.py can make the "did the pinned index change?" decision by
|
||||
# an EXACT string compare instead of inferring it from the wheel version tag
|
||||
# (which cannot encode the AMD per-arch gfx family). The path and format MUST match
|
||||
# install_python_stack.py, setup.ps1 and install.ps1:
|
||||
# <venv_dir>/.unsloth-torch-index (single line = the resolved index URL)
|
||||
_TORCH_INDEX_MARKER_NAME=".unsloth-torch-index"
|
||||
|
||||
# Normalise a wheel index URL for exact marker/pin comparison: trim whitespace,
|
||||
# strip ALL trailing slashes, lowercase ONLY the final path segment (the leaf).
|
||||
# Mirrors _normalize_index_url in install_python_stack.py / setup.ps1 / install.ps1.
|
||||
_normalize_index_url() {
|
||||
_n_url="$1"
|
||||
# Trim leading/trailing whitespace.
|
||||
_n_url="${_n_url#"${_n_url%%[![:space:]]*}"}"; _n_url="${_n_url%"${_n_url##*[![:space:]]}"}"
|
||||
[ -n "$_n_url" ] || { printf '%s' ""; return; }
|
||||
# Strip all trailing slashes.
|
||||
while [ "${_n_url%/}" != "$_n_url" ]; do _n_url="${_n_url%/}"; done
|
||||
[ -n "$_n_url" ] || { printf '%s' ""; return; }
|
||||
case "$_n_url" in
|
||||
*/*)
|
||||
_n_head="${_n_url%/*}"
|
||||
_n_leaf="${_n_url##*/}"
|
||||
_n_leaf=$(printf '%s' "$_n_leaf" | tr '[:upper:]' '[:lower:]')
|
||||
printf '%s/%s' "$_n_head" "$_n_leaf"
|
||||
;;
|
||||
*)
|
||||
printf '%s' "$_n_url" | tr '[:upper:]' '[:lower:]'
|
||||
;;
|
||||
esac
|
||||
}
|
||||
|
||||
# Write the resolved torch --index-url ($2) into the marker under venv dir ($1),
|
||||
# atomically (temp file + mv). Best-effort: a write failure never aborts the
|
||||
# install (the repair path then falls back to the version-tag heuristics). A blank
|
||||
# URL is ignored (nothing meaningful to record).
|
||||
_write_torch_index_marker() {
|
||||
_wm_venv="$1"
|
||||
_wm_url="$2"
|
||||
[ -n "$_wm_venv" ] || return 0
|
||||
[ -d "$_wm_venv" ] || return 0
|
||||
_wm_url="${_wm_url#"${_wm_url%%[![:space:]]*}"}"; _wm_url="${_wm_url%"${_wm_url##*[![:space:]]}"}"
|
||||
[ -n "$_wm_url" ] || return 0
|
||||
_wm_marker="$_wm_venv/$_TORCH_INDEX_MARKER_NAME"
|
||||
_wm_tmp="$_wm_marker.$$.tmp"
|
||||
if printf '%s\n' "$_wm_url" > "$_wm_tmp" 2>/dev/null; then
|
||||
mv -f "$_wm_tmp" "$_wm_marker" 2>/dev/null || rm -f "$_wm_tmp" 2>/dev/null || true
|
||||
else
|
||||
rm -f "$_wm_tmp" 2>/dev/null || true
|
||||
fi
|
||||
return 0
|
||||
}
|
||||
|
||||
get_radeon_wheel_url() {
|
||||
# Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing
|
||||
# contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/,
|
||||
|
|
@ -3090,6 +3145,15 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
|
|||
fi
|
||||
fi
|
||||
|
||||
# ── Record the resolved torch wheel index (marker) ──
|
||||
# Torch is now fully resolved; write the exact --index-url used so `unsloth studio
|
||||
# update` (install_python_stack.py / setup.ps1) can detect a later pin change by an
|
||||
# exact string compare rather than the version-tag heuristic. Only when torch was
|
||||
# actually installed from a resolved index (skip --no-torch / no-URL fallback).
|
||||
if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
|
||||
_write_torch_index_marker "$VENV_DIR" "$TORCH_INDEX_URL"
|
||||
fi
|
||||
|
||||
# ── Run studio setup ──
|
||||
tauri_log "STEP" "Running Studio setup"
|
||||
# When --local, use the repo's own setup.sh directly.
|
||||
|
|
|
|||
|
|
@ -81,6 +81,16 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
|
|||
# stay bare, so an override to one of them must NOT be forced onto the 2.11 line.
|
||||
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
|
||||
|
||||
# The pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (verified against
|
||||
# download.pytorch.org): rocm7.2 -> torch 2.11.0 is the ONLY stable 2.11 rocm
|
||||
# index today (rocm7.1 -> 2.10.0, rocm6.4 -> 2.9.1; rocm7.3 / torch 2.12 do NOT
|
||||
# exist as stable). Do NOT treat an unknown newer rocm (rocm7.3, rocm8.0, ...) as
|
||||
# 2.11 speculatively -- that is exactly the mismatch bug tracked in the review.
|
||||
# MUST match the rocm leaf in the KNOWN-2.11 case in install.sh / setup.ps1 /
|
||||
# install.ps1 (rocm7.2 there too). Bump alongside those when a new stable rocm
|
||||
# index publishes torch 2.11+. Stored as (major, minor) tuples for exact compares.
|
||||
_ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)})
|
||||
|
||||
# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x).
|
||||
_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
|
||||
"rocm7.2": (
|
||||
|
|
@ -110,6 +120,127 @@ _PYTORCH_WHL_BASE = (
|
|||
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
|
||||
).rstrip("/")
|
||||
|
||||
# ── Torch-index marker ─────────────────────────────────────────────────────────
|
||||
# After a successful torch install/reinstall, record the exact wheel --index-url
|
||||
# used at a stable per-venv path. On `studio update`/repair the marker turns the
|
||||
# "did the pinned index change?" decision into an EXACT string compare instead of
|
||||
# inferring it from the wheel's +rocm/+cu version tag (which cannot encode the AMD
|
||||
# per-arch gfx family: two gfx 2.11 indexes both install +rocm7.13.0). This
|
||||
# dissolves the per-arch-switch (gfx1151 -> gfx120X-all) and custom-URL
|
||||
# (/simple, /current) cases the version-tag heuristics cannot see.
|
||||
#
|
||||
# The path/format MUST match install.sh, install.ps1 and setup.ps1:
|
||||
# <venv_prefix>/.unsloth-torch-index (single line = the resolved index URL)
|
||||
# Written atomically (temp file + os.replace). A missing/empty/corrupt marker is
|
||||
# treated as absent, so old venvs (and torch installed out-of-band) fall back to
|
||||
# the existing +rocm/version-tag heuristics -- backward compatibility is required.
|
||||
_TORCH_INDEX_MARKER_NAME = ".unsloth-torch-index"
|
||||
|
||||
|
||||
def _normalize_index_url(url: "str | None") -> "str | None":
|
||||
"""Canonicalise a wheel index URL for exact marker/pin comparison.
|
||||
|
||||
Trims surrounding whitespace, strips ALL trailing slashes, and lowercases only
|
||||
the FINAL path segment (the wheel-family leaf: cu128 / cpu / rocm7.2 / gfx1151 /
|
||||
gfx120X-all). The host part is left untouched (it may be case-sensitive on some
|
||||
mirrors); the leaf is lowercased so the canonical gfx120X-all (capital X) and
|
||||
AMD's lowercase pip leaf gfx120x-all compare equal. MUST match the same
|
||||
normalization in install.sh / setup.ps1 / install.ps1. Returns None for an
|
||||
empty/whitespace-only input. Pure function.
|
||||
"""
|
||||
if url is None:
|
||||
return None
|
||||
url = url.strip()
|
||||
if not url:
|
||||
return None
|
||||
url = url.rstrip("/")
|
||||
if not url:
|
||||
return None
|
||||
head, sep, leaf = url.rpartition("/")
|
||||
if sep:
|
||||
return f"{head}/{leaf.lower()}"
|
||||
return url.lower()
|
||||
|
||||
|
||||
def _torch_index_marker_path() -> Path:
|
||||
"""Path to the per-venv torch-index marker (see _TORCH_INDEX_MARKER_NAME).
|
||||
|
||||
Anchored at sys.prefix (the venv the installer targets via --python
|
||||
sys.executable), so it matches install.sh's $VENV_DIR and the PowerShell
|
||||
$VenvDir marker location.
|
||||
"""
|
||||
return Path(sys.prefix) / _TORCH_INDEX_MARKER_NAME
|
||||
|
||||
|
||||
def _read_torch_index_marker() -> "str | None":
|
||||
"""Return the recorded torch --index-url from the marker, else None.
|
||||
|
||||
None when the marker is missing, empty, or unreadable (corrupt/permission) --
|
||||
all treated as "no marker" so the caller falls back to the version-tag
|
||||
heuristics. The stored URL is returned VERBATIM (not normalized); callers
|
||||
normalize both sides before comparing.
|
||||
"""
|
||||
try:
|
||||
text = _torch_index_marker_path().read_text(encoding = "utf-8")
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
line = text.strip()
|
||||
return line or None
|
||||
|
||||
|
||||
def _write_torch_index_marker(index_url: "str | None") -> None:
|
||||
"""Record the resolved torch wheel --index-url at the per-venv marker path.
|
||||
|
||||
Best-effort and atomic (temp file in the same dir + os.replace). Never raises:
|
||||
a marker write failure must not abort an otherwise-successful install (the
|
||||
repair path then falls back to the heuristics, same as an old venv). A blank
|
||||
index_url is ignored (nothing meaningful to record).
|
||||
"""
|
||||
if not index_url or not index_url.strip():
|
||||
return
|
||||
marker = _torch_index_marker_path()
|
||||
payload = index_url.strip() + "\n"
|
||||
try:
|
||||
marker.parent.mkdir(parents = True, exist_ok = True)
|
||||
fd, tmp = tempfile.mkstemp(prefix = ".unsloth-torch-index.", dir = str(marker.parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding = "utf-8") as fh:
|
||||
fh.write(payload)
|
||||
os.replace(tmp, str(marker))
|
||||
except OSError:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
except OSError:
|
||||
# Non-fatal: fall back to writing directly, then give up silently.
|
||||
try:
|
||||
marker.write_text(payload, encoding = "utf-8")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _marker_pin_mismatch(pin_url: str) -> "bool | None":
|
||||
"""Compare an explicit torch-index pin against the recorded marker.
|
||||
|
||||
Returns:
|
||||
* True -> the pin resolves to a DIFFERENT index than the marker records
|
||||
(exact compare after _normalize_index_url) -> reinstall needed.
|
||||
* False -> the pin matches the marker exactly -> no reinstall (no loop).
|
||||
* None -> no usable marker (missing/empty/corrupt) -> the caller must fall
|
||||
back to the +rocm/version-tag heuristics (backward compatibility).
|
||||
|
||||
This is the EXACT signal the version-tag heuristics cannot provide: switching
|
||||
one 2.11 gfx pin to another (gfx1151 -> gfx120X-all) changes the pinned URL but
|
||||
not the installed +rocm7.13.0 wheel tag. Pure w.r.t. its arg (reads the marker).
|
||||
"""
|
||||
marker = _read_torch_index_marker()
|
||||
if marker is None:
|
||||
return None
|
||||
return _normalize_index_url(pin_url) != _normalize_index_url(marker)
|
||||
|
||||
|
||||
# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are
|
||||
# pinned to the torch<2.11 family rather than left bare: the install uses an
|
||||
# exclusive --index-url (no PyPI fallback), so a bare name could resolve a
|
||||
|
|
@ -1145,8 +1276,10 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
|
|||
# pin -> mismatch; otherwise mismatch only when the installed torch is 2.11+.
|
||||
return (not _inst_has_rocm) or _inst_is_211
|
||||
|
||||
# rocmX.Y pin.
|
||||
_pin_is_211 = _pin_ver >= (7, 2) if _pin_ver is not None else False
|
||||
# rocmX.Y pin. Only the KNOWN-2.11 rocm indexes are treated as the 2.11 line;
|
||||
# an unknown newer rocm (rocm7.3, rocm8.0, ...) is NOT floored speculatively
|
||||
# (it does not exist yet). Aligns with install.sh / setup.ps1 / install.ps1.
|
||||
_pin_is_211 = _pin_ver in _ROCM_KNOWN_TORCH211_VERSIONS if _pin_ver is not None else False
|
||||
if _pin_ver is not None and _inst_ver is not None:
|
||||
# Both ROCm versions readable: exact (major, minor) comparison. A generic
|
||||
# rocm7.2 pin over the AMD per-arch (+rocm7.13.x) wheel compares (7, 2) vs
|
||||
|
|
@ -1201,6 +1334,69 @@ def _explicit_cuda_torch_index_url() -> "str | None":
|
|||
return url if _is_cuda_family_leaf(leaf) else None
|
||||
|
||||
|
||||
def _explicit_unknown_family_torch_index_url() -> "str | None":
|
||||
"""The pinned index URL when its leaf names NO known torch family, else None.
|
||||
|
||||
A "known" family leaf is rocm* / gfx* / cpu / cuXXX -- the ones the dedicated
|
||||
_explicit_{rocm,cpu,cuda}_torch_index_url helpers already classify. Anything
|
||||
else (a private PEP 503 mirror ending in /simple, /current, /custom, ...) is an
|
||||
UNKNOWN family: the version-tag heuristics cannot infer whether it is stale, so
|
||||
the marker drives the decision and, when it differs (or is absent), the URL is
|
||||
reinstalled VERBATIM -- "URL wins verbatim". Matches the unknown-leaf branch in
|
||||
install.sh / setup.ps1 / install.ps1.
|
||||
"""
|
||||
url = _explicit_torch_index_url()
|
||||
if url is None:
|
||||
return None
|
||||
leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
|
||||
if leaf.startswith(("rocm", "gfx")) or leaf == "cpu" or _is_cuda_family_leaf(leaf):
|
||||
return None
|
||||
return url
|
||||
|
||||
|
||||
def _ensure_verbatim_torch_index() -> None:
|
||||
"""Reinstall torch/vision/audio VERBATIM from an explicit custom index pin.
|
||||
|
||||
Handles the "URL wins verbatim" case for an explicit UNSLOTH_TORCH_INDEX_URL
|
||||
(or _FAMILY) whose leaf names no known family (e.g. a private mirror ending in
|
||||
/simple or /current). The other _ensure_* helpers all return None for such a
|
||||
pin (it is neither rocm/gfx nor cpu nor cuXXX), so without this the pin would be
|
||||
silently ignored and the GPU-probed default index used instead.
|
||||
|
||||
Fires ONLY when the marker exists and records a DIFFERENT index than the pin
|
||||
(or after this reinstalls, the marker is rewritten to match). With NO marker it
|
||||
is a no-op: an old venv (or torch installed out-of-band) must not be blindly
|
||||
force-reinstalled from an unverified custom index -- backward compatibility.
|
||||
macOS/no-torch: skipped (no torch to repair). The install uses the pinned URL
|
||||
exclusively (--index-url) with bare specs so it "wins verbatim".
|
||||
"""
|
||||
if NO_TORCH or IS_MACOS:
|
||||
return
|
||||
pin = _explicit_unknown_family_torch_index_url()
|
||||
if pin is None:
|
||||
return
|
||||
_mismatch = _marker_pin_mismatch(pin)
|
||||
if _mismatch is not True:
|
||||
# None -> no marker (fall back / do nothing); False -> already this index.
|
||||
return
|
||||
print(
|
||||
f" explicit torch index pin ({pin}) differs from the recorded index -- "
|
||||
f"reinstalling torch verbatim from it"
|
||||
)
|
||||
pip_install(
|
||||
"torch (pinned custom index)",
|
||||
"--force-reinstall",
|
||||
"--no-cache-dir",
|
||||
"torch",
|
||||
"torchvision",
|
||||
"torchaudio",
|
||||
"--index-url",
|
||||
pin,
|
||||
constrain = False,
|
||||
)
|
||||
_write_torch_index_marker(pin)
|
||||
|
||||
|
||||
def _ensure_cuda_torch() -> None:
|
||||
"""Repair a venv whose torch is a ROCm build on an NVIDIA host.
|
||||
|
||||
|
|
@ -1324,6 +1520,7 @@ def _ensure_cuda_torch() -> None:
|
|||
index_url,
|
||||
constrain = False,
|
||||
)
|
||||
_write_torch_index_marker(index_url)
|
||||
|
||||
|
||||
def _ensure_cpu_torch() -> None:
|
||||
|
|
@ -1394,6 +1591,7 @@ def _ensure_cpu_torch() -> None:
|
|||
pin,
|
||||
constrain = False,
|
||||
)
|
||||
_write_torch_index_marker(pin)
|
||||
|
||||
|
||||
def _ensure_rocm_torch() -> None:
|
||||
|
|
@ -1508,6 +1706,7 @@ def _ensure_rocm_torch() -> None:
|
|||
"later to retry ROCm."
|
||||
)
|
||||
return
|
||||
_write_torch_index_marker(index_url)
|
||||
# ROCm torch is installed (or already was); flag it so later phases
|
||||
# do not overwrite it with the generic CPU torch wheel. BNB is a
|
||||
# separate dependency -- a BNB install failure must NOT roll back the
|
||||
|
|
@ -1612,13 +1811,22 @@ def _ensure_rocm_torch() -> None:
|
|||
# Without this, `studio update` with UNSLOTH_TORCH_INDEX_FAMILY=rocm7.2 (or a
|
||||
# gfx* URL) on a venv that already carries an OLDER ROCm build (+rocm6.4 /
|
||||
# +rocm7.1) short-circuits on has_hip_torch and never applies the override.
|
||||
# Compare exact +rocmX.Y versions when both are readable; otherwise (gfx pin,
|
||||
# or an unreadable installed version) fall back to the torch 2.11 line, which
|
||||
# is what distinguishes the gfx/rocm>=7.2 wheels from older ROCm. Matches the
|
||||
# stale-venv comparison in setup.ps1.
|
||||
#
|
||||
# Prefer the torch-index MARKER when present: an EXACT compare of the pinned
|
||||
# index against the index the last install recorded. This is the ONLY signal
|
||||
# that catches a per-arch switch between two 2.11 gfx indexes (gfx1151 ->
|
||||
# gfx120X-all) -- both install a +rocm7.13.0 wheel, so the version-tag
|
||||
# heuristic below sees no difference and would leave the old arch in place. A
|
||||
# matching marker also guarantees a correctly-pinned venv does NOT reinstall
|
||||
# (no loop). When there is NO marker (old venv, or torch installed out-of-band)
|
||||
# fall back to the +rocm/version-tag heuristic -- backward compatibility.
|
||||
_rocm_pin_mismatch = False
|
||||
if has_hip_torch and _rocm_pin is not None:
|
||||
_rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
|
||||
_marker_verdict = _marker_pin_mismatch(_rocm_pin)
|
||||
if _marker_verdict is None:
|
||||
_rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
|
||||
else:
|
||||
_rocm_pin_mismatch = _marker_verdict
|
||||
|
||||
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
|
||||
|
||||
|
|
@ -1694,6 +1902,7 @@ def _ensure_rocm_torch() -> None:
|
|||
index_url,
|
||||
constrain = False,
|
||||
)
|
||||
_write_torch_index_marker(index_url)
|
||||
rocm_torch_ready = True
|
||||
elif not has_hip_torch or _rocm_pin_mismatch:
|
||||
# Reinstall when torch is not ROCm yet, OR when a ROCm build is present but
|
||||
|
|
@ -1746,6 +1955,7 @@ def _ensure_rocm_torch() -> None:
|
|||
index_url,
|
||||
constrain = False,
|
||||
)
|
||||
_write_torch_index_marker(index_url)
|
||||
rocm_torch_ready = True
|
||||
|
||||
# Install bitsandbytes only when torch links against ROCm. Prefers the
|
||||
|
|
@ -2580,6 +2790,7 @@ def install_python_stack() -> int:
|
|||
_ensure_cuda_torch()
|
||||
_ensure_rocm_torch()
|
||||
_ensure_cpu_torch()
|
||||
_ensure_verbatim_torch_index()
|
||||
|
||||
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
|
||||
# version or unknown ROCm version).
|
||||
|
|
@ -2775,6 +2986,7 @@ def install_python_stack() -> int:
|
|||
_ensure_cuda_torch()
|
||||
_ensure_rocm_torch()
|
||||
_ensure_cpu_torch()
|
||||
_ensure_verbatim_torch_index()
|
||||
|
||||
# 14. Final check (silent; third-party conflicts are expected)
|
||||
subprocess.run(
|
||||
|
|
|
|||
130
studio/setup.ps1
130
studio/setup.ps1
|
|
@ -425,6 +425,81 @@ function Get-TorchIndexLeaf {
|
|||
return ($Url.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
|
||||
}
|
||||
|
||||
# ── Torch-index marker ───────────────────────────────────────────────────────
|
||||
# After a successful torch install this records the exact wheel --index-url used
|
||||
# at a stable per-venv path so a later `unsloth studio update` can detect a pin
|
||||
# change by an EXACT string compare rather than the wheel version-tag heuristic
|
||||
# (which cannot encode the AMD per-arch gfx family). Path/format MUST match
|
||||
# install.sh, install.ps1 and install_python_stack.py:
|
||||
# <VenvDir>\.unsloth-torch-index (single line = the resolved index URL)
|
||||
$TorchIndexMarkerName = ".unsloth-torch-index"
|
||||
|
||||
# Normalise a wheel index URL for exact marker/pin comparison: trim whitespace,
|
||||
# strip ALL trailing slashes, lowercase ONLY the final path segment (the leaf).
|
||||
# Mirrors _normalize_index_url in install.sh / install_python_stack.py / install.ps1.
|
||||
function Get-NormalizedIndexUrl {
|
||||
param([string]$Url)
|
||||
if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
|
||||
$u = $Url.Trim().TrimEnd('/')
|
||||
if ([string]::IsNullOrWhiteSpace($u)) { return $null }
|
||||
$idx = $u.LastIndexOf('/')
|
||||
if ($idx -lt 0) { return $u.ToLowerInvariant() }
|
||||
$head = $u.Substring(0, $idx)
|
||||
$leaf = $u.Substring($idx + 1).ToLowerInvariant()
|
||||
return "$head/$leaf"
|
||||
}
|
||||
|
||||
# Path to the per-venv torch-index marker.
|
||||
function Get-TorchIndexMarkerPath {
|
||||
param([string]$VenvDir)
|
||||
if ([string]::IsNullOrWhiteSpace($VenvDir)) { return $null }
|
||||
return Join-Path $VenvDir $TorchIndexMarkerName
|
||||
}
|
||||
|
||||
# Return the recorded torch --index-url from the marker, else $null (missing /
|
||||
# empty / unreadable -> "no marker", so the caller falls back to the heuristics).
|
||||
function Read-TorchIndexMarker {
|
||||
param([string]$VenvDir)
|
||||
$marker = Get-TorchIndexMarkerPath -VenvDir $VenvDir
|
||||
if (-not $marker) { return $null }
|
||||
if (-not (Test-Path -LiteralPath $marker -PathType Leaf)) { return $null }
|
||||
try {
|
||||
$line = (Get-Content -LiteralPath $marker -Raw -ErrorAction Stop).Trim()
|
||||
} catch { return $null }
|
||||
if ([string]::IsNullOrWhiteSpace($line)) { return $null }
|
||||
return $line
|
||||
}
|
||||
|
||||
# Record the resolved torch wheel --index-url at the marker path, atomically
|
||||
# (temp file + Move). Best-effort: a write failure never aborts the install. A
|
||||
# blank URL is ignored (nothing to record).
|
||||
function Write-TorchIndexMarker {
|
||||
param([string]$VenvDir, [string]$IndexUrl)
|
||||
if ([string]::IsNullOrWhiteSpace($VenvDir)) { return }
|
||||
if ([string]::IsNullOrWhiteSpace($IndexUrl)) { return }
|
||||
if (-not (Test-Path -LiteralPath $VenvDir -PathType Container)) { return }
|
||||
$marker = Get-TorchIndexMarkerPath -VenvDir $VenvDir
|
||||
$tmp = "$marker.$PID.tmp"
|
||||
try {
|
||||
# Write a single line, LF-terminated, no BOM (parity with the sh/py writers).
|
||||
[System.IO.File]::WriteAllText($tmp, ($IndexUrl.Trim() + "`n"), (New-Object System.Text.UTF8Encoding($false)))
|
||||
Move-Item -LiteralPath $tmp -Destination $marker -Force -ErrorAction Stop
|
||||
} catch {
|
||||
try { if (Test-Path -LiteralPath $tmp) { Remove-Item -LiteralPath $tmp -Force -ErrorAction SilentlyContinue } } catch {}
|
||||
}
|
||||
}
|
||||
|
||||
# Compare an explicit torch-index pin against the recorded marker.
|
||||
# $true -> pin resolves to a DIFFERENT index than the marker -> reinstall.
|
||||
# $false -> pin matches the marker exactly -> no reinstall (no loop).
|
||||
# $null -> no usable marker -> caller falls back to the version-tag heuristics.
|
||||
function Test-MarkerPinMismatch {
|
||||
param([string]$VenvDir, [string]$PinUrl)
|
||||
$marker = Read-TorchIndexMarker -VenvDir $VenvDir
|
||||
if ($null -eq $marker) { return $null }
|
||||
return (Get-NormalizedIndexUrl $PinUrl) -ne (Get-NormalizedIndexUrl $marker)
|
||||
}
|
||||
|
||||
# The AMD per-arch index leaves that need the torch 2.11 floor (the _grouped_mm
|
||||
# null-ptr bug lives in the <2.11 wheels for these arches). MUST match the
|
||||
# $_pinGfx211 allowlist in the install-spec path below (and install.ps1 /
|
||||
|
|
@ -436,6 +511,17 @@ function Test-RocmGfx211Leaf {
|
|||
return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf
|
||||
}
|
||||
|
||||
# The pytorch.org rocmX.Y versions KNOWN to ship torch 2.11 (verified against
|
||||
# download.pytorch.org): rocm7.2 is the ONLY stable 2.11 rocm index today. Do NOT
|
||||
# treat an unknown newer rocm (rocm7.3, rocm8.0, ...) as 2.11 speculatively -- it
|
||||
# does not exist yet, and that speculative floor is the mismatch bug being fixed.
|
||||
# MUST match _ROCM_KNOWN_TORCH211_VERSIONS (Python) and the rocm7.2 KNOWN-2.11 leaf
|
||||
# in install.sh / install.ps1. $Major / $Minor are integers.
|
||||
function Test-RocmKnown211Version {
|
||||
param([int]$Major, [int]$Minor)
|
||||
return ($Major -eq 7 -and $Minor -eq 2)
|
||||
}
|
||||
|
||||
# True only for a real CUDA wheel-family leaf: "cu" followed by digits (cu118,
|
||||
# cu126, cu128, cu130, ...). Mirrors install_python_stack.py::_is_cuda_family_leaf.
|
||||
# A bare -like 'cu*' wrongly matches arbitrary mirror leaves like "custom" /
|
||||
|
|
@ -502,7 +588,9 @@ function Get-RocmPinStaleTags {
|
|||
}
|
||||
$_pinNeeds211 = $false
|
||||
if ($_pinRocm.Success) {
|
||||
$_pinNeeds211 = ([int]$_pinRocm.Groups[1].Value -gt 7) -or ([int]$_pinRocm.Groups[1].Value -eq 7 -and [int]$_pinRocm.Groups[2].Value -ge 2)
|
||||
# Only the KNOWN-2.11 rocm indexes (rocm7.2) are on the 2.11 line; an unknown
|
||||
# newer rocm is NOT floored speculatively. Matches _ROCM_KNOWN_TORCH211_VERSIONS.
|
||||
$_pinNeeds211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
|
||||
}
|
||||
# Fallback (installed rocm version unreadable): compare on the 2.11 line, but an
|
||||
# untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> report it stale.
|
||||
|
|
@ -2656,8 +2744,26 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
|
|||
$_expectedKnown = $true
|
||||
if ($_pinnedIdx) {
|
||||
$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx
|
||||
# Torch-index marker: when the last install recorded an index, compare the
|
||||
# pin against it EXACTLY. This is the only signal that catches a per-arch
|
||||
# switch between two 2.11 gfx indexes (gfx1151 -> gfx120X-all -- both
|
||||
# install a +rocm7.13.0 wheel, so the tag heuristic sees no difference) and
|
||||
# a custom-URL change (/simple -> /current). $null = no usable marker ->
|
||||
# fall back to the version-tag heuristic (old venv; backward compatible).
|
||||
$_markerMismatch = Test-MarkerPinMismatch -VenvDir $VenvDir -PinUrl $_pinnedIdx
|
||||
if ($null -ne $_markerMismatch) {
|
||||
# Drive the rebuild decision purely off the marker compare.
|
||||
$_expectedKnown = $true
|
||||
if ($_markerMismatch) {
|
||||
$expectedTorchTag = "pinned:$_pinnedIdx"
|
||||
$installedTorchTag = "marker-mismatch"
|
||||
} else {
|
||||
$expectedTorchTag = "pinned:$_pinnedIdx"
|
||||
$installedTorchTag = "pinned:$_pinnedIdx"
|
||||
}
|
||||
}
|
||||
# cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds.
|
||||
if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') {
|
||||
elseif ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') {
|
||||
# Do NOT collapse a pinned ROCm/gfx leaf to a generic "rocm": that
|
||||
# would match any installed +rocm wheel and mask a pin change from
|
||||
# one ROCm family to another (e.g. rocm6.4 -> gfx1151, or rocm6.4
|
||||
|
|
@ -2678,8 +2784,9 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
|
|||
$expectedTorchTag = $_pinLeaf
|
||||
} else {
|
||||
# Custom index whose final segment is not a torch flavor (e.g. a
|
||||
# PEP 503 mirror ending in /simple). We cannot infer the flavor, so
|
||||
# trust the pinned URL and do not rebuild on a bogus tag comparison.
|
||||
# PEP 503 mirror ending in /simple) and no marker to compare against.
|
||||
# We cannot infer the flavor, so trust the pinned URL and do not
|
||||
# rebuild on a bogus tag comparison.
|
||||
$_expectedKnown = $false
|
||||
$expectedTorchTag = $installedTorchTag
|
||||
}
|
||||
|
|
@ -2995,7 +3102,10 @@ if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) {
|
|||
$_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl
|
||||
$_pinRocm211 = $false
|
||||
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') {
|
||||
$_pinRocm211 = ([int]$Matches[1] -gt 7) -or ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -ge 2)
|
||||
# Only KNOWN-2.11 rocm indexes (rocm7.2) get the 2.11 floor; do not floor an
|
||||
# unknown newer rocm speculatively. Matches install.sh's rocm7.2 KNOWN-2.11
|
||||
# leaf and Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS.
|
||||
$_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2])
|
||||
}
|
||||
# Only the gfx families the AMD arch map above pins to torch 2.11 need the
|
||||
# floor here (gfx120X-all, gfx1151, gfx1150 -- the _grouped_mm bug arches).
|
||||
|
|
@ -3117,6 +3227,16 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
|
|||
}
|
||||
}
|
||||
|
||||
# ── Record the resolved torch wheel index (marker) ──
|
||||
# Torch was just installed; write the exact --index-url used so a later `unsloth
|
||||
# studio update` can detect a pin change by an EXACT string compare (see the
|
||||
# stale-venv check above) instead of the version-tag heuristic. Reflects the actual
|
||||
# installed family: $ROCmIndexUrl when the ROCm path ran, else $TorchInstallIndexUrl
|
||||
# (which is the pinned/CUDA/CPU index, or the CPU index after a ROCm fallback).
|
||||
# Path/format matches install.sh, install.ps1 and install_python_stack.py.
|
||||
$MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchInstallIndexUrl }
|
||||
Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl
|
||||
|
||||
# No unsloth.exe rename needed. setup.ps1 runs *via* unsloth.exe, so renaming the
|
||||
# running launcher only ever failed (WinError 32) and printed a scary warning. It's
|
||||
# also unnecessary: install.ps1 sets SKIP_STUDIO_BASE=1 (base never reinstalled) and
|
||||
|
|
|
|||
|
|
@ -47,9 +47,9 @@ class TestNoTorchBackendAutoInInstallSh:
|
|||
def test_fallback_uses_torch_backend_auto(self):
|
||||
"""The fallback branch should use --torch-backend=auto as recovery."""
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"GPU detection failed" in text
|
||||
), "install.sh should have a fallback branch for when GPU detection fails"
|
||||
assert "GPU detection failed" in text, (
|
||||
"install.sh should have a fallback branch for when GPU detection fails"
|
||||
)
|
||||
|
||||
|
||||
class TestInstallShHasGpuDetection:
|
||||
|
|
@ -57,15 +57,15 @@ class TestInstallShHasGpuDetection:
|
|||
|
||||
def test_function_exists(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"get_torch_index_url()" in text
|
||||
), "install.sh is missing the get_torch_index_url() function"
|
||||
assert "get_torch_index_url()" in text, (
|
||||
"install.sh is missing the get_torch_index_url() function"
|
||||
)
|
||||
|
||||
def test_torch_index_url_assigned(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"TORCH_INDEX_URL=$(get_torch_index_url)" in text
|
||||
), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
|
||||
assert "TORCH_INDEX_URL=$(get_torch_index_url)" in text, (
|
||||
"install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
|
||||
)
|
||||
|
||||
|
||||
class TestCudaMappingParity:
|
||||
|
|
@ -133,15 +133,15 @@ class TestPyTorchMirrorEnvVar:
|
|||
|
||||
def test_install_sh_has_mirror_var(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"UNSLOTH_PYTORCH_MIRROR" in text
|
||||
), "install.sh should reference UNSLOTH_PYTORCH_MIRROR"
|
||||
assert "UNSLOTH_PYTORCH_MIRROR" in text, (
|
||||
"install.sh should reference UNSLOTH_PYTORCH_MIRROR"
|
||||
)
|
||||
|
||||
def test_install_ps1_has_mirror_var(self):
|
||||
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"UNSLOTH_PYTORCH_MIRROR" in text
|
||||
), "install.ps1 should reference UNSLOTH_PYTORCH_MIRROR"
|
||||
assert "UNSLOTH_PYTORCH_MIRROR" in text, (
|
||||
"install.ps1 should reference UNSLOTH_PYTORCH_MIRROR"
|
||||
)
|
||||
|
||||
|
||||
class TestUvBytecodeCompileTimeout:
|
||||
|
|
@ -167,21 +167,21 @@ class TestUvBytecodeCompileTimeout:
|
|||
|
||||
def test_install_sh_preserves_timeout_override(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text
|
||||
), "install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers"
|
||||
assert (
|
||||
"export UV_COMPILE_BYTECODE_TIMEOUT" in text
|
||||
), "install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses"
|
||||
assert ': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text, (
|
||||
"install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers"
|
||||
)
|
||||
assert "export UV_COMPILE_BYTECODE_TIMEOUT" in text, (
|
||||
"install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses"
|
||||
)
|
||||
|
||||
def test_install_ps1_preserves_timeout_override(self):
|
||||
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text
|
||||
), "install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides"
|
||||
assert (
|
||||
'$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
|
||||
), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
|
||||
assert "if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text, (
|
||||
"install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides"
|
||||
)
|
||||
assert '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text, (
|
||||
"install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
|
||||
)
|
||||
|
||||
|
||||
class TestTorchIndexOverrideParity:
|
||||
|
|
@ -207,9 +207,9 @@ class TestTorchIndexOverrideParity:
|
|||
# The AMD ROCm reroute must be skipped when the index is explicitly pinned,
|
||||
# so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
|
||||
text = path.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"TorchIndexPinned" in text
|
||||
), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
|
||||
assert "TorchIndexPinned" in text, (
|
||||
f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
|
||||
)
|
||||
|
||||
def test_cuda_pin_overrides_cvd_hide_gate(self):
|
||||
# A pinned cu* index skips ALL host-GPU probing (parity with install.sh's
|
||||
|
|
@ -226,9 +226,9 @@ class TestTorchIndexOverrideParity:
|
|||
"_ensure_cuda_torch should compute a CUDA-pin flag so the pin can "
|
||||
"override the CVD hide gate"
|
||||
)
|
||||
assert re.search(
|
||||
r"if not _cuda_pinned and _cvd is not None", body
|
||||
), "the CVD hide gate must be bypassed when a CUDA index is pinned"
|
||||
assert re.search(r"if not _cuda_pinned and _cvd is not None", body), (
|
||||
"the CVD hide gate must be bypassed when a CUDA index is pinned"
|
||||
)
|
||||
|
||||
def test_cpu_repair_pins_supported_torch_range(self):
|
||||
# The explicit-CPU repair must not install a bare torch trio: the /cpu
|
||||
|
|
@ -281,12 +281,12 @@ class TestGfx211AllowlistParity:
|
|||
# install-spec path must reuse it, so the stale check and install spec can
|
||||
# never disagree again.
|
||||
text = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"function Test-RocmGfx211Leaf" in text
|
||||
), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
|
||||
assert re.search(
|
||||
r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()
|
||||
), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
|
||||
assert "function Test-RocmGfx211Leaf" in text, (
|
||||
"setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
|
||||
)
|
||||
assert re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()), (
|
||||
"Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
|
||||
)
|
||||
assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, (
|
||||
"setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not "
|
||||
"re-hardcode the allowlist (they must not diverge)"
|
||||
|
|
@ -294,9 +294,9 @@ class TestGfx211AllowlistParity:
|
|||
|
||||
def test_stack_py_allowlist(self):
|
||||
text = STACK_PY.read_text(encoding = "utf-8").lower()
|
||||
assert (
|
||||
'"gfx120x-all", "gfx1151", "gfx1150"' in text
|
||||
), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
|
||||
assert '"gfx120x-all", "gfx1151", "gfx1150"' in text, (
|
||||
"install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
|
||||
)
|
||||
|
||||
|
||||
class TestCudaLeafDigitParity:
|
||||
|
|
@ -307,32 +307,32 @@ class TestCudaLeafDigitParity:
|
|||
|
||||
def test_stack_py_requires_cu_digit(self):
|
||||
text = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert re.search(
|
||||
r'r"\^cu\[0-9\]"', text
|
||||
), "install_python_stack.py _is_cuda_family_leaf must match ^cu[0-9]"
|
||||
assert re.search(r'r"\^cu\[0-9\]"', text), (
|
||||
"install_python_stack.py _is_cuda_family_leaf must match ^cu[0-9]"
|
||||
)
|
||||
|
||||
def test_setup_ps1_requires_cu_digit(self):
|
||||
text = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert re.search(
|
||||
r"'\^cu\[0-9\]'", text
|
||||
), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9], not a bare cu* glob"
|
||||
assert re.search(r"'\^cu\[0-9\]'", text), (
|
||||
"setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9], not a bare cu* glob"
|
||||
)
|
||||
# The stale-venv branch must go through the digit-guarded helper.
|
||||
assert (
|
||||
"Test-CudaFamilyLeaf $_pinLeaf" in text
|
||||
), "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf"
|
||||
assert "Test-CudaFamilyLeaf $_pinLeaf" in text, (
|
||||
"setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf"
|
||||
)
|
||||
|
||||
def test_install_ps1_requires_cu_digit_in_gpu_branch(self):
|
||||
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
assert re.search(
|
||||
r"'\^cu\[0-9\]'", text
|
||||
), "install.ps1 Get-TauriGpuBranch must require a digit after cu"
|
||||
assert re.search(r"'\^cu\[0-9\]'", text), (
|
||||
"install.ps1 Get-TauriGpuBranch must require a digit after cu"
|
||||
)
|
||||
|
||||
def test_install_sh_requires_cu_digit_in_gpu_branch(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
# The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*.
|
||||
assert re.search(
|
||||
r"cu\[0-9\]\*\)\s*echo \"cuda\"", text
|
||||
), "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*"
|
||||
assert re.search(r"cu\[0-9\]\*\)\s*echo \"cuda\"", text), (
|
||||
"install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*"
|
||||
)
|
||||
|
||||
def test_install_sh_backend_export_requires_cu_digit(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
|
|
@ -340,13 +340,13 @@ class TestCudaLeafDigitParity:
|
|||
# bare catch-all *) -> cuda would mis-brand /current, /custom mirror pins
|
||||
# as CUDA and make the stack skip ROCm repair on AMD hosts (comment #2's
|
||||
# bug via install.sh instead of standalone studio update).
|
||||
assert re.search(
|
||||
r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text
|
||||
), "install.sh backend export must brand cuda only on cu[0-9]*"
|
||||
assert re.search(r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text), (
|
||||
"install.sh backend export must brand cuda only on cu[0-9]*"
|
||||
)
|
||||
# An unknown leaf must NOT commit a cuda backend (it unsets instead).
|
||||
assert re.search(
|
||||
r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text
|
||||
), "install.sh backend export must unset (not force cuda) on an unknown leaf"
|
||||
assert re.search(r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text), (
|
||||
"install.sh backend export must unset (not force cuda) on an unknown leaf"
|
||||
)
|
||||
|
||||
def test_install_sh_lowercases_backend_leaf(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
|
|
@ -356,3 +356,108 @@ class TestCudaLeafDigitParity:
|
|||
r"_torch_index_leaf=\$\(printf '%s' \"\$_torch_index_leaf\" \| tr '\[:upper:\]' '\[:lower:\]'\)",
|
||||
text,
|
||||
), "install.sh must lowercase _torch_index_leaf before the gfx/rocm/cu case matches"
|
||||
|
||||
|
||||
class TestTorchIndexMarkerParity:
|
||||
"""All four installers must agree on the torch-index marker (PR #6692):
|
||||
the same filename, the same write points, and the same read helpers."""
|
||||
|
||||
MARKER = ".unsloth-torch-index"
|
||||
|
||||
def test_all_installers_use_same_marker_filename(self):
|
||||
# The exact marker filename must appear in every installer so bash / py /
|
||||
# ps write and read the same per-venv path.
|
||||
for path, label in (
|
||||
(INSTALL_SH, "install.sh"),
|
||||
(INSTALL_PS1, "install.ps1"),
|
||||
(SETUP_PS1, "setup.ps1"),
|
||||
(STACK_PY, "install_python_stack.py"),
|
||||
):
|
||||
text = path.read_text(encoding = "utf-8")
|
||||
assert self.MARKER in text, f"{label} must reference the marker '{self.MARKER}'"
|
||||
|
||||
def test_all_installers_write_the_marker(self):
|
||||
# install.sh + PowerShell use a _write_torch_index_marker / Write-TorchIndexMarker
|
||||
# helper; the Python stack calls _write_torch_index_marker at its install sites.
|
||||
assert "_write_torch_index_marker" in INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert "Write-TorchIndexMarker" in INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
assert "Write-TorchIndexMarker" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "_write_torch_index_marker" in STACK_PY.read_text(encoding = "utf-8")
|
||||
|
||||
def test_setup_ps1_and_python_read_the_marker(self):
|
||||
# The repair/update side (setup.ps1 stale check, python _ensure_rocm_torch)
|
||||
# must READ the marker to make the exact pin-change decision.
|
||||
assert "Read-TorchIndexMarker" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "Test-MarkerPinMismatch" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
stack = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert "_read_torch_index_marker" in stack
|
||||
assert "_marker_pin_mismatch" in stack
|
||||
|
||||
def test_all_installers_normalize_index_url(self):
|
||||
# The exact-compare normalization (trim, strip trailing slash, lowercase
|
||||
# leaf) must exist in every language so the compare is identical.
|
||||
assert "_normalize_index_url" in INSTALL_SH.read_text(encoding = "utf-8")
|
||||
assert "Get-NormalizedIndexUrl" in SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "_normalize_index_url" in STACK_PY.read_text(encoding = "utf-8")
|
||||
|
||||
def test_marker_written_atomically(self):
|
||||
# bash uses a temp file + mv; PowerShell a temp file + Move-Item; Python
|
||||
# tempfile + os.replace. Confirm the atomic-write intent in each.
|
||||
assert re.search(r"mv -f \"\$_wm_tmp\"", INSTALL_SH.read_text(encoding = "utf-8"))
|
||||
for ps in (INSTALL_PS1, SETUP_PS1):
|
||||
assert "Move-Item" in ps.read_text(encoding = "utf-8")
|
||||
assert "os.replace" in STACK_PY.read_text(encoding = "utf-8")
|
||||
|
||||
|
||||
class TestKnown211SetParity:
|
||||
"""The KNOWN-2.11 rocm/gfx set must be identical across all four installers:
|
||||
exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}.
|
||||
rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively."""
|
||||
|
||||
def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self):
|
||||
text = INSTALL_SH.read_text(encoding = "utf-8")
|
||||
# The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves.
|
||||
assert re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text), (
|
||||
"install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
|
||||
)
|
||||
# No speculative rocm7.3 anywhere.
|
||||
assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"
|
||||
|
||||
def test_python_known_211_versions_is_only_rocm72(self):
|
||||
text = STACK_PY.read_text(encoding = "utf-8")
|
||||
assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text
|
||||
# The frozenset literal is exactly {(7, 2)}.
|
||||
m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text)
|
||||
assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
|
||||
assert "(7, 2)" in m.group(1)
|
||||
assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1)
|
||||
|
||||
def test_setup_ps1_known_211_helper_is_only_rocm72(self):
|
||||
text = SETUP_PS1.read_text(encoding = "utf-8")
|
||||
assert "Test-RocmKnown211Version" in text
|
||||
# The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2).
|
||||
assert re.search(
|
||||
r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text
|
||||
), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2"
|
||||
|
||||
def test_install_ps1_pin_floor_is_only_rocm72(self):
|
||||
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
||||
# The pinned-ROCm install-spec floor must be Major -eq 7 -and Minor -eq 2,
|
||||
# not the speculative >= 2 that would floor a non-existent rocm7.3.
|
||||
assert re.search(
|
||||
r"\$_pinRocm211 = \(\[int\]\$Matches\[1\] -eq 7 -and \[int\]\$Matches\[2\] -eq 2\)",
|
||||
text,
|
||||
), "install.ps1 pinned-ROCm floor must be rocm7.2 only (no speculative >= 2)"
|
||||
|
||||
def test_gfx_allowlist_matches_across_installers(self):
|
||||
# The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each.
|
||||
gfx = ("gfx120x-all", "gfx1151", "gfx1150")
|
||||
for path, label in (
|
||||
(INSTALL_SH, "install.sh"),
|
||||
(INSTALL_PS1, "install.ps1"),
|
||||
(SETUP_PS1, "setup.ps1"),
|
||||
(STACK_PY, "install_python_stack.py"),
|
||||
):
|
||||
low = path.read_text(encoding = "utf-8").lower()
|
||||
for g in gfx:
|
||||
assert g in low, f"{label} missing gfx 2.11 allowlist member {g}"
|
||||
|
|
|
|||
92
tests/sh/test_torch_index_marker.sh
Executable file
92
tests/sh/test_torch_index_marker.sh
Executable file
|
|
@ -0,0 +1,92 @@
|
|||
#!/bin/bash
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
# Unit tests for install.sh's torch-index MARKER helpers (_normalize_index_url,
|
||||
# _write_torch_index_marker). These converge the ROCm/gfx pin-change detection
|
||||
# across install.sh / install_python_stack.py / setup.ps1 / install.ps1 by
|
||||
# recording the resolved wheel --index-url so a later update can compare exactly.
|
||||
# Helpers are extracted from install.sh and sourced (parity with test_torch_flavor.sh).
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
|
||||
PASS=0
|
||||
FAIL=0
|
||||
|
||||
_TORCH_INDEX_MARKER_NAME=".unsloth-torch-index"
|
||||
|
||||
# Extract the marker helpers from install.sh and source them.
|
||||
_FUNC_FILE=$(mktemp)
|
||||
{
|
||||
sed -n '/^_normalize_index_url()/,/^}/p' "$INSTALL_SH"
|
||||
echo ""
|
||||
sed -n '/^_write_torch_index_marker()/,/^}/p' "$INSTALL_SH"
|
||||
} > "$_FUNC_FILE"
|
||||
# shellcheck disable=SC1090
|
||||
. "$_FUNC_FILE"
|
||||
rm -f "$_FUNC_FILE"
|
||||
|
||||
assert_eq() {
|
||||
_label="$1"; _expected="$2"; _actual="$3"
|
||||
if [ "$_actual" = "$_expected" ]; then
|
||||
echo " PASS: $_label"; PASS=$((PASS + 1))
|
||||
else
|
||||
echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
|
||||
fi
|
||||
}
|
||||
|
||||
echo "=== _normalize_index_url ==="
|
||||
assert_eq "trailing slashes stripped + leaf lowered" \
|
||||
"https://repo.amd.com/rocm/whl/gfx120x-all" \
|
||||
"$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120X-all///')"
|
||||
assert_eq "whitespace trimmed" \
|
||||
"https://download.pytorch.org/whl/cu128" \
|
||||
"$(_normalize_index_url ' https://download.pytorch.org/whl/cu128 ')"
|
||||
assert_eq "host case preserved, only leaf lowered" \
|
||||
"https://Mirror.Local/simple" \
|
||||
"$(_normalize_index_url 'https://Mirror.Local/Simple/')"
|
||||
# gfx120X-all (capital X) and AMD's lowercase pip leaf normalise equal.
|
||||
assert_eq "gfx120X-all == gfx120x-all after normalize" \
|
||||
"$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120x-all')" \
|
||||
"$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120X-all')"
|
||||
assert_eq "empty -> empty" "" "$(_normalize_index_url ' ')"
|
||||
# rocm7.2 KNOWN-2.11 leaf normalises to itself.
|
||||
assert_eq "rocm7.2 unchanged" \
|
||||
"https://download.pytorch.org/whl/rocm7.2" \
|
||||
"$(_normalize_index_url 'https://download.pytorch.org/whl/rocm7.2/')"
|
||||
|
||||
echo "=== _write_torch_index_marker ==="
|
||||
_VD=$(mktemp -d)
|
||||
_write_torch_index_marker "$_VD" "https://download.pytorch.org/whl/rocm7.2"
|
||||
assert_eq "marker written verbatim (single line)" \
|
||||
"https://download.pytorch.org/whl/rocm7.2" \
|
||||
"$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)"
|
||||
|
||||
# Overwrite (per-arch switch gfx1151 -> gfx120X-all): the marker is replaced.
|
||||
_write_torch_index_marker "$_VD" "https://repo.amd.com/rocm/whl/gfx120X-all"
|
||||
assert_eq "marker overwritten on re-install" \
|
||||
"https://repo.amd.com/rocm/whl/gfx120X-all" \
|
||||
"$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)"
|
||||
|
||||
# Blank URL is ignored (nothing meaningful to record) -- existing marker kept.
|
||||
_write_torch_index_marker "$_VD" " "
|
||||
assert_eq "blank url leaves prior marker intact" \
|
||||
"https://repo.amd.com/rocm/whl/gfx120X-all" \
|
||||
"$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)"
|
||||
|
||||
# No stray temp files left behind by the atomic write.
|
||||
_leftover=$(find "$_VD" -maxdepth 1 -name "$_TORCH_INDEX_MARKER_NAME.*.tmp" 2>/dev/null | wc -l | tr -d ' ')
|
||||
assert_eq "no stray temp file left" "0" "$_leftover"
|
||||
|
||||
# Missing venv dir -> no-op, no error, no file created.
|
||||
_MISSING="$_VD/does_not_exist_dir"
|
||||
_write_torch_index_marker "$_MISSING" "https://x/cu128"
|
||||
assert_eq "missing venv dir -> no marker" \
|
||||
"absent" \
|
||||
"$([ -e "$_MISSING/$_TORCH_INDEX_MARKER_NAME" ] && echo present || echo absent)"
|
||||
|
||||
rm -rf "$_VD"
|
||||
|
||||
echo ""
|
||||
echo "Results: $PASS passed, $FAIL failed"
|
||||
[ "$FAIL" -eq 0 ]
|
||||
|
|
@ -87,6 +87,13 @@ def _run_cuda_repair(
|
|||
return smi_path
|
||||
return None
|
||||
|
||||
import tempfile as _tempfile
|
||||
|
||||
# Isolate the torch-index marker so _ensure_cuda_torch's post-reinstall marker
|
||||
# write lands in a throwaway dir, never the real venv (and never leaks state).
|
||||
_marker_dir = _tempfile.mkdtemp(prefix = "unsloth-marker-test-")
|
||||
_marker_path = Path(_marker_dir) / ".unsloth-torch-index"
|
||||
|
||||
with (
|
||||
patch.object(stack_mod, "_TORCH_BACKEND", backend),
|
||||
patch.object(stack_mod, "IS_MACOS", is_macos),
|
||||
|
|
@ -95,6 +102,7 @@ def _run_cuda_repair(
|
|||
patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia),
|
||||
patch.object(stack_mod.shutil, "which", side_effect = _which),
|
||||
patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)),
|
||||
patch.object(stack_mod, "_torch_index_marker_path", return_value = _marker_path),
|
||||
patch.object(stack_mod, "pip_install") as mock_pip,
|
||||
patch.object(
|
||||
stack_mod.subprocess,
|
||||
|
|
@ -351,12 +359,12 @@ class TestTorchBackendDerivationFromPin:
|
|||
def test_source_uses_helper_not_bare_startswith(self):
|
||||
# Guard against a regression back to elif _idx_leaf.startswith("cu").
|
||||
src = _STACK_PATH.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"elif _is_cuda_family_leaf(_idx_leaf):" in src
|
||||
), "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf"
|
||||
assert (
|
||||
'elif _idx_leaf.startswith("cu"):' not in src
|
||||
), "_TORCH_BACKEND derivation must not use a bare startswith('cu')"
|
||||
assert "elif _is_cuda_family_leaf(_idx_leaf):" in src, (
|
||||
"_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf"
|
||||
)
|
||||
assert 'elif _idx_leaf.startswith("cu"):' not in src, (
|
||||
"_TORCH_BACKEND derivation must not use a bare startswith('cu')"
|
||||
)
|
||||
|
||||
|
||||
# CUDA index ladder.
|
||||
|
|
|
|||
|
|
@ -557,6 +557,17 @@ class TestDetectRocmVersion:
|
|||
class TestEnsureRocmTorch:
|
||||
"""Verify ROCm torch reinstall logic."""
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _isolate_torch_index_marker(self, tmp_path):
|
||||
"""Point the torch-index marker at a per-test tmp path so _ensure_rocm_torch's
|
||||
marker writes never touch the real venv and stale markers never leak between
|
||||
tests. The file is ABSENT by default -> these tests exercise the no-marker
|
||||
fallback to the +rocm/version-tag heuristic (backward compatibility). The
|
||||
path is exposed as self._marker_path for tests that seed a marker."""
|
||||
self._marker_path = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = self._marker_path):
|
||||
yield
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
def test_no_rocm_skips(self, mock_nvidia, mock_pip):
|
||||
|
|
@ -814,9 +825,9 @@ class TestEnsureRocmTorch:
|
|||
_args = [str(a) for a in _call.args]
|
||||
if "--index-url" in _args:
|
||||
_url = _args[_args.index("--index-url") + 1]
|
||||
assert "rocm7.2" not in _url or "torch" not in " ".join(
|
||||
_args
|
||||
), "torch must not be reinstalled when the pin already matches"
|
||||
assert "rocm7.2" not in _url or "torch" not in " ".join(_args), (
|
||||
"torch must not be reinstalled when the pin already matches"
|
||||
)
|
||||
# A torch reinstall would pass torch>=... as a positional; assert none did.
|
||||
assert not any(
|
||||
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
|
||||
|
|
@ -993,6 +1004,267 @@ class TestEnsureRocmTorch:
|
|||
mock_pip.assert_not_called()
|
||||
|
||||
|
||||
# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692)
|
||||
|
||||
|
||||
class TestTorchIndexMarkerHelpers:
|
||||
"""Pure marker helpers: normalization, read/write round-trip, exact compare."""
|
||||
|
||||
def test_normalize_index_url_lowercases_only_leaf(self):
|
||||
f = stack_mod._normalize_index_url
|
||||
# Trailing slashes stripped; ONLY the final segment lowercased.
|
||||
assert f("https://repo.amd.com/rocm/whl/gfx120X-all///") == (
|
||||
"https://repo.amd.com/rocm/whl/gfx120x-all"
|
||||
)
|
||||
# Host case preserved (only the leaf is lowered).
|
||||
assert f("https://Mirror.Local/Simple/") == "https://Mirror.Local/simple"
|
||||
# Whitespace trimmed.
|
||||
assert f(" https://download.pytorch.org/whl/cu128 ") == (
|
||||
"https://download.pytorch.org/whl/cu128"
|
||||
)
|
||||
# gfx120X-all (capital X) and AMD's lowercase pip leaf compare equal.
|
||||
assert f("https://repo.amd.com/rocm/whl/gfx120X-all") == (
|
||||
f("https://repo.amd.com/rocm/whl/gfx120x-all")
|
||||
)
|
||||
# Empty / whitespace-only -> None.
|
||||
assert f(" ") is None
|
||||
assert f(None) is None
|
||||
|
||||
def test_marker_write_read_round_trip(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
stack_mod._write_torch_index_marker("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
# Recorded verbatim (single stripped line).
|
||||
assert marker.read_text().strip() == ("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
assert stack_mod._read_torch_index_marker() == ("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
|
||||
def test_marker_write_is_atomic_and_ignores_blank(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
# Blank / whitespace-only is ignored (nothing to record).
|
||||
stack_mod._write_torch_index_marker("")
|
||||
stack_mod._write_torch_index_marker(" ")
|
||||
assert not marker.exists()
|
||||
# No stray temp files left behind by the atomic write.
|
||||
stack_mod._write_torch_index_marker("https://x/cu128")
|
||||
leftovers = [p.name for p in tmp_path.iterdir() if p.name != ".unsloth-torch-index"]
|
||||
assert leftovers == []
|
||||
|
||||
def test_missing_or_corrupt_marker_reads_as_absent(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
# Absent.
|
||||
assert stack_mod._read_torch_index_marker() is None
|
||||
# Empty file -> absent.
|
||||
marker.write_text("")
|
||||
assert stack_mod._read_torch_index_marker() is None
|
||||
# Whitespace-only -> absent.
|
||||
marker.write_text(" \n")
|
||||
assert stack_mod._read_torch_index_marker() is None
|
||||
|
||||
def test_marker_pin_mismatch_exact_compare(self, tmp_path):
|
||||
marker = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker):
|
||||
# No marker -> None (caller falls back to the heuristic).
|
||||
assert stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx1151") is None
|
||||
# Marker gfx1151, pin gfx120X-all -> mismatch (True). This is the exact
|
||||
# per-arch switch the version-tag heuristic cannot see (#2543).
|
||||
marker.write_text("https://repo.amd.com/rocm/whl/gfx1151\n")
|
||||
assert (
|
||||
stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx120X-all") is True
|
||||
)
|
||||
# Marker matches the pin (case-insensitive leaf, trailing slash) -> False.
|
||||
assert stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx1151/") is False
|
||||
|
||||
def test_known_211_versions_only_rocm72(self):
|
||||
# KNOWN-2.11 rocm set is exactly {rocm7.2} -- rocm7.1/rocm7.3 are NOT in it.
|
||||
known = stack_mod._ROCM_KNOWN_TORCH211_VERSIONS
|
||||
assert (7, 2) in known
|
||||
assert (7, 1) not in known
|
||||
assert (7, 3) not in known
|
||||
assert (8, 0) not in known
|
||||
|
||||
def test_rocm_pin_family_mismatch_known_211_set(self):
|
||||
# The rocmX.Y unreadable-installed fallback uses the KNOWN-2.11 set, so a
|
||||
# (hypothetical) rocm7.3 pin is NOT treated as the 2.11 line speculatively.
|
||||
f = stack_mod._rocm_pin_family_mismatch
|
||||
base = "https://download.pytorch.org/whl"
|
||||
# rocm7.3 pin (unknown -> <2.11 line) over an unreadable-version +rocm wheel
|
||||
# that is <2.11: not a mismatch (both on the non-2.11 line).
|
||||
assert f(f"{base}/rocm7.3", "2.10.0+rocm") is False
|
||||
# rocm7.2 pin (KNOWN-2.11) over the same <2.11 unreadable wheel: mismatch.
|
||||
assert f(f"{base}/rocm7.2", "2.10.0+rocm") is True
|
||||
|
||||
|
||||
class TestEnsureRocmTorchMarker:
|
||||
"""_ensure_rocm_torch marker integration (the #2543 per-arch switch fix)."""
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _marker(self, tmp_path):
|
||||
self._marker_path = tmp_path / ".unsloth-torch-index"
|
||||
with patch.object(stack_mod, "_torch_index_marker_path", return_value = self._marker_path):
|
||||
yield
|
||||
|
||||
def _seed(self, url):
|
||||
self._marker_path.write_text(url + "\n")
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install_try", return_value = True)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
|
||||
def test_marker_gfx_switch_reinstalls_despite_same_wheel_tag(
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
|
||||
):
|
||||
"""Marker records gfx1151; user re-pins to gfx120X-all. Both indexes install
|
||||
a +rocm7.13.0 per-arch wheel, so the version-tag heuristic sees NO difference
|
||||
and would leave the old arch in place. The marker's exact compare catches it
|
||||
and reinstalls from the new gfx index (#2543)."""
|
||||
self._seed("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
# has_hip_torch True + an already-installed AMD per-arch (three-part) wheel:
|
||||
# _rocm_pin_family_mismatch(gfx120X-all, ...+rocm7.13.0) would return False.
|
||||
mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\n"
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"}
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch("subprocess.run", return_value = mock_probe):
|
||||
with patch.object(
|
||||
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
|
||||
):
|
||||
_ensure_rocm_torch()
|
||||
torch_call = str(mock_pip.call_args_list[0])
|
||||
assert "gfx120X-all" in torch_call or "gfx120x-all" in torch_call
|
||||
# Marker rewritten to the new index.
|
||||
assert "gfx120X-all" in self._marker_path.read_text()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install_try", return_value = True)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
|
||||
def test_marker_matches_pin_no_reinstall(
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
|
||||
):
|
||||
"""Marker matches the current pin exactly -> NO torch reinstall (no loop),
|
||||
even when the wheel tag would otherwise look ambiguous."""
|
||||
self._seed("https://repo.amd.com/rocm/whl/gfx1151")
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\n"
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch("subprocess.run", return_value = mock_probe):
|
||||
with patch.object(
|
||||
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
|
||||
):
|
||||
_ensure_rocm_torch()
|
||||
# No torch reinstall (marker matches). bnb-only calls may still happen.
|
||||
assert not any(
|
||||
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
|
||||
)
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install_try", return_value = True)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
|
||||
def test_no_marker_falls_back_to_heuristic_no_forced_reinstall(
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
|
||||
):
|
||||
"""NO marker + a healthy ROCm wheel that already satisfies the pin -> the
|
||||
heuristic (_rocm_pin_family_mismatch) decides, and a correct venv is NOT
|
||||
force-reinstalled (backward compatibility for old venvs)."""
|
||||
# No marker seeded (autouse fixture leaves the file absent).
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
# rocm6.4 pin over an installed +rocm6.4 wheel -> heuristic says no mismatch.
|
||||
mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
|
||||
env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm6.4"}
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch("subprocess.run", return_value = mock_probe):
|
||||
_ensure_rocm_torch()
|
||||
assert not any(
|
||||
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
|
||||
)
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_custom_url_reinstall_on_marker_mismatch(self, mock_pip):
|
||||
"""Marker records .../simple; user pins a custom .../current index (leaf is
|
||||
neither rocm/gfx/cu/cpu). _ensure_verbatim_torch_index reinstalls torch
|
||||
VERBATIM from that URL -- "URL wins verbatim" (#2544)."""
|
||||
self._seed("https://mirror.local/simple")
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current"}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
assert mock_pip.call_count == 1
|
||||
call = str(mock_pip.call_args_list[0])
|
||||
assert "https://mirror.local/current" in call
|
||||
assert "torch" in call
|
||||
# Marker rewritten to the pinned custom URL.
|
||||
assert "current" in self._marker_path.read_text()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_custom_url_no_marker_is_noop(self, mock_pip):
|
||||
"""No marker + a custom-URL pin -> _ensure_verbatim_torch_index does NOTHING
|
||||
(an old venv must not be blindly force-reinstalled from an unverified URL)."""
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current"}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_custom_url_matching_marker_no_reinstall(self, mock_pip):
|
||||
"""Marker matches the custom URL pin -> no reinstall (idempotent, no loop)."""
|
||||
self._seed("https://mirror.local/current")
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current/"}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "IS_WINDOWS", False)
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_verbatim_skips_known_family_pins(self, mock_pip):
|
||||
"""A known-family pin (rocm/gfx/cu/cpu) is NOT handled by the verbatim path --
|
||||
the dedicated _ensure_{rocm,cuda,cpu} helpers own those. cu128 stays CUDA."""
|
||||
self._seed("https://download.pytorch.org/whl/cu126")
|
||||
for pin in (
|
||||
"https://download.pytorch.org/whl/cu128",
|
||||
"https://download.pytorch.org/whl/rocm7.2",
|
||||
"https://repo.amd.com/rocm/whl/gfx1151",
|
||||
"https://download.pytorch.org/whl/cpu",
|
||||
):
|
||||
mock_pip.reset_mock()
|
||||
env = {"UNSLOTH_TORCH_INDEX_URL": pin}
|
||||
with patch.object(stack_mod, "NO_TORCH", False):
|
||||
with patch.object(stack_mod, "IS_MACOS", False):
|
||||
with patch.dict(stack_mod.os.environ, env, clear = False):
|
||||
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
|
||||
stack_mod._ensure_verbatim_torch_index()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
|
||||
# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
|
||||
|
||||
|
||||
|
|
@ -1022,9 +1294,9 @@ class TestHasRocmGpuKfdVendorGuard:
|
|||
|
||||
src = self._src()
|
||||
# Word boundary so "vendor_id 41098" doesn't match "vendor_id 4098".
|
||||
assert (
|
||||
_re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
|
||||
), "_has_rocm_gpu vendor_id check should use word boundary anchors"
|
||||
assert _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src, (
|
||||
"_has_rocm_gpu vendor_id check should use word boundary anchors"
|
||||
)
|
||||
|
||||
def test_sysfs_fallback_guarded_by_non_win32(self):
|
||||
"""KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32')."""
|
||||
|
|
@ -1034,9 +1306,9 @@ class TestHasRocmGpuKfdVendorGuard:
|
|||
def test_cpu_node_excluded(self):
|
||||
"""gpu_id == '0' must be excluded (CPU topology nodes)."""
|
||||
src = self._src()
|
||||
assert (
|
||||
'!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src
|
||||
), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
|
||||
assert '!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src, (
|
||||
"_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
|
||||
)
|
||||
|
||||
def test_install_sh_has_vendor_check(self):
|
||||
"""_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098."""
|
||||
|
|
@ -1069,12 +1341,12 @@ class TestHasRocmGpuKfdVendorGuard:
|
|||
func_start = source.find("_has_amd_rocm_gpu()")
|
||||
func_end = source.find("\n}", func_start)
|
||||
func_body = source[func_start:func_end]
|
||||
assert (
|
||||
"_has_usable_nvidia_gpu" in func_body
|
||||
), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
|
||||
assert (
|
||||
"return 1" in func_body
|
||||
), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
|
||||
assert "_has_usable_nvidia_gpu" in func_body, (
|
||||
"_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
|
||||
)
|
||||
assert "return 1" in func_body, (
|
||||
"_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
|
||||
)
|
||||
|
||||
def test_has_usable_nvidia_gpu_proc_fallback_present(self):
|
||||
"""`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback."""
|
||||
|
|
@ -1290,12 +1562,12 @@ class TestInstallShStructure:
|
|||
rocm_call = body.find("_has_amd_rocm_gpu")
|
||||
assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
|
||||
assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch"
|
||||
assert (
|
||||
rocm_call > no_nvidia_branch
|
||||
), "ROCm detection should sit inside the 'no NVIDIA' branch"
|
||||
assert (
|
||||
nvidia_call < no_nvidia_branch
|
||||
), "NVIDIA detection should run before the no-NVIDIA branch"
|
||||
assert rocm_call > no_nvidia_branch, (
|
||||
"ROCm detection should sit inside the 'no NVIDIA' branch"
|
||||
)
|
||||
assert nvidia_call < no_nvidia_branch, (
|
||||
"NVIDIA detection should run before the no-NVIDIA branch"
|
||||
)
|
||||
|
||||
def test_bitsandbytes_amd_install(self):
|
||||
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
|
||||
|
|
@ -1361,9 +1633,9 @@ class TestInstallShStructure:
|
|||
stripped = line.lstrip()
|
||||
if stripped.startswith("#"):
|
||||
continue
|
||||
assert (
|
||||
"((" not in line or "))" not in line or "$(()" in line
|
||||
), f"get_torch_index_url line {i} may use non-POSIX (( ))"
|
||||
assert "((" not in line or "))" not in line or "$(()" in line, (
|
||||
f"get_torch_index_url line {i} may use non-POSIX (( ))"
|
||||
)
|
||||
|
||||
def test_macos_returns_cpu_before_rocm_check(self):
|
||||
"""macOS should return CPU immediately (before any ROCm check)."""
|
||||
|
|
@ -1382,9 +1654,9 @@ class TestInstallShStructure:
|
|||
torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
|
||||
backend_pos = source.find("UNSLOTH_TORCH_BACKEND")
|
||||
assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh"
|
||||
assert (
|
||||
backend_pos > torch_url_pos
|
||||
), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
|
||||
assert backend_pos > torch_url_pos, (
|
||||
"UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
|
||||
)
|
||||
assert '"cuda"' in source[backend_pos : backend_pos + 500]
|
||||
assert '"rocm"' in source[backend_pos : backend_pos + 500]
|
||||
assert '"cpu"' in source[backend_pos : backend_pos + 500]
|
||||
|
|
@ -1398,12 +1670,12 @@ class TestInstallShStructure:
|
|||
func_start = source.find("_has_amd_rocm_gpu()")
|
||||
func_end = source.find("\n}", func_start)
|
||||
func_body = source[func_start:func_end]
|
||||
assert (
|
||||
"vendor_id" in func_body
|
||||
), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
|
||||
assert (
|
||||
"4098" in func_body
|
||||
), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
|
||||
assert "vendor_id" in func_body, (
|
||||
"_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
|
||||
)
|
||||
assert "4098" in func_body, (
|
||||
"_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
|
||||
)
|
||||
|
||||
def test_kfd_awk_resets_state_per_file(self):
|
||||
"""KFD sysfs awk must reset gpu/amd state per file (FNR==1) to avoid Ryzen+NVIDIA false positives."""
|
||||
|
|
@ -1428,9 +1700,9 @@ class TestInstallShStructure:
|
|||
"get_torch_index_url must use a _nvidia_detected flag (separate from "
|
||||
"_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
|
||||
)
|
||||
assert (
|
||||
'_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
|
||||
), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
|
||||
assert '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body, (
|
||||
"get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
|
||||
)
|
||||
|
||||
|
||||
# TEST: Live regression on current host (NVIDIA B200 expected)
|
||||
|
|
@ -2419,9 +2691,9 @@ class TestRuntimeBnbRocmSourceGuards:
|
|||
"""A failed redetect must not downgrade a persisted suffix to '72'."""
|
||||
for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH):
|
||||
source = path.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
'_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source
|
||||
), path.name
|
||||
assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source, (
|
||||
path.name
|
||||
)
|
||||
|
||||
def test_main_requires_found_rocm_dll(self):
|
||||
"""HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force
|
||||
|
|
@ -2870,9 +3142,9 @@ class TestStrixHaloGfxArchDetection:
|
|||
"""Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output."""
|
||||
for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH):
|
||||
source = path.read_text(encoding = "utf-8")
|
||||
assert (
|
||||
"gfx\\d+" in source or r"gfx\d+" in source
|
||||
), f"gfx arch regex not found in {path.name}"
|
||||
assert "gfx\\d+" in source or r"gfx\d+" in source, (
|
||||
f"gfx arch regex not found in {path.name}"
|
||||
)
|
||||
|
||||
|
||||
# TEST: HIP SDK tool path resolution via HIP_PATH / ROCM_PATH env vars
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ $tokens = $null; $errors = $null
|
|||
$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
|
||||
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
|
||||
|
||||
foreach ($name in @("Test-RocmGfx211Leaf", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) {
|
||||
foreach ($name in @("Test-RocmGfx211Leaf", "Test-RocmKnown211Version", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) {
|
||||
$fn = $ast.FindAll({ param($n)
|
||||
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
|
||||
}, $true)
|
||||
|
|
@ -85,6 +85,17 @@ Check "gfx90a pin + 2.10.0 (untagged) -> stale" (IsStale "gfx90a" "2.10.0
|
|||
Check "gfx120x-all pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx120x-all" "2.11.0+rocm7.2")
|
||||
Check "gfx120x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx120x-all" "2.10.0")
|
||||
|
||||
Write-Host "Test-RocmKnown211Version + KNOWN-2.11 fallback (rocm7.2 only; no speculative rocm7.3)"
|
||||
Check "rocm7.2 -> known 2.11" (Test-RocmKnown211Version -Major 7 -Minor 2)
|
||||
Check "rocm7.1 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 1))
|
||||
Check "rocm7.3 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 3))
|
||||
Check "rocm8.0 -> not known" (-not (Test-RocmKnown211Version -Major 8 -Minor 0))
|
||||
# Unreadable-installed fallback: a rocm7.3 pin (unknown -> <2.11 line) over a <2.11
|
||||
# +rocm wheel with an unreadable version is NOT stale (both on the non-2.11 line);
|
||||
# rocm7.2 (KNOWN-2.11) over the same wheel IS stale. This is the #2534 alignment.
|
||||
Check "rocm7.3 pin + 2.10.0+rocm (unreadable ver) -> not stale" (-not (IsStale "rocm7.3" "2.10.0+rocm"))
|
||||
Check "rocm7.2 pin + 2.10.0+rocm (unreadable ver) -> stale" (IsStale "rocm7.2" "2.10.0+rocm")
|
||||
|
||||
Write-Host ""
|
||||
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
|
||||
Write-Host "All checks passed" -ForegroundColor Green
|
||||
|
|
|
|||
101
tests/studio/test_torch_index_marker.ps1
Normal file
101
tests/studio/test_torch_index_marker.ps1
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
#!/usr/bin/env pwsh
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
# Unit tests for studio/setup.ps1's torch-index MARKER helpers
|
||||
# (Get-NormalizedIndexUrl, Read-TorchIndexMarker, Write-TorchIndexMarker,
|
||||
# Test-MarkerPinMismatch, Test-RocmKnown211Version). These converge the ROCm/gfx
|
||||
# pin-change detection across install.sh / install_python_stack.py / setup.ps1 /
|
||||
# install.ps1. Pure helpers, AST-extracted and run in-process -- no GPU/venv.
|
||||
# Run: pwsh -NoProfile -File tests/studio/test_torch_index_marker.ps1
|
||||
|
||||
$ErrorActionPreference = "Stop"
|
||||
$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
|
||||
$setupPath = (Resolve-Path $setupPath).Path
|
||||
|
||||
# Parse setup.ps1 (also a syntax gate) and extract the helpers.
|
||||
$tokens = $null; $errors = $null
|
||||
$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
|
||||
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
|
||||
|
||||
# The marker filename constant the helpers close over.
|
||||
$TorchIndexMarkerName = ".unsloth-torch-index"
|
||||
|
||||
foreach ($name in @(
|
||||
"Get-NormalizedIndexUrl", "Get-TorchIndexMarkerPath", "Read-TorchIndexMarker",
|
||||
"Write-TorchIndexMarker", "Test-MarkerPinMismatch", "Test-RocmKnown211Version"
|
||||
)) {
|
||||
$fn = $ast.FindAll({ param($n)
|
||||
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
|
||||
}, $true)
|
||||
if ($fn.Count -ne 1) { throw "expected exactly one $name in setup.ps1, found $($fn.Count)" }
|
||||
Invoke-Expression $fn[0].Extent.Text
|
||||
}
|
||||
|
||||
$failures = 0
|
||||
function Check($name, $cond) {
|
||||
if ($cond) { Write-Host " PASS $name" }
|
||||
else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
|
||||
}
|
||||
|
||||
Write-Host "Get-NormalizedIndexUrl (trim / strip trailing slash / lowercase leaf)"
|
||||
Check "trailing slashes + leaf lowered" `
|
||||
((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all///") -eq "https://repo.amd.com/rocm/whl/gfx120x-all")
|
||||
Check "whitespace trimmed" `
|
||||
((Get-NormalizedIndexUrl " https://download.pytorch.org/whl/cu128 ") -eq "https://download.pytorch.org/whl/cu128")
|
||||
Check "host case preserved, leaf lowered" `
|
||||
((Get-NormalizedIndexUrl "https://Mirror.Local/Simple/") -eq "https://Mirror.Local/simple")
|
||||
Check "gfx120X-all == gfx120x-all after normalize" `
|
||||
((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all") -eq (Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120x-all"))
|
||||
Check "empty -> null" ($null -eq (Get-NormalizedIndexUrl " "))
|
||||
|
||||
Write-Host "Write/Read-TorchIndexMarker (round trip, atomic, blank ignored)"
|
||||
$venv = Join-Path ([System.IO.Path]::GetTempPath()) ("unsloth-marker-" + [System.Guid]::NewGuid().ToString("N"))
|
||||
New-Item -ItemType Directory -Path $venv | Out-Null
|
||||
try {
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx1151"
|
||||
Check "marker written verbatim" `
|
||||
((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx1151")
|
||||
# Per-arch switch overwrites the marker.
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all"
|
||||
Check "marker overwritten on re-install" `
|
||||
((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx120X-all")
|
||||
# Blank URL ignored -- prior marker kept.
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl " "
|
||||
Check "blank url leaves prior marker intact" `
|
||||
((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx120X-all")
|
||||
# No stray temp file left behind.
|
||||
$tmpLeft = @(Get-ChildItem -LiteralPath $venv -Filter "$TorchIndexMarkerName.*.tmp" -ErrorAction SilentlyContinue).Count
|
||||
Check "no stray temp file left" ($tmpLeft -eq 0)
|
||||
|
||||
Write-Host "Test-MarkerPinMismatch (exact compare; null when no marker)"
|
||||
# Marker gfx120X-all now recorded. A gfx1151 pin differs -> mismatch (#2543).
|
||||
Check "gfx marker vs gfx1151 pin -> mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx1151") -eq $true)
|
||||
# Same pin (trailing slash, case) -> not a mismatch (no reinstall loop).
|
||||
Check "same pin (slash/case) -> no mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx120x-all/") -eq $false)
|
||||
# Custom URL change /simple -> /current is a mismatch (#2544).
|
||||
Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://mirror.local/simple"
|
||||
Check "custom /simple marker vs /current pin -> mismatch" `
|
||||
((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://mirror.local/current") -eq $true)
|
||||
|
||||
Write-Host "Read-TorchIndexMarker (missing / empty -> null)"
|
||||
Remove-Item -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Force
|
||||
Check "missing marker -> null" ($null -eq (Read-TorchIndexMarker -VenvDir $venv))
|
||||
Set-Content -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Value "" -NoNewline
|
||||
Check "empty marker -> null" ($null -eq (Read-TorchIndexMarker -VenvDir $venv))
|
||||
Check "no marker -> Test-MarkerPinMismatch null" `
|
||||
($null -eq (Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://x/rocm7.2"))
|
||||
} finally {
|
||||
Remove-Item -LiteralPath $venv -Recurse -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
Write-Host "Test-RocmKnown211Version (KNOWN-2.11 set == rocm7.2 only)"
|
||||
Check "rocm7.2 -> true" (Test-RocmKnown211Version -Major 7 -Minor 2)
|
||||
Check "rocm7.1 -> false" (-not (Test-RocmKnown211Version -Major 7 -Minor 1))
|
||||
Check "rocm7.3 -> false" (-not (Test-RocmKnown211Version -Major 7 -Minor 3))
|
||||
Check "rocm8.0 -> false" (-not (Test-RocmKnown211Version -Major 8 -Minor 0))
|
||||
|
||||
Write-Host ""
|
||||
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
|
||||
Write-Host "All checks passed" -ForegroundColor Green
|
||||
Loading…
Add table
Add a link
Reference in a new issue