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 2026-07-07 01:26:16 +00:00
commit cf3ee0d072
5 changed files with 140 additions and 16 deletions

View file

@ -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