From 753ef2f5cfd0c58ef82b5ee8bca8110f4f9807ca Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 02:42:36 +0000 Subject: [PATCH] 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. --- install.ps1 | 41 ++- install.sh | 64 ++++ studio/install_python_stack.py | 226 ++++++++++++- studio/setup.ps1 | 130 +++++++- tests/python/test_cross_platform_parity.py | 231 +++++++++---- tests/sh/test_torch_index_marker.sh | 92 ++++++ tests/studio/install/test_cuda_repair.py | 20 +- tests/studio/install/test_rocm_support.py | 356 ++++++++++++++++++--- tests/studio/test_setup_pin_stale.ps1 | 13 +- tests/studio/test_torch_index_marker.ps1 | 101 ++++++ 10 files changed, 1149 insertions(+), 125 deletions(-) create mode 100755 tests/sh/test_torch_index_marker.sh create mode 100644 tests/studio/test_torch_index_marker.ps1 diff --git a/install.ps1 b/install.ps1 index 0b0f5478af..fc005e2c9a 100644 --- a/install.ps1 +++ b/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: + # \.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__). diff --git a/install.sh b/install.sh index 2d533ad414..f748aee08c 100755 --- a/install.sh +++ b/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: +# /.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. diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 3a9e653c84..90aac9533b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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: +# /.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( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index b525bd98c2..4e28af3d93 100644 --- a/studio/setup.ps1 +++ b/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: +# \.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 diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 9514683ad6..c3d101ee30 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -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}" diff --git a/tests/sh/test_torch_index_marker.sh b/tests/sh/test_torch_index_marker.sh new file mode 100755 index 0000000000..c8f63a5dc7 --- /dev/null +++ b/tests/sh/test_torch_index_marker.sh @@ -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 ] diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index 9fe24d9fc5..cb37a29d28 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -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. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 8898b6e7b6..79c3c2d45f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -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 diff --git a/tests/studio/test_setup_pin_stale.ps1 b/tests/studio/test_setup_pin_stale.ps1 index 26c8e2382e..7335044987 100644 --- a/tests/studio/test_setup_pin_stale.ps1 +++ b/tests/studio/test_setup_pin_stale.ps1 @@ -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 diff --git a/tests/studio/test_torch_index_marker.ps1 b/tests/studio/test_torch_index_marker.ps1 new file mode 100644 index 0000000000..9cd12ad248 --- /dev/null +++ b/tests/studio/test_torch_index_marker.ps1 @@ -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