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

Studio: source-build + CPU last-resort recovery when no GPU prebuilt offloads

Browse source 
setup.sh/setup.ps1 now run install_llama_prebuilt.py --smoke-test on a freshly
source-built GPU binary and retry a CPU build if it loaded on CPU only. When a
GPU host's source build produces no binary, both scripts fall back to the CPU
prebuilt (--cpu-fallback) as a labelled last resort instead of leaving the host
without llama.cpp. setup.ps1 also guards an empty CUDA arch (no PTX-only binary,
#5854). The POSIX smoke-test exit code is captured set -e safe. Adds a fake
llama-server and an end-to-end spoof test that runs the real validate_server
against it with no GPU.
This commit is contained in:
danielhanchen 2026-06-01 15:53:12 +00:00
commit c85c62e456
6 changed files with 484 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

141
studio/setup.ps1
View file

@ -2952,39 +2952,53 @@ if (-not $NeedLlamaSourceBuild) {
$CmakeArgs += '-DLLAMA_CURL=OFF'
}
$CmakeArgs += '-DCMAKE_EXE_LINKER_FLAGS=/NODEFAULTLIB:LIBCMT'
# CUDA flags -- only if GPU available, otherwise explicitly disable
# CUDA flags -- only if GPU available, otherwise explicitly disable.
# $LlamaCudaBuild gates the post-build GPU smoke test and CUDA->CPU
# retry below.
$LlamaCudaBuild = $false
if ($HasNvidiaSmi -and $NvccPath) {
$CmakeArgs += '-DGGML_CUDA=ON'
# Accept a host MSVC newer than nvcc's whitelist; a fresh toolkit
# (e.g. CUDA 13.3) otherwise aborts with "#error -- unsupported
# Microsoft Visual Studio version!". Mirrors the Linux fix. Via env
# (covers the configure probe + build), after Refresh-Environment, idempotent.
$nvccAllowFlag = '-allow-unsupported-compiler'
if ([string]::IsNullOrEmpty($env:NVCC_PREPEND_FLAGS)) {
$env:NVCC_PREPEND_FLAGS = $nvccAllowFlag
} elseif ($env:NVCC_PREPEND_FLAGS -notlike "*$nvccAllowFlag*") {
$env:NVCC_PREPEND_FLAGS = "$($env:NVCC_PREPEND_FLAGS) $nvccAllowFlag"
}
substep "NVCC_PREPEND_FLAGS = $env:NVCC_PREPEND_FLAGS"
$CmakeArgs += "-DCUDAToolkit_ROOT=$CudaToolkitRoot"
$CmakeArgs += "-DCUDA_TOOLKIT_ROOT_DIR=$CudaToolkitRoot"
$CmakeArgs += "-DCMAKE_CUDA_COMPILER=$NvccPath"
# Resolve a concrete CUDA architecture FIRST. A CUDA build with no
# -DCMAKE_CUDA_ARCHITECTURES is PTX-only and can fail at runtime on a
# driver older than the toolkit ("the provided PTX was compiled with
# an unsupported toolchain", #5854). If we cannot resolve a supported
# arch, build CPU-only instead of shipping a silently broken binary.
$cudaArchFlag = $null
if ($CudaArch) {
# Validate nvcc actually supports this architecture
if (Test-NvccArchSupport -NvccExe $NvccPath -Arch $CudaArch) {
$CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$CudaArch"
$cudaArchFlag = "-DCMAKE_CUDA_ARCHITECTURES=$CudaArch"
} else {
# GPU arch too new for this toolkit -- fall back to highest supported.
# PTX forward-compatibility will JIT-compile for the actual GPU at runtime.
# GPU arch too new for this toolkit -- fall back to highest
# supported. PTX forward-compat will JIT for the real GPU.
$maxArch = Get-NvccMaxArch -NvccExe $NvccPath
if ($maxArch) {
$CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$maxArch"
$cudaArchFlag = "-DCMAKE_CUDA_ARCHITECTURES=$maxArch"
substep "GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" "Yellow"
substep "Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" "Yellow"
}
# else: omit flag entirely, let cmake pick defaults
}
}
if ($cudaArchFlag) {
$CmakeArgs += '-DGGML_CUDA=ON'
# Accept a host MSVC newer than nvcc's whitelist; a fresh toolkit
# (e.g. CUDA 13.3) otherwise aborts with "#error -- unsupported
# Microsoft Visual Studio version!". Via env (covers the configure
# probe + build), after Refresh-Environment, idempotent.
$nvccAllowFlag = '-allow-unsupported-compiler'
if ([string]::IsNullOrEmpty($env:NVCC_PREPEND_FLAGS)) {
$env:NVCC_PREPEND_FLAGS = $nvccAllowFlag
} elseif ($env:NVCC_PREPEND_FLAGS -notlike "*$nvccAllowFlag*") {
$env:NVCC_PREPEND_FLAGS = "$($env:NVCC_PREPEND_FLAGS) $nvccAllowFlag"
}
substep "NVCC_PREPEND_FLAGS = $env:NVCC_PREPEND_FLAGS"
$CmakeArgs += "-DCUDAToolkit_ROOT=$CudaToolkitRoot"
$CmakeArgs += "-DCUDA_TOOLKIT_ROOT_DIR=$CudaToolkitRoot"
$CmakeArgs += "-DCMAKE_CUDA_COMPILER=$NvccPath"
$CmakeArgs += $cudaArchFlag
$LlamaCudaBuild = $true
} else {
substep "Could not resolve a supported CUDA architecture for this GPU/toolkit; building CPU-only to avoid a PTX-only binary that fails at runtime (#5854)" "Yellow"
$CmakeArgs += '-DGGML_CUDA=OFF'
}
} else {
$CmakeArgs += '-DGGML_CUDA=OFF'
}
@ -3025,6 +3039,60 @@ if (-not $NeedLlamaSourceBuild) {
}
}
# -- Step C.5: GPU smoke test + CUDA->CPU fallback (#5807 / #5854) --
# A CUDA build whose runtime backend fails to initialize still links and
# serves HTTP 200, but only from CPU; and setup.ps1 previously had no CPU
# fallback when a CUDA build failed at all. Both gaps are closed here.
if ($LlamaCudaBuild -and $BuildOk) {
$builtServer = Join-Path $BuildDir "bin\Release\llama-server.exe"
if (-not (Test-Path -LiteralPath $builtServer)) {
$builtServer = Join-Path $BuildDir "bin\llama-server.exe"
}
if (Test-Path -LiteralPath $builtServer) {
Write-Host ""
Write-Host "--- GPU smoke test ---" -ForegroundColor Cyan
& python "$PSScriptRoot\install_llama_prebuilt.py" --smoke-test "$builtServer" --install-dir "$LlamaCppDir" 2>&1 | Out-String | Write-Host
$smokeExit = $LASTEXITCODE
if ($smokeExit -eq 2) {
substep "GPU build runs on CPU only (GPU backend failed to initialize)" "Yellow"
$BuildOk = $false
$FailedStep = "GPU smoke test (ran on CPU)"
} elseif ($smokeExit -ne 0) {
substep "GPU smoke test inconclusive (exit $smokeExit); keeping GPU build" "Yellow"
}
}
}
# If a CUDA build was attempted and configure/build/smoke-test left it
# unusable, retry once with CUDA disabled so the user still gets a working
# (if slower) CPU llama-server instead of nothing.
if ($LlamaCudaBuild -and -not $BuildOk) {
substep "CUDA build unusable at: $FailedStep; retrying CPU-only build..." "Yellow"
$CpuCmakeArgs = @($CmakeArgs | Where-Object {
$_ -ne '-DGGML_CUDA=ON' -and
$_ -notlike '-DCMAKE_CUDA_ARCHITECTURES=*' -and
$_ -notlike '-DCUDAToolkit_ROOT=*' -and
$_ -notlike '-DCUDA_TOOLKIT_ROOT_DIR=*' -and
$_ -notlike '-DCMAKE_CUDA_COMPILER=*'
})
$CpuCmakeArgs += '-DGGML_CUDA=OFF'
if (Test-Path -LiteralPath $BuildDir) { Remove-Item -LiteralPath $BuildDir -Recurse -Force }
$cpuConfigure = cmake @CpuCmakeArgs 2>&1 | Out-String
if ($LASTEXITCODE -eq 0) {
$cpuBuild = cmake --build $BuildDir --config Release --target llama-server -j $NumCpu 2>&1 | Out-String
if ($LASTEXITCODE -eq 0) {
$BuildOk = $true
$LlamaCudaBuild = $false
$FailedStep = $null
substep "CPU-only llama.cpp build succeeded" "Green"
} else {
Write-LlamaFailureLog -Output $cpuBuild
}
} else {
Write-LlamaFailureLog -Output $cpuConfigure
}
}
# -- Step D: Build llama-quantize (optional, best-effort) --
if ($BuildOk) {
Write-Host ""
@ -3083,6 +3151,35 @@ if (-not $NeedLlamaSourceBuild) {
}
}
# ─────────────────────────────────────────────
# Windows GPU: CPU prebuilt as a last resort
# ─────────────────────────────────────────────
# A GPU host reaches the source build only when no GPU prebuilt offloads (the
# CPU prebuilt is deliberately not offered to a GPU host so it does not short
# circuit the source build, #5807). If that build produced no binary (no CUDA
# toolkit, compile failure), install the CPU prebuilt via --cpu-fallback so the
# host still gets a working (if slower) llama-server instead of nothing.
if ($script:LlamaCppDegraded -and ($HasNvidiaSmi -or $HasROCm)) {
substep "GPU build unavailable; trying CPU prebuilt as a last resort..." "Yellow"
$lastResortArgs = @(
"$PSScriptRoot\install_llama_prebuilt.py",
"--install-dir", $OriginalLlamaCppDir,
"--llama-tag", $RequestedLlamaTag,
"--published-repo", $HelperReleaseRepo,
"--simple-policy",
"--cpu-fallback"
)
$prevEAPLast = $ErrorActionPreference
$ErrorActionPreference = "Continue"
& python @lastResortArgs 2>&1 | Out-String | Write-Host
$lastResortExit = $LASTEXITCODE
$ErrorActionPreference = $prevEAPLast
if ($lastResortExit -eq 0) {
step "llama.cpp" "CPU prebuilt installed (GPU unavailable; inference will run on CPU)" "Yellow"
$script:LlamaCppDegraded = $false
}
}
# ─────────────────────────────────────────────
# Footer
# ─────────────────────────────────────────────

81
studio/setup.sh
View file

@ -1340,6 +1340,56 @@ else
fi
fi
# Map an install_llama_prebuilt.py --smoke-test exit code to a decision
# token. 2 (EXIT_FALLBACK) = definitively GPU-intended but CPU-only ->
# rebuild CPU; 0 = offload confirmed; anything else (1/EXIT_ERROR,
# signals) = inconclusive -> keep the GPU build rather than downgrade on
# uncertain evidence. Tiny + side-effect free so
# tests/sh/test_llama_gpu_smoke.sh can exercise it.
_classify_smoke_exit() {
case "$1" in
0) echo "ok" ;;
2) echo "cpu_only" ;;
*) echo "inconclusive" ;;
esac
}
# Post-build GPU smoke test (#5807 / #5854): a GPU build whose runtime
# backend fails to initialize still links and serves HTTP 200, but only
# from CPU. Confirm the fresh binary actually offloads to the GPU; if it
# ran on CPU only, retry a CPU build so the user gets a working (if
# slower) llama-server instead of a silently CPU-only "GPU" build.
# _gpu_fallback_label is empty for a pure CPU build (nothing to verify).
if [ "$BUILD_OK" = true ]; then
_SMOKE_LABEL="$(_gpu_fallback_label)"
if [ -n "$_SMOKE_LABEL" ] && [ -f "$_BUILD_TMP/build/bin/llama-server" ]; then
# if/else keeps set -e from aborting before we read the code.
if python "$SCRIPT_DIR/install_llama_prebuilt.py" \
--smoke-test "$_BUILD_TMP/build/bin/llama-server" \
--install-dir "$_BUILD_TMP" > "$_BUILD_TMP/gpu-smoke.log" 2>&1; then
_SMOKE_RC=0
else
_SMOKE_RC=$?
fi
case "$(_classify_smoke_exit "$_SMOKE_RC")" in
cpu_only)
substep "$_SMOKE_LABEL build runs on CPU only; retrying CPU build..." "$C_WARN"
rm -rf "$_BUILD_TMP/build"
if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then
_BUILD_DESC="building (CPU fallback after $_SMOKE_LABEL smoke test ran on CPU)"
GPU_BACKEND=""
run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
else
BUILD_OK=false
fi
;;
inconclusive)
substep "GPU smoke test inconclusive (exit $_SMOKE_RC); keeping $_SMOKE_LABEL build" "$C_WARN"
;;
esac
fi
fi
if [ "$BUILD_OK" = true ]; then
run_quiet_no_exit "build llama-quantize" cmake --build "$_BUILD_TMP/build" --config Release --target llama-quantize -j"$NCPU" || true
fi
@ -1373,27 +1423,36 @@ 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.
# ── Linux GPU: CPU prebuilt as a last resort ──
# A Linux GPU host reaches a source build when no GPU prebuilt offloads (the
# CPU prebuilt is deliberately not offered to a GPU host so it does not short
# circuit the source build -- #5807). If that build produced no binary (no
# toolkit, compile failure), install a CPU prebuilt instead of leaving the host
# without llama.cpp. arm64 has no CUDA prebuilt anywhere, so it always lands
# here on a degraded GPU build; x86_64 only when its source build also failed.
# Repo: ggml-org ships the arm64 CPU tarball; x86_64 uses the same published
# repo as the primary path (the unslothai fork carries linux-x64-cpu).
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=(
&& [ "$_LINUX_HAS_GPU" = true ]; then
if [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; then
_LASTRESORT_CPU_REPO="ggml-org/llama.cpp"
else
_LASTRESORT_CPU_REPO="$_HELPER_RELEASE_REPO"
fi
substep "GPU build unavailable; trying $_LASTRESORT_CPU_REPO CPU prebuilt as a last resort..."
_LASTRESORT_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"
--published-repo "$_LASTRESORT_CPU_REPO"
--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"
if run_quiet_no_exit "CPU prebuilt (last resort)" "${_LASTRESORT_CPU_CMD[@]}"; then
step "llama.cpp" "CPU prebuilt installed (GPU unavailable; inference will run on CPU)" "$C_WARN"
_LLAMA_CPP_DEGRADED=false
print_installed_llama_prebuilt_release "$LLAMA_CPP_DIR"
fi

1
tests/run_all.sh
View file

@ -9,6 +9,7 @@ sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
sh "$TESTS_DIR/sh/test_torch_constraint.sh"
sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
sh "$TESTS_DIR/sh/test_llama_gpu_smoke.sh"
echo ""
echo "=== Python tests ==="

59
tests/sh/test_llama_gpu_smoke.sh Executable file
View file

@ -0,0 +1,59 @@
#!/bin/bash
# Unit tests for _classify_smoke_exit() from studio/setup.sh.
#
# After a llama.cpp source build, setup.sh runs install_llama_prebuilt.py
# --smoke-test against the fresh binary and maps its exit code to a decision:
# 2 (EXIT_FALLBACK) -> "cpu_only" : GPU was requested but the model ran
# on CPU -> rebuild CPU (#5807 / #5854).
# 0 -> "ok" : GPU offload confirmed -> keep build.
# 1 / signals / etc -> "inconclusive" : could not validate -> keep GPU build
# (never downgrade on uncertainty).
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh"
PASS=0
FAIL=0
# Extract just the helper (same approach as test_nvcc_meets_llama_minimum.sh).
# The function and its closing brace sit at 8-space indent inside setup.sh.
_FUNC_FILE=$(mktemp)
sed -n '/ _classify_smoke_exit() {/,/^ }/p' "$SETUP_SH" > "$_FUNC_FILE"
assert_eq() {
_label="$1"; _expected="$2"; _actual="$3"
if [ "$_actual" = "$_expected" ]; then
echo " PASS: $_label"
PASS=$((PASS + 1))
else
echo " FAIL: $_label (expected '$_expected', got '$_actual')"
FAIL=$((FAIL + 1))
fi
}
run_classify() {
bash -c ". '$_FUNC_FILE'; _classify_smoke_exit '$1'"
}
echo "=== test_llama_gpu_smoke (_classify_smoke_exit) ==="
assert_eq "exit 0 -> ok" "ok" "$(run_classify 0)"
assert_eq "exit 2 -> cpu_only" "cpu_only" "$(run_classify 2)"
assert_eq "exit 1 -> inconclusive" "inconclusive" "$(run_classify 1)"
assert_eq "exit 3 -> inconclusive" "inconclusive" "$(run_classify 3)"
assert_eq "exit 137 -> inconclusive" "inconclusive" "$(run_classify 137)"
# Sanity: the extracted function is non-empty and well-formed.
if [ -s "$_FUNC_FILE" ] && grep -q 'cpu_only' "$_FUNC_FILE"; then
echo " PASS: _classify_smoke_exit extracted from setup.sh"
PASS=$((PASS + 1))
else
echo " FAIL: _classify_smoke_exit could not be extracted from setup.sh"
FAIL=$((FAIL + 1))
fi
rm -f "$_FUNC_FILE"
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ] || exit 1

106
tests/studio/install/fake_llama_server.py Normal file
View file

@ -0,0 +1,106 @@
#!/usr/bin/env python3
"""A fake llama-server for GPU-offload validation tests (no GPU required).
It accepts the arguments install_llama_prebuilt.py's validate_server passes
(``-m`` / ``--host`` / ``--port`` / ``--n-gpu-layers`` / ...), prints a canned
llama.cpp startup log chosen by the ``FAKE_LLAMA_MODE`` env var, then serves
HTTP 200 from ``/completion`` until killed. This lets CI exercise the real
validate_server subprocess + HTTP + log-classifier path on GPU-less Windows /
macOS / Linux runners: the same binary "starts and serves 200" while its log
says CPU-only or GPU, which is exactly the #5807 / #5830 situation.
FAKE_LLAMA_MODE (default "cuda"):
cuda device_info enumerates CUDA0 (GPU offload confirmed)
cuda_buffer older "CUDA0 model buffer size" + offloaded 33/33 lines
cpu device_info enumerates only CPU (the silent CPU fallback)
offloaded_zero "offloaded 0/33 layers to GPU" (definite CPU-only)
no_signal a log with no offload evidence (validator must not reject)
"""
import os
import sys
from http.server import BaseHTTPRequestHandler, HTTPServer
LOGS = {
"cuda": (
"0.00 I device_info:\n"
"0.01 I - CUDA0 : NVIDIA GeForce RTX 5070 (12282 MiB, 11000 MiB free)\n"
"0.01 I - CPU : Generic CPU (32000 MiB free)\n"
"0.01 I system_info: n_threads = 8 | CUDA : ARCHS = 1200 | CPU : AVX2 = 1\n"
"0.02 I srv llama_server: model loaded\n"
),
"cuda_buffer": (
"load_tensors: offloaded 33/33 layers to GPU\n"
"load_tensors: CUDA0 model buffer size = 21000.0 MiB\n"
"load_tensors: CPU_Mapped model buffer size = 0.6 MiB\n"
"srv llama_server: model loaded\n"
),
"cpu": (
"0.00 I device_info:\n"
"0.00 I - CPU : Generic CPU (32000 MiB free)\n"
"0.00 I system_info: n_threads = 8 | CPU : AVX2 = 1\n"
"0.01 I srv llama_server: model loaded\n"
),
"offloaded_zero": (
"load_tensors: offloaded 0/33 layers to GPU\n"
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB\n"
"srv llama_server: model loaded\n"
),
"no_signal": (
"INFO [main] starting server\n"
"load_tensors: file format = GGUF V3\n"
"srv llama_server: model loaded\n"
),
}
def _arg(name, default=None):
argv = sys.argv
for i, token in enumerate(argv):
if token == name and i + 1 < len(argv):
return argv[i + 1]
if token.startswith(name + "="):
return token.split("=", 1)[1]
return default
class _Handler(BaseHTTPRequestHandler):
def _ok(self):
body = b'{"content": "x", "tokens_predicted": 1}'
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def do_POST(self):
length = int(self.headers.get("Content-Length", 0) or 0)
if length:
self.rfile.read(length)
self._ok()
def do_GET(self):
self._ok()
def log_message(self, *args):
pass # keep stdout clean for the classifier
def main() -> int:
mode = os.environ.get("FAKE_LLAMA_MODE", "cuda")
sys.stdout.write(LOGS.get(mode, LOGS["cuda"]))
sys.stdout.flush()
host = _arg("--host", "127.0.0.1")
port = int(_arg("--port", "8080"))
server = HTTPServer((host, port), _Handler)
try:
server.serve_forever()
except KeyboardInterrupt:
pass
return 0
if __name__ == "__main__":
raise SystemExit(main())

129
tests/studio/install/test_gpu_offload_spoof.py Normal file
View file

@ -0,0 +1,129 @@
"""End-to-end GPU-offload spoof: run the real validate_server (real subprocess
+ real HTTP + real log classifier) against a fake llama-server, no GPU needed.
Unlike test_validate_server_gpu_offload.py (which mocks subprocess/urlopen),
this launches an actual process that "starts and serves HTTP 200" while its log
reports CPU-only or GPU offload, reproducing #5807 / #5830 end to end. POSIX
only: validate_server execs the binary path directly, which needs a shebang
wrapper; the Windows equivalent runs in the studio-gpu-offload-smoke workflow
via a .bat shim.
"""
import importlib.util
import os
import stat
import sys
from pathlib import Path
import pytest
if sys.platform == "win32":
pytest.skip("POSIX-only (Windows covered by the spoof CI workflow)", allow_module_level = True)
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
FAKE_SERVER = Path(__file__).resolve().parent / "fake_llama_server.py"
SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt_e2e", MODULE_PATH)
M = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = M
SPEC.loader.exec_module(M)
HostInfo = M.HostInfo
def linux_cuda_host(**overrides):
defaults = dict(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = "/usr/bin/nvidia-smi",
driver_cuda_version = (13, 0),
compute_caps = ["120"],
visible_cuda_devices = None,
has_physical_nvidia = True,
has_usable_nvidia = True,
has_rocm = False,
)
defaults.update(overrides)
return HostInfo(**defaults)
@pytest.fixture
def fake_server_binary(tmp_path):
"""A `llama-server` that execs the fake server, so validate_server runs it
exactly as it would a real prebuilt binary."""
binary = tmp_path / "llama-server"
binary.write_text(
"#!/bin/sh\n"
f'exec "{sys.executable}" "{FAKE_SERVER}" "$@"\n'
)
binary.chmod(binary.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
return binary
def _validate(binary, tmp_path, host, install_kind, mode, monkeypatch):
monkeypatch.setenv("FAKE_LLAMA_MODE", mode)
probe = tmp_path / "probe.gguf"
probe.write_bytes(b"GGUF\x00fake")
M.validate_server(binary, probe, host, tmp_path, install_kind = install_kind)
def test_cpu_only_binary_tagged_cuda_is_rejected(fake_server_binary, tmp_path, monkeypatch):
# The #5807 case: a binary that serves 200 but loaded the model on CPU.
with pytest.raises(M.GpuOffloadFailure):
_validate(fake_server_binary, tmp_path, linux_cuda_host(), "linux-cuda", "cpu", monkeypatch)
def test_offloaded_zero_binary_tagged_cuda_is_rejected(fake_server_binary, tmp_path, monkeypatch):
with pytest.raises(M.GpuOffloadFailure):
_validate(
fake_server_binary, tmp_path, linux_cuda_host(), "linux-cuda", "offloaded_zero", monkeypatch
)
def test_gpu_binary_tagged_cuda_passes(fake_server_binary, tmp_path, monkeypatch):
_validate(fake_server_binary, tmp_path, linux_cuda_host(), "linux-cuda", "cuda", monkeypatch)
def test_gpu_buffer_format_passes(fake_server_binary, tmp_path, monkeypatch):
_validate(fake_server_binary, tmp_path, linux_cuda_host(), "linux-cuda", "cuda_buffer", monkeypatch)
def test_cpu_only_binary_tagged_cpu_is_accepted(fake_server_binary, tmp_path, monkeypatch):
# A linux-cpu bundle is the intentional fallback; never GPU-gated.
_validate(fake_server_binary, tmp_path, linux_cuda_host(), "linux-cpu", "cpu", monkeypatch)
def test_no_signal_binary_tagged_cuda_is_accepted(fake_server_binary, tmp_path, monkeypatch):
# No offload evidence -> conservative: do not reject on no signal.
_validate(fake_server_binary, tmp_path, linux_cuda_host(), "linux-cuda", "no_signal", monkeypatch)
def test_smoke_test_cli_exit_codes(fake_server_binary, tmp_path, monkeypatch):
# The contract setup.sh / setup.ps1 depend on, exercised end to end through
# the real --smoke-test CLI: CPU-only -> 2, GPU -> 0.
probe = tmp_path / "probe.gguf"
probe.write_bytes(b"GGUF\x00fake")
monkeypatch.setattr(M, "detect_host", lambda: linux_cuda_host())
def run(mode):
monkeypatch.setenv("FAKE_LLAMA_MODE", mode)
monkeypatch.setattr(
sys,
"argv",
[
"install_llama_prebuilt.py",
"--smoke-test", str(fake_server_binary),
"--probe", str(probe),
"--install-kind", "linux-cuda",
],
)
return M.main()
assert run("cpu") == M.EXIT_FALLBACK
assert run("cuda") == M.EXIT_SUCCESS