diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 347f469d3..211543e31 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3697,12 +3697,22 @@ class LlamaCppBackend: hf_repo: str, hf_variant: Optional[str] = None, hf_token: Optional[str] = None, + force: bool = False, + allow_smaller_fallback: bool = True, + cancel_event: Optional[threading.Event] = None, ) -> str: """Download GGUF file(s) from HuggingFace. Returns local path. Runs WITHOUT self._lock so unload_model() can set _cancel_event at any time; checks it between each shard download. + + ``force`` re-fetches even when a (possibly stale) blob is cached. + ``allow_smaller_fallback=False`` raises on low disk instead of silently + switching to a smaller quant. ``cancel_event`` overrides + ``self._cancel_event`` so an update can use a private event without + touching the shared one; defaults to the shared event. """ + cancel_event = cancel_event if cancel_event is not None else self._cancel_event try: import huggingface_hub # noqa: F401 -- presence check only except ImportError: @@ -3768,21 +3778,22 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - for p in path_infos: - if not p.size: - continue - try: - cached_path = try_to_load_from_cache(hf_repo, p.path) - except Exception: - cached_path = None - if isinstance(cached_path, str) and os.path.exists(cached_path): + if not force: + for p in path_infos: + if not p.size: + continue try: - on_disk = os.path.getsize(cached_path) - except OSError: - on_disk = 0 - # Satisfied only when the full blob is present. - if on_disk >= p.size: - already_cached_bytes += p.size + cached_path = try_to_load_from_cache(hf_repo, p.path) + except Exception: + cached_path = None + if isinstance(cached_path, str) and os.path.exists(cached_path): + try: + on_disk = os.path.getsize(cached_path) + except OSError: + on_disk = 0 + # Satisfied only when the full blob is present. + if on_disk >= p.size: + already_cached_bytes += p.size total_download_bytes = max(0, total_bytes - already_cached_bytes) @@ -3805,6 +3816,13 @@ class LlamaCppBackend: ) if total_download_bytes > free_bytes: + if not allow_smaller_fallback: + # Update path: never silently switch to a smaller quant; + # surface the disk shortfall for the requested variant. + raise RuntimeError( + f"Not enough disk space to download {gguf_filename}. " + f"Only {free_gb:.1f} GB free in {cache_dir}" + ) smaller = self._find_smallest_fitting_variant( hf_repo, free_bytes, @@ -3845,7 +3863,7 @@ class LlamaCppBackend: ) logger.info(f"Resolving GGUF: {gguf_label}") try: - if self._cancel_event.is_set(): + if cancel_event.is_set(): raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. @@ -3853,18 +3871,20 @@ class LlamaCppBackend: hf_repo, gguf_filename, hf_token, - cancel_event = self._cancel_event, + cancel_event = cancel_event, on_status = lambda m: logger.info(m), + force_download = force, ) for shard in gguf_extra_shards: - if self._cancel_event.is_set(): + if cancel_event.is_set(): raise RuntimeError("Cancelled") logger.info(f"Resolving GGUF shard: {shard}") hf_hub_download_with_xet_fallback( hf_repo, shard, hf_token, - cancel_event = self._cancel_event, + cancel_event = cancel_event, + force_download = force, ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): @@ -3887,6 +3907,7 @@ class LlamaCppBackend: hf_token: Optional[str], pick: Callable[[list[str]], Optional[str]], label: str, + cancel_event: Optional[threading.Event] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. @@ -3894,8 +3915,10 @@ class LlamaCppBackend: (offline, same fallback as _download_gguf), then hf_hub_download. Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. + ``cancel_event`` overrides ``self._cancel_event`` (defaults to it). """ - if self._cancel_event.is_set(): + cancel_event = cancel_event if cancel_event is not None else self._cancel_event + if cancel_event.is_set(): return None target: Optional[str] = None @@ -3904,7 +3927,7 @@ class LlamaCppBackend: # Retry a transient listing blip; permanent repo/auth errors and offline # mode are not retried (offline raises at once -> fall through to cache). for attempt in range(3): - if self._cancel_event.is_set(): + if cancel_event.is_set(): return None try: target = pick(list_repo_files(hf_repo, token = hf_token)) @@ -3923,7 +3946,7 @@ class LlamaCppBackend: f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}" ) if attempt < 2: - self._cancel_event.wait(2**attempt) + cancel_event.wait(2**attempt) if target is None: try: @@ -3937,7 +3960,7 @@ class LlamaCppBackend: except Exception as e: logger.debug(f"Offline cache lookup for {label} failed: {e}") - if target is None or self._cancel_event.is_set(): + if target is None or cancel_event.is_set(): return None try: @@ -3947,7 +3970,7 @@ class LlamaCppBackend: hf_repo, target, hf_token, - cancel_event = self._cancel_event, + cancel_event = cancel_event, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") @@ -3958,11 +3981,13 @@ class LlamaCppBackend: *, hf_repo: str, hf_token: Optional[str] = None, + cancel_event: Optional[threading.Event] = None, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local - path, or None if none exists. + path, or None if none exists. ``cancel_event`` overrides + ``self._cancel_event`` (defaults to it). """ def _pick_mmproj(candidates: list[str]) -> Optional[str]: @@ -3983,6 +4008,7 @@ class LlamaCppBackend: hf_token = hf_token, pick = _pick_mmproj, label = "mmproj", + cancel_event = cancel_event, ) def _download_mtp( diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index 44ff545e7..ef95efe2f 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -27,6 +27,9 @@ class GgufVariantDetail(BaseModel): downloaded: bool = Field( False, description = "Whether this variant is already in the local HF cache" ) + update_available: bool = Field( + False, description = "Whether a newer main GGUF blob is available on Hugging Face" + ) partial: bool = Field( False, description = "Whether this variant has an in-progress (.incomplete) blob in cache", diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index a961a6ae9..a27e4860e 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -39,8 +39,10 @@ from hub.services.models.common import ( logger = get_logger(__name__) -_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict() -_repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict() +_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = ( + OrderedDict() +) +_repo_size_neg_cache: "OrderedDict[tuple[str, str, str], float]" = OrderedDict() _REPO_SIZE_CACHE_MAX = 256 _REPO_SIZE_POS_TTL = 60.0 _REPO_SIZE_NEG_TTL = 60.0 @@ -52,7 +54,7 @@ def get_repo_snapshot_metadata_cached( repo_id: str, hf_token: Optional[str] = None ) -> tuple[int, frozenset[str]]: token_fp = hf_cache_scan.token_fingerprint(hf_token) - cache_key = (repo_id, token_fp) + cache_key = (repo_id, token_fp, "snapshot") with _repo_size_cache_lock: cached = _repo_size_cache.get(cache_key) if cached is not None: @@ -119,6 +121,52 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _cached_repo_file_name(file_obj) -> str: + file_path = getattr(file_obj, "file_path", None) + if file_path: + try: + path = Path(file_path) + parts = path.parts + snapshots_idx = max(i for i, part in enumerate(parts) if part == "snapshots") + if len(parts) > snapshots_idx + 2: + return Path(*parts[snapshots_idx + 2 :]).as_posix() + except Exception: + pass + return str(getattr(file_obj, "file_name", "")).replace("\\", "/") + + +def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]: + """Map each cached GGUF file's repo-relative name to the SET of its local + blob hashes across all cached revisions. + + HF names each local cache blob FILE by the file's etag (lfs.sha256 else + blob_id), so a local file's blob hash == ``Path(blob_path).name``. An updated + repo keeps BOTH the old and new revision snapshots until HF garbage-collects + them, so the same file resolves to several blobs; collecting them ALL (not + just the first one seen, since ``repo_info.revisions`` is a frozenset and + yields them in arbitrary order) lets the remote-vs-local diff treat the file + as current when the remote (``main``) blob is present in any cached revision. + Mirrors the ``cached_blob_ids`` membership test in routes/models.py. + + By default this keeps the historical MAIN-GGUF-only behavior. GGUF update + checks opt into companions so a shared mmproj/MTP blob can be compared too. + """ + blob_map: dict[str, set[str]] = {} + for revision in repo_info.revisions: + for f in revision.files: + if include_companions: + if not _is_gguf_filename(f.file_name): + continue + elif not _is_main_gguf_filename(f.file_name): + continue + blob_path = getattr(f, "blob_path", None) + if not blob_path: + continue + name = _cached_repo_file_name(f) + blob_map.setdefault(name, set()).add(Path(blob_path).name) + return blob_map + + def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool: if existing is None: return True diff --git a/studio/backend/hub/services/models/deletion.py b/studio/backend/hub/services/models/deletion.py index 916bb9d4f..e7c54fc75 100644 --- a/studio/backend/hub/services/models/deletion.py +++ b/studio/backend/hub/services/models/deletion.py @@ -153,6 +153,30 @@ def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, l return removed, failures +def _remove_empty_snapshot_dirs(target_repos: list) -> tuple[int, list[str]]: + removed = 0 + failures: list[str] = [] + for target_repo in target_repos: + repo_path = getattr(target_repo, "repo_path", None) + if not repo_path: + continue + snapshots = Path(repo_path) / "snapshots" + if not snapshots.is_dir(): + continue + try: + snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()] + except OSError: + continue + for snap in snap_dirs: + try: + snap.rmdir() + removed += 1 + except OSError as e: + if e.errno != errno.ENOTEMPTY: + failures.append(f"{snap.name}: {e}") + return removed, failures + + def _delete_gguf_variant_from_repos( repo_id: str, variant: str, @@ -255,6 +279,9 @@ def _delete_gguf_variant_from_repos( state_purged = download_manifest.purge_state("model", repo_id, variant) # Reclaim the empty quant folder so it stops 404ing on delete. removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant) + removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos) + removed_dirs += removed_snap_dirs + dir_failures.extend(snap_dir_failures) if dir_failures: raise HTTPException( status_code = 409, @@ -284,6 +311,181 @@ def _delete_gguf_variant_from_repos( return {"status": "deleted", "repo_id": repo_id, "variant": variant} +def reclaim_replaced_gguf_variant( + repo_id: str, + variant: str, + keep_main_hashes: frozenset[str], + hf_token: Optional[str] = None, +) -> dict: + """Prune stale main-GGUF files for a variant after a replacement verified. + + This is intentionally narrower than user-driven delete: it removes only + same-variant main files whose local blob hash is not in *keep_main_hashes*, + then unlinks their blobs only if no remaining snapshot references them. + Shared companions and sibling variants are left intact. + """ + if not keep_main_hashes: + logger.info( + "Skipping stale GGUF reclaim for %s [%s]: current main hashes unresolved", + repo_id, + variant, + ) + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "unresolved_hashes", + } + if not _is_valid_repo_id(repo_id) or not _is_valid_gguf_variant(variant): + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "invalid_target", + } + + failures: list[str] = [] + removed_snapshots = 0 + deleted_blobs = 0 + deleted_bytes = 0 + variant_key = variant.lower() + + try: + cache_scans = cache_inventory.all_hf_cache_scans() + except Exception as e: + logger.warning( + "Skipping stale GGUF reclaim for %s [%s]: cache scan failed: %s", + repo_id, + variant, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "scan_failed", + } + + candidate_repos = [ + repo_info + for hf_cache in cache_scans + for repo_info in hf_cache.repos + if str(getattr(repo_info, "repo_type", "")) == "model" + and str(getattr(repo_info, "repo_id", "")).lower() == repo_id.lower() + ] + try: + matched_repo_ids = resolve_destructive_repo_ids( + repo_id, + [str(getattr(repo_info, "repo_id", "")) for repo_info in candidate_repos], + noun = "models", + ) + except HTTPException as e: + detail = getattr(e, "detail", str(e)) + logger.warning( + "Skipping stale GGUF reclaim for %s [%s]: %s", + repo_id, + variant, + download_registry.scrub_secrets(str(detail), hf_token = hf_token), + ) + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "ambiguous_repo", + } + target_repos = [ + repo_info + for repo_info in candidate_repos + if str(getattr(repo_info, "repo_id", "")) in matched_repo_ids + ] + + for target_repo in target_repos: + repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None + stale_matches: list[tuple[Path, Optional[Path], str]] = [] + matches = _repo_file_matches( + target_repo, + lambda name: _is_main_gguf_filename(name) + and extract_quant_label(name).lower() == variant_key, + ) + for snap, blob, name in matches: + blob_hash = _blob_hash_from_path(blob) if blob is not None else None + if blob_hash is None or blob_hash in keep_main_hashes: + continue + stale_matches.append((snap, blob, name)) + + if not stale_matches: + continue + + for snap, _blob, name in stale_matches: + try: + if _path_exists_or_symlink(snap): + snap.unlink() + removed_snapshots += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + ref_counts = _snapshot_blob_reference_counts(repo_dir) + seen_blobs: set[Path] = set() + for _snap, blob, name in stale_matches: + if blob is None: + continue + try: + blob_key = blob.resolve() + except OSError: + blob_key = blob + if blob_key in seen_blobs: + continue + seen_blobs.add(blob_key) + if ref_counts.get(blob_key, 0) > 0: + continue + try: + if blob.exists(): + deleted_bytes += blob.stat().st_size + blob.unlink() + deleted_blobs += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + removed_dirs = 0 + dir_failures: list[str] = [] + if target_repos: + removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant) + removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos) + removed_dirs += removed_snap_dirs + dir_failures.extend(snap_dir_failures) + failures.extend(dir_failures) + + if failures: + logger.warning( + "Stale GGUF reclaim for %s [%s] left %d failure(s): %s", + repo_id, + variant, + len(failures), + "; ".join(failures[:3]), + ) + + if removed_snapshots or deleted_blobs or removed_dirs: + cache_inventory.invalidate_hf_cache_scans() + logger.info( + "Reclaimed stale GGUF %s [%s]: snapshots=%d blobs=%d dirs=%d freed=%.1f MB", + repo_id, + variant, + removed_snapshots, + deleted_blobs, + removed_dirs, + deleted_bytes / (1024 * 1024), + ) + + return { + "status": "reclaimed", + "repo_id": repo_id, + "variant": variant, + "removed_snapshots": removed_snapshots, + "deleted_blobs": deleted_blobs, + "removed_dirs": removed_dirs, + } + + def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool: """True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``.""" rid = repo_id.lower() diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py index 0c0e6f925..0147bba19 100644 --- a/studio/backend/hub/services/models/gguf_variants.py +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -291,6 +291,75 @@ def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]: return hf_cache_scan.partial_transport_for("model", repo_id, variant) +def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str]]]: + """Map quant -> repo-relative expected GGUF filename -> cached blob hashes. + + Shared companions are copied into each main-quant bucket so update checks can + detect mmproj/MTP-only upstream changes without a separate remote call. + """ + result: dict[str, dict[str, set[str]]] = {} + companion_blobs: dict[str, set[str]] = {} + try: + from hub.services.models import cache_inventory + scans = cache_inventory.all_hf_cache_scans() + except Exception as e: + logger.warning("Failed to scan local GGUF blobs for %s: %s", repo_id, e) + return result + + target_lower = repo_id.lower() + for hf_cache in scans: + for repo_info in hf_cache.repos: + if str(getattr(repo_info, "repo_type", "")) != "model": + continue + if str(getattr(repo_info, "repo_id", "")).lower() != target_lower: + continue + for path, hashes in cache_inventory._repo_gguf_blob_map( + repo_info, + include_companions = True, + ).items(): + normalized = str(path).replace("\\", "/") + if not hashes: + continue + if _is_mmproj_filename(normalized) or _is_mtp_drafter_path(normalized): + companion_blobs.setdefault(normalized, set()).update( + str(blob) for blob in hashes if blob + ) + continue + quant = extract_quant_label(normalized).lower() + if is_big_endian_gguf_path(normalized, quant): + continue + bucket = result.setdefault(quant, {}).setdefault(normalized, set()) + bucket.update(str(blob) for blob in hashes if blob) + if companion_blobs: + for local_blobs in result.values(): + for path, hashes in companion_blobs.items(): + local_blobs.setdefault(path, set()).update(hashes) + return result + + +def _variant_update_available_from_requirement( + local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str +) -> bool: + if requirement is None or not local_blobs: + return False + local_by_posix = {path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()} + for expected in requirement.expected_files: + path = str(expected.path).replace("\\", "/") + if not ( + is_main_gguf_variant_path(path, variant) + or _is_mmproj_filename(path) + or _is_mtp_drafter_path(path) + ): + continue + remote_blob = expected.sha256 + if not remote_blob: + continue + local_set = local_by_posix.get(path) + if not local_set or remote_blob not in local_set: + return True + return False + + def delete_variant_incomplete_blobs_result( repo_id: str, variant: str, @@ -657,9 +726,12 @@ async def get_gguf_variants_response( _partial_transport_for_variant(repo_id, variant.quant), ) + local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id) + def _variant_detail(v) -> GgufVariantDetail: is_partial = v.quant in partial_quants requirement = requirements_by_quant.get(v.quant.lower()) + downloaded = _is_fully_downloaded(v) and not is_partial return GgufVariantDetail( filename = v.filename, quant = v.quant, @@ -668,7 +740,13 @@ async def get_gguf_variants_response( download_size_bytes = ( requirement.download_size_bytes if requirement is not None else v.size_bytes ), - downloaded = _is_fully_downloaded(v) and not is_partial, + downloaded = downloaded, + update_available = downloaded + and _variant_update_available_from_requirement( + local_blobs_by_quant.get(v.quant.lower(), {}), + requirement, + v.quant, + ), partial = is_partial, partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None), ) diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 1eb7042e4..e22aaba28 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -1632,6 +1632,34 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp ) +def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path): + """A verified GGUF update can prune an older snapshot and make that old + directory the newest by mtime. The variant is still complete when another + snapshot satisfies its manifest.""" + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + repo_dir = tmp_path / "cache" / "models--Org--Repo" + old_snapshot = repo_dir / "snapshots" / "old" + new_snapshot = repo_dir / "snapshots" / "new" + old_snapshot.mkdir(parents = True) + new_snapshot.mkdir(parents = True) + (old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling") + (new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new") + assert download_manifest.write_manifest( + "model", + "Org/Repo", + "Q4_K_M", + [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)], + "http", + ) + + assert not inventory_scan.is_variant_partial( + "Org/Repo", + "Q4_K_M", + snapshot_dir = old_snapshot, + repo_cache_dir = repo_dir, + ) + + def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path): async def _run_inline(fn, *args, **kwargs): return fn(*args, **kwargs) diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 40d5a3254..2abdb0fb7 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -36,7 +36,10 @@ def sibling_sha256(sibling) -> Optional[str]: value = lfs.get("sha256") else: value = getattr(lfs, "sha256", None) - return value if isinstance(value, str) and value else None + if isinstance(value, str) and value: + return value + blob_id = getattr(sibling, "blob_id", None) + return blob_id if isinstance(blob_id, str) and blob_id else None def sibling_size(sibling) -> int: diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index 0f7ce6fe3..57ad7f665 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -387,9 +387,55 @@ def _manifest_partial( ) if resolved is None: return True + if repo_type == "model" and variant is not None: + if download_manifest.verify_against_disk(manifest, resolved).ok: + return False + for candidate in _manifest_snapshot_dirs(repo_type, repo_id, repo_cache_dir): + if candidate == resolved: + continue + if download_manifest.verify_against_disk(manifest, candidate).ok: + return False + return True return not download_manifest.verify_against_disk(manifest, resolved).ok +def _manifest_snapshot_dirs( + repo_type: RepoType, + repo_id: str, + repo_cache_dir: Optional[Path] = None, +) -> list[Path]: + repo_dirs = ( + [repo_cache_dir] + if repo_cache_dir is not None + else list(iter_repo_cache_dirs(repo_type, repo_id)) + ) + snapshots: list[Path] = [] + seen: set[str] = set() + for repo_dir in repo_dirs: + if repo_dir is None: + continue + snapshots_dir = repo_dir / "snapshots" + try: + if not snapshots_dir.is_dir(): + continue + entries = list(snapshots_dir.iterdir()) + except OSError: + continue + for entry in entries: + try: + if not entry.is_dir(): + continue + resolved = entry.resolve() + except OSError: + continue + key = str(resolved) + if key in seen: + continue + seen.add(key) + snapshots.append(resolved) + return snapshots + + def is_snapshot_partial( repo_type: RepoType, repo_id: str, diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index 42a8ca52b..e45357d31 100644 --- a/studio/backend/hub/workers/hf_download.py +++ b/studio/backend/hub/workers/hf_download.py @@ -653,6 +653,21 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod snapshot_path, metadata_unavailable = metadata_unavailable, ) + if plan is not None: + try: + from hub.services.models.deletion import reclaim_replaced_gguf_variant + reclaim_replaced_gguf_variant( + repo_id, + variant, + plan.main_hashes, + hf_token, + ) + except Exception as e: + print( + f"Verified GGUF update for {repo_id} [{variant}], but stale-cache " + f"reclaim failed ({type(e).__name__}: {e})", + file = sys.stderr, + ) def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None: diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 20dea5ec1..54e88fed5 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -136,9 +136,13 @@ class GgufVariantDetail(BaseModel): filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')") quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')") size_bytes: int = Field(0, description = "File size in bytes") + download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant") downloaded: bool = Field( False, description = "Whether this variant is already in the local HF cache" ) + update_available: bool = Field( + False, description = "Whether a newer version of this variant is available on HF" + ) class GgufVariantsResponse(BaseModel): diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index e22f65751..7c75e8522 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -12,6 +12,7 @@ import sys import uuid from pathlib import Path from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query +from pydantic import BaseModel from typing import List, Optional import structlog from loggers import get_logger @@ -22,10 +23,27 @@ import re as _re _VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") +class CachedModelRepo(BaseModel): + repo_id: str + size_bytes: int + last_modified: Optional[float] = None + + +class CachedModelsResponse(BaseModel): + cached: List[CachedModelRepo] + + def _is_valid_repo_id(repo_id: str) -> bool: return bool(_VALID_REPO_ID.fullmatch(repo_id)) +def _normalize_hf_token(hf_token) -> Optional[str]: + if not isinstance(hf_token, str): + return None + token = hf_token.strip() + return token or None + + def _safe_is_dir(path) -> bool: """``Path.is_dir()`` returning ``False`` instead of raising. @@ -74,6 +92,7 @@ if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token try: from utils.models import ( @@ -2594,109 +2613,41 @@ async def get_gguf_variants( ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')" ), hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"), + hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - """List GGUF quantization variants for a HF repo or local directory. - - Returns all variants with file sizes, vision support, and the - recommended default. - """ + """List GGUF quantization variants for a HF repo or local directory.""" try: - from utils.models.model_config import is_local_path, list_local_gguf_variants + hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token) + from hub.services.models import gguf_variants as hub_gguf_variants - # Local directory path — scan filesystem. - if is_local_path(repo_id): - variants, has_vision = list_local_gguf_variants(repo_id) - - filenames = [v.filename for v in variants] - best = _pick_best_gguf(filenames) - default_variant = _extract_quant_label(best) if best else None - - return GgufVariantsResponse( - repo_id = repo_id, - variants = [ - GgufVariantDetail( - filename = v.filename, - quant = v.quant, - size_bytes = v.size_bytes, - downloaded = True, # all local variants are downloaded - ) - for v in variants - ], - has_vision = has_vision, - default_variant = default_variant, - context_length = _read_native_context_length(repo_id, is_local = True), - ) - - # Remote HuggingFace repo — query HF API. - variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token) - - filenames = [v.filename for v in variants] - best = _pick_best_gguf(filenames) - default_variant = _extract_quant_label(best) if best else None - - # Per-snapshot so a split GGUF's shards must all sit in one snapshot; - # mmproj adapters are excluded so they can't inflate a quant's bytes. - cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = [] - try: - from huggingface_hub import constants as hf_constants - - if not _is_valid_repo_id(repo_id): - raise ValueError(f"Invalid repo_id format: {repo_id}") - - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"models--{repo_id.replace('/', '--')}".lower() - for entry in cache_dir.iterdir(): - if entry.name.lower() == target: - snapshots = entry / "snapshots" - if snapshots.is_dir(): - for snap in snapshots.iterdir(): - by_quant: dict[str, int] = {} - for f in _iter_gguf_paths(snap): - if _is_mmproj_filename(f.name): - continue - try: - size = f.stat().st_size - except OSError: - continue # broken symlink / unreadable: skip - rel = f.relative_to(snap).as_posix() - q = _extract_quant_label(rel) - if _is_big_endian_gguf_path(rel, q): - continue - q = q.lower() - by_quant[q] = by_quant.get(q, 0) + size - if by_quant: - cached_bytes_by_quant_per_snapshot.append(by_quant) - break - except Exception: - pass - - def _is_fully_downloaded(variant) -> bool: - if variant.size_bytes == 0: - return False - # Complete within one snapshot (tolerance for symlink size jitter). - quant = variant.quant.lower() - return any( - by_quant.get(quant, 0) >= variant.size_bytes * 0.99 - for by_quant in cached_bytes_by_quant_per_snapshot - ) + response = await hub_gguf_variants.get_gguf_variants_response( + repo_id, + hf_token = hf_token, + ) + local = is_local_path(repo_id) return GgufVariantsResponse( - repo_id = repo_id, + repo_id = response.repo_id, variants = [ GgufVariantDetail( filename = v.filename, quant = v.quant, size_bytes = v.size_bytes, - downloaded = _is_fully_downloaded(v), + download_size_bytes = int( + getattr(v, "download_size_bytes", v.size_bytes) or v.size_bytes + ), + downloaded = bool(v.downloaded), + update_available = bool(getattr(v, "update_available", False)), ) - for v in variants + for v in response.variants ], - has_vision = has_vision, - default_variant = default_variant, - context_length = _read_native_context_length(repo_id, is_local = False), + has_vision = response.has_vision, + default_variant = response.default_variant, + context_length = _read_native_context_length(repo_id, is_local = local), ) - + except HTTPException: + raise except Exception as e: logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True) raise HTTPException( @@ -3123,10 +3074,14 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): return {"cached": []} -@router.get("/cached-models") -async def list_cached_models(current_subject: str = Depends(get_current_subject)): +@router.get("/cached-models", response_model = CachedModelsResponse) +async def list_cached_models( + current_subject: str = Depends(get_current_subject), + hf_token: Optional[str] = Depends(get_hf_token), +): """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" _WEIGHT_EXTENSIONS = (".safetensors", ".bin") + hf_token = _normalize_hf_token(hf_token) try: cache_scans = _all_hf_cache_scans() @@ -3147,20 +3102,16 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) ) if total_size == 0: continue - has_weights = any( - f.file_name.endswith(_WEIGHT_EXTENSIONS) + weight_files = [ + f for rev in repo_info.revisions for f in rev.files - ) - if not has_weights: + if f.file_name.endswith(_WEIGHT_EXTENSIONS) + ] + if not weight_files: continue last_modified = max( - ( - _blob_mtime(f) - for rev in repo_info.revisions - for f in rev.files - if f.file_name.endswith(_WEIGHT_EXTENSIONS) - ), + (_blob_mtime(f) for f in weight_files), default = 0.0, ) key = repo_id.lower() @@ -3182,9 +3133,12 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") continue - # Newest download first; stable repo_id tie-break for equal/missing mtimes. + + rows = list(seen_lower.values()) + # Local-only list path: update checks are GGUF-only and happen lazily + # when a repo's variants are viewed. cached = sorted( - seen_lower.values(), + rows, key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()), ) return {"cached": cached} diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index b2ead305b..d4a7cae20 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -20,6 +20,7 @@ if "structlog" not in sys.modules: ) import routes.models as models_route +from hub.services.models import gguf_variants as GV def _repo( @@ -527,21 +528,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa """The per-quant 'downloaded' flag is driven by the real weight file in a single snapshot; an mmproj vision adapter (matching a quant label) must not make that quant appear downloaded.""" - import huggingface_hub.constants as hf_constants - variants = [ - SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000), - SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000), + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 10_000, + ), + SimpleNamespace( + filename = "model-F16.gguf", + quant = "F16", + display_label = None, + size_bytes = 20_000, + ), ] monkeypatch.setattr( - models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True) + GV, + "list_gguf_variants", + lambda repo_id, hf_token = None: (variants, True, []), ) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16" + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -555,21 +567,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): - import huggingface_hub.constants as hf_constants - siblings = [ SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100), SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10), ] monkeypatch.setattr( - "huggingface_hub.model_info", - lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings), + GV, + "list_gguf_variants", + lambda repo_id, hf_token = None: ( + [ + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 10, + ) + ], + False, + siblings, + ), ) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -583,19 +606,25 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path): - import huggingface_hub.constants as hf_constants - variants = [ - SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10), + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 10, + ), ] monkeypatch.setattr( - models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False) + GV, + "list_gguf_variants", + lambda repo_id, hf_token = None: (variants, False, []), ) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) result = asyncio.run( models_route.get_gguf_variants( diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 39ecebd32..9e40fbf50 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -206,6 +206,7 @@ class _FakeAttempt: interval, grace_period, on_status, + force_download = False, ): self.calls.append( _types.SimpleNamespace( diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py new file mode 100644 index 000000000..9cf2a62c3 --- /dev/null +++ b/studio/backend/tests/test_model_update_robustness.py @@ -0,0 +1,483 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +"""Tests for model-update detection and the GGUF force-download helper. + +Covers: + * GGUF variant listing computes update_available from the already-fetched + sibling metadata instead of a second Hub call. + * hf_hub_download_with_xet_fallback(force_download=True) bypasses the + try_to_load_from_cache cache-first early-return. + +The cache "Update" action now runs through the download manager as a normal +managed download (so it shows in the Downloads panel with progress + cancel), +so the old POST /api/models/update endpoint and its tests are gone. Update +*detection* — the "Update available" cue — is still exercised here. +""" + +import asyncio +import sys +import types +from types import SimpleNamespace + +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *a, **k: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, get_logger = lambda *a, **k: _DummyLogger() + ) + +import pytest +from hub.services.models import cache_inventory as CI +from hub.services.models import deletion as D +from hub.services.models import gguf_variants as GV + + +def _variants(): + return [ + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 1000, + ), + SimpleNamespace( + filename = "model-Q8_0.gguf", + quant = "Q8_0", + display_label = None, + size_bytes = 2000, + ), + ] + + +def _seed_cache(tmp_path, repo_id, blob_ids, gguf_files): + repo = tmp_path / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / ("a" * 40) + snap.mkdir(parents = True, exist_ok = True) + for name, size in gguf_files.items(): + (snap / name).write_bytes(b"\0" * size) + blobs = repo / "blobs" + blobs.mkdir(exist_ok = True) + for b in blob_ids: + (blobs / b).write_bytes(b"x") + return repo, snap, blobs + + +@pytest.fixture +def patch_hub_gguf(monkeypatch): + """Patch GGUF listing and cache scans for sibling-derived update checks.""" + + def _sibling( + path: str, + size: int, + sha = None, + *, + lfs_dict = False, + blob_id = None, + ): + if lfs_dict: + lfs = {"sha256": sha} if sha else {} + else: + lfs = SimpleNamespace(sha256 = sha) if sha else None + return SimpleNamespace(rfilename = path, size = size, lfs = lfs, blob_id = blob_id) + + def _repo_info(repo_id: str, repo_path, files: list[tuple[str, str]]): + return SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = name, + blob_path = str(repo_path / "blobs" / blob), + ) + for name, blob in files + ] + ) + ], + ) + + def _apply(tmp_path, repo_id: str, *, local_blob: str, remote_sibling): + with GV._VARIANT_HASH_LOCK: + GV._VARIANT_HASH_CACHE.clear() + GV._VARIANT_REQUIREMENT_CACHE.clear() + GV._VARIANT_REQUIREMENT_NEG_CACHE.clear() + repo, snap, _blobs = _seed_cache( + tmp_path, + repo_id, + blob_ids = [local_blob], + gguf_files = {"model-Q4_K_M.gguf": 1000}, + ) + monkeypatch.setattr( + GV, + "list_gguf_variants", + lambda r, hf_token = None: (_variants(), False, [remote_sibling]), + raising = True, + ) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [ + SimpleNamespace( + repos = [ + _repo_info( + repo_id, + repo, + [("model-Q4_K_M.gguf", local_blob)], + ) + ] + ) + ], + ) + + return SimpleNamespace(apply = _apply, sibling = _sibling) + + +def _call(coro): + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + +# ── GGUF variant update detection ─────────────────────────────── + + +def test_variant_update_check_missing_remote_blob_id_is_not_phantom_update( + tmp_path, patch_hub_gguf +): + """Missing sha/blob metadata is unknown, not update_available=True.""" + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "oldsha", + remote_sibling = patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, None), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + assert len(resp.variants) == 2 + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + assert q4.downloaded is True + assert q4.update_available is False + + +def test_variant_update_check_detects_update_from_existing_siblings(tmp_path, patch_hub_gguf): + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "oldsha", + remote_sibling = patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "NEWsha"), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + assert q4.update_available is True + + +def test_variant_update_check_no_update_when_blob_matches(tmp_path, patch_hub_gguf): + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "samesha", + remote_sibling = patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "samesha"), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + assert q4.update_available is False + + +@pytest.mark.parametrize( + ("companion_path", "has_vision"), + [ + ("mmproj-F16.gguf", True), + ("mtp-drafter-Q8_0.gguf", False), + ], +) +def test_variant_update_check_detects_companion_only_update( + monkeypatch, tmp_path, patch_hub_gguf, companion_path, has_vision +): + repo_id = "unsloth/gemma-4-GGUF" + with GV._VARIANT_HASH_LOCK: + GV._VARIANT_HASH_CACHE.clear() + GV._VARIANT_REQUIREMENT_CACHE.clear() + GV._VARIANT_REQUIREMENT_NEG_CACHE.clear() + repo, snap, _blobs = _seed_cache( + tmp_path, + repo_id, + blob_ids = ["mainsha", "old-companion"], + gguf_files = { + "model-Q4_K_M.gguf": 1000, + companion_path: 100, + }, + ) + siblings = [ + patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "mainsha"), + patch_hub_gguf.sibling(companion_path, 100, "new-companion"), + ] + monkeypatch.setattr( + GV, + "list_gguf_variants", + lambda r, hf_token = None: (_variants(), has_vision, siblings), + raising = True, + ) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [ + SimpleNamespace( + repos = [ + SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + blob_path = str(repo / "blobs" / "mainsha"), + ), + SimpleNamespace( + file_name = companion_path, + blob_path = str(repo / "blobs" / "old-companion"), + ), + ] + ) + ], + ) + ] + ) + ], + ) + + resp = _call(GV.get_gguf_variants_response(repo_id)) + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + + assert q4.downloaded is True + assert q4.update_available is True + + +def test_variant_update_check_accepts_lfs_dict_and_blob_id_fallback(tmp_path, patch_hub_gguf): + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "dictsha", + remote_sibling = patch_hub_gguf.sibling( + "model-Q4_K_M.gguf", + 1000, + "dictsha", + lfs_dict = True, + ), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + assert next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "blobid", + remote_sibling = patch_hub_gguf.sibling( + "model-Q4_K_M.gguf", + 1000, + None, + blob_id = "blobid", + ), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + assert next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + + +def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): + repo_path = tmp_path / "models--Org--SafeTensorRepo" + repo = SimpleNamespace( + repo_id = "Org/SafeTensorRepo", + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "config.json", + size_on_disk = 10, + blob_path = None, + ), + SimpleNamespace( + file_name = "model.safetensors", + size_on_disk = 100, + blob_path = str(repo_path / "blobs" / "modelsha"), + ), + ] + ) + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + monkeypatch.setattr( + CI.hf_cache_scan, + "is_snapshot_partial", + lambda *args, **kwargs: False, + ) + + rows = CI._scan_cached_models() + + assert len(rows) == 1 + assert rows[0]["repo_id"] == "Org/SafeTensorRepo" + assert rows[0]["model_format"] == "safetensors" + assert rows[0]["size_bytes"] == 100 + + +# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── + + +def test_force_download_bypasses_cache_first_early_return(monkeypatch): + """force_download=True skips the try_to_load_from_cache early-return and + proceeds to the real download path; force_download=False returns the cached + path without ever attempting a download (X2/F2).""" + import huggingface_hub as hf + import utils.hf_xet_fallback as X + + cached_path = "/cache/blob/cached.gguf" + + # Pretend the blob IS cached on disk (try_to_load_from_cache is imported + # inside the function from huggingface_hub, and os.path.exists must agree). + monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False) + monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False) + + attempts = [] + + def fake_attempt(repo_id, filename, token, **kwargs): + attempts.append( + {"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")} + ) + return ("ok", "/freshly/downloaded/path") + + monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True) + + # force_download=False: cache-first early-return, no download attempt. + out = X.hf_hub_download_with_xet_fallback( + "unsloth/repo", "model.gguf", token = None, force_download = False + ) + assert out == cached_path + assert attempts == [] # never reached the real download + + # force_download=True: bypass the early-return, run the real download. + out2 = X.hf_hub_download_with_xet_fallback( + "unsloth/repo", "model.gguf", token = None, force_download = True + ) + assert out2 == "/freshly/downloaded/path" + assert len(attempts) == 1 + assert attempts[0]["force"] is True + + +# ── multi-revision GGUF blob comparison and update reclaim ── +# +# Regression for the phantom "Update available" cue that lingered AFTER a model +# was already updated. A re-download leaves BOTH the old and new revision +# snapshots in the HF cache, so the same gguf file resolves to several blobs. +# The local collection must keep ALL of them (a set per file), and stale hashes +# must be pruned only after the replacement revision verifies. + + +def _rev(*files): + return SimpleNamespace( + files = [SimpleNamespace(file_name = name, blob_path = f"/blobs/{blob}") for name, blob in files] + ) + + +def test_repo_gguf_blob_map_collects_all_revision_blobs(): + """Every cached revision's blob for a gguf file is kept as a set, not + collapsed to one arbitrary blob.""" + repo_info = SimpleNamespace( + revisions = [ + _rev(("lfm2-350m-q4_k_m.gguf", "OLDsha")), + _rev(("lfm2-350m-q4_k_m.gguf", "NEWsha")), + ] + ) + assert CI._repo_gguf_blob_map(repo_info) == {"lfm2-350m-q4_k_m.gguf": {"OLDsha", "NEWsha"}} + + +def test_reclaim_replaced_gguf_variant_prunes_old_revision_only(monkeypatch, tmp_path): + """After a verified update, stale same-variant files/blobs are removed while + the freshly downloaded hash and sibling variants remain cached.""" + repo_id = "org/repo-GGUF" + repo_path = tmp_path / "models--org--repo-GGUF" + old_snap = repo_path / "snapshots" / ("a" * 40) / "model-Q4_K_M.gguf" + new_snap = repo_path / "snapshots" / ("b" * 40) / "model-Q4_K_M.gguf" + sibling_snap = repo_path / "snapshots" / ("b" * 40) / "model-Q8_0.gguf" + old_blob = repo_path / "blobs" / "OLDsha" + new_blob = repo_path / "blobs" / "NEWsha" + sibling_blob = repo_path / "blobs" / "Q8sha" + for path, payload in ( + (old_snap, b"old"), + (new_snap, b"new"), + (sibling_snap, b"sibling"), + (old_blob, b"old-blob"), + (new_blob, b"new-blob"), + (sibling_blob, b"sibling-blob"), + ): + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(payload) + + repo_info = SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + file_path = str(old_snap), + blob_path = str(old_blob), + ) + ] + ), + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + file_path = str(new_snap), + blob_path = str(new_blob), + ), + SimpleNamespace( + file_name = "model-Q8_0.gguf", + file_path = str(sibling_snap), + blob_path = str(sibling_blob), + ), + ] + ), + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo_info])], + ) + invalidated = [] + monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: invalidated.append(True)) + + result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"NEWsha"})) + + assert result["removed_snapshots"] == 1 + assert result["deleted_blobs"] == 1 + assert result["removed_dirs"] == 1 + assert old_snap.exists() is False + assert old_snap.parent.exists() is False + assert old_blob.exists() is False + assert new_snap.exists() is True + assert new_blob.exists() is True + assert sibling_snap.exists() is True + assert sibling_blob.exists() is True + assert invalidated == [True] diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index a6ba69fff..15961ac03 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -161,6 +161,7 @@ def _download_child_entry( repo_type: str, disable_xet: bool, result_queue: Any, + force_download: bool = False, ) -> None: """Spawn-child entrypoint: download one file and report the result. @@ -211,6 +212,7 @@ def _download_child_entry( filename = filename, repo_type = repo_type, token = token, + force_download = force_download, ) result_queue.put({"ok": True, "path": path}) except BaseException as e: # noqa: BLE001 - report every failure to the parent @@ -264,6 +266,7 @@ def _run_download_attempt( interval: float, grace_period: float, on_status: Optional[Callable[[str], None]], + force_download: bool = False, ) -> tuple[str, Optional[str]]: """Run one download in a spawn child supervised by the no-progress watchdog. @@ -280,6 +283,7 @@ def _run_download_attempt( repo_type = repo_type, disable_xet = disable_xet, result_queue = result_queue, + force_download = force_download, ), daemon = True, ) @@ -345,21 +349,28 @@ def hf_hub_download_with_xet_fallback( interval: float = DEFAULT_HEARTBEAT_INTERVAL, grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, + force_download: bool = False, ) -> str: """Download a single file with Xet primary and HTTP as a stall-only fallback. Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if *cancel_event* is set, re-raises a deterministic child error unchanged (no fallback), and raises ``DownloadStallError`` only if BOTH transports stall. + + When *force_download* is True the cache-first early-return is skipped and the + flag is threaded to ``hf_hub_download`` so a newer remote blob is re-fetched + even if an older blob is already cached. """ # Finalized blob already cached: return it with no child and no network. - try: - from huggingface_hub import try_to_load_from_cache - cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) - if isinstance(cached, str) and os.path.exists(cached): - return cached - except Exception as e: - logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) + # Skipped when force_download is set so an update re-fetches a newer blob. + if not force_download: + try: + from huggingface_hub import try_to_load_from_cache + cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) + if isinstance(cached, str) and os.path.exists(cached): + return cached + except Exception as e: + logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) if cancel_event is not None and cancel_event.is_set(): raise RuntimeError("Cancelled") @@ -386,6 +397,7 @@ def hf_hub_download_with_xet_fallback( interval = interval, grace_period = grace_period, on_status = on_status, + force_download = force_download, ) if kind == "ok": diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx index d6afa30e0..4de96d364 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx @@ -1,16 +1,7 @@ // 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 { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, -} from "@/components/ui/alert-dialog"; +import { DeleteConfirmDialog } from "@/features/hub/catalog/download-card"; import { cn } from "@/lib/utils"; import { Delete02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -22,7 +13,6 @@ interface ModelDeleteActionProps { title: string; description: ReactNode; successMessage: string; - loadingLabel?: string; buttonClassName?: string; iconClassName?: string; disabled?: boolean; @@ -35,7 +25,6 @@ export function ModelDeleteAction({ title, description, successMessage, - loadingLabel = "Deleting...", buttonClassName, iconClassName, disabled = false, @@ -85,33 +74,17 @@ export function ModelDeleteAction({ /> - { if (!nextOpen && deleting) return; setOpen(nextOpen); }} - > - - - {title} - {description} - - - No - { - e.preventDefault(); - handleConfirm(); - }} - > - {deleting ? loadingLabel : "Yes"} - - - - + title={title} + description={description} + deleting={deleting} + onConfirm={() => void handleConfirm()} + /> ); } diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx new file mode 100644 index 000000000..d00c81232 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx @@ -0,0 +1,106 @@ +// 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 { subscribeJobListeners } from "@/features/hub/download-manager"; +import { UpdateConfirmDialog } from "@/features/hub/catalog/download-card"; +import { ggufVariantsMatch } from "@/features/hub/lib/model-identity"; +import { cn } from "@/lib/utils"; +import { RefreshCw } from "lucide-react"; +import { useCallback, useEffect, useRef, useState, type ReactNode } from "react"; +import { toast } from "sonner"; + +interface ModelUpdateActionProps { + ariaLabel: string; + title: string; + description: ReactNode; + /** Repo + variant the update targets, so this action can refresh its caller + * (clearing the "update available" cue) when the matching managed download + * completes. `variant` is null for full-model (safetensors / MLX) rows. */ + repoId: string; + variant?: string | null; + buttonClassName?: string; + iconClassName?: string; + disabled?: boolean; + /** Starts the update — which now runs as a managed download. Resolves once the + * download has been handed to the download manager, NOT when it finishes. */ + onConfirm: () => Promise | void; + /** Fired when THIS repo+variant's managed update actually completes. */ + onUpdated?: () => void; +} + +export function ModelUpdateAction({ + ariaLabel, + title, + description, + repoId, + variant = null, + buttonClassName, + iconClassName, + disabled = false, + onConfirm, + onUpdated, +}: ModelUpdateActionProps) { + const [open, setOpen] = useState(false); + + // The update is a managed download (it surfaces in the global Downloads panel + // with progress + cancel). When this exact repo+variant finishes, refresh the + // caller so the "update available" cue clears once the new revision is on + // disk. A ref keeps the subscription stable across renders without resubscribing. + const onUpdatedRef = useRef(onUpdated); + onUpdatedRef.current = onUpdated; + useEffect(() => { + return subscribeJobListeners("model", repoId, { + onComplete: (completedVariant) => { + const matches = variant + ? ggufVariantsMatch(completedVariant, variant) + : !completedVariant; + if (matches) onUpdatedRef.current?.(); + }, + }); + }, [repoId, variant]); + + const handleConfirm = useCallback(() => { + // Start the background re-download and close the dialog immediately; the + // Downloads panel owns progress + cancel from here. Only a failure to START + // surfaces a toast — a failed download reports itself in the panel. + void Promise.resolve() + .then(onConfirm) + .catch((err) => { + toast.error( + err instanceof Error ? err.message : "Failed to start update", + ); + }); + setOpen(false); + }, [onConfirm]); + + return ( + <> + + + + + ); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 4a79904e3..e11a08f8a 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -23,18 +23,19 @@ import { listScanFolders, removeScanFolder, } from "@/features/chat/api/chat-api"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import type { CachedGgufRepo, CachedModelRepo, LocalModelInfo, } from "@/features/chat/api/chat-api"; -import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import type { GgufVariantDetail } from "@/features/chat/types/api"; import { DotTag } from "@/features/hub/catalog/dot-tag"; import { type HubOption, HubOptionMenu, } from "@/features/hub/catalog/hub-option-menu"; +import { TransportConflictDialog } from "@/features/hub/catalog/transport-conflict-dialog"; import { TrainIcon } from "@/features/hub/components/train-icon"; import { useHubInfiniteScroll } from "@/features/hub/hooks/use-hub-infinite-scroll"; import { @@ -46,6 +47,11 @@ import { useOnlineStatus } from "@/features/hub/hooks/use-online-status"; import { isHiddenModelId } from "@/features/hub/lib/hidden-models"; import { classifyUnslothSupport } from "@/features/hub/lib/unsloth-support"; import { useHfTokenStore } from "@/features/hub/stores/hf-token-store"; +import { + downloadManager, + jobKeyOf, + useDownloadManagerStore, +} from "@/features/hub/download-manager"; import { useDebouncedValue, useGpuInfo } from "@/hooks"; import { extractParamLabel } from "@/lib/model-size"; import { toast } from "@/lib/toast"; @@ -85,6 +91,7 @@ import { hasAnyCapability, } from "./model-capabilities"; import { ModelDeleteAction } from "./model-delete-action"; +import { ModelUpdateAction } from "./model-update-action"; import { ModelLoadSettingsAction } from "./model-load-settings-action"; import { type ModelLoadTimes, @@ -616,20 +623,26 @@ function normalizeGgufVariantsResponse(res: { }; } +function ggufVariantExpectedBytes(variant: GgufVariantDetail): number { + const downloadBytes = variant.download_size_bytes; + return typeof downloadBytes === "number" && + Number.isFinite(downloadBytes) && + downloadBytes > 0 + ? downloadBytes + : variant.size_bytes; +} + function GgufVariantExpander({ repoId, onSelect, gpuGb, systemRamGb, + hfToken, parentOptionKey, onNavigatePastStart, onNavigatePastEnd, - onDeleteVariant, sourceOverride, - deleteVariantTitle = "Delete cached model?", - renderDeleteVariantDescription, - getDeleteVariantSuccessMessage, - deleteDisabled = false, + variantActions, onDevice = false, onHasVision, }: { @@ -637,21 +650,42 @@ function GgufVariantExpander({ onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; gpuGb?: number; systemRamGb?: number; + /** HF token threaded into the variant fetch so private/gated repos resolve + * their GGUF variants (and update badges). */ + hfToken?: string; parentOptionKey?: string; onNavigatePastStart?: () => void; onNavigatePastEnd?: () => void; - onDeleteVariant?: (quant: string) => Promise | void; sourceOverride?: ModelSelectorChangeMeta["source"]; - deleteVariantTitle?: string; - renderDeleteVariantDescription?: (quant: string) => ReactNode; - getDeleteVariantSuccessMessage?: (quant: string) => string; - deleteDisabled?: boolean; + /** Update/delete actions for cached variant rows. Omitted by browse-only + * expanders (Recommended, etc.) that don't manage on-disk variants. */ + variantActions?: { + onUpdate?: (quant: string, expectedBytes: number) => Promise | void; + updateTitle?: string; + renderUpdateDescription?: (quant: string) => ReactNode; + getUpdateSuccessMessage?: (quant: string) => string; + updateDisabled?: boolean; + onDelete?: (quant: string) => Promise | void; + deleteTitle?: string; + renderDeleteDescription?: (quant: string) => ReactNode; + getDeleteSuccessMessage?: (quant: string) => string; + deleteDisabled?: boolean; + }; /** On Device rows honor the Show all quantizations setting; Recommended and * other browse lists always show every quant. */ onDevice?: boolean; /** Report GGUF vision support up so the parent row can badge it. */ onHasVision?: (hasVision: boolean) => void; }) { + const onUpdateVariant = variantActions?.onUpdate; + const updateVariantTitle = variantActions?.updateTitle ?? "Update cached model?"; + const renderUpdateVariantDescription = variantActions?.renderUpdateDescription; + const updateDisabled = variantActions?.updateDisabled ?? false; + const onDeleteVariant = variantActions?.onDelete; + const deleteVariantTitle = variantActions?.deleteTitle ?? "Delete cached model?"; + const renderDeleteVariantDescription = variantActions?.renderDeleteDescription; + const getDeleteVariantSuccessMessage = variantActions?.getDeleteSuccessMessage; + const deleteDisabled = variantActions?.deleteDisabled ?? false; const [variants, setVariants] = useState(null); const [defaultVariant, setDefaultVariant] = useState(null); const [hasVision, setHasVision] = useState(false); @@ -659,13 +693,14 @@ function GgufVariantExpander({ const [nativeContext, setNativeContext] = useState(null); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); + const [refreshKey, setRefreshKey] = useState(0); useEffect(() => { let canceled = false; setLoading(true); setError(null); - listGgufVariants(repoId) + listGgufVariants(repoId, hfToken) .then((res) => { if (canceled) return; const normalized = normalizeGgufVariantsResponse(res); @@ -688,7 +723,7 @@ function GgufVariantExpander({ return () => { canceled = true; }; - }, [repoId]); + }, [repoId, refreshKey, hfToken]); // Covers Unix absolute (/), Windows drive (C:\, D:/), UNC (\\server), relative (./, ../), tilde (~/) const isLocalPath = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test( @@ -870,6 +905,7 @@ function GgufVariantExpander({ const fit = getGgufFit(v.size_bytes); const oom = fit === "oom"; const tight = fit === "tight"; + const expectedBytes = ggufVariantExpectedBytes(v); const keyBase = `${repoId}:${v.filename}`; const variantOptionKey = makeModelOptionKey("gguf-variant", keyBase); return ( @@ -878,7 +914,7 @@ function GgufVariantExpander({ type="button" {...variantList.getOptionProps(variantOptionKey, false)} onClick={() => - handleVariantClick(v.quant, v.downloaded, v.size_bytes) + handleVariantClick(v.quant, v.downloaded, expectedBytes) } className={cn( "flex min-w-0 flex-1 items-center justify-between gap-2 rounded-full px-2 py-1 text-left text-sm transition-colors hover:bg-[#ececec] focus-visible:bg-[#ececec] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring/45 dark:hover:bg-[var(--sidebar-accent)] dark:focus-visible:bg-[var(--sidebar-accent)]", @@ -891,9 +927,16 @@ function GgufVariantExpander({ {v.quant} {v.downloaded ? ( - - downloaded - + <> + + downloaded + + {v.update_available ? ( + + update available + + ): null} + ) : v.quant === effectiveRecommended ? ( recommended @@ -916,6 +959,29 @@ function GgufVariantExpander({ + {v.downloaded && v.update_available && onUpdateVariant && ( + + This will update{" "} + + {repoId} ({v.quant}) + {"."} + + ) + } + repoId={repoId} + variant={v.quant} + buttonClassName="p-1" + iconClassName="size-3" + disabled={updateDisabled} + onConfirm={() => onUpdateVariant(v.quant, expectedBytes)} + onUpdated={() => setRefreshKey((key) => key + 1)} + /> + )} {v.downloaded && ( void; }) { const gpu = useGpuInfo(); + // The currently-loaded/running model id. We read params.checkpoint from the + // runtime store (backend-mirrored from /api/inference/status.active_model, see + // chat-runtime-store) rather than the dropdown `isSelected` highlight (which is + // just `value === repo_id` and can reflect a staged, not-yet-loaded pick). Used + // to disable the cached-row update action for the model that's live in memory. + const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint); // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). const loadTimes = useModelLoadTimes(value); // Fade the list's top edge once scrolled, and its bottom edge while more // rows sit below the fold. const [listScrolled, setListScrolled] = useState(false); const [listMoreBelow, setListMoreBelow] = useState(false); + const hfToken = useHfTokenStore((s) => s.token); const [query, setQuery] = useState(""); const debouncedQuery = useDebouncedValue(query); // Shared Hub search stack (the same hooks the Hub page uses) so the picker // and Hub run one implementation. Scoped to unsloth like the old listing. const online = useOnlineStatus(); - const accessToken = useHfTokenStore((s) => s.token) || undefined; + const accessToken = hfToken || undefined; // Recommended section: a live unsloth listing sorted by the dropdown. The // same sort drives the search results so the dropdown works while searching. const [recommendedSort, setRecommendedSort] = @@ -1358,6 +1431,28 @@ export function HubModelPicker({ const alreadyCached = _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0; const [cachedReady, setCachedReady] = useState(alreadyCached); + const [updateConflictKey, setUpdateConflictKey] = useState( + null, + ); + const updateTransportConflict = useDownloadManagerStore((state) => + updateConflictKey + ? (state.conflicts[updateConflictKey]?.info ?? null) + : null, + ); + const cancelUpdateConflict = useCallback(() => { + if (updateConflictKey) downloadManager.cancelConflict(updateConflictKey); + setUpdateConflictKey(null); + }, [updateConflictKey]); + const resumeUpdateConflict = useCallback(() => { + if (!updateConflictKey) return; + downloadManager.resumeConflict(updateConflictKey); + setUpdateConflictKey(null); + }, [updateConflictKey]); + const restartUpdateConflict = useCallback(() => { + if (!updateConflictKey) return; + downloadManager.restartConflict(updateConflictKey); + setUpdateConflictKey(null); + }, [updateConflictKey]); // LM Studio local models -- module-level cache, same pattern as above. const [lmStudioModels, setLmStudioModels] = @@ -1480,14 +1575,42 @@ export function HubModelPicker({ setCachedGguf(v); }) .catch(() => {}); - listCachedModels() + listCachedModels(hfToken || undefined) .then((v) => { _cachedModelsCache = v; setCachedModels(v); }) .catch(() => {}); refreshLocalModelsList(); - }, [refreshLocalModelsList]); + }, [hfToken, refreshLocalModelsList]); + + // Updates run as MANAGED downloads (they show in the global Downloads panel + // with manifest-based progress + a working Cancel), instead of a blocking + // call. The worker re-resolves `main` and pulls only changed blobs, so the + // cached copy stays usable until the new revision lands. The row's + // ModelUpdateAction refreshes the list when this repo+variant completes. + const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => { + return downloadManager + .requestStart({ + kind: "model", + repoId, + variant, + expectedBytes, + }) + .then((outcome) => { + if (outcome === "conflict") { + setUpdateConflictKey(jobKeyOf("model", repoId, variant)); + } else if (outcome === "error") { + throw new Error("Failed to start update"); + } + }); + }, []); + + const updateGgufVariant = useCallback( + (repoId: string, quant: string, expectedBytes: number) => + startManagedUpdate(repoId, quant, expectedBytes), + [startManagedUpdate], + ); useEffect(() => { // Always refresh LM Studio + custom folder models (not gated by alreadyCached). @@ -1512,14 +1635,14 @@ export function HubModelPicker({ }) .catch(() => {}) .finally(check); - listCachedModels() + listCachedModels(hfToken || undefined) .then((v) => { _cachedModelsCache = v; setCachedModels(v); }) .catch(() => {}) .finally(check); - }, [refreshLocalModelsList, refreshScanFolders]); + }, [hfToken, refreshLocalModelsList, refreshScanFolders]); // Hide downloaded models from the recommended list. Case-insensitive // since the HF cache lowercases repo IDs. @@ -2317,14 +2440,21 @@ export function HubModelPicker({ onDevice={true} onHasVision={(v) => reportVision(c.repo_id, v)} onSelect={onSelect} + hfToken={hfToken || undefined} parentOptionKey={optionKey} onNavigatePastStart={() => hubModelList.focusOption(optionKey)} onNavigatePastEnd={() => hubModelList.moveFocus(optionKey, "next")} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={gpu.systemRamAvailableGb || undefined} - onDeleteVariant={async (quant) => { - await deleteCachedModel(c.repo_id, quant); - refreshCachedLists(); + variantActions={{ + onUpdate: (quant, expectedBytes) => + updateGgufVariant(c.repo_id, quant, expectedBytes), + // Can't update the model that's live in memory under itself. + updateDisabled: loadedModelId === c.repo_id, + onDelete: async (quant) => { + await deleteCachedModel(c.repo_id, quant); + refreshCachedLists(); + }, }} /> )} @@ -2383,7 +2513,8 @@ export function HubModelPicker({ }; return ( -
+ <> +
{/* A small right inset shortens the search bar so Search Hub lands on the last dropdown's right edge (none on the wider Connected box). */}
hubModelList.focusOption(optionKey) @@ -3250,9 +3382,11 @@ export function HubModelPicker({ gpu.available ? gpu.memoryTotalGb : undefined } systemRamGb={gpu.systemRamAvailableGb || undefined} - onDeleteVariant={async (quant) => { - await deleteCachedModel(id, quant); - refreshCachedLists(); + variantActions={{ + onDelete: async (quant) => { + await deleteCachedModel(id, quant); + refreshCachedLists(); + }, }} /> )} @@ -3324,6 +3458,7 @@ export function HubModelPicker({ hubModelList.focusOption(optionKey) @@ -3335,9 +3470,11 @@ export function HubModelPicker({ gpu.available ? gpu.memoryTotalGb : undefined } systemRamGb={gpu.systemRamAvailableGb || undefined} - onDeleteVariant={async (quant) => { - await deleteCachedModel(id, quant); - refreshCachedLists(); + variantActions={{ + onDelete: async (quant) => { + await deleteCachedModel(id, quant); + refreshCachedLists(); + }, }} /> )} @@ -3411,6 +3548,7 @@ export function HubModelPicker({ hubModelList.focusOption(optionKey) @@ -3422,9 +3560,11 @@ export function HubModelPicker({ gpu.available ? gpu.memoryTotalGb : undefined } systemRamGb={gpu.systemRamAvailableGb || undefined} - onDeleteVariant={async (quant) => { - await deleteCachedModel(id, quant); - refreshCachedLists(); + variantActions={{ + onDelete: async (quant) => { + await deleteCachedModel(id, quant); + refreshCachedLists(); + }, }} /> )} @@ -3459,7 +3599,14 @@ export function HubModelPicker({
) : null} -
+
+ + ); } @@ -3611,22 +3758,21 @@ function FineTunedRows({ gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={gpu.systemRamAvailableGb || undefined} sourceOverride={isExportedGguf ? "exported" : undefined} - deleteVariantTitle="Delete exported GGUF variant?" - renderDeleteVariantDescription={(quant) => ( - <> - This will remove{" "} - - {adapter.name} ({quant}) - {" "} - from disk. This cannot be undone. - - )} - getDeleteVariantSuccessMessage={(quant) => - `Deleted ${adapter.name} ${quant}` - } - deleteDisabled={deleteDisabled} - onDeleteVariant={ - isExportedGguf + variantActions={{ + deleteTitle: "Delete exported GGUF variant?", + renderDeleteDescription: (quant) => ( + <> + This will remove{" "} + + {adapter.name} ({quant}) + {" "} + from disk. This cannot be undone. + + ), + getDeleteSuccessMessage: (quant) => + `Deleted ${adapter.name} ${quant}`, + deleteDisabled: deleteDisabled, + onDelete: isExportedGguf ? async (quant) => { await deleteFineTunedModel({ modelPath: adapter.id, @@ -3639,8 +3785,8 @@ function FineTunedRows({ ggufVariant: quant, }); } - : undefined - } + : undefined, + }} /> )} diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 15066caa8..8195d8952 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { authFetch } from "@/features/auth"; +import { hubTokenHeader } from "@/features/hub/lib/hub-token-header"; import { consumeNativePathToken } from "@/features/native-intents/api"; import { formatFastApiDetail } from "@/lib/format-fastapi-error"; import type { @@ -313,8 +314,12 @@ export interface CachedModelRepo { last_modified?: number; } -export async function listCachedModels(): Promise { - const response = await authFetch("/api/models/cached-models"); +export async function listCachedModels( + hfToken?: string | null, +): Promise { + const response = await authFetch("/api/models/cached-models", { + headers: hubTokenHeader(hfToken), + }); const data = await parseJsonOrThrow<{ cached: CachedModelRepo[] }>(response); return data.cached; } @@ -778,8 +783,9 @@ export async function listGgufVariants( hfToken?: string, ): Promise { const params = new URLSearchParams({ repo_id: repoId }); - if (hfToken) params.set("hf_token", hfToken); - const response = await authFetch(`/api/models/gguf-variants?${params}`); + const response = await authFetch(`/api/models/gguf-variants?${params}`, { + headers: hubTokenHeader(hfToken), + }); return parseJsonOrThrow(response); } diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 412f35d1a..94a7e1c56 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -84,7 +84,9 @@ export interface GgufVariantDetail { filename: string; quant: string; size_bytes: number; + download_size_bytes?: number; downloaded?: boolean; + update_available?: boolean; } export interface GgufVariantsResponse { diff --git a/studio/frontend/src/features/hub/catalog/download-card.tsx b/studio/frontend/src/features/hub/catalog/download-card.tsx index 5083fb2f8..17177936a 100644 --- a/studio/frontend/src/features/hub/catalog/download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/download-card.tsx @@ -19,7 +19,11 @@ import { TooltipTrigger, } from "@/components/ui/tooltip"; import { cn } from "@/lib/utils"; -import { Delete02Icon, Download01Icon } from "@hugeicons/core-free-icons"; +import { + ArrowReloadHorizontalIcon, + Delete02Icon, + Download01Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { DownloadProgressBar, @@ -111,6 +115,73 @@ export function CardDeleteButton({ ); } +export function CardUpdateButton({ + label, + onClick, + emphasized = false, +}: { + label: string; + onClick: () => void; + /** When true (a newer revision is available) the control becomes a prominent + * labeled amber pill instead of the quiet hover-revealed icon — the + * "update available" cue. Amber is the established status tone in this surface + * (the older-cache hint banner), kept tinted (never solid) per the design + * system's feedback-color convention, so the one emerald accent stays scarce. */ + emphasized?: boolean; +}) { + if (emphasized) { + return ( + + + + + + A newer version is available on Hugging Face + + + ); + } + return ( + + + + + + Update from Hugging Face + + + ); +} + /** Download / Cancel / Resume button for the safetensors and dataset cards. */ export function DownloadActionButton({ downloading, @@ -213,3 +284,44 @@ export function DeleteConfirmDialog({ ); } + +/** Confirmation dialog shared by the model and quantization update flows. */ +export function UpdateConfirmDialog({ + open, + onOpenChange, + title, + description, + updating, + onConfirm, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + title: string; + description: ReactNode; + updating: boolean; + onConfirm: () => void; +}) { + return ( + + + + {title} + {description} + + + Cancel + { + e.preventDefault(); + onConfirm(); + }} + > + {updating ? "Updating…" : "Update"} + + + + + ); +} diff --git a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx index d089835c3..9e4f6593b 100644 --- a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx @@ -17,7 +17,10 @@ import { useDownloadManagerStore, useRepoDownload, } from "../download-manager"; -import { type GgufVariantDetail, deleteCachedModel } from "../inventory"; +import { + type GgufVariantDetail, + deleteCachedModel, +} from "../inventory"; import { formatBytes } from "../lib/format"; import { type GgufFitClass, classifyGgufFit } from "../lib/gguf-fit"; import { HUB_GGUF_RUN_ACTIONS_VISIBLE } from "../lib/hub-feature-flags"; @@ -27,7 +30,9 @@ import { } from "../lib/model-identity"; import { cn } from "@/lib/utils"; import { useHfTokenStore } from "../stores/hf-token-store"; +import { useOnlineStatus } from "../hooks/use-online-status"; import { + ArrowReloadHorizontalIcon, Delete02Icon, Download01Icon, InformationCircleIcon, @@ -56,6 +61,7 @@ import { CardDivider, DeleteConfirmDialog, DownloadCard, + UpdateConfirmDialog, } from "./download-card"; import { activeDownloadState, @@ -160,9 +166,13 @@ function QuantBadge({ {quant} ) : ( + // Trigger quant label is the row's primary identity and is short + // (e.g. "Q4_K_M"); keep it `shrink-0` + `whitespace-nowrap` so it never + // collapses to "q…" when the Update/Run actions crowd the row. The info + // group's `overflow-hidden` sacrifices the trailing status tags instead. @@ -173,7 +183,7 @@ function QuantBadge({ className={cn("size-3.5 shrink-0", meta.iconClassName)} /> )} - {quant} + {quant} ); if (!showFit || tooltipMode === "none") return inner; @@ -288,6 +298,11 @@ const GgufVariantMenuRow = memo(function GgufVariantMenuRow({ : "hover:bg-foreground/[0.05] dark:hover:bg-foreground/[0.06]", )} > + {/* Status (On device / Loaded / Partial) sits beside the quant on the + left so the model's identity reads as one unit; only the size pins + right. No per-row "GGUF" tag: every row here is a GGUF quant and the + trigger already labels it, so repeating it only stole the room the + quant label needs (it would otherwise truncate to "q…"). */} - - {item.downloaded && ( )} @@ -319,7 +332,8 @@ const GgufVariantMenuRow = memo(function GgufVariantMenuRow({ )} - + + {item.downloadSizeLabel} @@ -377,6 +391,7 @@ export function GgufDownloadCard({ onChange?: () => void; }) { const hfToken = useHfTokenStore((s) => s.token); + const online = useOnlineStatus(); const localVariantPath = cachePath?.trim() || null; const { variants, loading, error, refreshError, refresh } = useGgufVariantFetchState({ @@ -394,6 +409,7 @@ export function GgufDownloadCard({ selectedQuantState.repoId === repoId ? selectedQuantState.quant : null; const [open, setOpen] = useState(false); const [deleteTarget, setDeleteTarget] = useState(null); + const [updateTarget, setUpdateTarget] = useState(null); const [completedVariantKeys, setCompletedVariantKeys] = useState< ReadonlySet >(() => new Set()); @@ -416,7 +432,7 @@ export function GgufDownloadCard({ if (completedVariantKeys.size === 0) return withLive; return withLive.map((v) => completedVariantKeys.has(normalizeGgufVariantIdentity(v.quant)) - ? { ...v, downloaded: true, partial: false } + ? { ...v, downloaded: true, partial: false, update_available: false } : v, ); }, [completedVariantKeys, liveVariantStates, rawSortedVariants]); @@ -444,6 +460,7 @@ export function GgufDownloadCard({ setCompletedVariantKeys((prev) => prev.has(key) ? prev : new Set(prev).add(key), ); + void refresh(); }, }); const progress = job.progress; @@ -529,6 +546,8 @@ export function GgufDownloadCard({ const selectedDownloadSizeLabel = selected ? formatBytes(ggufVariantDownloadSizeBytes(selected)) : null; + const updateAvailable = + selected?.downloaded === true && selected.update_available === true; const selectedVariantKey = selectedQuant ? normalizeGgufVariantIdentity(selectedQuant) : null; @@ -587,6 +606,35 @@ export function GgufDownloadCard({ setDeleteTarget(null); }, }); + const updateTargetVariant = + updateTarget && sortedVariants + ? sortedVariants.find((v) => ggufVariantsMatch(v.quant, updateTarget)) + : null; + const updateTargetLabel = updateTargetVariant + ? ggufVariantDisplayLabel(updateTargetVariant) + : updateTarget; + // Confirm → close the dialog and run the re-download as a MANAGED download, so + // it surfaces in the "Downloading N items" panel with correct manifest-based + // progress and a working Cancel — the same UX as any other download — instead + // of a bespoke modal/toast. The worker re-resolves `main` and pulls only the + // changed blobs, so the cached version stays intact (and runnable) until the + // new revision lands. Completion refreshes the variant list, whose metadata + // carries the "Update available" cue. + const handleConfirmUpdate = useCallback(() => { + if (!updateTarget) return; + const variant = updateTarget; + const expectedBytes = + updateTargetVariant?.download_size_bytes ?? + updateTargetVariant?.size_bytes ?? + 0; + setUpdateTarget(null); + void downloadManager.requestStart({ + kind: "model", + repoId, + variant, + expectedBytes, + }); + }, [updateTarget, updateTargetVariant, repoId]); const variantListUnavailable = !sortedVariants || sortedVariants.length === 0; const showVariantLoadingState = loading && variantListUnavailable; @@ -640,24 +688,44 @@ export function GgufDownloadCard({ job={job} progress={downloadingThisVariant ? progress : null} dialogs={ - { - if (!o && !deleting) setDeleteTarget(null); - }} - title="Delete quantization?" - deleting={deleting} - onConfirm={() => void runDelete()} - description={ - <> - This will remove{" "} - - {repoId} ({deleteTargetLabel}) - {" "} - from disk. You can re-download it later. - - } - /> + <> + { + if (!o && !deleting) setDeleteTarget(null); + }} + title="Delete quantization?" + deleting={deleting} + onConfirm={() => void runDelete()} + description={ + <> + This will remove{" "} + + {repoId} ({deleteTargetLabel}) + {" "} + from disk. You can re-download it later. + + } + /> + { + if (!o) setUpdateTarget(null); + }} + title="Update quantization?" + updating={false} + onConfirm={handleConfirmUpdate} + description={ + <> + This will re-download the latest version of{" "} + + {repoId} ({updateTargetLabel}) + {" "} + from Hugging Face. + + } + /> + } > @@ -668,21 +736,24 @@ export function GgufDownloadCard({ e.preventDefault(); setOpen((o) => !o); }} - className="hub-menu-trigger flex h-9 min-w-0 flex-1 cursor-pointer items-center gap-2.5 rounded-full px-3 text-left transition-colors hover:bg-foreground/[0.04] data-[state=open]:bg-foreground/[0.06] dark:hover:bg-white/[0.04] dark:data-[state=open]:bg-white/[0.06]" + className="hub-menu-trigger flex h-9 min-w-0 flex-1 cursor-pointer items-center gap-2 rounded-full px-3 text-left transition-colors hover:bg-foreground/[0.04] data-[state=open]:bg-foreground/[0.06] dark:hover:bg-white/[0.04] dark:data-[state=open]:bg-white/[0.06]" > - {selected ? ( - - ) : ( - - Select quantization - - )} - + {/* Quant label + status tags travel together as one left-aligned + group so the fit-info icon never floats orphaned from its tags; + only the chevron pins right, the standard select affordance. */} + + {selected ? ( + + ) : ( + + Select quantization + + )} {selected?.downloaded && ( + {selectedDownloadSizeLabel} )} - +
{variantMenuItems.map((item) => { @@ -760,6 +831,25 @@ export function GgufDownloadCard({ {!isGgufRunCta && } + {selected?.downloaded && + online && + updateAvailable && + !selectedIsActive && + !downloadingThisVariant && ( + + )} +