diff --git a/install.ps1 b/install.ps1 index 148672b7cc..3bf9317960 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2234,9 +2234,21 @@ exit 0 return $installed } + # Track whether THIS run actually installed or repaired torch, so the marker + # write below reflects the real wheel source. A migrated venv that keeps its + # existing torch (no reinstall) must NOT rewrite its marker to the newly + # requested pin, or a later update compares the new pin against a marker that + # already matches and skips the reinstall the pin needs (e.g. a per-arch + # gfx1151 -> gfx120X-all switch, identical +rocm tag). Fresh installs below + # always install torch; the shared flavor repair re-lands it when wrong. + $_TorchInstalledThisRun = $true + if ($_Migrated) { # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # in the new venv location, while preserving existing torch/CUDA. Torch is + # preserved unless the shared flavor repair below re-lands it, so the marker + # must stay as the migrated venv recorded it for the preserved case. + $_TorchInstalledThisRun = $false Write-TauriLog "STEP" "Installing unsloth" substep "upgrading unsloth in migrated environment..." if ($SkipTorch) { @@ -2422,6 +2434,9 @@ exit 0 return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) } $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + # torch was re-landed from $ROCmIndexUrl, so record it (the + # gfx-switch case keeps the same rocm flavor and never reaches here). + $_TorchInstalledThisRun = $true } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" @@ -2431,6 +2446,7 @@ exit 0 return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) } $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + $_TorchInstalledThisRun = $true # torch re-landed from $TorchIndexUrl } } # Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly. @@ -2450,7 +2466,7 @@ exit 0 # $ROCmIndexUrl when the ROCm path ran, the CPU fallback index when a pinned # ROCm install failed over to a CPU base, else the CUDA/CPU/pinned $TorchIndexUrl. # Skipped for --no-torch (nothing installed). Matches install.sh / setup.ps1. - if (-not $SkipTorch) { + if ((-not $SkipTorch) -and $_TorchInstalledThisRun) { $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } elseif ($RocmCpuFallbackIndexUrl) { $RocmCpuFallbackIndexUrl } else { $TorchIndexUrl } Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl } diff --git a/install.sh b/install.sh index ea77ecd4d9..c16648aa73 100755 --- a/install.sh +++ b/install.sh @@ -2831,9 +2831,19 @@ esac # ── Install unsloth directly into the venv (no activation needed) ── tauri_log "STEP" "Installing PyTorch" _VENV_PY="$VENV_DIR/bin/python" +# Track whether THIS run actually installed or repaired torch, so the marker at the +# end reflects the real wheel source. A migrated venv that keeps its existing torch +# (no reinstall) must NOT rewrite its marker to the newly requested pin, or a later +# `unsloth studio update` compares the new pin against a marker that already matches +# and skips the reinstall the pin needs (e.g. a per-arch gfx1151 -> gfx120X-all +# switch, identical +rocm tag). Fresh installs below always install torch. +_TORCH_INSTALLED_THIS_RUN=true if [ "$_MIGRATED" = true ]; then # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # in the new venv location, while preserving existing torch/CUDA. Torch is + # preserved (not reinstalled) unless the ROCm repair below fires, so the + # marker must stay as the migrated venv recorded it for the preserved case. + _TORCH_INSTALLED_THIS_RUN=false substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps (current @@ -2880,6 +2890,9 @@ if [ "$_MIGRATED" = true ]; then "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall + # torch was actually reinstalled from $TORCH_INDEX_URL now, so the + # marker should record it (the preserved-torch case above must not). + _TORCH_INSTALLED_THIS_RUN=true fi ;; esac @@ -3158,6 +3171,10 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio + # torch was re-landed from $TORCH_INDEX_URL, so record it even on a + # migrated venv whose flavor was genuinely wrong (the gfx-switch case + # keeps the same rocm flavor and does NOT reach here, so its marker stays). + _TORCH_INSTALLED_THIS_RUN=true _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") @@ -3176,11 +3193,13 @@ fi # 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). +# actually installed from a resolved index (skip --no-torch / no-URL fallback) AND +# actually installed/repaired this run (a migrated venv that kept its existing torch +# leaves the old marker so a later update can still detect a pin change). # Reflects the actual source: the Radeon --find-links path sets # _TORCH_MARKER_INDEX_URL to its repo.radeon.com base; every other path falls back # to $TORCH_INDEX_URL (the CUDA/CPU/ROCm/pinned index it installed from). -if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then +if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ] && [ "$_TORCH_INSTALLED_THIS_RUN" = true ]; then _write_torch_index_marker "$VENV_DIR" "${_TORCH_MARKER_INDEX_URL:-$TORCH_INDEX_URL}" fi diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index b68aebb328..eeef85e51f 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1225,13 +1225,25 @@ def _explicit_torch_index_url() -> "str | None": return None +def _is_pip_rocm_family_leaf(leaf: str) -> bool: + """True when a lowercased index leaf names a pip --index-url ROCm family: an + actual rocm... leaf (download.pytorch.org/whl/rocm7.2) or a repo.amd.com + per-arch gfx leaf (gfx120x-all). A Radeon find-links directory leaf + (repo.radeon.com/.../rocm-rel-7.2.1, which install.sh records in the marker) + starts with "rocm" but is NOT a pip index -- it must route to the verbatim/marker + path, not a --index-url reinstall that fails against a find-links listing. + Mirrors install.sh's rocm[0-9]* / setup.ps1's ^(rocm[0-9]|gfx) gate. Pure function. + """ + return bool(re.match(r"^rocm\d", leaf)) or leaf.startswith("gfx") + + def _explicit_rocm_torch_index_url() -> "str | None": - """The pinned wheel index URL when it names a ROCm family (rocm*/gfx*), else None.""" + """The pinned wheel index URL when it names a pip ROCm family (rocm/gfx*), else None.""" url = _explicit_torch_index_url() if url is None: return None leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() - return url if leaf.startswith(("rocm", "gfx")) else None + return url if _is_pip_rocm_family_leaf(leaf) else None def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool: @@ -1365,7 +1377,7 @@ def _explicit_unknown_family_torch_index_url() -> "str | None": 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): + if _is_pip_rocm_family_leaf(leaf) or leaf == "cpu" or _is_cuda_family_leaf(leaf): return None return url @@ -1596,13 +1608,22 @@ def _ensure_cpu_torch() -> None: _lines = [ line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() ] - if not _lines or _lines[-1] != "gpu": - return # already CPU (or unreadable) -- nothing to repair + if not _lines: + return # unreadable -- the base install step handles a missing torch + if _lines[-1] != "gpu": + # torch is already a CPU build. Normally nothing to repair, BUT a standalone + # update may change the CPU index URL itself (official /cpu -> a private + # UNSLOTH_PYTORCH_MIRROR /cpu) with the same +cpu tag. The tag cannot see the + # host change, so consult the exact-URL marker and reinstall from the new pin + # only when it records a DIFFERENT index -- mirroring the CUDA/ROCm + # same-family-URL handling. No marker (or a matching one) -> leave it alone. + if _marker_pin_mismatch(pin) is not True: + return + _why = "the pinned CPU index URL differs from the recorded marker" + else: + _why = "torch is a GPU build but an explicit CPU index is pinned" - print( - f" torch is a GPU build but an explicit CPU index is pinned -- " - f"reinstalling CPU torch from {pin}" - ) + print(f" {_why} -- reinstalling CPU torch from {pin}") # Pin to the supported torch<2.11 family (same bounds as the CUDA/ROCm repair # specs). The /cpu index now also serves torch 2.11+, so a bare trio off the # exclusive --index-url could resolve outside the supported range or drag in diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 45b819f98f..e878b48775 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -952,6 +952,73 @@ class TestEnsureRocmTorch: assert f(f"{amd}/gfx110X-all", "2.10.0") is True assert f(f"{amd}/gfx90a", "2.10.0") is True + def test_radeon_url_not_classified_as_pip_rocm_family(self): + """A repo.radeon.com find-links directory (leaf rocm-rel-7.2.1, which + install.sh records in the marker) starts with "rocm" but is NOT a pip + --index-url ROCm family: it must route to the verbatim/marker path, not a + --index-url reinstall that fails against a find-links listing (Codex P2).""" + leaf_f = stack_mod._is_pip_rocm_family_leaf + # Real pip ROCm families (download.pytorch.org/whl/rocmX.Y, repo.amd.com gfx). + assert leaf_f("rocm7.2") is True + assert leaf_f("rocm6.4") is True + assert leaf_f("gfx120x-all") is True + assert leaf_f("gfx1151") is True + # A Radeon find-links dir leaf, a custom mirror, cpu and cuda are NOT pip rocm. + assert leaf_f("rocm-rel-7.2.1") is False + assert leaf_f("simple") is False + assert leaf_f("current") is False + assert leaf_f("cpu") is False + assert leaf_f("cu128") is False + + radeon = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1" + pip_rocm = "https://download.pytorch.org/whl/rocm7.2" + amd_gfx = "https://repo.amd.com/rocm/whl/gfx120X-all" + + def _classify(url, fn): + with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + return fn() + + rocm_fn = stack_mod._explicit_rocm_torch_index_url + unk_fn = stack_mod._explicit_unknown_family_torch_index_url + # Real pip rocm/gfx pins ARE a ROCm family (reinstallable via --index-url) and + # are NOT "unknown". + assert _classify(pip_rocm, rocm_fn) == pip_rocm + assert _classify(amd_gfx, rocm_fn) == amd_gfx + assert _classify(pip_rocm, unk_fn) is None + assert _classify(amd_gfx, unk_fn) is None + # The Radeon find-links URL is NOT a pip ROCm family (so _ensure_rocm_torch + # skips it instead of a failing --index-url reinstall) and IS unknown, so it + # routes to the verbatim/marker path -- a no-op when the marker already matches + # (the finding's "leave the matching marker alone" scenario). + assert _classify(radeon, rocm_fn) is None + assert _classify(radeon, unk_fn) == radeon + + @patch.object(stack_mod, "_write_torch_index_marker") + @patch.object(stack_mod, "pip_install") + def test_ensure_cpu_torch_honors_marker_url_change(self, mock_pip, mock_marker): + """_ensure_cpu_torch: a CPU venv whose marker records a DIFFERENT /cpu index + (official -> a private UNSLOTH_PYTORCH_MIRROR /cpu, same +cpu tag) must + reinstall from the new pin -- the tag cannot see the host change, so the marker + drives it (Codex P2). A matching marker (or none) leaves CPU torch alone.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"cpu\n" # torch is already a CPU build + env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/cpu"} + + def _run(mismatch): + mock_pip.reset_mock() + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("subprocess.run", return_value = mock_probe): + with patch.object(stack_mod, "_marker_pin_mismatch", return_value = mismatch): + stack_mod._ensure_cpu_torch() + return mock_pip.called + + assert _run(True) is True # marker records a different /cpu index -> reinstall + assert _run(False) is False # marker matches the pin -> no reinstall (no loop) + assert _run(None) is False # no usable marker -> no blind reinstall + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") diff --git a/tests/studio/test_torch_index_marker.ps1 b/tests/studio/test_torch_index_marker.ps1 index 9cd12ad248..516f3f6f90 100644 --- a/tests/studio/test_torch_index_marker.ps1 +++ b/tests/studio/test_torch_index_marker.ps1 @@ -21,8 +21,9 @@ if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has $TorchIndexMarkerName = ".unsloth-torch-index" foreach ($name in @( - "Get-NormalizedIndexUrl", "Get-TorchIndexMarkerPath", "Read-TorchIndexMarker", - "Write-TorchIndexMarker", "Test-MarkerPinMismatch", "Test-RocmKnown211Version" + "Get-NormalizedFamilyLeaf", "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