install: finish pinned ROCm/CUDA edge cases on Windows + repair path
Follow-ups to the previous round: - studio/setup.ps1: a pinned gfx*/rocm>=7.2 index now routes through the ROCm install path with the 2.11 floor + companions (it previously fell through to the CUDA branch with bare torch/torchvision/torchaudio against the ROCm index). The CPU/CUDA fallback index is forced to the CPU wheel index when a ROCm index is active, so a failed pinned-ROCm install does not retry the ROCm mirror. - studio/setup.ps1: the stale-venv check no longer treats an unrecognized pinned URL leaf (e.g. a PEP 503 mirror ending in /simple) as a torch flavor tag, which was marking a correct venv stale; cu*/cpu/rocm/gfx leaves are still compared. - install.ps1: the post-failure CPU fallback uses an explicit CPU index instead of , which for a pinned ROCm index was the ROCm mirror itself (so the 'fallback' just retried the failing index and aborted the installer). - install_python_stack.py: _ensure_cuda_torch now also reinstalls when the venv's CUDA family differs from a pinned one (installed cu126 vs pinned cu128), not only CPU->CUDA; the probe reports the installed cuXXX tag for the comparison.
This commit is contained in:
parent
d2d5f90175
commit
4a5baba8ce
3 changed files with 63 additions and 15 deletions
|
|
@ -1053,11 +1053,13 @@ def _ensure_cuda_torch() -> None:
|
|||
sys.executable,
|
||||
"-c",
|
||||
(
|
||||
"import torch; "
|
||||
"import torch, re; "
|
||||
"hip = getattr(torch.version, 'hip', '') or ''; "
|
||||
"cuda = getattr(torch.version, 'cuda', '') or ''; "
|
||||
"ver = getattr(torch, '__version__', '').lower(); "
|
||||
"print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
|
||||
"m = re.search(r'\\+(cu\\d+)', ver); "
|
||||
"marker = 'hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'); "
|
||||
"print(marker + '|' + (m.group(1) if m else ''))"
|
||||
),
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
|
|
@ -1075,22 +1077,26 @@ def _ensure_cuda_torch() -> None:
|
|||
]
|
||||
if not _marker_lines:
|
||||
return
|
||||
_marker = _marker_lines[-1]
|
||||
_marker, _, _installed_cu = _marker_lines[-1].partition("|")
|
||||
# Reinstall CUDA torch when the venv carries a ROCm build on an NVIDIA host
|
||||
# (the poisoning signature), OR when an explicit CUDA index is pinned but the
|
||||
# venv still has a CPU wheel. The latter is the headless CPU-venv-to-CUDA
|
||||
# cross-install (`studio update` with UNSLOTH_TORCH_INDEX_FAMILY=cu128): the
|
||||
# update path preserves torch rather than preinstalling it from install.sh, so
|
||||
# without this the explicit CUDA pin stays ineffective. A healthy CUDA torch,
|
||||
# or a CPU wheel with no CUDA pin, is deliberate and left alone.
|
||||
# (the poisoning signature), or when an explicit CUDA index is pinned but the
|
||||
# venv has the wrong family -- a CPU wheel, or a different cuXXX than pinned.
|
||||
# This covers the headless cross-install (`studio update` with
|
||||
# UNSLOTH_TORCH_INDEX_FAMILY=cu128): the update path preserves torch rather
|
||||
# than preinstalling it from install.sh, so without this an explicit CUDA pin
|
||||
# stays ineffective. A healthy CUDA torch matching the pin, or a CPU wheel
|
||||
# with no CUDA pin, is deliberate and left alone.
|
||||
_pin = _explicit_torch_index_url()
|
||||
_pinned_cuda = bool(_pin) and _pin.rstrip("/").rsplit("/", 1)[-1].lower().startswith("cu")
|
||||
_pin_leaf = _pin.rstrip("/").rsplit("/", 1)[-1].lower() if _pin else ""
|
||||
_pinned_cuda = _pin_leaf.startswith("cu")
|
||||
if _marker == "hip":
|
||||
_why = "torch is a ROCm build on an NVIDIA host"
|
||||
elif _marker == "cpu" and _pinned_cuda:
|
||||
_why = "torch is a CPU build but an explicit CUDA index is pinned"
|
||||
elif _marker == "cuda" and _pinned_cuda and _installed_cu and _installed_cu != _pin_leaf:
|
||||
_why = f"torch is {_installed_cu} but the pinned CUDA index is {_pin_leaf}"
|
||||
else:
|
||||
return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
|
||||
return # healthy CUDA torch matching the pin, or a deliberate CPU wheel
|
||||
|
||||
index_url = _detect_cuda_torch_index_url()
|
||||
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue