diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 705ba708d2..f7d4b1fa21 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -6704,6 +6704,44 @@ def validate_prebuilt_attempts( raise PrebuiltFallback("no prebuilt bundle passed validation") +def existing_gpu_install_offloads( + install_dir: Path, + host: HostInfo, + plan: InstallReleasePlan, + probe_path: Path, +) -> bool: + """For a GPU plan whose metadata matches the existing install, smoke-test + the already-installed llama-server's offload before reusing it. Returns True + to keep the existing install (non-GPU kind, passed, or inconclusive) and + False on a definite CPU-only load so the caller reinstalls. This closes the + "reinstall/restart keeps the silently CPU-only binary" path (#5807): without + it, a metadata match short-circuits the new offload validation entirely.""" + choice = plan.attempts[0] + if choice.install_kind not in _GPU_INSTALL_KINDS: + return True + server_name = "llama-server.exe" if host.is_windows else "llama-server" + try: + server_path = discover_installed_executable(install_dir, server_name) + except PrebuiltFallback: + return True + try: + validate_server( + server_path, + probe_path, + host, + install_dir, + runtime_line = choice.runtime_line, + install_kind = choice.install_kind, + ) + except GpuOffloadFailure: + return False + except PrebuiltFallback: + # Inconclusive (failed to start, etc.): keep the existing install rather + # than reinstall on uncertain evidence. + return True + return True + + def install_prebuilt( install_dir: Path, llama_tag: str, @@ -6747,8 +6785,12 @@ def install_prebuilt( published_repo, published_release_tag, ) - if release_plans and existing_install_matches_plan( - install_dir, host, release_plans[0] + # Non-GPU match: keep the fast path (no probe download). A GPU match + # must still be smoke-tested below, so it does not short-circuit here. + if ( + release_plans + and existing_install_matches_plan(install_dir, host, release_plans[0]) + and release_plans[0].attempts[0].install_kind not in _GPU_INSTALL_KINDS ): current = release_plans[0] log( @@ -6765,9 +6807,27 @@ def install_prebuilt( release_count = len(release_plans) for release_index, plan in enumerate(release_plans): choice = plan.attempts[0] - if existing_install_matches_plan(install_dir, host, plan): + # A metadata match used to skip straight to reuse; now a + # matching GPU install is smoke-tested first so a previously + # installed CPU-only "GPU" binary is rebuilt instead of kept + # forever across reruns/restarts (#5807). + matched = existing_install_matches_plan(install_dir, host, plan) + existing_install_dir: Path | None = install_dir + if matched and not existing_gpu_install_offloads( + install_dir, host, plan, probe_path + ): log( - "existing llama.cpp install already matches fallback release " + "existing GPU llama.cpp install served a completion but " + "loaded the model on CPU; reinstalling to recover GPU " + "offload (unslothai/unsloth#5807)" + ) + matched = False + # Don't let validate_prebuilt_attempts short-circuit on + # the same bad install we just rejected. + existing_install_dir = None + if matched: + log( + "existing llama.cpp install already matches release " f"{plan.release_tag} upstream_tag={plan.llama_tag}; skipping reinstall" ) return @@ -6788,7 +6848,7 @@ def install_prebuilt( release_tag = plan.release_tag, approved_checksums = plan.approved_checksums, initial_fallback_used = release_index > 0, - existing_install_dir = install_dir, + existing_install_dir = existing_install_dir, ) except ExistingInstallSatisfied: return diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 3a4e9d4ca3..da4729f135 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2633,6 +2633,16 @@ if (Test-Path -LiteralPath $LlamaServerBin) { Write-Host " Existing llama-server was built with CUDA but no GPU detected -- rebuilding" -ForegroundColor Yellow $NeedRebuild = $true } + # The cache check catches a CPU build on a GPU host; this catches a CUDA + # build that still loads on CPU at runtime (PTX / runtime-init failure, + # #5807). Smoke-test before reusing it on a GPU host. + if (-not $NeedRebuild -and $HasNvidiaSmi -and $cachedCuda) { + & python "$PSScriptRoot\install_llama_prebuilt.py" --smoke-test "$LlamaServerBin" --install-dir "$LlamaCppDir" --install-kind "windows-cuda" 2>&1 | Out-String | Write-Host + if ($LASTEXITCODE -eq 2) { + Write-Host " Existing CUDA llama-server runs on CPU only -- rebuilding" -ForegroundColor Yellow + $NeedRebuild = $true + } + } } } diff --git a/studio/setup.sh b/studio/setup.sh index 10aa831b72..a16488c0b8 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -919,12 +919,41 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && \ [ -z "$_LLAMA_PR" ] && \ [ -x "$LLAMA_CPP_DIR/build/bin/llama-server" ] && \ [ -x "$LLAMA_CPP_DIR/build/bin/llama-quantize" ]; then - step "llama.cpp" "existing source build found; skipping rebuild" - ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize" - if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then - : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true + _REUSE_SOURCE=true + # On a GPU host, smoke-test the existing source binary first so a stale + # CPU-only build (e.g. an earlier no-toolkit fallback) is rebuilt instead of + # silently reused on a GPU host (#5807). Exit 2 = ran on CPU -> rebuild; + # anything else keeps the build (never downgrade on uncertain evidence). + _REUSE_KIND="" + if [ "$_HOST_SYSTEM" = "Darwin" ] && { [ "$_HOST_MACHINE" = "arm64" ] || [ "$_HOST_MACHINE" = "aarch64" ]; }; then + _REUSE_KIND="macos-arm64" + elif command -v nvidia-smi >/dev/null 2>&1; then + _REUSE_KIND="linux-cuda" + elif [ "$_LINUX_HAS_GPU" = true ]; then + _REUSE_KIND="linux-rocm" + fi + if [ -n "$_REUSE_KIND" ]; then + if python "$SCRIPT_DIR/install_llama_prebuilt.py" \ + --smoke-test "$LLAMA_CPP_DIR/build/bin/llama-server" \ + --install-dir "$LLAMA_CPP_DIR" \ + --install-kind "$_REUSE_KIND" > "$LLAMA_CPP_DIR/gpu-smoke-existing.log" 2>&1; then + _REUSE_RC=0 + else + _REUSE_RC=$? + fi + if [ "$_REUSE_RC" -eq 2 ]; then + substep "existing source build runs on CPU only; rebuilding for GPU..." "$C_WARN" + _REUSE_SOURCE=false + fi + fi + if [ "$_REUSE_SOURCE" = true ]; then + step "llama.cpp" "existing source build found; skipping rebuild" + ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize" + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true + fi + _NEED_LLAMA_SOURCE_BUILD=false fi - _NEED_LLAMA_SOURCE_BUILD=false fi # ── 8. WSL: pre-install GGUF build dependencies for fallback source builds ── diff --git a/tests/studio/install/test_validate_server_gpu_offload.py b/tests/studio/install/test_validate_server_gpu_offload.py index 6b0afc9de0..cc64d2560d 100644 --- a/tests/studio/install/test_validate_server_gpu_offload.py +++ b/tests/studio/install/test_validate_server_gpu_offload.py @@ -499,3 +499,84 @@ def test_main_smoke_exit_error_on_inconclusive(monkeypatch): rc = _run_main_smoke(monkeypatch, impl) assert rc == M.EXIT_ERROR # 1 -> setup scripts keep the GPU build + + +# -- 6. existing_gpu_install_offloads: re-validate a matching install (#5807) -- + + +def _gpu_plan(install_kind = "linux-cuda"): + choice = M.AssetChoice( + repo = "unslothai/llama.cpp", + tag = "b9001", + name = "app-b9001-linux-x64-cuda13-newer.tar.gz", + url = "https://example.com/x", + source_label = "published", + install_kind = install_kind, + ) + return M.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "rel", + attempts = [choice], + approved_checksums = M.ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", release_tag = "rel", upstream_tag = "b9001", + artifacts = {}, + ), + ) + + +def _with_server(tmp_path): + server = tmp_path / "llama-server" + server.write_text("#!/bin/sh\n") + probe = tmp_path / "probe.gguf" + probe.write_bytes(b"GGUF") + return probe + + +def test_existing_cpu_kind_install_is_kept(tmp_path): + # A non-GPU existing install is never offload-gated. + probe = _with_server(tmp_path) + assert M.existing_gpu_install_offloads( + tmp_path, nvidia_host(), _gpu_plan("linux-cpu"), probe + ) is True + + +def test_existing_gpu_install_cpu_only_triggers_reinstall(monkeypatch, tmp_path): + probe = _with_server(tmp_path) + + def fake_validate(*a, **k): + raise M.GpuOffloadFailure("loaded the model entirely on CPU") + + monkeypatch.setattr(M, "validate_server", fake_validate) + assert M.existing_gpu_install_offloads( + tmp_path, nvidia_host(), _gpu_plan("linux-cuda"), probe + ) is False + + +def test_existing_gpu_install_offloading_is_kept(monkeypatch, tmp_path): + probe = _with_server(tmp_path) + monkeypatch.setattr(M, "validate_server", lambda *a, **k: None) + assert M.existing_gpu_install_offloads( + tmp_path, nvidia_host(), _gpu_plan("linux-cuda"), probe + ) is True + + +def test_existing_gpu_install_inconclusive_is_kept(monkeypatch, tmp_path): + probe = _with_server(tmp_path) + + def fake_validate(*a, **k): + raise PrebuiltFallback("llama-server exited during startup") + + monkeypatch.setattr(M, "validate_server", fake_validate) + assert M.existing_gpu_install_offloads( + tmp_path, nvidia_host(), _gpu_plan("linux-cuda"), probe + ) is True + + +def test_existing_gpu_install_no_binary_is_kept(tmp_path): + # No llama-server present -> let the normal flow reinstall, don't crash. + probe = tmp_path / "probe.gguf" + probe.write_bytes(b"GGUF") + assert M.existing_gpu_install_offloads( + tmp_path, nvidia_host(), _gpu_plan("linux-cuda"), probe + ) is True