From 9d39aeec2b505541e6be6ef44c8bddb28cc17802 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 12 Jun 2026 14:09:41 +0000
Subject: [PATCH] studio: honor UNSLOTH_TORCH_INDEX_FAMILY in CUDA repair path,
 assert torch CUDA family at studio image build

_detect_cuda_torch_index_url now respects the explicit family override
before probing nvidia-smi, matching install.sh get_torch_index_url and
install.ps1 Get-TorchIndexUrl. Without it, a GPU-less environment falls
back to cu126 wheels which lack sm_100/sm_120 kernels and break training
on Blackwell. ROCm repair path is intentionally unchanged.

Dockerfile.studio now fails the build if the Studio venv torch local
version tag does not match the pinned TORCH_FAMILY, so a studio ref whose
installer ignores the override can never ship a silently wrong image.
Metadata-only check so QEMU arm64 builds do not need to load torch.
---
 docker/Dockerfile.studio       | 6 ++++++
 studio/install_python_stack.py | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio
index 40dfbcd69f..fc77902f15 100644
--- a/docker/Dockerfile.studio
+++ b/docker/Dockerfile.studio
@@ -101,6 +101,12 @@ RUN set -eux \
  && UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \
     UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \
     bash install.sh --local \
+    # Fail loud if the Studio venv torch missed the pinned CUDA family (an
+    # install.sh that ignores UNSLOTH_TORCH_INDEX_FAMILY falls back to
+    # nvidia-smi probing, which cannot work at build time and lands on cu126
+    # wheels with no sm_100/sm_120 kernels). metadata check only: importing
+    # torch needs native libs, which QEMU arm64 builds cannot load.
+ && "${UNSLOTH_STUDIO_HOME}/unsloth_studio/bin/python" -c "from importlib.metadata import version; v = version('torch'); assert v.endswith('+${TORCH_FAMILY}'), 'Studio venv torch ' + v + ' does not match ${TORCH_FAMILY}'; print('Studio venv torch', v)" \
  && rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" /root/.cache \
  && if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \
         for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index e540aac305..8f7c2d2c5a 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -816,6 +816,12 @@ def _detect_cuda_torch_index_url() -> str:
     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).
     """
+    # Explicit override (parity with install.sh / install.ps1):
+    # UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel
+    # index when probing is wrong or impossible (no GPU at build time, CI).
+    family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY")
+    if family:
+        return f"{_PYTORCH_WHL_BASE}/{family}"
     exe = shutil.which("nvidia-smi")
     if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
         exe = "/usr/bin/nvidia-smi"