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

install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override CUDA wheel detection

Browse source 
get_torch_index_url (and the studio-update mirror _detect_cuda_torch_index_url)
chose the torch wheel family solely by probing the host GPU, with no override.
In a headless / container / CI build the host driver is visible via the
/proc/driver/nvidia/gpus fallback but nvidia-smi cannot report a CUDA version,
so the function fell back to its cu126 default and installed the wrong wheels
(e.g. a cu128 image got cu126 torch).

Add an explicit override checked before any probing, in both the shell installer
and the Python studio-update path:
  - UNSLOTH_TORCH_INDEX_URL   full index URL, used verbatim (wins)
  - UNSLOTH_TORCH_INDEX_FAMILY family (cpu, cu128, rocm6.4, ...) appended to the
                               mirror base (UNSLOTH_PYTORCH_MIRROR still honoured)

This matches how the published GPU images select CUDA -- vLLM and SGLang take the
CUDA version from an explicit build ARG rather than detecting it, and the Unsloth
Docker base image already pins the cu128 index directly. Desktop installs are
unchanged: with no override set, detection runs exactly as before.

Adds test_get_torch_index_url.sh cases for the override (family, full URL,
precedence, mirror base, trailing-slash strip, empty-ignored).
This commit is contained in:
Daniel Han 2026-06-26 04:59:32 +00:00
commit 138c0949b4
3 changed files with 64 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

17
studio/install_python_stack.py
View file

@ -933,11 +933,20 @@ def _detect_cuda_torch_index_url() -> str:
Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update`
repairs to the same wheel family a fresh `curl | sh` install would pick.
Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the
legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings.
Defaults to cu126 when nvidia-smi is missing or the version is unreadable
(e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
Honours the same explicit overrides first (UNSLOTH_TORCH_INDEX_URL /
UNSLOTH_TORCH_INDEX_FAMILY) so a headless / container / CI install never lets
the host GPU decide the wheel family. Otherwise probes nvidia-smi (PATH, then
/usr/bin/nvidia-smi) and parses both the legacy "CUDA Version:" and the newer
"CUDA UMD Version:" spellings. Defaults to cu126 when nvidia-smi is missing or
the version is unreadable (e.g. NVIDIA detected only via the
/proc/driver/nvidia/gpus fallback).
"""
_override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
if _override_url:
return _override_url.rstrip("/")
_override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
if _override_family:
return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}"
exe = shutil.which("nvidia-smi")
if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
exe = "/usr/bin/nvidia-smi"