diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 3aae6f4209..5559cc0404 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -55,6 +55,9 @@ class LlamaUpdateStatusResponse(BaseModel): source_build: bool = Field( False, description = "True when there is no marker (source build) but a prebuilt is offered." ) + update_size_bytes: Optional[int] = Field( + None, description = "Download size of the prebuilt Update would fetch, in bytes." + ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index f90c4ba0e7..2a2e113585 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -520,3 +520,114 @@ def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path) assert info["latest_tag"] == "b9596-mix-aaa" assert info["behind"] is True assert info["stale"] is True + + +# update_download_size_bytes (banner download-size lookup). + + +def _patch_assets(monkeypatch, mapping): + """Stub latest_release_assets with a per-repo {asset_name: size} lookup.""" + monkeypatch.setattr( + fr, + "latest_release_assets", + lambda repo, *, force_refresh = False: mapping.get(repo), + ) + + +def test_update_size_unsloth_prebuilt_exact_match(monkeypatch): + # The unsloth fork's own bundle (app--): the want= exact match + # on app-- wins. + marker = { + "asset": "app-b9190-linux-x64-cuda13-newer.tar.gz", + "published_repo": "unslothai/llama.cpp", + } + _patch_assets( + monkeypatch, + { + "unslothai/llama.cpp": { + "app-b9300-linux-x64-cuda13-newer.tar.gz": 123_456_789, + "app-b9300-windows-x64-cuda13-newer.zip": 999, + } + }, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 123_456_789 + + +def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch): + # macOS bundles use the upstream-style llama--bin-macos-*, matched via the + # endswith fallback in the publish repo. + marker = { + "asset": "llama-b9190-bin-macos-arm64.tar.gz", + "published_repo": "unslothai/llama.cpp", + } + _patch_assets( + monkeypatch, + {"unslothai/llama.cpp": {"llama-b9300-bin-macos-arm64.tar.gz": 55_000_000}}, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 55_000_000 + + +def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch): + # #6338 P2: ggml-org ubuntu-* prebuilt lives in binary_repo, not the fork + # publish repo. The size must still resolve. + marker = { + "asset": "llama-b9190-bin-ubuntu-x64.tar.gz", + "published_repo": "unslothai/llama.cpp", + "binary_repo": "ggml-org/llama.cpp", + } + _patch_assets( + monkeypatch, + { + "unslothai/llama.cpp": {"app-b9300-linux-x64-cuda13-newer.tar.gz": 1}, + "ggml-org/llama.cpp": { + "llama-b9673-bin-ubuntu-x64.tar.gz": 42_000_000, + "llama-b9673-bin-ubuntu-vulkan-x64.tar.gz": 7, + }, + }, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 42_000_000 + + +def test_update_size_upstream_windows_uses_binary_repo(monkeypatch): + # Regression (#6338 P2): the Windows upstream CPU prebuilt uses a win-* token. + marker = { + "asset": "llama-b9190-bin-win-cpu-x64.zip", + "published_repo": "unslothai/llama.cpp", + "binary_repo": "ggml-org/llama.cpp", + } + _patch_assets( + monkeypatch, + {"ggml-org/llama.cpp": {"llama-b9673-bin-win-cpu-x64.zip": 33_000_000}}, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 33_000_000 + + +def test_update_size_no_matching_asset_fails_open(monkeypatch): + # A ROCm version drift (installed 6.4 vs latest 7.2) leaves no suffix match; + # the helper fails open to None rather than guessing a wrong artifact. + marker = { + "asset": "llama-b9190-bin-ubuntu-rocm-6.4-x64.tar.gz", + "published_repo": "unslothai/llama.cpp", + "binary_repo": "ggml-org/llama.cpp", + } + _patch_assets( + monkeypatch, + {"ggml-org/llama.cpp": {"llama-b9673-bin-ubuntu-rocm-7.2-x64.tar.gz": 9}}, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") is None + + +def test_update_size_missing_inputs_fail_open(monkeypatch): + _patch_assets( + monkeypatch, + {"unslothai/llama.cpp": {"app-b9300-linux-x64-cpu.tar.gz": 5}}, + ) + # No marker, no latest tag, or no asset string -> None (never raise). + assert fr.update_download_size_bytes(None, "b9300", "unslothai/llama.cpp") is None + assert ( + fr.update_download_size_bytes( + {"asset": "app-b9190-linux-x64-cpu.tar.gz"}, None, "unslothai/llama.cpp" + ) + is None + ) + assert fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") is None diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index c42c3e8a79..828d439e0a 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -901,3 +901,49 @@ def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path): res = upd.start_update() assert res["started"] is False assert res["reason"] == "up_to_date" + + +def test_status_update_available_includes_size(monkeypatch, tmp_path): + # Marker (prebuilt) update path attaches the download size of the asset the + # banner would fetch. + binary = _write_install(tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + freshness, + "latest_release_assets", + lambda repo, *, force_refresh = False: { + "app-b9518-linux-x64-cuda13-newer.tar.gz": 88_000_000 + }, + ) + st = upd.get_update_status(force_refresh = True) + assert st["update_available"] is True + assert st["update_size_bytes"] == 88_000_000 + + +def test_status_source_build_includes_update_size(monkeypatch, tmp_path): + # #6338 P3: a source build offered a prebuilt must carry the asset size too. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker -> source build + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt( + monkeypatch, + repo = "unslothai/llama.cpp", + release_tag = "b9585", + asset = "app-b9585-linux-x64-cpu.tar.gz", + ) + monkeypatch.setattr(upd, "_installed_build_number", lambda b: None) + monkeypatch.setattr( + upd, + "latest_release_assets", + lambda repo, *, force_refresh = False: ( + {"app-b9585-linux-x64-cpu.tar.gz": 77_000_000} + if repo == "unslothai/llama.cpp" + else None + ), + ) + st = upd.get_update_status() + assert st["source_build"] is True + assert st["update_available"] is True + assert st["update_size_bytes"] == 77_000_000 diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index 7f1d0c8a42..91979e3e0f 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -281,6 +281,7 @@ def test_kill_process_records_timestamp_on_actual_kill(): backend = LlamaCppBackend.__new__(LlamaCppBackend) backend._process = None backend._healthy = False + backend._stats_logger = None # _kill_process stops it in finally backend._stdout_thread = None backend._llama_log_fh = None backend._last_kill_monotonic = 0.0 diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py index bf0c4b731f..333f710b44 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -78,6 +78,27 @@ def test_status_response_exposes_source_build(): rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1}) +def test_status_response_exposes_update_size_bytes(): + payload = { + "supported": True, + "update_available": True, + "stale": False, + "installed_tag": "b9493", + "latest_tag": "b9518", + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "update_size_bytes": 123_456_789, + "job": {"state": "idle"}, + } + model = rl.LlamaUpdateStatusResponse(**payload) + assert model.model_dump()["update_size_bytes"] == 123_456_789 + # Omitted -> defaults to None (the offline / no-matching-asset case). + without = {k: v for k, v in payload.items() if k != "update_size_bytes"} + assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None + + def test_status_handler_runs_off_event_loop(monkeypatch): seen = {} diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index 87d0d2ec01..7d077bfa3b 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -33,6 +33,8 @@ _INSTALL_MARKER_NAME = "UNSLOTH_PREBUILT_INFO.json" _marker_cache: dict[str, Optional[dict]] = {} _release_memo: dict[str, tuple[float, Optional[str]]] = {} +# Newest-release asset sizes (name -> bytes), memoized like the tag (24h TTL). +_assets_memo: dict[str, tuple[float, dict[str, int]]] = {} def _cache_dir() -> Path: @@ -180,6 +182,113 @@ def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optio return latest +def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]: + """Asset name -> size (bytes) for the newest published release of `repo`, + selected exactly like _fetch_latest_release_tag. None on any failure.""" + import urllib.error + import urllib.request + + url = f"https://api.github.com/repos/{repo}/releases?per_page=30" + headers = { + "Accept": "application/vnd.github+json", + "User-Agent": "unsloth-studio-freshness-check", + } + token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, headers = headers) + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + data = json.loads(resp.read().decode("utf-8")) + except ( + urllib.error.URLError, + urllib.error.HTTPError, + OSError, + json.JSONDecodeError, + ) as exc: + logger.debug("freshness asset fetch failed", repo = repo, error = str(exc)) + return None + if not isinstance(data, list): + return None + published = [ + r + for r in data + if isinstance(r, dict) + and not r.get("draft") + and not r.get("prerelease") + and isinstance(r.get("tag_name"), str) + and r.get("tag_name") + ] + if not published: + return None + newest = max(published, key = lambda r: r.get("published_at") or "") + assets: dict[str, int] = {} + for a in newest.get("assets") or []: + name, size = a.get("name"), a.get("size") + if isinstance(name, str) and isinstance(size, int): + assets[name] = size + return assets + + +def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]: + """Newest-release asset sizes for `repo`, memoized (24h TTL). None when + offline and never fetched. In-memory only -- a restart simply re-fetches.""" + if not repo: + return None + now = time.time() + if not force_refresh: + memo = _assets_memo.get(repo) + if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS: + return memo[1] + assets = _fetch_latest_release_assets(repo) + if assets is None: + memo = _assets_memo.get(repo) + return memo[1] if memo else None + _assets_memo[repo] = (now, assets) + return assets + + +def update_download_size_bytes( + marker: Optional[dict], + latest_tag: Optional[str], + repo: Optional[str], + *, + force_refresh: bool = False, +) -> Optional[int]: + """Download size of the latest-release asset matching this host's installed + bundle (same platform/arch/runtime suffix as the installed asset). None when + there is no marker asset, the latest assets can't be read, or no match.""" + if not marker or not latest_tag or not repo: + return None + installed_asset = marker.get("asset") + if not isinstance(installed_asset, str): + return None + # Tag-independent platform suffix: accept the fork's "app-*" bundles and the + # upstream ggml-org "ubuntu-*"/"win-*" prebuilts ("windows" before "win"). + m = re.search(r"-((?:linux|ubuntu|windows|win|macos|darwin)-.*)$", installed_asset) + if not m: + return None + suffix = m.group(1) + # Upstream ubuntu/win assets live in the marker's binary_repo, not the fork + # publish repo; try the publish repo first, then it. + repos = [repo] + binary_repo = marker.get("binary_repo") + if isinstance(binary_repo, str) and binary_repo and binary_repo != repo: + repos.append(binary_repo) + want = f"app-{latest_tag}-{suffix}" + for r in repos: + assets = latest_release_assets(r, force_refresh = force_refresh) + if not assets: + continue + if want in assets: + return assets[want] + # Tag formatting can vary (mix suffixes); fall back to the platform suffix. + for name, size in assets.items(): + if name.endswith(suffix): + return size + return None + + def _parse_installed_at(value: object) -> Optional[datetime]: if not isinstance(value, str) or not value: return None @@ -313,6 +422,7 @@ def reset_caches(*, drop_disk: bool = False) -> None: open (off) instead of pointing at the just-replaced build.""" _marker_cache.clear() _release_memo.clear() + _assets_memo.clear() if drop_disk: import shutil diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index c2c6c67432..22ffc58b24 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -37,9 +37,11 @@ from utils.llama_cpp_freshness import ( _INSTALL_MARKER_NAME, check_prebuilt_freshness, latest_published_release, + latest_release_assets, parse_base_build, read_install_marker, reset_caches, + update_download_size_bytes, ) logger = structlog.get_logger(__name__) @@ -291,6 +293,18 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: update_available = False # Display the mix tag when that's what makes it newer; otherwise the base. latest = release_tag if latest_is_mix else base_tag + # Size of the resolved prebuilt, so source builds show it like the marker + # path. Fails open to None (offline / asset absent from the release). + update_size_bytes = None + if update_available: + asset_name = res.get("asset") + if isinstance(asset_name, str) and asset_name: + try: + assets = latest_release_assets(res.get("repo"), force_refresh = force_refresh) + if assets: + update_size_bytes = assets.get(asset_name) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: source-build size lookup failed", error = str(exc)) with _job_lock: job = dict(_job) return { @@ -303,6 +317,7 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: "installed_at_utc": None, "age_days": None, "source_build": True, + "update_size_bytes": update_size_bytes, "job": job, } @@ -345,6 +360,20 @@ def get_update_status(*, force_refresh: bool = False) -> dict: # (see llama_cpp_freshness.is_behind). update_available = bool(freshness.get("has_marker") and freshness.get("behind")) + # Size of the prebuilt that Update would download, for the banner. Only when + # an update is offered; fails open to None (offline / no matching asset). + update_size_bytes = None + if update_available: + try: + update_size_bytes = update_download_size_bytes( + marker, + latest, + freshness.get("published_repo") or repo, + force_refresh = force_refresh, + ) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: size lookup failed", error = str(exc)) + with _job_lock: job = dict(_job) @@ -358,6 +387,7 @@ def get_update_status(*, force_refresh: bool = False) -> dict: "installed_at_utc": freshness.get("installed_at_utc"), "age_days": freshness.get("age_days"), "source_build": False, + "update_size_bytes": update_size_bytes, "job": job, } diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 8c8fcc3646..0383d8a150 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -7,10 +7,7 @@ import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; import { toast } from "@/lib/toast"; import { cn } from "@/lib/utils"; import { Download } from "lucide-react"; -import { AnimatePresence, motion } from "motion/react"; import { type ReactElement, useEffect, useRef, useState } from "react"; - -const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1]; // Backend progress is coarse (5% steps, ~0.9 max) and the extract tail emits no // signal. Creep toward this cap so the bar keeps moving rather than freezing. const RUNNING_CAP = 0.95; @@ -114,6 +111,12 @@ export function LlamaUpdateBanner({ const show = visible && status != null && (status.update_available || applying); + const sizeBytes = status?.update_size_bytes ?? null; + // Round to whole MB; these prebuilts are hundreds of MB. + const sizeLabel = + sizeBytes && sizeBytes > 0 + ? `${Math.round(sizeBytes / (1024 * 1024))} MB` + : null; const updateProgress = status?.job.progress ?? null; const jobSucceeded = status?.job.state === "success"; // Drives the bar so it animates continuously; aria reports the real value. @@ -123,113 +126,110 @@ export function LlamaUpdateBanner({ jobSucceeded, ); - return ( - - {show ? ( - -
- {applying ? null : ( - - )} - -
-