From dab0b77673abd28789fbbd5fba09dcf41402cea1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 10 Jun 2026 10:04:26 -0700 Subject: [PATCH] Studio: in-app Update llama.cpp button to install the latest prebuilt (#6097) Adds an in-app "Update llama.cpp" banner and button to Unsloth Studio. When the installed prebuilt is behind the latest published release, a non-invasive banner appears; clicking Update downloads the latest prebuilt for this host and swaps it in place in the background, with no restart. Detection reuses the freshness check from #5529. The update re-runs install_llama_prebuilt.py the same way setup.sh and setup.ps1 do after #5963: it forwards the published repo and the AMD gfx target derived from the install marker, and does not pass the removed --simple-policy or the arm64-only --cpu-fallback. While the installer swaps binaries the backend enters a maintenance state (flag set under the serial load lock, active server unloaded) so a concurrent load cannot start a server from a half-swapped binary; the next load uses the new build. The banner also handles refused responses and jobs started in another tab so it never sticks on "Updating...". Verified end to end on an NVIDIA B200: installed b9493, detected the update, applied it, and confirmed the binary at the same path advanced to b9585 in the same process. Hermetic backend tests and the frontend type-check pass. --- studio/backend/core/inference/llama_cpp.py | 7 + studio/backend/main.py | 2 + studio/backend/routes/llama.py | 75 +++ studio/backend/tests/test_desktop_auth.py | 4 + studio/backend/tests/test_llama_cpp_update.py | 444 ++++++++++++++++++ studio/backend/utils/llama_cpp_update.py | 317 +++++++++++++ studio/frontend/src/app/provider.tsx | 18 +- .../src/components/llama-update-banner.tsx | 109 +++++ .../src/hooks/use-llama-update-check.ts | 236 ++++++++++ 9 files changed, 1211 insertions(+), 1 deletion(-) create mode 100644 studio/backend/routes/llama.py create mode 100644 studio/backend/tests/test_llama_cpp_update.py create mode 100644 studio/backend/utils/llama_cpp_update.py create mode 100644 studio/frontend/src/components/llama-update-banner.tsx create mode 100644 studio/frontend/src/hooks/use-llama-update-check.ts diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3c582a695f..e608a4b552 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -719,6 +719,9 @@ class LlamaCppBackend: # Wraps load_model() end-to-end so concurrent loads serialise and never # coexist as two llama-server processes (#5401). self._serial_load_lock = threading.Lock() + # Set by the in-app updater while it swaps prebuilt binaries; load_model() + # rejects fast so no server starts from a half-swapped binary. + self._llama_update_in_progress = False # Last extra_args / requested n_ctx, preserved across unload so the chat # UI's /unload+/load Apply path can inherit them (#5401). # ``_extra_args_source`` records the (model_identifier, hf_variant) the @@ -2806,6 +2809,10 @@ class LlamaCppBackend: # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. with self._serial_load_lock: + # In-app update swapping binaries: refuse fast (set under this lock, + # so any in-flight load has drained) instead of using a half-swapped one. + if getattr(self, "_llama_update_in_progress", False): + raise RuntimeError("llama.cpp is updating; try again in a moment.") # Duplicate /load that raced past the route check: do nothing if the # live server already satisfies this request. if self._already_in_target_state( diff --git a/studio/backend/main.py b/studio/backend/main.py index 74a7dbc8c5..dc3d802b78 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -266,6 +266,7 @@ from routes import ( training_history_router, training_router, ) +from routes.llama import router as llama_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, @@ -804,6 +805,7 @@ app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"]) app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"]) app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) +app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"]) app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"]) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py new file mode 100644 index 0000000000..11cc4c798e --- /dev/null +++ b/studio/backend/routes/llama.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""llama.cpp prebuilt update endpoints. + +GET /api/llama/update-status -> is a newer prebuilt available + job state +POST /api/llama/update -> download + atomically swap to the latest + +Detection reuses utils.llama_cpp_freshness; the swap reuses +install_llama_prebuilt.py via utils.llama_cpp_update. Both fail open so the UI +never blocks on a missing marker / offline GitHub. +""" + +from __future__ import annotations + +from typing import Optional + +from fastapi import APIRouter, Depends, Query +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject +from utils.llama_cpp_update import get_update_status, start_update + +router = APIRouter() + + +class LlamaUpdateJob(BaseModel): + state: str = Field("idle", description = "idle | running | success | error") + message: str = "" + from_tag: Optional[str] = None + to_tag: Optional[str] = None + error: Optional[str] = None + started_at: Optional[str] = None + finished_at: Optional[str] = None + + +class LlamaUpdateStatusResponse(BaseModel): + supported: bool = Field( + False, + description = "True when the install came from an Unsloth prebuilt (has a marker).", + ) + update_available: bool = Field(False, description = "True when installed_tag != latest_tag.") + stale: bool = Field( + False, description = "Update available AND install older than the staleness threshold." + ) + installed_tag: Optional[str] = None + latest_tag: Optional[str] = None + published_repo: Optional[str] = None + installed_at_utc: Optional[str] = None + age_days: Optional[int] = None + job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) + + +class LlamaUpdateActionResponse(BaseModel): + started: bool + reason: Optional[str] = None + message: Optional[str] = None + job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) + + +@router.get("/update-status", response_model = LlamaUpdateStatusResponse) +async def llama_update_status( + force_refresh: bool = Query( + False, description = "Bypass the 24h release cache for an explicit check." + ), + current_subject: str = Depends(get_current_subject), +) -> LlamaUpdateStatusResponse: + return LlamaUpdateStatusResponse(**get_update_status(force_refresh = force_refresh)) + + +@router.post("/update", response_model = LlamaUpdateActionResponse) +async def llama_update( + current_subject: str = Depends(get_current_subject), +) -> LlamaUpdateActionResponse: + return LlamaUpdateActionResponse(**start_update()) diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index 50bd0e89d7..b4b39ba1a9 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -419,6 +419,8 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): routes_module.__path__ = [] settings_module = ModuleType("routes.settings") settings_module.router = APIRouter() + llama_module = ModuleType("routes.llama") + llama_module.router = APIRouter() prompts_module = ModuleType("routes.prompts") prompts_module.router = APIRouter() @@ -440,9 +442,11 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): }.items(): setattr(routes_module, name, router) routes_module.settings = settings_module + routes_module.llama = llama_module monkeypatch.setitem(sys.modules, "routes", routes_module) monkeypatch.setitem(sys.modules, "routes.settings", settings_module) + monkeypatch.setitem(sys.modules, "routes.llama", llama_module) monkeypatch.setitem(sys.modules, "routes.prompts", prompts_module) import studio.backend.main as backend_main diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py new file mode 100644 index 0000000000..d2c854cdee --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -0,0 +1,444 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic tests for the in-app llama.cpp update orchestration. + +No network, no real install: the GitHub release lookup and the installer +subprocess are both monkeypatched. Verifies detection (update_available) and +the apply flow (job lifecycle, installer invocation, post-swap re-read). +""" + +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path +from types import ModuleType + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import utils.llama_cpp_freshness as freshness # noqa: E402 +import utils.llama_cpp_update as upd # noqa: E402 + +MARKER = "UNSLOTH_PREBUILT_INFO.json" + + +def _write_install( + dir_: Path, + tag: str, + repo: str = "unslothai/llama.cpp", + asset: str | None = None, +) -> str: + """Create a fake prebuilt install tree and return the llama-server path. + + ``asset`` is the bundle filename recorded in the marker; omit it to model an + older marker that predates asset-based ROCm forwarding (backward compat).""" + bin_dir = dir_ / "build" / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + binary = bin_dir / "llama-server" + binary.write_text("#!/bin/sh\necho stub\n") + marker = { + "tag": tag, + "release_tag": tag, + "published_repo": repo, + "installed_at_utc": "2020-01-01T00:00:00Z", + "bundle_profile": "cuda13-newer", + "runtime_line": "cuda13", + } + if asset is not None: + marker["asset"] = asset + (dir_ / MARKER).write_text(json.dumps(marker)) + return str(binary) + + +@pytest.fixture(autouse = True) +def _clean_state(monkeypatch): + freshness.reset_caches() + upd._reset_job_for_tests() + # 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() + + +def test_status_no_marker(monkeypatch, tmp_path): + 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)) + 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_update_available(monkeypatch, tmp_path): + binary = _write_install(tmp_path, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + st = upd.get_update_status(force_refresh = True) + assert st["supported"] is True + assert st["installed_tag"] == "b9493" + assert st["latest_tag"] == "b9518" + assert st["update_available"] is True + + +def test_status_up_to_date(monkeypatch, tmp_path): + binary = _write_install(tmp_path, "b9518") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + st = upd.get_update_status(force_refresh = True) + assert st["installed_tag"] == "b9518" + assert st["latest_tag"] == "b9518" + assert st["update_available"] is False + + +def test_start_update_no_marker_refuses(monkeypatch, tmp_path): + binary = tmp_path / "llama-server" + binary.write_text("stub") # no marker + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "no_prebuilt_marker" + + +def test_start_update_happy_path(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + captured["cmd"] = cmd + # Simulate the installer writing a new marker with the latest tag. + _write_install(install_dir, "b9518") + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + + res = upd.start_update() + assert res["started"] is True + assert res["job"]["from_tag"] == "b9493" + + # Wait for the background worker. + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert job["to_tag"] == "b9518" + # Installer was invoked with the resolved install dir + latest + repo. + assert "--install-dir" in captured["cmd"] + assert str(install_dir) in captured["cmd"] + assert "--llama-tag" in captured["cmd"] and "latest" in captured["cmd"] + assert "unslothai/llama.cpp" in captured["cmd"] + + +def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + + class _Proc: + returncode = 2 + stdout = "" + stderr = "boom: network error" + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "error" + assert "boom" in (job["error"] or "") + + +# --- installer-argument construction (mirrors the post-#5963 setup scripts) --- + + +def test_rocm_install_args_lemonade_gfx(): + # Lemonade HIP app bundle: gfx family lives in the asset name. + assert upd._rocm_install_args("app-b9585-linux-x64-rocm-gfx110X.tar.gz") == [ + "--rocm-gfx", + "gfx110x", + ] + assert upd._rocm_install_args("app-b9585-windows-x64-rocm-gfx1150.zip") == [ + "--rocm-gfx", + "gfx1150", + ] + + +def test_rocm_install_args_fork_version_bundle(): + # Fork ROCm bundles encode a ROCm version, not a gfx -> forward --has-rocm. + assert upd._rocm_install_args("llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz") == ["--has-rocm"] + + +def test_rocm_install_args_windows_hip(): + assert upd._rocm_install_args("llama-b9334-bin-win-hip-radeon-x64.zip") == ["--has-rocm"] + + +def test_rocm_install_args_non_rocm_and_missing(): + assert upd._rocm_install_args("llama-b9334-bin-ubuntu-x64.tar.gz") == [] + assert upd._rocm_install_args("app-b9585-linux-x64-cuda13.tar.gz") == [] + assert upd._rocm_install_args(None) == [] + + +def _capture_install_cmd( + monkeypatch, + tmp_path, + *, + tag = "b9493", + repo = "unslothai/llama.cpp", + asset = None, + latest = "b9518", +) -> list: + """Run start_update() with the installer subprocess stubbed; return the argv.""" + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, tag, repo = repo, asset = asset) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest) + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + captured["cmd"] = list(cmd) + _write_install(install_dir, latest, repo = repo, asset = asset) + 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) + return captured.get("cmd", []) + + +def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path): + cmd = _capture_install_cmd( + monkeypatch, tmp_path, asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz" + ) + assert "--rocm-gfx" in cmd + assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" + assert "--has-rocm" not in cmd + assert "--cpu-fallback" not in cmd + assert "--simple-policy" not in cmd + assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd + + +def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): + cmd = _capture_install_cmd( + monkeypatch, tmp_path, asset = "llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz" + ) + assert "--has-rocm" in cmd + assert "--rocm-gfx" not in cmd + + +def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): + # CPU installs come from ggml-org. Re-running into the same install-dir/repo + # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU + # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + cmd = _capture_install_cmd( + monkeypatch, + tmp_path, + repo = "ggml-org/llama.cpp", + asset = "llama-b9334-bin-ubuntu-x64.tar.gz", + ) + assert "--cpu-fallback" not in cmd + assert "--rocm-gfx" not in cmd + assert "--has-rocm" not in cmd + assert "--simple-policy" not in cmd + assert "--published-repo" in cmd and "ggml-org/llama.cpp" in cmd + + +def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tmp_path): + # Marker without an asset field (older install): no ROCm flags, no crash, and + # never the obsolete --simple-policy that #5963 removed from setup. + cmd = _capture_install_cmd(monkeypatch, tmp_path, asset = None) + assert "--simple-policy" not in cmd + assert "--rocm-gfx" not in cmd + assert "--has-rocm" not in cmd + assert "--cpu-fallback" not in cmd + + +# --- refusal + maintenance-state coordination --- + + +def test_start_update_already_running_refuses(monkeypatch, tmp_path): + binary = _write_install(tmp_path / "llama.cpp", "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + with upd._job_lock: + upd._job.update(state = upd._JOB_RUNNING) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "already_running" + + +def test_start_update_installer_missing_refuses(monkeypatch, tmp_path): + binary = _write_install(tmp_path / "llama.cpp", "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: None) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "installer_missing" + + +class _FakeBackend: + """Minimal stand-in for LlamaCppBackend's update-coordination surface.""" + + def __init__(self): + import threading + + self._serial_load_lock = threading.Lock() + self._llama_update_in_progress = False + self.is_active = True + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + +def _inject_backend(monkeypatch, backend): + routes_pkg = ModuleType("routes") + routes_pkg.__path__ = [] + inference_mod = ModuleType("routes.inference") + inference_mod.get_llama_cpp_backend = lambda: backend + monkeypatch.setitem(sys.modules, "routes", routes_pkg) + monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) + + +def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + backend = _FakeBackend() + _inject_backend(monkeypatch, backend) + + seen = {} + + class _Proc: + returncode = 0 + stdout = "ok" + stderr = "" + + def _fake_run(cmd, **kwargs): + # The maintenance flag must be set while the installer runs. + seen["flag_during_install"] = backend._llama_update_in_progress + _write_install(install_dir, "b9518") + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + + assert backend.unloaded is True + assert seen.get("flag_during_install") is True + # Cleared in the finally so model loads work again after the swap. + assert backend._llama_update_in_progress is False + + +def test_update_clears_maintenance_flag_on_installer_failure(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + + backend = _FakeBackend() + _inject_backend(monkeypatch, backend) + + class _Proc: + returncode = 1 + stdout = "" + stderr = "boom" + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + assert upd.get_update_status()["job"]["state"] == "error" + assert backend._llama_update_in_progress is False + + +def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _raise(): + raise RuntimeError("no backend") + + inference_mod = ModuleType("routes.inference") + inference_mod.get_llama_cpp_backend = lambda: _raise() + routes_pkg = ModuleType("routes") + routes_pkg.__path__ = [] + monkeypatch.setitem(sys.modules, "routes", routes_pkg) + monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) + + class _Proc: + returncode = 0 + stdout = "ok" + stderr = "" + + def _fake_run(cmd, **kwargs): + _write_install(install_dir, "b9518") + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py new file mode 100644 index 0000000000..19946b7966 --- /dev/null +++ b/studio/backend/utils/llama_cpp_update.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""In-app llama.cpp prebuilt update. + +Builds on utils.llama_cpp_freshness (which detects whether a newer prebuilt +release exists) and adds the *apply* half: run install_llama_prebuilt.py to +download the newest bundle for this host and atomically swap it in place, so +the next model load uses it. + +Design notes: +- Detection is delegated to check_prebuilt_freshness(). We surface an + ``update_available`` flag (installed_tag != latest_tag) which is laxer than + freshness' ``stale`` (which additionally requires the install to be >= 3 days + old). The UI shows the "Update llama.cpp" affordance on update_available. +- The install is slow (download + extract + validate), so it runs on a daemon + thread; callers poll get_update_status() for the job state. +- Everything fails open: a missing marker / offline GitHub / source build just + reports update_available=False and never blocks the app. +""" + +from __future__ import annotations + +import os +import re +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Optional + +import structlog + +from utils.llama_cpp_freshness import ( + _INSTALL_MARKER_NAME, + check_prebuilt_freshness, + latest_published_release, + read_install_marker, + reset_caches, +) + +logger = structlog.get_logger(__name__) + +DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" +_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate + +# Background job state. Single in-flight update at a time, guarded by _job_lock. +_JOB_IDLE = "idle" +_JOB_RUNNING = "running" +_JOB_SUCCESS = "success" +_JOB_ERROR = "error" + +_job_lock = threading.Lock() +_job: dict = { + "state": _JOB_IDLE, + "message": "", + "from_tag": None, + "to_tag": None, + "error": None, + "started_at": None, + "finished_at": None, +} + + +def _utcnow() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + + +def _find_binary() -> Optional[str]: + """Locate the active llama-server binary via the inference backend's own + resolver, so update targets exactly what Studio runs. Lazy import keeps the + heavy inference module off this module's import path.""" + try: + from core.inference.llama_cpp import LlamaCppBackend + return LlamaCppBackend._find_llama_server_binary() + except Exception as exc: # pragma: no cover - defensive + logger.debug("llama update: binary discovery failed", error = str(exc)) + return None + + +def _install_dir_for(binary_path: Optional[str]) -> Optional[Path]: + """The directory holding UNSLOTH_PREBUILT_INFO.json -- i.e. the install root + install_llama_prebuilt.py wrote and the one we re-install into. Walks up from + the binary the same way read_install_marker() does.""" + if not binary_path: + return None + p = Path(binary_path) + for parent in p.parents[:5]: + if (parent / _INSTALL_MARKER_NAME).is_file(): + return parent + return None + + +def _installer_script() -> Optional[Path]: + """Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then + searches up from this file for both ``/install_llama_prebuilt.py`` and + ``/studio/install_llama_prebuilt.py`` so it works in the dev tree and + in an installed Studio layout.""" + env = os.environ.get("UNSLOTH_LLAMA_INSTALLER") + if env and Path(env).is_file(): + return Path(env) + here = Path(__file__).resolve() + for up in here.parents: + for cand in (up / "install_llama_prebuilt.py", up / "studio" / "install_llama_prebuilt.py"): + if cand.is_file(): + return cand + return None + + +def get_update_status(*, force_refresh: bool = False) -> dict: + """Report whether a newer prebuilt exists plus the current job state. + + force_refresh bypasses the 24h release cache for an explicit "check now". + """ + binary = _find_binary() + marker = read_install_marker(binary) + repo = (marker or {}).get("published_repo") or DEFAULT_PUBLISHED_REPO + + if force_refresh and repo: + # Prime the cache so the freshness read below sees the newest tag. + try: + latest_published_release(repo, force_refresh = True) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: force refresh failed", error = str(exc)) + + freshness = check_prebuilt_freshness(binary) + installed = freshness.get("installed_tag") + latest = freshness.get("latest_tag") + update_available = bool( + freshness.get("has_marker") and installed and latest and installed != latest + ) + + with _job_lock: + job = dict(_job) + + return { + "supported": bool(freshness.get("has_marker")), + "update_available": update_available, + "stale": bool(freshness.get("stale")), + "installed_tag": installed, + "latest_tag": latest, + "published_repo": freshness.get("published_repo") or repo, + "installed_at_utc": freshness.get("installed_at_utc"), + "age_days": freshness.get("age_days"), + "job": job, + } + + +def _rocm_install_args(asset: Optional[str]) -> list[str]: + """Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh. + The installer probe can miss the gfx arch on amd-smi-only hosts; lemonade + bundles carry the family in the name (rocm-gfx110X), fork bundles only rocm/hip.""" + if not asset: + return [] + low = asset.lower() + if "rocm" not in low and "hip" not in low: + return [] + gfx = re.search(r"-gfx[0-9a-z]+", low) + if gfx: + # _normalize_forwarded_gfx accepts the family form (gfx110x -> gfx110X). + return ["--rocm-gfx", gfx.group(0).lstrip("-")] + return ["--has-rocm"] + + +def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path) -> None: + """Worker: put the backend into a maintenance state, run the installer for + the latest prebuilt, then refresh caches so the next load uses the new build.""" + backend = None + model_was_active = False + try: + # Maintenance state so no load starts a server from the half-swapped binary + # (and the old binary is freed for the swap). Fails open without a backend. + try: + from routes.inference import get_llama_cpp_backend + backend = get_llama_cpp_backend() + except Exception as exc: + logger.debug( + "llama update: backend unavailable, skipping load coordination", error = str(exc) + ) + backend = None + + if backend is not None: + try: + with backend._serial_load_lock: + backend._llama_update_in_progress = True + # is_active covers the loading/unhealthy window is_loaded misses + # (a live process also locks the exe on Windows during the swap). + if getattr(backend, "is_active", False): + model_was_active = True + backend.unload_model() + except Exception as exc: + logger.debug("llama update: load coordination failed", error = str(exc)) + + cmd = [ + sys.executable, + str(script), + "--install-dir", + str(install_dir), + "--llama-tag", + "latest", + "--published-repo", + repo, + ] + cmd.extend(_rocm_install_args(asset)) + logger.info("llama update: installing", cmd = " ".join(cmd)) + proc = subprocess.run( + cmd, + capture_output = True, + text = True, + timeout = _INSTALL_TIMEOUT_SECONDS, + ) + if proc.returncode != 0: + tail = (proc.stderr or proc.stdout or "").strip()[-1500:] + raise RuntimeError(f"installer exited {proc.returncode}: {tail or 'no output'}") + + # New UNSLOTH_PREBUILT_INFO.json is on disk; drop caches so the next + # status read reflects the freshly installed tag. + reset_caches() + new_marker = read_install_marker(_find_binary()) + new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag") + + with _job_lock: + _job.update( + state = _JOB_SUCCESS, + message = ( + f"Updated llama.cpp to {new_tag}." + + (" Reload your model to use it." if model_was_active else "") + ), + to_tag = new_tag, + error = None, + finished_at = _utcnow(), + ) + logger.info("llama update: success", to_tag = new_tag) + except Exception as exc: + logger.warning("llama update: failed", error = str(exc)) + with _job_lock: + _job.update( + state = _JOB_ERROR, + message = "llama.cpp update failed.", + error = str(exc), + finished_at = _utcnow(), + ) + finally: + # Lift the maintenance state so model loads work again, success or not. + if backend is not None: + try: + backend._llama_update_in_progress = False + except Exception: # pragma: no cover - defensive + pass + + +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 { + "started": False, + "reason": "installer_missing", + "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") + + with _job_lock: + if _job["state"] == _JOB_RUNNING: + return {"started": False, "reason": "already_running", "job": dict(_job)} + _job.update( + state = _JOB_RUNNING, + message = "Downloading and installing the latest llama.cpp prebuilt...", + from_tag = from_tag, + to_tag = None, + error = None, + started_at = _utcnow(), + finished_at = None, + ) + job_snapshot = dict(_job) + + thread = threading.Thread( + target = _run_update, + args = (install_dir, repo, asset, script), + name = "llama-cpp-update", + daemon = True, + ) + thread.start() + return {"started": True, "reason": None, "job": job_snapshot} + + +def _reset_job_for_tests() -> None: + """Test-only: return the job tracker to idle.""" + with _job_lock: + _job.update( + state = _JOB_IDLE, + message = "", + from_tag = None, + to_tag = None, + error = None, + started_at = None, + finished_at = None, + ) diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 25dbfdc780..91d99e238e 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -11,6 +11,7 @@ import { import { Toaster } from "@/components/ui/sonner"; import { TooltipProvider } from "@/components/ui/tooltip"; import { WebUpdateBanner } from "@/components/web/update-banner"; +import { LlamaUpdateBanner } from "@/components/llama-update-banner"; import { DownloadManagerPanel } from "@/features/hub/download-manager"; import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth"; import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain"; @@ -259,6 +260,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { {children} + ); } @@ -294,7 +296,18 @@ function TauriWrapper({ children }: { children: ReactNode }) { /> ); - if (!shouldUseCustomWindowTitlebar()) return content; + if (!shouldUseCustomWindowTitlebar()) { + // macOS desktop uses the native titlebar and returns here before the + // custom-titlebar branch, so mount the updater banner on this path too. + return ( + <> + {content} + + + ); + } const showSidebarSurface = showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); @@ -305,6 +318,9 @@ function TauriWrapper({ children }: { children: ReactNode }) {
{content}
+ ); } diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx new file mode 100644 index 0000000000..8919c24cc7 --- /dev/null +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; +import { toast } from "@/lib/toast"; +import { AnimatePresence, motion } from "motion/react"; +import { type ReactElement } from "react"; + +const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1]; + +interface LlamaUpdateBannerProps { + enabled?: boolean; +} + +/** + * Non-invasive "Update llama.cpp" affordance. Appears bottom-right when a newer + * prebuilt is available, fades on its own after ~10s, and re-surfaces hourly. + * Clicking Update swaps the prebuilt in place via POST /api/llama/update. + */ +export function LlamaUpdateBanner({ + enabled = true, +}: LlamaUpdateBannerProps): ReactElement | null { + const { status, visible, applying, apply, dismiss } = useLlamaUpdateCheck({ + enabled, + }); + + async function handleUpdate() { + const result = await apply(); + if (result?.ok) { + toast.success( + `llama.cpp updated to ${result.tag ?? "the latest build"}. Reload your model to use it.`, + ); + } else if (result) { + toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`); + } + } + + const show = visible && status != null && (status.update_available || applying); + + return ( + + {show ? ( + +
+ {!applying ? ( + + ) : null} + +
+ +

+ {applying ? "Updating llama.cpp..." : "New llama.cpp prebuilt"} +

+
+

+ {status?.installed_tag ?? "unknown"} →{" "} + + {status?.latest_tag ?? ""} + +

+ +
+ +
+
+ + ) : null} + + ); +} diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts new file mode 100644 index 0000000000..801576534c --- /dev/null +++ b/studio/frontend/src/hooks/use-llama-update-check.ts @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch, getAuthToken } from "@/features/auth"; +import { useCallback, useEffect, useRef, useState } from "react"; + +// First check shortly after load, then re-surface as an hourly reminder. +const FIRST_CHECK_DELAY_MS = 8000; +const REMINDER_INTERVAL_MS = 60 * 60 * 1000; // ~1 hour +// The banner fades on its own after this long (non-invasive). It re-appears on +// the next hourly reminder while an update is still available. +const AUTO_HIDE_MS = 10000; +// While an update is applying, poll the job state at this cadence. +const JOB_POLL_INTERVAL_MS = 3000; + +export interface LlamaUpdateJob { + state: "idle" | "running" | "success" | "error"; + message: string; + from_tag: string | null; + to_tag: string | null; + error: string | null; +} + +export interface LlamaUpdateStatus { + supported: boolean; + update_available: boolean; + installed_tag: string | null; + latest_tag: string | null; + job: LlamaUpdateJob; +} + +function parseStatus(value: unknown): LlamaUpdateStatus | null { + if (!value || typeof value !== "object") return null; + const s = value as Record; + const job = (s.job ?? {}) as Record; + return { + supported: s.supported === true, + update_available: s.update_available === true, + installed_tag: typeof s.installed_tag === "string" ? s.installed_tag : null, + latest_tag: typeof s.latest_tag === "string" ? s.latest_tag : null, + job: { + state: (job.state as LlamaUpdateJob["state"]) ?? "idle", + message: typeof job.message === "string" ? job.message : "", + from_tag: typeof job.from_tag === "string" ? job.from_tag : null, + to_tag: typeof job.to_tag === "string" ? job.to_tag : null, + error: typeof job.error === "string" ? job.error : null, + }, + }; +} + +async function fetchStatus(forceRefresh = false): Promise { + if (!getAuthToken()) return null; + try { + const res = await authFetch( + `/api/llama/update-status${forceRefresh ? "?force_refresh=true" : ""}`, + ); + if (!res.ok) return null; + return parseStatus(await res.json()); + } catch { + return null; + } +} + +interface UseLlamaUpdateCheckOptions { + enabled?: boolean; +} + +export interface LlamaApplyResult { + ok: boolean; + tag?: string | null; + error?: string | null; +} + +/** + * Polls the backend for a newer llama.cpp prebuilt. When one exists, `visible` + * is true for a 10s window (fades on its own), and re-surfaces every ~hour as a + * gentle reminder. `apply()` triggers the in-place swap and tracks the job. + */ +export function useLlamaUpdateCheck({ enabled = true }: UseLlamaUpdateCheckOptions = {}) { + const [status, setStatus] = useState(null); + const [visible, setVisible] = useState(false); + const [applying, setApplying] = useState(false); + const hideTimer = useRef | null>(null); + const pollTimer = useRef | null>(null); + + const clearHideTimer = useCallback(() => { + if (hideTimer.current) { + clearTimeout(hideTimer.current); + hideTimer.current = null; + } + }, []); + + const armAutoHide = useCallback(() => { + clearHideTimer(); + hideTimer.current = setTimeout(() => setVisible(false), AUTO_HIDE_MS); + }, [clearHideTimer]); + + const clearPollTimer = useCallback(() => { + if (pollTimer.current) { + clearInterval(pollTimer.current); + pollTimer.current = null; + } + }, []); + + // Poll the job to completion. Shared by apply() and surfaceIfAvailable() so a + // job is tracked once whoever noticed it; onDone resolves with the result. + const startJobPoll = useCallback( + (onDone?: (result: LlamaApplyResult) => void) => { + clearPollTimer(); + pollTimer.current = setInterval(async () => { + const s = await fetchStatus(); + if (!s) return; + setStatus(s); + if (s.job.state === "running") return; + clearPollTimer(); + setApplying(false); + if (s.job.state === "success") { + setVisible(false); + onDone?.({ ok: true, tag: s.job.to_tag }); + } else if (s.job.state === "error") { + armAutoHide(); + onDone?.({ ok: false, error: s.job.error }); + } else { + // idle without a terminal result (job reset): stop so the banner + // does not stick on "Updating...". + armAutoHide(); + onDone?.({ ok: false, error: "update did not complete" }); + } + }, JOB_POLL_INTERVAL_MS); + }, + [armAutoHide, clearPollTimer], + ); + + // Surface the banner for the auto-hide window when an update is available. + const surfaceIfAvailable = useCallback( + (next: LlamaUpdateStatus | null) => { + if (!next) return; + setStatus(next); + if (next.job.state === "running") { + // Swap in progress (e.g. another tab): keep the banner up and track the + // job so "Updating..." clears when it finishes instead of sticking. + setApplying(true); + setVisible(true); + clearHideTimer(); + if (!pollTimer.current) startJobPoll(); + return; + } + if (next.update_available) { + setVisible(true); + armAutoHide(); + } + }, + [armAutoHide, clearHideTimer, startJobPoll], + ); + + useEffect(() => { + if (!enabled) return; + let canceled = false; + + const firstTimer = setTimeout(() => { + fetchStatus().then((s) => { + if (!canceled) surfaceIfAvailable(s); + }); + }, FIRST_CHECK_DELAY_MS); + + const reminder = setInterval(() => { + fetchStatus().then((s) => { + if (!canceled) surfaceIfAvailable(s); + }); + }, REMINDER_INTERVAL_MS); + + return () => { + canceled = true; + clearTimeout(firstTimer); + clearInterval(reminder); + clearHideTimer(); + clearPollTimer(); + }; + }, [enabled, surfaceIfAvailable, clearHideTimer, clearPollTimer]); + + const dismiss = useCallback(() => { + clearHideTimer(); + setVisible(false); + }, [clearHideTimer]); + + const apply = useCallback(async (): Promise => { + if (applying) return { ok: false, error: "already running" }; + setApplying(true); + setVisible(true); + clearHideTimer(); + let action: { + started?: boolean; + reason?: string | null; + message?: string | null; + } | null = null; + try { + const res = await authFetch("/api/llama/update", { method: "POST" }); + if (!res.ok) { + setApplying(false); + armAutoHide(); + return { ok: false, error: `HTTP ${res.status}` }; + } + try { + action = await res.json(); + } catch { + action = null; + } + } catch (e) { + setApplying(false); + armAutoHide(); + return { ok: false, error: String(e) }; + } + + // 200 without a started job (no marker / installer missing) leaves it idle, + // so surface the reason instead of polling forever. already_running is the + // exception: a job is in flight, so track it to completion below. + if (action && action.started === false && action.reason !== "already_running") { + setApplying(false); + armAutoHide(); + return { + ok: false, + error: action.message ?? action.reason ?? "update was not started", + }; + } + + return await new Promise((resolve) => startJobPoll(resolve)); + }, [applying, armAutoHide, clearHideTimer, startJobPoll]); + + return { + status: enabled ? status : null, + visible: enabled && visible, + applying, + apply, + dismiss, + }; +}