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:
parent
257ca41753
commit
5c93ffd450
3 changed files with 67 additions and 8 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue