From e8ac40fa5b3ba7c34a3dd9c8c9b490a5deee6286 Mon Sep 17 00:00:00 2001
From: danielhanchen <danielhanchen@gmail.com>
Date: Fri, 12 Jun 2026 05:06:25 +0000
Subject: [PATCH] docker/studio: deterministic Studio install inside the image
 build

Two failures from the first in-image Studio install, both rooted in
install.sh probing the build host:

1. setup.sh aborted on the pre-linked llama.cpp dir: 'already exists and
   is not marked as a Studio-owned llama.cpp install'. The dir is the
   image's baked prebuilt, provisioned exclusively for Studio, so write
   the .unsloth-studio-owned marker next to the binaries.

2. With no GPU and no nvidia-smi in the build container, install.sh fell
   back to cu126 torch wheels for the Studio venv (and would pick cpu
   wheels on a CI runner without /proc/driver/nvidia), so the published
   image's Studio venv would depend on which host built it and could not
   train on Blackwell. get_torch_index_url now honours an explicit
   UNSLOTH_TORCH_INDEX_FAMILY override naming the index leaf (cu128,
   cu130, rocm7.2, cpu, ...). The resolved family flows into
   UNSLOTH_TORCH_BACKEND, which install_python_stack.py already consumes,
   so the whole downstream chain follows the pin. Dockerfile.studio sets
   cu128 on amd64 and cu130 on arm64 (DGX Spark / Grace).
---
 docker/Dockerfile.studio | 26 +++++++++++++++++++++++---
 install.sh               | 10 ++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio
index d72c6f3a4f..8afd9e9819 100644
--- a/docker/Dockerfile.studio
+++ b/docker/Dockerfile.studio
@@ -33,6 +33,7 @@ FROM ${BASE_IMAGE}
 # that pins BASE_IMAGE to a digest should pin this too (same UNSLOTH_REF as
 # the base) so the published image is reproducible against a known ref.
 ARG UNSLOTH_STUDIO_REF=main
+ARG TARGETARCH
 
 USER root
 ENV UNSLOTH_STUDIO_HOME=/opt/unsloth-studio \
@@ -55,18 +56,37 @@ RUN apt-get update \
 # The llama.cpp symlink BEFORE install.sh points Studio's prebuilt dir at
 # the bundle already baked into the base image (validated, sha256-checked,
 # UNSLOTH_PREBUILT_INFO.json present), so the installer's prebuilt step
-# recognises it and skips a second ~400MB download.
+# recognises it and skips a second ~400MB download. The
+# .unsloth-studio-owned marker satisfies setup.sh's ownership assertion for
+# custom STUDIO_HOMEs -- the dir IS provisioned exclusively for Studio.
+#
+# UNSLOTH_TORCH_INDEX_FAMILY pins the torch wheel index for the Studio
+# venv: at build time there is no GPU and no nvidia-smi, so install.sh's
+# probing would land on cpu or cu126 wheels depending on which host built
+# the image. The image targets CUDA: cu128 on amd64 (Turing..Blackwell,
+# same line as the base venv), cu130 on arm64 (DGX Spark / Grace, the
+# aarch64 CUDA wheel line).
+#
 # fetch+checkout FETCH_HEAD instead of `clone --branch` because the CI
 # pipeline passes a commit SHA as the ref (clone --branch only accepts
 # branch/tag names).
-RUN mkdir -p "${UNSLOTH_STUDIO_HOME}" \
+RUN set -eux \
+ && case "${TARGETARCH:-amd64}" in \
+        amd64) TORCH_FAMILY="cu128" ;; \
+        arm64) TORCH_FAMILY="cu130" ;; \
+        *)     echo "ERROR: unsupported TARGETARCH=${TARGETARCH}" >&2; exit 1 ;; \
+    esac \
+ && mkdir -p "${UNSLOTH_STUDIO_HOME}" \
  && ln -s /opt/unsloth/llama.cpp "${UNSLOTH_STUDIO_HOME}/llama.cpp" \
+ && touch /opt/unsloth/llama.cpp/.unsloth-studio-owned \
  && git init -q "${UNSLOTH_STUDIO_HOME}/src" \
  && cd "${UNSLOTH_STUDIO_HOME}/src" \
  && git remote add origin https://github.com/unslothai/unsloth \
  && git fetch -q --depth 1 origin "${UNSLOTH_STUDIO_REF}" \
  && git checkout -q FETCH_HEAD \
- && UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" bash install.sh --local \
+ && UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \
+    UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \
+    bash install.sh --local \
  && rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" /root/.cache
 
 COPY supervisord.conf /etc/supervisor/supervisord.conf
diff --git a/install.sh b/install.sh
index 532ac61bc0..c7be27942e 100755
--- a/install.sh
+++ b/install.sh
@@ -1807,6 +1807,16 @@ _has_usable_nvidia_gpu() {
 get_torch_index_url() {
     _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
     _base="${_base%/}"
+    # Explicit pin for hosts where probing is impossible or must not happen
+    # (Docker image builds, CI runners). Names the index path leaf directly:
+    # UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|rocm7.2|cpu|...
+    # The Blackwell Docker image build uses this: at build time there is no
+    # GPU and no nvidia-smi, but the image targets CUDA, so probing would
+    # land on the cpu (CI) or cu126 (GPU build hosts leak /proc/driver/nvidia
+    # but not nvidia-smi) wheels depending on which host built the image.
+    if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
+        echo "$_base/${UNSLOTH_TORCH_INDEX_FAMILY}"; return
+    fi
     # macOS: always CPU (no CUDA support)
     case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
     # Try nvidia-smi -- require the binary to actually list a usable GPU.