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, + }; +}