From 37ce5077f757335f707c16e74d20d390d9100dba Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 29 Jun 2026 05:35:08 +0000
Subject: [PATCH] torch-index override: classify CUDA pin by leaf; trim blank
 shell overrides

_ensure_cuda_torch only overrode the NVIDIA-presence gate for *any* pinned index,
so a non-CUDA mirror URL (or a ROCm/CPU pin) on a non-NVIDIA host with ROCm torch
could force a CUDA reinstall over a working ROCm venv. Add
_explicit_cuda_torch_index_url() (leaf cu*), matching the ROCm/CPU helpers, and
gate on it instead.

install.sh::get_torch_index_url treated a whitespace-only UNSLOTH_TORCH_INDEX_URL
/ _FAMILY as authoritative (yielding an invalid index), unlike the Python .strip()
and PowerShell IsNullOrWhiteSpace paths; trim leading/trailing whitespace first.
---
 install.sh                     | 13 +++++++++----
 studio/install_python_stack.py | 17 ++++++++++++++++-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/install.sh b/install.sh
index 964a715f22..995a4622cf 100755
--- a/install.sh
+++ b/install.sh
@@ -2001,16 +2001,21 @@ get_torch_index_url() {
     # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); UNSLOTH_TORCH_INDEX_FAMILY
     # is the convenience form (cpu, cu124, cu126, cu128, cu130, rocm6.4, ...)
     # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured.
-    if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then
+    # Trim leading/trailing whitespace so a whitespace-only value is treated as
+    # unset (parity with the Python .strip() and PowerShell IsNullOrWhiteSpace
+    # paths); otherwise "   " would pass -n and yield an invalid index URL.
+    _url="${UNSLOTH_TORCH_INDEX_URL:-}"
+    _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}"
+    if [ -n "$_url" ]; then
         # Strip ALL trailing slashes (match the Python side's .rstrip("/") and the
         # Strix mirror handling below) -- a double/triple-slash URL 404s on strict
         # pip proxies (artifactory, sonatype).
-        _url="${UNSLOTH_TORCH_INDEX_URL}"
         while [ "${_url%/}" != "$_url" ]; do _url="${_url%/}"; done
         echo "$_url"; return
     fi
-    if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
-        _family="${UNSLOTH_TORCH_INDEX_FAMILY}"
+    _family="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
+    _family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}"
+    if [ -n "$_family" ]; then
         while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done
         while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done
         echo "$_base/$_family"; return
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 9dc7cdd02f..2bbf2329dd 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -1018,6 +1018,21 @@ def _explicit_cpu_torch_index_url() -> "str | None":
     return url if leaf == "cpu" else None
 
 
+def _explicit_cuda_torch_index_url() -> "str | None":
+    """The pinned wheel index URL when it names a CUDA family (leaf cu*), else None.
+
+    Mirrors _explicit_rocm/cpu_torch_index_url so _ensure_cuda_torch only treats a
+    *CUDA* pin as authority to override the NVIDIA-presence gate. An arbitrary
+    mirror URL (or a ROCm/CPU pin) must not force a CUDA reinstall over a working
+    ROCm/CPU venv on a non-NVIDIA host.
+    """
+    url = _explicit_torch_index_url()
+    if url is None:
+        return None
+    leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
+    return url if leaf.startswith("cu") else None
+
+
 def _ensure_cuda_torch() -> None:
     """Repair a venv whose torch is a ROCm build on an NVIDIA host.
 
@@ -1053,7 +1068,7 @@ def _ensure_cuda_torch() -> None:
     # An explicit CUDA wheel-index pin (headless / container / CI cross-install)
     # commits to CUDA wheels regardless of whether a GPU is visible here, so it
     # overrides the GPU-presence gate.
-    if not _has_usable_nvidia_gpu() and _explicit_torch_index_url() is None:
+    if not _has_usable_nvidia_gpu() and _explicit_cuda_torch_index_url() is None:
         return
 
     # Classify the installed torch: "hip" (ROCm build -- the poisoning