diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 11cc4c798e..b30c0f5aea 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -13,6 +13,7 @@ never blocks on a missing marker / offline GitHub. from __future__ import annotations +import asyncio from typing import Optional from fastapi import APIRouter, Depends, Query @@ -48,6 +49,9 @@ class LlamaUpdateStatusResponse(BaseModel): published_repo: Optional[str] = None installed_at_utc: Optional[str] = None age_days: Optional[int] = None + source_build: bool = Field( + False, description = "True when there is no marker (source build) but a prebuilt is offered." + ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) @@ -65,11 +69,14 @@ async def llama_update_status( ), current_subject: str = Depends(get_current_subject), ) -> LlamaUpdateStatusResponse: - return LlamaUpdateStatusResponse(**get_update_status(force_refresh = force_refresh)) + # Off the event loop: detection may probe the host and read GitHub. + status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh) + return LlamaUpdateStatusResponse(**status) @router.post("/update", response_model = LlamaUpdateActionResponse) async def llama_update( current_subject: str = Depends(get_current_subject), ) -> LlamaUpdateActionResponse: - return LlamaUpdateActionResponse(**start_update()) + action = await asyncio.to_thread(start_update) + return LlamaUpdateActionResponse(**action) diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py new file mode 100644 index 0000000000..ede5629664 --- /dev/null +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. + +These back the in-app update for source-build (markerless) installs: the backend +asks the installer whether an official prebuilt exists for this host without +downloading. Network and host detection are stubbed; no GPU or internet needed. +""" + +from __future__ import annotations + +import importlib +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_studio = Path(__file__).resolve().parent.parent.parent +if str(_studio) not in sys.path: + sys.path.insert(0, str(_studio)) + +ilp = importlib.import_module("install_llama_prebuilt") + +if not hasattr(ilp, "published_repo_for_host") or not hasattr( + ilp, "resolve_simple_install_release_plans" +): + pytest.skip("PR symbols not present - check branch", allow_module_level = True) + +FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp +UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp + + +def _host(**kw): + base = dict( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = False, + is_macos = False, + is_x86_64 = False, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = False, + rocm_gfx_target = None, + macos_version = None, + ) + base.update(kw) + return ilp.HostInfo(**base) + + +def test_published_repo_for_host(): + # CPU-only Linux (x64 and arm64) -> ggml-org upstream. + assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) + == UPSTREAM + ) + # GPU Linux -> fork. + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) + == FORK + ) + assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK + # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). + assert ( + ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) + == UPSTREAM + ) + # GPU Windows -> fork. + assert ( + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) + ) + == FORK + ) + # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). + assert ( + ilp.published_repo_for_host( + _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") + ) + == FORK + ) + # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). + assert ( + ilp.published_repo_for_host( + _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True + ) + == FORK + ) + # The tooling hint is Linux-only: Windows CPU stays on ggml-org. + assert ( + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True), + linux_amd_tooling_present = True, + ) + == UPSTREAM + ) + + +def _run_resolve(monkeypatch, capsys, plans_or_exc): + monkeypatch.setattr( + ilp, + "detect_host", + lambda: _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"), + ) + + def _resolver(tag, host, repo, published_release_tag): + if isinstance(plans_or_exc, Exception): + raise plans_or_exc + return ("b9585", plans_or_exc) + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + ) + rc = ilp.main() + assert rc == ilp.EXIT_SUCCESS + return json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + + +def test_resolve_prebuilt_available(monkeypatch, capsys): + plan = SimpleNamespace( + release_tag = "b9585", + llama_tag = "b9585", + attempts = [ + SimpleNamespace(name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64") + ], + ) + out = _run_resolve(monkeypatch, capsys, [plan]) + assert out["prebuilt_available"] is True + assert out["repo"] == FORK + assert out["release_tag"] == "b9585" + assert out["asset"] == "llama-b9585-bin-macos-arm64.tar.gz" + assert out["install_kind"] == "macos-arm64" + + +def test_resolve_prebuilt_unavailable(monkeypatch, capsys): + out = _run_resolve(monkeypatch, capsys, ilp.PrebuiltFallback("no macOS asset")) + assert out["prebuilt_available"] is False + assert out["repo"] == FORK + + +def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): + # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the + # fork so a HIP source build is not offered an upstream CPU prebuilt. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") + seen = {} + + def _resolver(tag, host, repo, published_release_tag): + seen["repo"] = repo + raise ilp.PrebuiltFallback("no asset") + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + assert seen["repo"] == FORK + assert out["repo"] == FORK diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index d2c854cdee..05a107377c 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -60,24 +60,151 @@ def _write_install( def _clean_state(monkeypatch): freshness.reset_caches() upd._reset_job_for_tests() + upd._resolve_memo.clear() + # Deterministic markerless paths: no host-pinned binary, no custom dir. + monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) # Never hit the network in these tests. monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) yield freshness.reset_caches() upd._reset_job_for_tests() + upd._resolve_memo.clear() -def test_status_no_marker(monkeypatch, tmp_path): +def _no_prebuilt(monkeypatch): + """Stub the host prebuilt probe to 'none available' (no source-build offer).""" + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: None) + + +def _prebuilt( + monkeypatch, + *, + repo = "unslothai/llama.cpp", + release_tag = "b9585", + llama_tag = None, + asset = None, +): + """Stub the host prebuilt probe to report an available prebuilt.""" + payload = { + "prebuilt_available": True, + "repo": repo, + "release_tag": release_tag, + "llama_tag": llama_tag or release_tag, + "asset": asset or f"llama-{release_tag}-bin-macos-arm64.tar.gz", + "install_kind": "macos-arm64", + } + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: payload) + + +def test_status_no_marker_no_prebuilt(monkeypatch, tmp_path): + # No marker AND no prebuilt available for the host -> unsupported (the genuine + # source-build-with-nothing-to-offer case). binary = tmp_path / "build" / "bin" / "llama-server" binary.parent.mkdir(parents = True) binary.write_text("stub") # no marker file alongside monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _no_prebuilt(monkeypatch) st = upd.get_update_status() assert st["supported"] is False assert st["update_available"] is False assert st["installed_tag"] is None +def test_status_source_build_offers_prebuilt(monkeypatch, tmp_path): + # Markerless source build with a prebuilt now available for the host: surface + # the update. Unknown installed version (source build) is treated as behind. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "b9585") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: None) + st = upd.get_update_status() + assert st["supported"] is True + assert st["update_available"] is True + assert st["source_build"] is True + assert st["latest_tag"] == "b9585" + assert st["published_repo"] == "unslothai/llama.cpp" + + +def test_status_source_build_compares_llama_tag(monkeypatch, tmp_path): + # release_tag may be a fork wrapper (v1.0); compare/display the upstream + # llama_tag (b9457) so a source build is not wrongly judged newer. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "v1.0", llama_tag = "b9457") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9000) + st = upd.get_update_status() + assert st["latest_tag"] == "b9457" # not the wrapper tag + assert st["update_available"] is True # 9000 < 9457 + + +def test_status_source_build_pinned_binary_not_offered(monkeypatch, tmp_path): + # LLAMA_SERVER_PATH pins a custom binary outside any llama.cpp dir; an apply + # could not take effect, so the button must not surface. + binary = tmp_path / "custom" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch) + st = upd.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + + +def test_llama_install_root_pinned_returns_none(monkeypatch, tmp_path): + binary = tmp_path / "custom" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + assert upd._llama_install_root(str(binary)) is None + + +def test_status_source_build_suppressed_when_newer(monkeypatch, tmp_path): + # A source build already newer than the latest prebuilt is not nagged. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "b9518") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) + st = upd.get_update_status() + assert st["supported"] is True + assert st["update_available"] is False + assert st["installed_tag"] == "b9600" + + +def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path): + # While the updater swaps the tree, status polls must not exec the binary + # being replaced (on Windows that exec can fail the installer's os.replace); + # the 3s poller only consumes job progress. + binary = tmp_path / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + probes = {"resolve": 0, "version": 0} + + def _count_resolve(*, force_refresh = False): + probes["resolve"] += 1 + return None + + def _count_version(b): + probes["version"] += 1 + return None + + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", _count_resolve) + monkeypatch.setattr(upd, "_installed_build_number", _count_version) + with upd._job_lock: + upd._job["state"] = upd._JOB_RUNNING + st = upd.get_update_status() + assert st["job"]["state"] == "running" + assert probes == {"resolve": 0, "version": 0} + + def test_status_update_available(monkeypatch, tmp_path): binary = _write_install(tmp_path, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) @@ -99,13 +226,62 @@ def test_status_up_to_date(monkeypatch, tmp_path): assert st["update_available"] is False -def test_start_update_no_marker_refuses(monkeypatch, tmp_path): +def test_start_update_no_marker_no_prebuilt_refuses(monkeypatch, tmp_path): binary = tmp_path / "llama-server" binary.write_text("stub") # no marker monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _no_prebuilt(monkeypatch) res = upd.start_update() assert res["started"] is False - assert res["reason"] == "no_prebuilt_marker" + assert res["reason"] == "no_prebuilt_available" + + +def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path): + # Markerless install + available prebuilt: install in place into the resolved + # root, with the asset-derived ROCm forwarding and the resolved repo. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _prebuilt( + monkeypatch, repo = "unslothai/llama.cpp", asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz" + ) + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + cmd = list(cmd) + # Status polls probe `llama-server --version`; keep the installer argv. + if "--version" in cmd: + return _Proc() + captured["cmd"] = cmd + _write_install(install_dir, "b9585") # installer writes the marker + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + + res = upd.start_update() + assert res["started"] is True, res + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + cmd = captured["cmd"] + assert "--install-dir" in cmd and str(install_dir) in cmd + assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd + assert "--llama-tag" in cmd and "latest" in cmd + assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" + assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd def test_start_update_happy_path(monkeypatch, tmp_path): @@ -123,6 +299,10 @@ def test_start_update_happy_path(monkeypatch, tmp_path): stderr = "" def _fake_run(cmd, **kwargs): + cmd = list(cmd) + # Status polls probe `llama-server --version`; keep the installer argv. + if "--version" in cmd: + return _Proc() captured["cmd"] = cmd # Simulate the installer writing a new marker with the latest tag. _write_install(install_dir, "b9518") @@ -229,7 +409,11 @@ def _capture_install_cmd( stderr = "" def _fake_run(cmd, **kwargs): - captured["cmd"] = list(cmd) + cmd = list(cmd) + # Status polls probe `llama-server --version`; keep the installer argv. + if "--version" in cmd: + return _Proc() + captured["cmd"] = cmd _write_install(install_dir, latest, repo = repo, asset = asset) return _Proc() @@ -442,3 +626,136 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): break time.sleep(0.05) assert job["state"] == "success", job + + +# --- markerless helper units --- + + +def test_resolve_prebuilt_parses_and_caches(monkeypatch, tmp_path): + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + calls = {"n": 0} + + class _Proc: + returncode = 0 + # stderr noise plus the JSON line on stdout (installer logs to stderr). + stdout = ( + '{"prebuilt_available": true, "repo": "unslothai/llama.cpp", "release_tag": "b9585"}' + ) + stderr = "[llama-prebuilt] some log\n" + + def _fake_run(cmd, **kwargs): + calls["n"] += 1 + assert "--resolve-prebuilt" in cmd + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + res = upd._resolve_prebuilt_for_host() + assert res["prebuilt_available"] is True and res["release_tag"] == "b9585" + # Second call is memoized (no second subprocess). + upd._resolve_prebuilt_for_host() + assert calls["n"] == 1 + + +def test_resolve_prebuilt_fails_open(monkeypatch, tmp_path): + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + + def _boom(cmd, **kwargs): + raise OSError("subprocess failed") + + monkeypatch.setattr(upd.subprocess, "run", _boom) + assert upd._resolve_prebuilt_for_host() is None + # Failures are not cached: a later success is observed. + + class _Proc: + returncode = 0 + stdout = '{"prebuilt_available": false}' + stderr = "" + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + assert upd._resolve_prebuilt_for_host() == {"prebuilt_available": False} + + +def test_installed_build_number(monkeypatch): + def _ver(text): + class _Proc: + returncode = 0 + stdout = "" + stderr = text + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + return upd._installed_build_number("/bin/llama-server") + + assert _ver("version: 9585 (abc1234)\nbuilt with clang\n") == 9585 + assert _ver("version: 1 (deadbee)\n") is None # source build without tags + assert _ver("no version here") is None + assert upd._installed_build_number(None) is None + + +def test_llama_install_root_finds_llama_cpp_ancestor(monkeypatch, tmp_path): + root = tmp_path / "llama.cpp" + binary = root / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) == root + + +def test_llama_install_root_unmanaged_path_returns_none(monkeypatch, tmp_path): + # A binary on PATH (no marker, no env pin, no llama.cpp ancestor) is foreign: + # installing elsewhere would not replace it, so report no manageable root. + binary = tmp_path / "usr" / "local" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) is None + + +def test_llama_install_root_unsloth_env_dir(monkeypatch, tmp_path): + # UNSLOTH_LLAMA_CPP_PATH dir holding the active binary is the managed root. + root = tmp_path / "vendor" / "llama" + binary = root / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(root)) + assert upd._llama_install_root(str(binary)) == root + + +def test_llama_install_root_ignores_inactive_env_root(monkeypatch, tmp_path): + # UNSLOTH_LLAMA_CPP_PATH set but the active binary is not under it: do not + # target the stale env root, resolve from the binary's own llama.cpp tree. + inactive = tmp_path / "custom-empty" + inactive.mkdir() + active = tmp_path / "llama.cpp" + binary = active / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(inactive)) + assert upd._llama_install_root(str(binary)) == active + + +def test_llama_install_root_refuses_pinned_checkout_under_llama_cpp(monkeypatch, tmp_path): + # The LLAMA_SERVER_PATH pin guard must run before the ancestor scan, or a + # user's own llama.cpp checkout could be handed to the installer. + root = tmp_path / "my-project" / "llama.cpp" + binary = root / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) is None + + +def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path): + # A direct POST on a source build already newer than the prebuilt must not + # downgrade it; start_update mirrors the detection suppression. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _prebuilt(monkeypatch, release_tag = "b9518") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "up_to_date" diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py new file mode 100644 index 0000000000..bf0c4b731f --- /dev/null +++ b/studio/backend/tests/test_llama_route.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""routes/llama.py: the source_build field is exposed and the handlers run the +(now subprocess-touching) detection off the event loop via a worker thread. + +The route file is loaded standalone with a stubbed auth dependency so the test +does not pull the whole routes package (matplotlib-heavy training router) and +works in a minimal env. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import threading +import types +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +pytest.importorskip("fastapi") + + +def _load_route(): + # Prefer the real auth module; stub it only in minimal envs where its + # deps are absent. Stubs are popped after the load so they never leak + # into sys.modules for the rest of the suite. + stubbed = [] + try: + import auth.authentication # noqa: F401 + except Exception: + auth_pkg = types.ModuleType("auth") + auth_pkg.__path__ = [] + auth_mod = types.ModuleType("auth.authentication") + auth_mod.get_current_subject = lambda: "test" + for name, stub in (("auth", auth_pkg), ("auth.authentication", auth_mod)): + if name not in sys.modules: + sys.modules[name] = stub + stubbed.append(name) + try: + spec = importlib.util.spec_from_file_location( + "llama_route_under_test", str(_BACKEND / "routes" / "llama.py") + ) + mod = importlib.util.module_from_spec(spec) + sys.modules["llama_route_under_test"] = mod # so pydantic resolves forward refs + spec.loader.exec_module(mod) + return mod + finally: + for name in stubbed: + sys.modules.pop(name, None) + + +rl = _load_route() + + +def test_status_response_exposes_source_build(): + payload = { + "supported": True, + "update_available": True, + "stale": False, + "installed_tag": None, + "latest_tag": "b9585", + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": None, + "age_days": None, + "source_build": True, + "job": {"state": "idle"}, + } + model = rl.LlamaUpdateStatusResponse(**payload) + assert model.model_dump()["source_build"] is True + # Extra/unknown keys must not crash the response model. + rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1}) + + +def test_status_handler_runs_off_event_loop(monkeypatch): + seen = {} + + def fake_status(force_refresh = False): + seen["thread"] = threading.current_thread() + return { + "supported": True, + "update_available": True, + "source_build": True, + "latest_tag": "b9585", + "job": {"state": "idle"}, + } + + monkeypatch.setattr(rl, "get_update_status", fake_status) + out = asyncio.run(rl.llama_update_status(force_refresh = False, current_subject = "t")) + assert out.source_build is True + # Detection ran in a worker thread, not the event-loop thread. + assert seen["thread"] is not threading.main_thread() + + +def test_update_handler_runs_off_event_loop(monkeypatch): + seen = {} + + def fake_start(): + seen["thread"] = threading.current_thread() + return {"started": True, "reason": None, "job": {"state": "running"}} + + monkeypatch.setattr(rl, "start_update", fake_start) + out = asyncio.run(rl.llama_update(current_subject = "t")) + assert out.started is True + assert seen["thread"] is not threading.main_thread() diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 19946b7966..b138b29af3 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -21,6 +21,7 @@ Design notes: from __future__ import annotations +import json import os import re import subprocess @@ -108,6 +109,143 @@ def _installer_script() -> Optional[Path]: return None +# Markerless (source-build) installs have no UNSLOTH_PREBUILT_INFO.json, so we +# ask the installer whether an official prebuilt now exists for this host. Memo +# is 24h; only successful answers are cached so a network blip retries. +_RESOLVE_TTL_SECONDS = 24 * 60 * 60 +_resolve_memo: dict = {} + + +def _resolve_prebuilt_for_host(*, force_refresh: bool = False) -> Optional[dict]: + """Run install_llama_prebuilt.py --resolve-prebuilt (no download) and return + {prebuilt_available, repo, release_tag, llama_tag, asset, install_kind} or + None. Fail-open: any error -> None so a source build never blocks the app.""" + now = time.time() + if not force_refresh and _resolve_memo: + if now - _resolve_memo.get("at", 0.0) < _RESOLVE_TTL_SECONDS: + return _resolve_memo.get("value") + script = _installer_script() + if script is None: + return None + value: Optional[dict] = None + try: + proc = subprocess.run( + [ + sys.executable, + str(script), + "--resolve-prebuilt", + "latest", + "--output-format", + "json", + ], + capture_output = True, + text = True, + timeout = 60, + ) + out = (proc.stdout or "").strip() + if proc.returncode == 0 and out: + parsed = json.loads(out.splitlines()[-1]) + if isinstance(parsed, dict): + value = parsed + except Exception as exc: # pragma: no cover - subprocess/json defensive + logger.debug("llama update: resolve-prebuilt failed", error = str(exc)) + value = None + if value is not None: # cache real answers; let failures retry next poll + _resolve_memo.update(at = now, value = value) + return value + + +def _installed_build_number(binary: Optional[str]) -> Optional[int]: + """Best-effort build number from ``llama-server --version`` (e.g. + 'version: 9585 (abc)'). None when unparseable or <= 1: a source build with + no git tags reports 'version: 1', which we treat as unknown (offer update).""" + if not binary: + return None + try: + proc = subprocess.run([binary, "--version"], capture_output = True, text = True, timeout = 20) + except Exception: # pragma: no cover - defensive + return None + m = re.search(r"version:\s*(\d+)", (proc.stderr or "") + (proc.stdout or "")) + if not m: + return None + n = int(m.group(1)) + return n if n > 1 else None + + +def _is_under(path: Path, root: Path) -> bool: + try: + p, r = path.resolve(), root.resolve() + except (OSError, ValueError): + p, r = path, root + return p == r or r in p.parents + + +def _llama_install_root(binary: Optional[str]) -> Optional[Path]: + """The Studio-managed llama.cpp root the active binary lives under, or None + when the binary is unmanaged. Installing anywhere the active binary is not + would not replace what _find_llama_server_binary runs (which prefers a pinned + LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we + refuse rather than silently install into an inactive or foreign tree.""" + marked = _install_dir_for(binary) + if marked is not None: + return marked + if not binary: + return None + # LLAMA_SERVER_PATH is an explicit user pin that always wins in discovery; + # never auto-replace its tree (even a user's own llama.cpp checkout). + if os.environ.get("LLAMA_SERVER_PATH"): + return None + p = Path(binary) + env = os.environ.get("UNSLOTH_LLAMA_CPP_PATH") + if env and _is_under(p, Path(env)): + return Path(env) + for parent in p.parents: + if parent.name == "llama.cpp": + return parent + # PATH / system / custom install: not a managed tree, so do not offer. + return None + + +def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: + """Update status for a markerless (source-build) install: offer the official + prebuilt when one exists for this host and is newer than the installed + binary. None -> caller falls through to the no-marker default (unsupported).""" + res = _resolve_prebuilt_for_host(force_refresh = force_refresh) + if not res or not res.get("prebuilt_available"): + return None + # llama_tag is the upstream build (bNNNN, what --version reports); release_tag + # can be a fork wrapper tag, so compare/display against llama_tag. + latest = res.get("llama_tag") or res.get("release_tag") + if not latest: + return None + # No resolvable install root (e.g. a pinned LLAMA_SERVER_PATH we cannot + # manage) means an apply would not take effect, so do not offer. + if _llama_install_root(binary) is None: + return None + installed_build = _installed_build_number(binary) + m = re.search(r"(\d+)", latest) + latest_build = int(m.group(1)) if m else None + # Suppress only when the source build is reliably newer/equal; unknown + # version (the involuntary source-build case) is treated as behind. + update_available = ( + installed_build is None or latest_build is None or installed_build < latest_build + ) + with _job_lock: + job = dict(_job) + return { + "supported": True, + "update_available": update_available, + "stale": False, + "installed_tag": (f"b{installed_build}" if installed_build else None), + "latest_tag": latest, + "published_repo": res.get("repo"), + "installed_at_utc": None, + "age_days": None, + "source_build": True, + "job": job, + } + + def get_update_status(*, force_refresh: bool = False) -> dict: """Report whether a newer prebuilt exists plus the current job state. @@ -115,6 +253,20 @@ def get_update_status(*, force_refresh: bool = False) -> dict: """ binary = _find_binary() marker = read_install_marker(binary) + + with _job_lock: + job_running = _job["state"] == _JOB_RUNNING + + # No marker = source build / custom path. Offer the official prebuilt if one + # now exists for this host (this is why macOS source builds showed no button). + # Skipped while the updater swaps the tree: each 3s poll would exec the + # half-replaced binary (on Windows that exec can make the installer's + # os.replace fail) and the poller only consumes job progress. + if marker is None and binary is not None and not job_running: + src = _source_build_status(binary, force_refresh = force_refresh) + if src is not None: + return src + repo = (marker or {}).get("published_repo") or DEFAULT_PUBLISHED_REPO if force_refresh and repo: @@ -143,6 +295,7 @@ def get_update_status(*, force_refresh: bool = False) -> dict: "published_repo": freshness.get("published_repo") or repo, "installed_at_utc": freshness.get("installed_at_utc"), "age_days": freshness.get("age_days"), + "source_build": False, "job": job, } @@ -254,19 +407,7 @@ def start_update() -> dict: """Kick off a background update. Idempotent: a second call while one is running returns the in-flight job rather than starting another.""" binary = _find_binary() - install_dir = _install_dir_for(binary) marker = read_install_marker(binary) - if install_dir is None or not marker: - return { - "started": False, - "reason": "no_prebuilt_marker", - "message": ( - "This llama.cpp install was not provisioned from an Unsloth " - "prebuilt (source build or custom path); in-app update is " - "unavailable." - ), - "job": get_update_status()["job"], - } script = _installer_script() if script is None: return { @@ -275,9 +416,47 @@ def start_update() -> dict: "message": "install_llama_prebuilt.py could not be located.", "job": get_update_status()["job"], } - repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO - from_tag = marker.get("tag") or marker.get("release_tag") - asset = marker.get("asset") + + if marker: + install_dir = _install_dir_for(binary) + repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO + from_tag = marker.get("tag") or marker.get("release_tag") + asset = marker.get("asset") + else: + # Source build / custom path: only proceed when the same detection logic + # would offer the update (prebuilt exists, install is behind, root is + # manageable), so a direct POST cannot downgrade a newer source build. + src = _source_build_status(binary, force_refresh = True) if binary else None + if src is None: + return { + "started": False, + "reason": "no_prebuilt_available", + "message": ( + "No official llama.cpp prebuilt is available for this host, " + "so the source build cannot be swapped automatically." + ), + "job": get_update_status()["job"], + } + if not src.get("update_available"): + return { + "started": False, + "reason": "up_to_date", + "message": "The installed llama.cpp build is already at or newer than the latest prebuilt.", + "job": get_update_status()["job"], + } + res = _resolve_prebuilt_for_host() + install_dir = _llama_install_root(binary) + repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO + from_tag = None + asset = (res or {}).get("asset") + + if install_dir is None: + return { + "started": False, + "reason": "no_install_dir", + "message": "Could not determine the llama.cpp install directory.", + "job": get_update_status()["job"], + } with _job_lock: if _job["state"] == _JOB_RUNNING: diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 822a32ddb8..f9425f5c4f 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -3045,6 +3045,21 @@ def _apply_host_overrides( return host +def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str: + """The release repo setup.sh / setup.ps1 pick for this host: macOS always the + fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows + -> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA + or ROCm) -> the fork. linux_amd_tooling_present mirrors setup.sh routing Linux + hosts that expose AMD tooling (rocminfo/amd-smi/hipconfig/hipinfo) to the fork + even when the probe cannot confirm an active GPU. Mirrors the shell routing.""" + if host.is_macos: + return DEFAULT_PUBLISHED_REPO + has_gpu = ( + host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present) + ) + return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO + + def pick_windows_cuda_runtime(host: HostInfo) -> str | None: if not host.driver_cuda_version: return None @@ -6690,6 +6705,16 @@ def parse_args() -> argparse.Namespace: const = "latest", help = ("Resolve the source-build fallback plan."), ) + resolve_group.add_argument( + "--resolve-prebuilt", + nargs = "?", + const = "latest", + help = ( + "Report whether an official prebuilt exists for this host without " + "downloading. Picks the host's published repo when --published-repo " + "is left at the default. Use --output-format json." + ), + ) parser.add_argument( "--output-format", choices = ("plain", "json"), @@ -6774,6 +6799,46 @@ def main() -> int: ) return EXIT_SUCCESS + if args.resolve_prebuilt is not None: + # Host-aware "is a prebuilt available" probe, no download. A default repo + # means "pick the repo for this host"; PrebuiltFallback == source build. + host = _apply_host_overrides( + detect_host(), + override_has_rocm = args.has_rocm, + override_rocm_gfx = args.rocm_gfx, + force_cpu = args.cpu_fallback, + ) + # setup.sh routes Linux hosts with AMD tooling to the fork even when no GPU + # is probed; mirror that so a HIP source build is not offered a CPU prebuilt. + amd_tooling = host.is_linux and any( + shutil.which(t) for t in ("rocminfo", "amd-smi", "hipconfig", "hipinfo") + ) + repo = ( + published_repo_for_host(host, linux_amd_tooling_present = amd_tooling) + if args.published_repo == DEFAULT_PUBLISHED_REPO + else args.published_repo + ) + try: + _requested, plans = resolve_simple_install_release_plans( + args.resolve_prebuilt, host, repo, args.published_release_tag or "" + ) + choice = plans[0].attempts[0] if plans and plans[0].attempts else None + if choice is None: + payload = {"prebuilt_available": False, "repo": repo} + else: + payload = { + "prebuilt_available": True, + "repo": repo, + "release_tag": plans[0].release_tag, + "llama_tag": plans[0].llama_tag, + "asset": choice.name, + "install_kind": choice.install_kind, + } + except PrebuiltFallback: + payload = {"prebuilt_available": False, "repo": repo} + emit_resolver_output(payload, output_format = args.output_format) + return EXIT_SUCCESS + if not args.install_dir: raise SystemExit( "install_llama_prebuilt.py: --install-dir is required unless --resolve-llama-tag, --resolve-install-tag, or --resolve-source-build is used"