From 172d9d1c8e98b136d2fa124bf7f76ece39ee5335 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 1 Jun 2026 06:35:43 -0700 Subject: [PATCH] Studio: source-build arm64 Linux GPU hosts, with a CPU prebuilt fallback (#5924) * 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> --- studio/install_llama_prebuilt.py | 32 ++++ studio/setup.sh | 26 +++ tests/studio/install/test_selection_logic.py | 162 +++++++++++++++++++ 3 files changed, 220 insertions(+) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index b0af64fc8d..ca26a112bc 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -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 diff --git a/studio/setup.sh b/studio/setup.sh index 3a149d465f..77013f0e6d 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -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 "" diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 8cb53d1aeb..30ad3dd88b 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -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