install: fix 3 torch-index marker edge cases (CPU mirror pin, Radeon leaf, migrated venv)
Addresses three review findings on the torch-index override path:
1. CPU index URL change on an already-CPU venv. _ensure_cpu_torch returned
early whenever torch was already a CPU build, so a standalone update that
moved the pin (official /cpu -> a private UNSLOTH_PYTORCH_MIRROR /cpu, same
+cpu tag) never reinstalled. It now consults the exact-URL marker and
reinstalls only when _marker_pin_mismatch reports a different index,
mirroring the CUDA/ROCm same-family handling. A matching marker (or none)
still leaves CPU torch untouched, so there is no reinstall loop.
2. Radeon find-links directory misclassified as a pip ROCm family. A
repo.radeon.com/.../rocm-rel-7.2.1 leaf starts with "rocm" but is a
find-links listing, not a pip --index-url. The old startswith(("rocm",
"gfx")) test routed it into a --index-url reinstall that fails against
find-links. New _is_pip_rocm_family_leaf gates on ^rocm\d / gfx (matching
install.sh's rocm[0-9]* and setup.ps1's ^(rocm[0-9]|gfx)), so a Radeon URL
routes to the verbatim/marker path instead.
3. Migrated venv rewriting its marker to a pin it did not install. install.sh
and install.ps1 write the marker unconditionally, so a migration that
preserves existing torch recorded the newly requested pin and a later
update then found a matching marker and skipped the reinstall the pin
needs (e.g. a per-arch gfx1151 -> gfx120X-all switch, identical +rocm tag).
Both now track _TORCH_INSTALLED_THIS_RUN and write the marker only when
torch was actually installed or repaired this run.
Also add Get-NormalizedFamilyLeaf to the setup.ps1 helper-extraction list in
test_torch_index_marker.ps1 (it was added to setup.ps1 and the shell test in an
earlier round but missed here) and add two unit tests covering findings 1 and 2.
This commit is contained in:
Daniel Han
parent
5ed073755c
commit
cf3ee0d072
5 changed files with 140 additions and 16 deletions
20
install.ps1
20
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
|
||||
}
|
||||
|
|
|
|||
25
install.sh
25
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
|
||||
|
||||
|
|
|
|||
|
|
@ -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<digit>... 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<d>/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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue