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

Studio: route Linux hosts to the fork only when a usable GPU is present

Browse source
This commit is contained in:
oobabooga 2026-06-09 22:47:00 -03:00
commit 82f993e593
2 changed files with 33 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

27
studio/setup.sh
View file

@ -670,6 +670,7 @@ fi
# ── GPU detection summary (mirrors setup.ps1 step "gpu" block) ──
_setup_amd_detected=false
_setup_nvidia_usable=false
_setup_gfx_all=""
_setup_mkt=""
if command -v rocminfo >/dev/null 2>&1 && \
@ -690,6 +691,7 @@ fi
if command -v nvidia-smi >/dev/null 2>&1 && \
nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
_setup_nvidia_usable=true
step "gpu" "NVIDIA GPU detected"
elif [ "$_setup_amd_detected" = true ]; then
_setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
@ -767,12 +769,25 @@ _HOST_MACHINE="$(uname -m 2>/dev/null || true)"
# CPU-only Linux (x86_64 and arm64) routes there; GPU Linux, Windows and macOS
# use unslothai.
_LINUX_HAS_GPU=false
for _GPU_TOOL in nvidia-smi rocminfo amd-smi hipconfig hipinfo; do
if command -v "$_GPU_TOOL" >/dev/null 2>&1; then
_LINUX_HAS_GPU=true
break
fi
done
# Route to the fork only for a usable GPU. NVIDIA counts only when a device is
# actually enumerated (_setup_nvidia_usable, from the nvidia-smi -L probe above)
# AND not hidden via CUDA_VISIBLE_DEVICES=-1 -- mirroring install_llama_prebuilt.py's
# has_usable_nvidia. Mere nvidia-smi presence (CPU-only CUDA-toolkit containers,
# broken drivers) or a hidden GPU therefore takes the ggml-org CPU prebuilt
# instead of a slow source build. AMD is deliberately left on tooling presence,
# not usability: an unusable NVIDIA host has a good CPU prebuilt to fall back to,
# whereas tightening AMD would regress ROCm hosts exposing only hipconfig/hipinfo
# into an unnecessary CPU build.
if [ "$_setup_nvidia_usable" = true ] && [ "${CUDA_VISIBLE_DEVICES:-}" != "-1" ]; then
_LINUX_HAS_GPU=true
else
for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do
if command -v "$_GPU_TOOL" >/dev/null 2>&1; then
_LINUX_HAS_GPU=true
break
fi
done
fi
if [ "$_HOST_SYSTEM" = "Linux" ] \
&& [ "$_HOST_MACHINE" = "x86_64" ] \

12
tests/studio/install/test_pr4562_bugfixes.py
View file

@ -694,6 +694,18 @@ class TestSourceCodePatterns:
assert "_HELPER_RELEASE_REPO}/releases/latest" not in content
assert "ggml-org/llama.cpp/releases/latest" not in content
def test_setup_sh_routes_to_fork_only_on_usable_gpu(self):
"""Linux fork-vs-ggml routing must gate NVIDIA on actual GPU usability,
not mere nvidia-smi presence, so CPU-only / hidden-GPU hosts (e.g.
CUDA_VISIBLE_DEVICES=-1) get the ggml CPU prebuilt instead of a source
build. Guards against a silent revert to the old presence-only loop."""
content = SETUP_SH.read_text()
assert '[ "$_setup_nvidia_usable" = true ]' in content
assert "CUDA_VISIBLE_DEVICES" in content
# nvidia-smi must NOT be back in the bare presence loop.
assert "for _GPU_TOOL in nvidia-smi" not in content
assert "for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo" in content
def test_setup_sh_reports_installed_prebuilt_release(self):
"""Shell wrapper should report the installed prebuilt release from metadata."""
content = SETUP_SH.read_text()