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>
This commit is contained in:
Daniel Han
GitHub
parent
f695fbd0fe
commit
172d9d1c8e
3 changed files with 220 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue