install: honor custom pins and repair pinned venvs in place

Four follow-ups to the torch-index marker work:

- install_python_stack.py: _ensure_cuda_torch/_ensure_rocm_torch now bail when an
  explicit custom-index pin names no known torch family, so a verbatim URL override
  (a private/simple mirror) is not clobbered by auto-detected CUDA/ROCm wheels
  before _ensure_verbatim_torch_index applies it.

- install_python_stack.py: the ROCm marker is additive, not a substitute -- a
  matching marker still runs the family/version check so a wheel swapped after the
  marker was written is caught. Mirrors setup.ps1.

- setup.ps1: a stale venv under an explicit pin, whose torch still imports, is
  repaired in place (force-reinstall torch from the pin in the dependency pass)
  instead of wiped. The wipe path only delegates to install.ps1, so on a direct
  update it stranded the user at "Virtual environment not found" instead of
  applying the new pin. A broken venv or unpinned drift still wipes/delegates.

- install.ps1: when a pinned ROCm install fails over to a CPU base, the marker now
  records the CPU index actually used instead of the ROCm pin, so the next managed
  setup does not see CPU torch under a ROCm pin and abort as stale.
This commit is contained in:
danielhanchen 2026-07-06 05:33:36 +00:00
commit 5c93ffd450
3 changed files with 67 additions and 8 deletions

View file

@ -1414,6 +1414,11 @@ def _ensure_cuda_torch() -> None:
# (or any unrecognised value) are deliberate and must not be overridden.
if _TORCH_BACKEND not in ("", "cuda"):
return
# An explicit custom-index pin whose leaf names no known torch family wins
# VERBATIM (_ensure_verbatim_torch_index applies it); do not override it with
# auto-detected CUDA wheels here.
if _explicit_unknown_family_torch_index_url() is not None:
return
# No CUDA torch on macOS; Windows venv/torch lifecycle is owned by
# install.ps1 (and the KFD poisoning bug is Linux-only), so skip both.
if IS_MACOS or IS_WINDOWS or NO_TORCH:
@ -1611,6 +1616,14 @@ def _ensure_rocm_torch() -> None:
# different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
if _TORCH_BACKEND in ("cuda", "cpu"):
return
# An explicit custom-index pin whose leaf names no known torch family (a private
# PEP 503 mirror, /simple, /current, ...) wins VERBATIM -- _ensure_verbatim_torch_index()
# applies it. Never override such a pin with the auto-detected ROCm index here: the
# user chose that index (a valid cross-install / custom-mirror case), and installing
# ROCm wheels over it before the verbatim pass runs would leave the URL override
# unhonored (the marker then matches the pin, so verbatim skips the restore).
if _explicit_unknown_family_torch_index_url() is not None:
return
# setup.ps1 sets this after installing AMD wheels; skip the probe only when
# torch is actually importable as ROCm. If the venv was wiped between runs,
# the stale env-var would suppress a needed reinstall.
@ -1823,10 +1836,19 @@ def _ensure_rocm_torch() -> None:
_rocm_pin_mismatch = False
if has_hip_torch and _rocm_pin is not None:
_marker_verdict = _marker_pin_mismatch(_rocm_pin)
if _marker_verdict is None:
_rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
if _marker_verdict is True:
# Marker records a DIFFERENT index than the pin -> reinstall. This is the
# only signal that catches a per-arch gfx switch (gfx1151 -> gfx120X-all,
# both +rocm7.13.0) the version-tag heuristic below cannot see.
_rocm_pin_mismatch = True
else:
_rocm_pin_mismatch = _marker_verdict
# Marker matches OR is absent: the marker is an ADDITIONAL rebuild signal,
# not a substitute for validating the installed wheel. Still run the
# family/version check so a stale wheel (torch swapped after the marker was
# written -- e.g. marker records gfx1151 but the venv now carries generic
# +rocm7.2 or an older +rocm6.4) is caught. Mirrors setup.ps1, which keeps
# the flavor check alongside the marker for this reason.
_rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch