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

Studio: source-build arm64 Linux GPU hosts, with a CPU prebuilt fallback (#5924)

Browse source 
* Studio: fall back to source build for arm64 Linux GPU hosts

setup.sh routes a Linux host with any GPU tool to the unslothai/llama.cpp fork,
which publishes only linux-x64 bundles. On an arm64 host with a GPU (GH200,
GB200, DGX Spark) the resolver then selected an x64 CUDA bundle, which cannot run
on aarch64. Routing those hosts to ggml-org instead would install a CPU-only
arm64 build, silently losing the GPU.

Guard resolve_simple_install_release_plans so an arm64 Linux host on the fork
raises PrebuiltFallback before any release is fetched, letting setup.sh do a
source build that actually targets the GPU. x86_64 hosts and arm64 CPU hosts
(which route to ggml-org) are unaffected.

Add tests covering the arm64 fork raise, the x86_64 pass-through, and the arm64
CPU ggml-org path.

* Studio: install ggml-org arm64 CPU prebuilt when the arm64 GPU source build fails

Per review of #5924: arm64 Linux GPU hosts have no CUDA prebuilt anywhere (the
unslothai fork is x64 only, ggml-org ships no Linux CUDA build), so they source
build for the GPU. If that build produces no binary, the host was left without
llama.cpp.

Add a --cpu-fallback flag to install_llama_prebuilt.py that drops the host GPU
attributes so the CPU prebuilt for the host arch is selected (a GPU host cannot
otherwise pick the CPU bundle). setup.sh calls it against ggml-org as a last
resort for arm64 Linux when the source build degraded, installing the
ubuntu-arm64 CPU build instead of leaving the host with no llama.cpp.

Add tests: force_cpu drops GPU attrs before planning, a CPU-forced arm64 host
selects the ggml-org ubuntu-arm64 bundle, and setup.sh wires the fallback.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-01 06:35:43 -07:00 • committed by GitHub
commit 172d9d1c8e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 220 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

32
studio/install_llama_prebuilt.py
View file

@ -1661,6 +1661,15 @@ def resolve_simple_install_release_plans(
) -> tuple[str, list[InstallReleasePlan]]:
repo = published_repo or DEFAULT_PUBLISHED_REPO
requested_tag = normalized_requested_llama_tag(llama_tag)
# The unslothai/llama.cpp fork ships only linux-x64 bundles. An arm64 Linux
# host with a GPU (GH200/GB200/DGX Spark) routes here; it must not install an
# x64 binary, so fall back to a source build that targets the GPU rather than
# selecting the wrong arch (or silently dropping to a CPU arm64 build).
if host.is_linux and not host.is_x86_64 and repo == DEFAULT_PUBLISHED_REPO:
raise PrebuiltFallback(
f"{repo} ships only linux-x64 prebuilts; "
f"{host.machine or 'non-x64'} Linux falls back to source build"
)
allow_older_release_fallback = (
requested_tag == "latest" and not published_release_tag
)
@ -6451,10 +6460,22 @@ def install_prebuilt(
*,
simple_policy: bool = False,
override_has_rocm: bool = False,
force_cpu: bool = False,
) -> None:
host = detect_host()
if override_has_rocm and not host.has_rocm:
host = dataclasses_replace(host, has_rocm = True)
if force_cpu:
# Explicit CPU fallback: drop GPU attributes so the CPU prebuilt for this
# OS/arch is selected. setup.sh uses this for arm64 Linux GPU hosts whose
# source build failed, where no arm64 CUDA prebuilt exists anywhere.
host = dataclasses_replace(
host,
has_usable_nvidia = False,
has_physical_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
)
choice: AssetChoice | None = None
try:
with install_lock(install_lock_path(install_dir)):
@ -6599,6 +6620,16 @@ def parse_args() -> argparse.Namespace:
"so the HIP llama.cpp prebuilt is selected even when hipinfo is not on PATH."
),
)
parser.add_argument(
"--cpu-fallback",
action = "store_true",
default = False,
help = (
"Select the CPU prebuilt for this OS/arch even when a GPU is present. "
"setup.sh uses this as a last resort for arm64 Linux GPU hosts whose "
"source build failed (no arm64 CUDA prebuilt exists anywhere)."
),
)
resolve_group = parser.add_mutually_exclusive_group()
resolve_group.add_argument(
"--resolve-llama-tag",
@ -6720,6 +6751,7 @@ def main() -> int:
published_release_tag = args.published_release_tag or "",
simple_policy = args.simple_policy,
override_has_rocm = args.has_rocm,
force_cpu = args.cpu_fallback,
)
return EXIT_SUCCESS

26
studio/setup.sh
View file

@ -1363,6 +1363,32 @@ else
}
fi # end _SKIP_GGUF_BUILD check
# ── arm64 Linux GPU: CPU prebuilt as a last resort ──
# arm64 Linux with a GPU has no CUDA prebuilt anywhere (the unslothai fork is
# x64 only; ggml-org ships no Linux CUDA build), so it source-builds for the
# GPU above. If that produced no binary, install ggml-org's arm64 CPU prebuilt
# instead of leaving the host without llama.cpp.
if [ "$_LLAMA_CPP_DEGRADED" = true ] \
&& [ "$_HOST_SYSTEM" = "Linux" ] \
&& { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; }; then
substep "GPU source build unavailable; trying ggml-org arm64 CPU prebuilt..."
_ARM64_CPU_CMD=(
python "$SCRIPT_DIR/install_llama_prebuilt.py"
--install-dir "$LLAMA_CPP_DIR"
--llama-tag "$_REQUESTED_LLAMA_TAG"
--published-repo "ggml-org/llama.cpp"
--simple-policy
--cpu-fallback
)
# Trust the installer's exit code: it validates the server before exiting 0,
# the same signal the primary prebuilt path above relies on.
if run_quiet_no_exit "arm64 CPU prebuilt" "${_ARM64_CPU_CMD[@]}"; then
step "llama.cpp" "arm64 CPU prebuilt installed (GPU build unavailable)" "$C_WARN"
_LLAMA_CPP_DEGRADED=false
print_installed_llama_prebuilt_release "$LLAMA_CPP_DIR"
fi
fi
# ── Footer ──
if [ "$_LLAMA_ONLY" = "1" ]; then
echo ""

162
tests/studio/install/test_selection_logic.py
View file

@ -2762,3 +2762,165 @@ class TestResolveSimpleMacosPin:
assert plans[0].release_tag == "b9442"
# No pin: the iterator was asked for latest, not a specific tag.
assert calls[0][2] == "latest"
# ===========================================================================
# Linux arm64 + GPU must not install the x64-only fork bundle
# ===========================================================================
class TestLinuxArm64ForkFallsBackToSource:
"""The unslothai/llama.cpp fork ships only linux-x64 bundles. An arm64
Linux host with a GPU (GH200/GB200/DGX Spark) routes to the fork and must
fall back to a source build instead of selecting an x64 binary."""
def test_arm64_nvidia_fork_raises_before_fetching_releases(self, monkeypatch):
# Guard fires before any release is fetched: poison the iterator to prove
# it is never called.
def _boom(*_a, **_k):
raise AssertionError("iterator must not run for arm64 fork hosts")
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", _boom
)
host = make_host(system = "Linux", machine = "aarch64")
with pytest.raises(PrebuiltFallback, match = "linux-x64 prebuilts"):
resolve_simple_install_release_plans(
"latest", host, "unslothai/llama.cpp", ""
)
def test_x86_64_fork_is_not_blocked_by_the_arch_guard(self, monkeypatch):
# x64 host must pass the guard and reach the iterator (here empty, so it
# raises the generic message, not the arch one).
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"iter_release_payloads_by_time",
lambda *_a, **_k: iter(()),
)
host = make_host(system = "Linux", machine = "x86_64")
with pytest.raises(PrebuiltFallback) as exc:
resolve_simple_install_release_plans(
"latest", host, "unslothai/llama.cpp", ""
)
assert "linux-x64 prebuilts" not in str(exc.value)
def test_arm64_cpu_on_ggml_org_is_not_blocked(self, monkeypatch):
# CPU-only arm64 routes to ggml-org (not the fork), so the guard must not
# fire; it reaches the iterator (empty here -> generic message).
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"iter_release_payloads_by_time",
lambda *_a, **_k: iter(()),
)
host = make_host(
system = "Linux",
machine = "aarch64",
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
has_physical_nvidia = False,
has_usable_nvidia = False,
)
with pytest.raises(PrebuiltFallback) as exc:
resolve_simple_install_release_plans(
"latest", host, "ggml-org/llama.cpp", ""
)
assert "linux-x64 prebuilts" not in str(exc.value)
# ===========================================================================
# arm64 Linux GPU: CPU prebuilt fallback after a failed source build (--cpu-fallback)
# ===========================================================================
class TestCpuFallback:
"""--cpu-fallback drops GPU attributes so the CPU prebuilt for the host's
OS/arch is selected, letting an arm64 GPU host install ggml-org's arm64 CPU
build as a last resort when its source build produced no binary."""
_SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
def _arm64_nvidia(self):
return make_host(
system = "Linux",
machine = "aarch64",
driver_cuda_version = (13, 0),
compute_caps = ["90"],
has_physical_nvidia = True,
has_usable_nvidia = True,
)
def test_force_cpu_drops_gpu_attrs_before_planning(self, monkeypatch, tmp_path):
captured = {}
def _capture(llama_tag, host, *a, **k):
captured["host"] = host
raise PrebuiltFallback("stop after capture")
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", self._arm64_nvidia)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"resolve_simple_install_release_plans",
_capture,
)
# install_prebuilt exits EXIT_FALLBACK on PrebuiltFallback; we only care
# about the host it handed to the resolver before that.
with pytest.raises(SystemExit):
INSTALL_LLAMA_PREBUILT.install_prebuilt(
install_dir = tmp_path / "llama",
llama_tag = "latest",
published_repo = "ggml-org/llama.cpp",
published_release_tag = "",
simple_policy = True,
force_cpu = True,
)
host = captured["host"]
assert host.has_usable_nvidia is False
assert host.has_physical_nvidia is False
assert host.has_rocm is False
# Arch is preserved so the arm64 CPU bundle (not x64) is chosen.
assert host.is_arm64 is True
def test_cpu_forced_arm64_selects_ubuntu_arm64(self):
tag = "b9444"
release = {
"tag_name": tag,
"assets": [
{
"name": f"llama-{tag}-bin-ubuntu-arm64.tar.gz",
"browser_download_url": f"https://x/llama-{tag}-bin-ubuntu-arm64.tar.gz",
},
{
"name": f"llama-{tag}-bin-ubuntu-x64.tar.gz",
"browser_download_url": f"https://x/llama-{tag}-bin-ubuntu-x64.tar.gz",
},
],
}
# A GPU arm64 host cannot pick the CPU arm64 bundle on its own.
with pytest.raises(PrebuiltFallback):
direct_upstream_release_plan(
release, self._arm64_nvidia(), "ggml-org/llama.cpp", "latest"
)
# force_cpu drops the GPU attributes, so the CPU arm64 bundle is selected.
cpu_host = make_host(
system = "Linux",
machine = "aarch64",
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
has_physical_nvidia = False,
has_usable_nvidia = False,
)
plan = direct_upstream_release_plan(
release, cpu_host, "ggml-org/llama.cpp", "latest"
)
assert plan.attempts[0].install_kind == "linux-arm64"
assert plan.attempts[0].name == f"llama-{tag}-bin-ubuntu-arm64.tar.gz"
def test_setup_sh_has_arm64_cpu_prebuilt_fallback(self):
source = self._SETUP_SH.read_text(encoding = "utf-8")
assert "--cpu-fallback" in source
# Fallback targets ggml-org (the only repo with an arm64 Linux build) and
# is gated on a degraded source build for arm64.
assert "ggml-org/llama.cpp" in source
assert "_LLAMA_CPP_DEGRADED" in source