ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

install: make the torch-index override authoritative across ROCm paths

Browse source 
Address review feedback on the override added in this PR so a pinned index is
honoured everywhere, not just in get_torch_index_url:

- Skip the WSL ROCm bootstrap (root privilege + large downloads, probes
  /dev/dxg) when UNSLOTH_TORCH_INDEX_URL / _FAMILY is set; it previously ran
  before the override was consulted.
- Skip the Radeon/Strix rerouting (which re-probes the GPU and overwrites the
  resolved URL with repo.radeon.com / repo.amd.com) when the index is pinned, so
  an explicit ROCm override (e.g. UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) is kept.
- install_python_stack.py: derive _TORCH_BACKEND from the override when
  UNSLOTH_TORCH_BACKEND is unset (standalone studio update), so _ensure_rocm_torch
  / _ensure_cuda_torch repair to the requested family instead of re-detecting.
- Strip ALL leading/trailing slashes in the shell override to match the Python
  side (avoids 404s on strict pip proxies).

Adds test cases for double-slash and leading/trailing-slash overrides.
This commit is contained in:
Daniel Han 2026-06-26 05:28:36 +00:00
commit b02a609af5
3 changed files with 50 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

17
studio/install_python_stack.py
View file

@ -1429,6 +1429,23 @@ NO_TORCH = _infer_no_torch()
# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
# (standalone `unsloth studio update` runs, where we re-detect normally).
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
# When install.sh did not run (standalone `unsloth studio update`) but the caller
# pinned the wheel index explicitly, derive the backend from that override so the
# CUDA/ROCm repair helpers honour it instead of re-probing the GPU and possibly
# reinstalling a different family. Classify on the final URL/family segment,
# mirroring install.sh's UNSLOTH_TORCH_BACKEND case.
if not _TORCH_BACKEND:
_idx_override = (
os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
or os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
)
_idx_leaf = _idx_override.rstrip("/").rsplit("/", 1)[-1].lower()
if _idx_leaf.startswith(("rocm", "gfx")):
_TORCH_BACKEND = "rocm"
elif _idx_leaf == "cpu":
_TORCH_BACKEND = "cpu"
elif _idx_leaf.startswith("cu"):
_TORCH_BACKEND = "cuda"
def _torch_step_label(suffix: str) -> str: