diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 528c059fbc..8163caa639 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -432,6 +432,14 @@ def render_native_template( if native_tpl is None: # A LoRA adapter's native template lives on the base model, not the adapter id. template_source = model_info.get("base_model") or active_model_name + # A local-only (background) load resolved its weights to a local + # snapshot; the template reload must read the SAME files. For a + # non-LoRA model prefer the stored load path over the repo id, and + # either way pass local_files_only so a cache miss fails the fallback + # instead of downloading tokenizer files mid-generation. + local_files_only = bool(model_info.get("local_files_only", False)) + if local_files_only and not model_info.get("base_model"): + template_source = model_info.get("model_path") or template_source # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can # instantiate its class (the stored flag already covers template_source). trust_remote_code = bool(model_info.get("trust_remote_code", False)) @@ -441,6 +449,7 @@ def render_native_template( template_source, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, + local_files_only = local_files_only, ) native_tpl = nt.chat_template or False except Exception as exc: diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 563a6732a1..439a286f37 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -322,6 +322,7 @@ class InferenceBackend: hf_token: Optional[str] = None, trust_remote_code: bool = False, gpu_ids: Optional[list[int]] = None, + local_files_only: bool = False, ) -> bool: """Load any model: base, LoRA adapter, text, or vision.""" # Keep the token so the native-template fallback can fetch a @@ -364,6 +365,9 @@ class InferenceBackend: "trust_remote_code": trust_remote_code, "is_vision": config.is_vision, "is_lora": config.is_lora, + # Local-only loads: generation-time repo fallbacks (native + # template reload) must also resolve from cache, not the Hub. + "local_files_only": local_files_only, "is_audio": config.is_audio, "audio_type": config.audio_type, "has_audio_input": config.has_audio_input, @@ -562,7 +566,11 @@ class InferenceBackend: isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor") ): # LoRA adapters: use base model. Local merged exports: read base from export_metadata.json. - processor_source = config.base_model if config.is_lora else config.identifier + # Non-LoRA: config.path, not config.identifier. They are the + # same for ordinary loads, but a local-only load rewrites + # path to the cached snapshot and this fallback must stay + # on those local files instead of refetching by repo id. + processor_source = config.base_model if config.is_lora else config.path if not config.is_lora and config.is_local: _meta_path = Path(config.path) / "export_metadata.json" try: @@ -578,10 +586,14 @@ class InferenceBackend: ) from transformers import AutoProcessor + # Local-only loads: a LoRA base or export_metadata base_model + # is a Hub repo id; resolve it from cache or fail the load + # (candidate failover) instead of downloading the processor. processor = AutoProcessor.from_pretrained( processor_source, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, + local_files_only = local_files_only, ) logger.info(f"Loaded {type(processor).__name__} from {processor_source}") diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index ff966e8446..113a786db3 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -515,29 +515,84 @@ def _hf_env_offline() -> bool: return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} +def _hub_offline_env_truthy() -> bool: + """HF_HUB_OFFLINE specifically is truthy. huggingface_hub does not honor + TRANSFORMERS_OFFLINE, so a forced local-only guard may only no-op when the + Hub flag itself is already set.""" + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + +# Overlapping offline guards share one env override: the count tracks active +# guards and the saved values are restored only when the LAST guard exits, so +# a request finishing early cannot re-enable network for one still running. +_OFFLINE_GUARD_LOCK = threading.Lock() +_OFFLINE_GUARD_STATE: dict = {"count": 0, "hub_prev": None, "transformers_prev": None} + + @contextlib.contextmanager -def _hf_offline_if_dns_dead(): +def _hf_offline_if_dns_dead(force: bool = False): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; restores env on exit so a transient hiccup can't quarantine the process. - No-op if the user already set it.""" - if "HF_HUB_OFFLINE" in os.environ: + No-op when the user already set it to a truthy value. ``force`` skips the + DNS probe and goes offline unconditionally (local-only background loads + resolve metadata from the cache without any network), overriding even an + explicitly falsy HF_HUB_OFFLINE=0 for the block and restoring it after. + Guards are refcounted so overlapping loads/validations each keep offline + until the last one exits.""" + entered = False # this guard joined or created the env override + owner = False # this guard created it (first in) + with _OFFLINE_GUARD_LOCK: + if _OFFLINE_GUARD_STATE["count"] > 0: + if force: + # Join the active override so the env survives until every + # local-only guard exits, whichever request finishes first. + _OFFLINE_GUARD_STATE["count"] += 1 + entered = True + # An ordinary guard never joins: its block tolerates either env + # state (it would have run online), so extending the forced + # window would only widen the exposure of concurrent online + # work to the process-global override. It no-ops instead. + elif _hf_env_offline() and (not force or _hub_offline_env_truthy()): + # User-set truthy env (count is 0): already offline, nothing to + # arrange or restore. A forced guard only trusts HF_HUB_OFFLINE + # itself: TRANSFORMERS_OFFLINE=1 alone leaves hub API calls (e.g. + # get_paths_info) online, so force still installs both flags. + pass + elif not force and "HF_HUB_OFFLINE" in os.environ: + # A user-pinned falsy value stays authoritative for ordinary loads. + pass + elif force or _probe_dns_dead(): + _OFFLINE_GUARD_STATE["hub_prev"] = os.environ.get("HF_HUB_OFFLINE") + _OFFLINE_GUARD_STATE["transformers_prev"] = os.environ.get("TRANSFORMERS_OFFLINE") + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + _OFFLINE_GUARD_STATE["count"] = 1 + entered = True + owner = True + if not entered: yield False return - if not _probe_dns_dead(): - yield False - return - - transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ - os.environ["HF_HUB_OFFLINE"] = "1" - if not transformers_was_set: - os.environ["TRANSFORMERS_OFFLINE"] = "1" - logger.warning("huggingface.co unreachable; using local HF cache for this load.") + if owner: + if force: + logger.info("Local-only load: forcing HF offline for this block.") + else: + logger.warning("huggingface.co unreachable; using local HF cache for this load.") try: yield True finally: - os.environ.pop("HF_HUB_OFFLINE", None) - if not transformers_was_set: - os.environ.pop("TRANSFORMERS_OFFLINE", None) + with _OFFLINE_GUARD_LOCK: + _OFFLINE_GUARD_STATE["count"] -= 1 + if _OFFLINE_GUARD_STATE["count"] == 0: + for key, prev in ( + ("HF_HUB_OFFLINE", _OFFLINE_GUARD_STATE["hub_prev"]), + ("TRANSFORMERS_OFFLINE", _OFFLINE_GUARD_STATE["transformers_prev"]), + ): + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + _OFFLINE_GUARD_STATE["hub_prev"] = None + _OFFLINE_GUARD_STATE["transformers_prev"] = None try: @@ -5273,6 +5328,7 @@ class LlamaCppBackend: force: bool = False, allow_smaller_fallback: bool = True, cancel_event: Optional[threading.Event] = None, + local_files_only: bool = False, ) -> str: """Download GGUF file(s) from HuggingFace. Returns local path. @@ -5310,16 +5366,19 @@ class LlamaCppBackend: gguf_filename = None gguf_extra_shards: list[str] = [] if hf_variant: - try: - from huggingface_hub import list_repo_files + # Local-only loads resolve from the cache alone: the live listing + # is both network and unnecessary for files already on disk. + if not local_files_only: + try: + from huggingface_hub import list_repo_files - files = list_repo_files(hf_repo, token = hf_token) - gguf_files = _gguf_files_for_variant(files, hf_variant) - if gguf_files: - gguf_filename = gguf_files[0] - gguf_extra_shards = _gguf_extra_shards(gguf_files, gguf_filename) - except Exception as e: - logger.warning(f"Could not list repo files: {e}") + files = list_repo_files(hf_repo, token = hf_token) + gguf_files = _gguf_files_for_variant(files, hf_variant) + if gguf_files: + gguf_filename = gguf_files[0] + gguf_extra_shards = _gguf_extra_shards(gguf_files, gguf_filename) + except Exception as e: + logger.warning(f"Could not list repo files: {e}") # Fall back to the local cache when the repo listing is unavailable. if not gguf_filename: @@ -5343,10 +5402,16 @@ class LlamaCppBackend: # Resolve by variant so a newer revision's filename does not hide # the complete older copy. Size-check against that older snapshot's # own revision when its metadata remains available. + # Local-only loads skip the remote size check outright: the + # hub API behind it (get_paths_info) performs no offline-mode + # check and HF_HUB_OFFLINE is baked into hub constants at + # import, so only a per-call skip reliably keeps this off the + # network. A truncated cache then fails the load into + # candidate failover instead of being caught up front. cached_main = cached_gguf_for_load( hf_repo, hf_variant, - verify_sizes = True, + verify_sizes = not local_files_only, hf_token = hf_token, ) else: @@ -5354,13 +5419,24 @@ class LlamaCppBackend: cached_main = ( candidate[0] if candidate is not None - and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + and ( + local_files_only + or _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + ) else None ) if cached_main is not None: logger.info(f"Reusing cached GGUF: {cached_main}") return cached_main + # A local-only load reaching this point has no complete cached copy; + # fail closed instead of downloading (background loads never fetch). + if local_files_only: + raise RuntimeError( + f"GGUF '{hf_repo}' ({hf_variant or 'default'}) is not fully " + "available in the local cache; select it explicitly to download it." + ) + # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards try: @@ -5523,6 +5599,7 @@ class LlamaCppBackend: label: str, cancel_event: Optional[threading.Event] = None, near_path: Optional[str] = None, + local_files_only: bool = False, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. @@ -5544,6 +5621,12 @@ class LlamaCppBackend: logger.info("Reusing cached %s: %s", label, cached) return cached + # Background loads never download: without a cached copy the model + # runs without the optional companion instead of fetching it. + if local_files_only: + logger.info("Skipping %s fetch (local-only load)", label) + return None + from utils.hf_cache_settings import get_hf_cache_paths companion_cache_dir = _hub_cache_dir_for_snapshot_path(near_path) or str( @@ -5632,6 +5715,7 @@ class LlamaCppBackend: hf_token: Optional[str] = None, cancel_event: Optional[threading.Event] = None, near_path: Optional[str] = None, + local_files_only: bool = False, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. @@ -5648,6 +5732,7 @@ class LlamaCppBackend: label = "mmproj", cancel_event = cancel_event, near_path = near_path, + local_files_only = local_files_only, ) def _cached_repo_mtp_drafter( @@ -5689,6 +5774,7 @@ class LlamaCppBackend: hf_repo: str, hf_token: Optional[str] = None, near_path: Optional[str] = None, + local_files_only: bool = False, ) -> Optional[str]: """Download the separate MTP drafter (speculative head) from a GGUF repo. @@ -5722,7 +5808,7 @@ class LlamaCppBackend: # fetched). Online, _download_companion_gguf/hf_hub_download reuse the # current cached file and refetch a changed one, so skip the probe here # rather than pair new weights with a stale draft. - if _hf_env_offline(): + if local_files_only or _hf_env_offline(): cached = self._cached_repo_mtp_drafter( hf_repo, cache_dir = _hub_cache_dir_for_snapshot_path(near_path), @@ -5737,6 +5823,7 @@ class LlamaCppBackend: pick = _pick_mtp, label = "MTP drafter", near_path = near_path, + local_files_only = local_files_only, ) def _resolve_launch_mmproj_path( @@ -6433,6 +6520,9 @@ class LlamaCppBackend: hf_repo: Optional[str] = None, hf_variant: Optional[str] = None, hf_token: Optional[str] = None, + # Background auto-loads: resolve every file (main quant, mmproj, MTP + # drafter) from the local cache only and never download. + local_files_only: bool = False, # Common model_identifier: str, is_vision: bool = False, @@ -6470,6 +6560,7 @@ class LlamaCppBackend: "gguf_path": gguf_path, "mmproj_path": mmproj_path, "mtp_draft_path": mtp_draft_path, + "local_files_only": local_files_only, "hf_repo": hf_repo, "hf_variant": hf_variant, "hf_token": hf_token, @@ -6578,11 +6669,15 @@ class LlamaCppBackend: hf_repo, ) hf_repo = _resolved_repo - with _hf_offline_if_dns_dead(): + # Same local-only policy as the Phase 2 download below: this + # preflight runs FIRST, so without the flag it would be the + # download path a background load slips through. + with _hf_offline_if_dns_dead(force = local_files_only): _preflight_model_path = self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, hf_token = hf_token, + local_files_only = local_files_only, ) self._reject_vulkan_diffusion_gpu_ids_before_teardown( _preflight_model_path, @@ -6620,11 +6715,14 @@ class LlamaCppBackend: hf_repo, ) hf_repo = _resolved_repo - with _hf_offline_if_dns_dead(): + # Forced offline under local-only so even the cache-size + # verification (get_paths_info) stays off the network. + with _hf_offline_if_dns_dead(force = local_files_only): model_path = _preflight_model_path or self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, hf_token = hf_token, + local_files_only = local_files_only, ) # Auto-download mmproj for vision models unless opted out. if is_vision and not mmproj_path and not extra_args_disable_mmproj(extra_args): @@ -6632,6 +6730,7 @@ class LlamaCppBackend: hf_repo = hf_repo, hf_token = hf_token, near_path = model_path, + local_files_only = local_files_only, ) # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head @@ -6650,6 +6749,7 @@ class LlamaCppBackend: hf_repo = hf_repo, hf_token = hf_token, near_path = model_path, + local_files_only = local_files_only, ) elif gguf_path: if not Path(gguf_path).is_file(): diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index d19c67a01a..b3372026b7 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -524,6 +524,9 @@ class MLXInferenceBackend: dtype = None, parallel_mode = None, distributed_group = None, + # Accepted for worker parity; the load runs under the worker's forced + # offline env when set, which is what enforces local-only here. + local_files_only: bool = False, ) -> bool: import mlx.core as mx @@ -603,8 +606,13 @@ class MLXInferenceBackend: else: load_kwargs["tensor_group"] = distributed_group + # Registry identity stays the repo id (model_name); the LOAD source + # honors config.path so a route-resolved local snapshot (local-only + # loads against a moved live cache) is read instead of re-resolving + # the id through the import-time cache location. + load_source = getattr(config, "path", None) or model_name model, tokenizer_or_processor = FastMLXModel.from_pretrained( - model_name, + load_source, **load_kwargs, ) @@ -627,6 +635,10 @@ class MLXInferenceBackend: "hf_token": hf_token, # Per-model trust_remote_code reused by the native-template reload (matches transformers). "trust_remote_code": trust_remote_code, + # Local-only loads: generation-time repo fallbacks (native + # template reload) must also resolve from cache, not the Hub. + "local_files_only": local_files_only, + "model_path": load_source, "model": self._model, "tokenizer": self._tokenizer, "processor": self._processor, diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 616384386d..4a213d591e 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -973,6 +973,7 @@ class InferenceOrchestrator: subject: Optional[str] = None, tensor_parallel: bool = False, mlx_distributed: bool = False, + local_files_only: bool = False, ) -> bool: """Load a model for inference. @@ -985,7 +986,24 @@ class InferenceOrchestrator: self.loading_models.add(model_name) try: - needed_major = "5" if needs_transformers_5(model_name) else "4" + # Parent-side metadata preflight stays offline for local-only + # loads: the tier probe reads model configs and GPU sizing calls + # hf model_info. The guard closes BEFORE the spawn below, so the + # worker does not inherit HF_HUB_OFFLINE for its whole lifetime + # (its own bootstrap guard scopes the load and restores for + # generation-time fetches). + import contextlib + + from core.inference.llama_cpp import _hf_offline_if_dns_dead + + def _preflight_offline(): + # Ordinary loads keep their current (unguarded) behavior. + if local_files_only: + return _hf_offline_if_dns_dead(force = True) + return contextlib.nullcontext() + + with _preflight_offline(): + needed_major = "5" if needs_transformers_5(model_name) else "4" # Build config dict for subprocess sub_config = { @@ -1000,16 +1018,25 @@ class InferenceOrchestrator: "gpu_ids": gpu_ids, "tensor_parallel": bool(tensor_parallel), "mlx_distributed": bool(mlx_distributed), + "local_files_only": bool(local_files_only), + # Route-resolved local snapshot: the worker rebuilds its + # ModelConfig from model_name, which would lose the rewrite. + "local_snapshot_path": ( + getattr(config, "path", None) + if local_files_only and getattr(config, "path", None) != model_name + else None + ), "mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline") if mlx_distributed else None, } - resolved_gpu_ids, gpu_selection = prepare_gpu_selection( - gpu_ids, - model_name = model_name, - hf_token = hf_token, - load_in_4bit = load_in_4bit, - ) + with _preflight_offline(): + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + gpu_ids, + model_name = model_name, + hf_token = hf_token, + load_in_4bit = load_in_4bit, + ) sub_config["resolved_gpu_ids"] = resolved_gpu_ids sub_config["gpu_selection"] = gpu_selection # Parent-detected backend for the worker's apply_gpu_ids(). diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 254eda40a3..c188295d28 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker. from __future__ import annotations import base64 +import contextlib import json from loggers import get_logger import os @@ -95,6 +96,28 @@ def _clean_token(value: str | None) -> str | None: return value if value and value.strip() else None +@contextlib.contextmanager +def _local_only_offline_env(enabled: bool): + """Force HF offline for a local-only background load. The worker is a + fresh per-load process, so scoping env here covers config resolution, + security gates, and every from_pretrained in the load, then restores it so + generation-time fetches (e.g. the chat template fallback) still work.""" + if not enabled: + yield + return + prev = {key: os.environ.get(key) for key in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")} + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + try: + yield + finally: + for key, value in prev.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + def _build_model_config(config: dict): """Build a ModelConfig from the config dict.""" from utils.models import ModelConfig @@ -107,6 +130,12 @@ def _build_model_config(config: dict): ) if not mc: raise ValueError(f"Invalid model identifier: {model_name}") + # The route resolves a cached repo id to its local snapshot for local-only + # loads; from_identifier here rebuilds path as the repo id, losing that + # rewrite across the process boundary, so re-apply it. + snapshot_override = config.get("local_snapshot_path") + if snapshot_override: + mc.path = snapshot_override return mc @@ -175,19 +204,48 @@ def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool: return load_in_4bit -def _ensure_ssm_kernels(targets: list, resp_queue: Any) -> bool: +def _ensure_ssm_kernels( + targets: list, + resp_queue: Any, + local_files_only: bool = False, +) -> bool: """Install the SSM kernels the given model(s) lazy-import in from_pretrained; no-op for non-SSM models, idempotent. Returns True on success; on a fatal mamba-ssm failure sends a 'loaded' failure response and returns False. Call BEFORE importing transformers, which snapshots its optional-backend gates at import (a later install may not be picked up). + Under ``local_files_only`` nothing is ever installed: kernels already present are used, + a missing fatal kernel fails the load into candidate failover, and the optional + causal-conv1d fast path is skipped (its torch fallback covers it). """ try: - from utils.ssm_runtime import ensure_ssm_runtime + from utils.ssm_runtime import ensure_ssm_runtime, model_is_ssm except Exception as exc: logger.debug("ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc) return True _ssm_status = lambda m: _send_response(resp_queue, {"type": "status", "message": m}) + if local_files_only: + import importlib.util + for ssm_target in dict.fromkeys(t for t in targets if t): + try: + needs_mamba = model_is_ssm(ssm_target) + except Exception: + needs_mamba = False + if needs_mamba and importlib.util.find_spec("mamba_ssm") is None: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + "This model needs the mamba-ssm kernel, which is not " + "installed; select the model explicitly to install it." + ), + "error_kind": "ssm_runtime_install_failed", + }, + ) + return False + return True try: for ssm_target in dict.fromkeys(t for t in targets if t): ensure_ssm_runtime(ssm_target, status_cb = _ssm_status) @@ -285,7 +343,13 @@ def _run_security_gates( def _handle_load(backend, config: dict, resp_queue: Any) -> None: """Handle a load command: load a model into the backend.""" + _offline_guard = contextlib.ExitStack() try: + # Local-only loads run the WHOLE load offline: config resolution, + # security gates, and from_pretrained can otherwise all reach the Hub. + _offline_guard.enter_context( + _local_only_offline_env(bool(config.get("local_files_only", False))) + ) mc = _build_model_config(config) hf_token = _clean_token(config.get("hf_token")) @@ -335,7 +399,11 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: str(mc.base_model) if (mc.is_lora and getattr(mc, "base_model", None)) else None ) ssm_targets = [ssm_probe_identifier(config["model_name"], _ssm_base)] - if not _ensure_ssm_kernels(ssm_targets, resp_queue): + if not _ensure_ssm_kernels( + ssm_targets, + resp_queue, + local_files_only = bool(config.get("local_files_only", False)), + ): return # Heartbeat keeps the orchestrator's inactivity deadline alive during slow @@ -363,6 +431,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: "hf_token": hf_token, "trust_remote_code": trust_remote_code, "gpu_ids": config.get("resolved_gpu_ids"), + "local_files_only": bool(config.get("local_files_only", False)), } if getattr(backend, "device", None) == "mlx": load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode") @@ -435,6 +504,8 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: "stack": traceback.format_exc(limit = 20), }, ) + finally: + _offline_guard.close() def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool: @@ -796,6 +867,16 @@ def run_inference_process( apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) + # Local-only background loads keep the ENTIRE bootstrap offline: base + # resolution, transformers activation, security gates, kernel probes and + # the initial load can all reach the Hub otherwise. Closed before the + # command loop so generation-time fetches (e.g. the chat template + # fallback) still work; error-path returns end the process anyway. + _bootstrap_offline = contextlib.ExitStack() + _bootstrap_offline.enter_context( + _local_only_offline_env(bool(config.get("local_files_only", False))) + ) + model_name = config["model_name"] # ── 0. MLX fast-path — skip torch/transformers ── @@ -857,6 +938,7 @@ def run_inference_process( return # Enter the same command loop as the GPU path. + _bootstrap_offline.close() logger.info("MLX inference subprocess ready, entering command loop") while True: try: @@ -1015,7 +1097,11 @@ def run_inference_process( from utils.ssm_runtime import ssm_probe_identifier _ssm_targets = [ssm_probe_identifier(model_name, _base)] - if not _ensure_ssm_kernels(_ssm_targets, resp_queue): + if not _ensure_ssm_kernels( + _ssm_targets, + resp_queue, + local_files_only = bool(config.get("local_files_only", False)), + ): return # ── 2. Import ML libraries (fresh in this clean process) ── @@ -1077,6 +1163,8 @@ def run_inference_process( ) return + _bootstrap_offline.close() + # ── 4. Command loop — process commands until shutdown ── # cancel_event is an mp.Event the parent can set anytime to cancel # generation instantly (no queue polling needed). diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index ca0f4658a3..581d6dc90d 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -188,6 +188,13 @@ class CachedModelRepo(CachedRepoBase): pipeline_tag: Optional[str] = None library_name: Optional[str] = None tags: Optional[List[str]] = None + snapshot_size_bytes: Optional[int] = Field( + None, + description = ( + "Weight bytes of the newest cached snapshot only (what a load " + "resolves); size_bytes sums blobs across every cached revision." + ), + ) class CachedModelsResponse(BaseModel): diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 807ec70991..3155e36608 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -327,7 +327,7 @@ def _scan_cached_gguf() -> list[dict]: continue repo_id = repo_info.repo_id repo_path = Path(repo_info.repo_path) - snapshot_path = _cached_model_snapshot_path(repo_path) + snapshot_path = _cached_gguf_repo_snapshot_path(repo_path) total_size = _repo_gguf_size_bytes(repo_info) has_variant_state, variant_state_size = _gguf_variant_state_summary( repo_id, @@ -415,16 +415,35 @@ async def list_cached_gguf_response(hf_token: Optional[str] = None): class _CachedNonGgufPayload(NamedTuple): size_bytes: int + snapshot_size_bytes: int has_runnable_weights: bool model_format: ModelFormat last_modified: float +def _snapshot_dir_mtime(revision) -> float: + """mtime of a revision's snapshot dir; the same signal latest_snapshot_dir + (and therefore the load-side snapshot resolution) selects by.""" + snapshot_path = getattr(revision, "snapshot_path", None) + if not snapshot_path: + return 0.0 + try: + return float(Path(snapshot_path).stat().st_mtime) + except OSError: + return 0.0 + + def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: all_weight_blobs: dict[str, tuple[int, float]] = {} adapter_blobs: dict[str, tuple[int, float]] = {} safetensors_blobs: dict[str, tuple[int, float]] = {} checkpoint_blobs: dict[str, tuple[int, float]] = {} + # Per-revision selected-format byte sums, so the row can also report the + # size of the ONE snapshot a load resolves (newest by snapshot-dir mtime) + # rather than only the all-revisions total. + rev_category_sizes: dict[str, dict[str, int]] = {} + rev_snapshot_mtimes: dict[str, float] = {} + rev_has_config: dict[str, bool] = {} has_config = False has_adapter_config = False has_adapter_weights = False @@ -433,7 +452,7 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: has_checkpoint = False def _record_blob( - target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str + target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str, category: str ) -> None: blob_path = getattr(file_obj, "blob_path", None) size = int(file_obj.size_on_disk or 0) @@ -441,9 +460,13 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: value = (size, _blob_mtime(file_obj)) target[key] = value all_weight_blobs[key] = value + per_rev = rev_category_sizes.setdefault(rev_id, {}) + per_rev[category] = per_rev.get(category, 0) + size + per_rev["all"] = per_rev.get("all", 0) + size for revision in repo_info.revisions: rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) + rev_snapshot_mtimes[rev_id] = _snapshot_dir_mtime(revision) for f in revision.files: file_name = str(f.file_name) lower = file_name.lower() @@ -452,6 +475,7 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: continue if name == "config.json": has_config = True + rev_has_config[rev_id] = True continue if name == "adapter_config.json": has_adapter_config = True @@ -461,15 +485,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: is_checkpoint = _is_checkpoint_weight_name(name) if is_adapter: has_adapter_weights = True - _record_blob(adapter_blobs, f, rev_id, file_name) + _record_blob(adapter_blobs, f, rev_id, file_name, "adapter") if is_safetensors: has_safetensors = True if _is_transformers_safetensors_weight_name(name): has_transformers_safetensors = True - _record_blob(safetensors_blobs, f, rev_id, file_name) + _record_blob(safetensors_blobs, f, rev_id, file_name, "safetensors") if is_checkpoint: has_checkpoint = True - _record_blob(checkpoint_blobs, f, rev_id, file_name) + _record_blob(checkpoint_blobs, f, rev_id, file_name, "checkpoint") model_format = ( _classify_non_gguf_model_format( @@ -485,22 +509,134 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: ) if model_format == "adapter": selected_blobs = adapter_blobs + selected_category = "adapter" elif model_format == "safetensors": selected_blobs = safetensors_blobs + selected_category = "safetensors" elif model_format == "checkpoint": selected_blobs = checkpoint_blobs + selected_category = "checkpoint" else: selected_blobs = all_weight_blobs + selected_category = "all" + + size_bytes = sum(size for size, _mtime in selected_blobs.values()) + # The one snapshot a load resolves (newest snapshot-dir mtime) among the + # revisions holding selected-format weights. For safetensors rows the + # snapshot resolvers additionally require config.json in the SAME + # revision, so sizing must use that predicate too: a weight-only newer + # revision would otherwise be sized while the older complete revision is + # what actually loads. Falls back to weight-only revisions, then to the + # all-revisions total. + weight_revs = [ + rev_id + for rev_id, sizes in rev_category_sizes.items() + if sizes.get(selected_category, 0) > 0 + ] + if selected_category == "safetensors": + complete_revs = [rev_id for rev_id in weight_revs if rev_has_config.get(rev_id, False)] + if complete_revs: + weight_revs = complete_revs + if weight_revs: + newest_rev = max(weight_revs, key = lambda rev_id: rev_snapshot_mtimes.get(rev_id, 0.0)) + snapshot_size_bytes = rev_category_sizes[newest_rev].get(selected_category, 0) + else: + snapshot_size_bytes = size_bytes return _CachedNonGgufPayload( - size_bytes = sum(size for size, _mtime in selected_blobs.values()), + size_bytes = size_bytes, + snapshot_size_bytes = snapshot_size_bytes, has_runnable_weights = model_format != "unknown", model_format = model_format, last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0), ) +def _newest_snapshot_where(repo_path: Path, loadable) -> Optional[Path]: + """Newest snapshots/* dir (by mtime) whose entry names satisfy *loadable*. + + Inactive-cache rows emit the selected snapshot path as their load_id, + which the load consumes directly (it bypasses the repo-id resolver), so + the snapshot must actually hold the format the row was classified from. + """ + snapshots = repo_path / "snapshots" + try: + revisions = [entry for entry in snapshots.iterdir() if entry.is_dir()] + except OSError: + return None + + def _ok(rev: Path) -> bool: + try: + names = [entry.name for entry in rev.iterdir()] + except OSError: + return False + return loadable(names) + + candidates = [rev for rev in revisions if _ok(rev)] + if not candidates: + return None + try: + return max(candidates, key = lambda rev: rev.stat().st_mtime).resolve() + except OSError: + return None + + +def _weightful_snapshot_path(repo_path: Path) -> Optional[Path]: + """Newest snapshot holding config.json plus a safetensors weight (the + non-GGUF resolvers' complete-revision predicate).""" + return _newest_snapshot_where( + repo_path, + lambda names: "config.json" in names and any(n.endswith(".safetensors") for n in names), + ) + + +def _gguf_snapshot_path(repo_path: Path) -> Optional[Path]: + """Newest snapshot holding a GGUF file (top level or one folder deep, + where multi-quant repos keep per-quant subfolders).""" + direct = _newest_snapshot_where( + repo_path, + lambda names: any(_is_gguf_filename(n.lower()) for n in names), + ) + if direct is not None: + return direct + snapshots = repo_path / "snapshots" + try: + revisions = [entry for entry in snapshots.iterdir() if entry.is_dir()] + except OSError: + return None + + def _has_nested_gguf(rev: Path) -> bool: + try: + return any(True for _ in rev.glob("*/*.gguf")) + except OSError: + return False + + candidates = [rev for rev in revisions if _has_nested_gguf(rev)] + if not candidates: + return None + try: + return max(candidates, key = lambda rev: rev.stat().st_mtime).resolve() + except OSError: + return None + + def _cached_model_snapshot_path(repo_path: Path) -> Optional[Path]: + weightful = _weightful_snapshot_path(repo_path) + if weightful is not None: + return weightful + resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_path) + if not resolved: + return None + path = Path(resolved) + return path if path.is_dir() else None + + +def _cached_gguf_repo_snapshot_path(repo_path: Path) -> Optional[Path]: + """Snapshot path for a GGUF row: a safetensors-bearing revision of a + mixed repo must not become the GGUF row's load target.""" + gguf_snapshot = _gguf_snapshot_path(repo_path) + if gguf_snapshot is not None: + return gguf_snapshot resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_path) if not resolved: return None @@ -636,6 +772,7 @@ def _scan_cached_models() -> list[dict]: row = { "repo_id": repo_id, "size_bytes": payload.size_bytes, + "snapshot_size_bytes": payload.snapshot_size_bytes, "cache_path": str(repo_info.repo_path), "partial": snapshot_partial, "partial_transport": ( diff --git a/studio/backend/hub/utils/local_snapshot.py b/studio/backend/hub/utils/local_snapshot.py new file mode 100644 index 0000000000..33c8d2e59e --- /dev/null +++ b/studio/backend/hub/utils/local_snapshot.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local-only snapshot resolution for background model loads.""" + +import os +from typing import Optional + + +def _snapshot_has_weights(rev: str) -> bool: + """Whether a snapshot dir holds at least one safetensors weight file. + + Background loads only target safetensors-format rows (adapters and + pickle checkpoints are excluded upstream), so this is the runnable-weight + signal that made the inventory row eligible. + """ + try: + return any(name.endswith(".safetensors") for name in os.listdir(rev)) + except OSError: + return False + + +def _snapshot_dir_fallback(repo_id: str, cache_dir: Optional[str]) -> Optional[str]: + """Newest snapshots/* dir (by mtime) holding a config.json, preferring + revisions that also hold weights. + + This mirrors the inventory scanner's selection (newest revision that + actually carries the inventoried weights), so the load targets the + snapshot that made the row eligible: a newest metadata-only revision must + not shadow an older complete one. It also covers revision-only layouts + (refs/ missing or pruned) that snapshot_download(local_files_only = True) + cannot map through refs/main. + """ + if cache_dir is None: + try: + from huggingface_hub.constants import HF_HUB_CACHE + cache_dir = HF_HUB_CACHE + except Exception: + return None + folder = os.path.join(cache_dir, "models--" + repo_id.replace("/", "--"), "snapshots") + try: + revisions = [ + os.path.join(folder, name) + for name in os.listdir(folder) + if os.path.isdir(os.path.join(folder, name)) + ] + except OSError: + return None + candidates = [rev for rev in revisions if os.path.isfile(os.path.join(rev, "config.json"))] + if not candidates: + return None + weightful = [rev for rev in candidates if _snapshot_has_weights(rev)] + return max(weightful or candidates, key = os.path.getmtime) + + +def resolve_local_snapshot_path( + repo_id: str, + hf_token: Optional[str] = None, + cache_dir: Optional[str] = None, +) -> Optional[str]: + """Resolve a Hub repo id to its snapshot directory in the local HF cache + without any network access; None when the repo is not cached. + + The newest snapshot dir (by mtime) holding a config.json is preferred: + that is the same selection the inventory scanner surfaces, so the load + targets the revision that made the row eligible. refs/main can lag it + when a non-main commit was downloaded later, so consulting refs first + could load an older revision or fail on its missing files. + ``snapshot_download(local_files_only = True)`` stays as the fallback for + layouts the directory scan cannot interpret. Either way resolution never + touches the network, and a missing-file snapshot still resolves so the + subsequent weight load fails instead of downloading the gaps, which is + the fail-closed behavior background loads need. + """ + resolved = _snapshot_dir_fallback(repo_id, cache_dir) + if resolved is not None: + return resolved + try: + from huggingface_hub import snapshot_download + return snapshot_download( + repo_id, + local_files_only = True, + token = hf_token or None, + cache_dir = cache_dir or None, + ) + except Exception: + return None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index add3228a28..29230beb02 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -40,6 +40,13 @@ class LoadRequest(BaseModel): gguf_variant: Optional[str] = Field( None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" ) + local_files_only: bool = Field( + False, + description = ( + "Resolve a cached repo id against the local snapshot only and " + "never download missing files (background auto-loads)." + ), + ) trust_remote_code: bool = Field( False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", @@ -249,6 +256,14 @@ class ValidateModelRequest(BaseModel): "delegate fitting to llama.cpp, while explicit layers are user-owned." ), ) + local_files_only: bool = Field( + False, + description = ( + "Background auto-loads: resolve model metadata from the local " + "cache only and reject candidates whose runtime would pull " + "remote auxiliaries (e.g. audio codec models)." + ), + ) include_context_length: bool = Field( False, description = "Also read the native context length from the local GGUF header. " diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1e6959472c..474b2a5f03 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5045,8 +5045,9 @@ async def _load_model_impl( # is_lora auto-detected from adapter_config.json on disk/HF. # DNS-probe wrap so offline loads skip 30-60s of soft-failed network - # checks before the worker starts. - with _hf_offline_if_dns_dead(): + # checks before the worker starts. Local-only loads force offline so + # metadata resolution itself cannot reach the Hub. + with _hf_offline_if_dns_dead(force = request.local_files_only): config = ModelConfig.from_identifier( model_id = model_identifier, hf_token = request.hf_token, @@ -5059,6 +5060,51 @@ async def _load_model_impl( detail = f"Invalid model identifier: {model_log_label}", ) + # A background auto-load promises on-device-only resolution, but a + # cache populated outside Studio can pass the partial check while + # missing shards, and from_pretrained on a repo id would download the + # gaps. Rewriting the load path to the locally resolved snapshot keeps + # the registry identity (config.identifier) intact while forcing the + # weight load to the files actually on disk, so an incomplete cache + # fails over to the next candidate instead of fetching. + from utils.paths import is_local_path + + # Python audio runtimes fetch codec auxiliaries (SNAC/DAC/BiCodec repos) + # at load time regardless of where the weights live, so a local-only + # load of a non-GGUF audio model still hits the network. Refuse it. + if request.local_files_only and not config.is_gguf and getattr(config, "is_audio", False): + raise HTTPException( + status_code = 409, + detail = ( + f"Model '{model_log_label}' needs audio codec downloads; " + "select it explicitly to load it." + ), + ) + + if request.local_files_only and not config.is_gguf and not is_local_path(config.path): + from hub.utils.local_snapshot import resolve_local_snapshot_path + from utils.hf_cache_settings import get_hf_cache_paths + + # Resolve against the LIVE Studio-managed hub cache: Studio can + # move the cache at runtime without rewriting Hugging Face's + # import-time env constants, so the helper's default would search + # the stale location and 409 models the inventory just found. + local_snapshot = await asyncio.to_thread( + resolve_local_snapshot_path, + config.path, + request.hf_token, + str(get_hf_cache_paths().hub_cache), + ) + if local_snapshot is None: + raise HTTPException( + status_code = 409, + detail = ( + f"Model '{model_log_label}' is not available on device; " + "select it explicitly to download it." + ), + ) + config.path = local_snapshot + # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None @@ -5089,7 +5135,13 @@ async def _load_model_impl( # to match. Off-loop: tier resolution reads configs. if effective_load_in_4bit and not config.is_gguf: from utils.transformers_version import latest_tier_active_for - if await asyncio.to_thread(latest_tier_active_for, config.identifier, request.hf_token): + + # Local-only loads probe the tier against the resolved snapshot: + # a repo id would reach _remote_lora_base's raw HTTP request and + # Hub config reads, while a local path resolves from config.json + # on disk (non-canonical ids skip the remote adapter probe). + _tier_target = config.path if request.local_files_only else config.identifier + if await asyncio.to_thread(latest_tier_active_for, _tier_target, request.hf_token): effective_load_in_4bit = False logger.info( f"Latest-transformers sidecar active for '{model_log_label}' - " @@ -5111,10 +5163,17 @@ async def _load_model_impl( # Apply the training coexistence policy before the unload step below # frees the resident model. Off-loop: the default-mode guard does sync work. + # Local-only non-GGUF loads size against the resolved snapshot so the + # guard's memory estimation reads local files instead of hf model_info + # (which performs no offline-mode check). GGUF sizing already reads + # the cached file under its own local-only handling. + _guard_identifier = ( + config.path if request.local_files_only and not config.is_gguf else model_identifier + ) await asyncio.to_thread( _guard_chat_load_against_training, config, - model_identifier = model_identifier, + model_identifier = _guard_identifier, hf_token = request.hf_token, load_in_4bit = effective_load_in_4bit, max_seq_length = request.max_seq_length, @@ -5201,6 +5260,10 @@ async def _load_model_impl( hf_repo = config.gguf_hf_repo, hf_variant = config.gguf_variant, hf_token = request.hf_token, + # Background auto-loads never download: main quant resolves + # cache-only and the optional mmproj/MTP companions are + # cached-or-skipped instead of fetched. + local_files_only = request.local_files_only, ) else: # Local mode: llama-server loads via -m @@ -5414,6 +5477,9 @@ async def _load_model_impl( approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, gpu_ids = effective_gpu_ids, subject = current_subject, + # Threads the rewritten snapshot path and forced-offline load into + # the worker subprocess, which rebuilds its own ModelConfig. + local_files_only = request.local_files_only, ) if not success: @@ -5652,10 +5718,19 @@ async def validate_model( native_grant_backed = False model_log_label = request.model_path + # Local-only validation (background auto-loads) covers the WHOLE metadata + # and security preflight with a forced-offline env, not just the identifier + # probe: the upgrade, remote-code, and file-security helpers below all read + # Hub configs and would otherwise refetch them for a cache-only candidate. + import contextlib + + _local_only_offline = contextlib.ExitStack() try: model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "validate-model") ) + if request.local_files_only: + _local_only_offline.enter_context(_hf_offline_if_dns_dead(force = True)) config = ModelConfig.from_identifier( model_id = model_identifier, hf_token = request.hf_token, @@ -5668,6 +5743,22 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) + # Same audio gate as /load: non-GGUF audio models pull codec repos at + # load time, so a local-only candidate is rejected here before the + # frontend burns a load attempt on it. + if ( + request.local_files_only + and not getattr(config, "is_gguf", False) + and getattr(config, "is_audio", False) + ): + raise HTTPException( + status_code = 409, + detail = ( + f"Model '{model_log_label}' needs audio codec downloads; " + "select it explicitly to load it." + ), + ) + # Apply the same training coexistence policy as /load before the frontend # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None @@ -5903,6 +5994,8 @@ async def validate_model( status_code = 400, detail = "Invalid model", ) + finally: + _local_only_offline.close() # studio_router only: admin action, kept off the OpenAI-compatible /v1 mount. diff --git a/studio/backend/tests/test_local_snapshot_resolution.py b/studio/backend/tests/test_local_snapshot_resolution.py new file mode 100644 index 0000000000..3aeed0ffb7 --- /dev/null +++ b/studio/backend/tests/test_local_snapshot_resolution.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local-only snapshot resolution for background auto-loads. + +A cache populated outside Studio (no download manifest) passes the partial +check while missing shard files, and ``from_pretrained`` on a repo id would +download the gaps. Background loads therefore rewrite the load path to the +LOCALLY resolved snapshot: resolution never touches the network, an uncached +repo resolves to None (409 upstream), and an incomplete snapshot still +resolves so the weight load fails on the missing files instead of fetching +them. No GPU or network required. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +pytest.importorskip("huggingface_hub") + +from hub.utils.local_snapshot import resolve_local_snapshot_path + + +_REV = "0123456789abcdef0123456789abcdef01234567" + + +def _build_cached_repo( + cache_dir: Path, + repo_id: str, + files: dict[str, str], + with_refs: bool = True, + rev: str = _REV, +) -> Path: + """Lay out a minimal HF hub cache entry the way huggingface_hub expects: + ``models--org--name/refs/main`` pointing at a snapshot directory. + ``with_refs = False`` builds the revision-only layout (pruned or foreign + caches) the inventory scanner accepts.""" + repo_dir = cache_dir / f"models--{repo_id.replace('/', '--')}" + snapshot = repo_dir / "snapshots" / rev + snapshot.mkdir(parents = True) + if with_refs: + (repo_dir / "refs").mkdir(exist_ok = True) + (repo_dir / "refs" / "main").write_text(rev) + for name, content in files.items(): + (snapshot / name).write_text(content) + return snapshot + + +def test_cached_repo_resolves_to_its_snapshot_dir(tmp_path): + snapshot = _build_cached_repo( + tmp_path, + "org/tiny-model", + {"config.json": "{}", "model.safetensors": "weights"}, + ) + resolved = resolve_local_snapshot_path("org/tiny-model", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == snapshot.resolve() + + +def test_incomplete_snapshot_still_resolves_locally(tmp_path): + """Missing shards must not block resolution: the local path is what makes + the subsequent weight load fail closed instead of downloading.""" + snapshot = _build_cached_repo( + tmp_path, + "org/half-downloaded", + { + "config.json": "{}", + "model-00001-of-00002.safetensors": "first shard only", + }, + ) + resolved = resolve_local_snapshot_path("org/half-downloaded", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == snapshot.resolve() + + +def test_uncached_repo_resolves_to_none(tmp_path): + assert resolve_local_snapshot_path("org/never-downloaded", cache_dir = str(tmp_path)) is None + + +def test_newest_snapshot_preferred_over_refs_main(tmp_path): + """A newer snapshot downloaded at an explicit revision outranks the older + refs/main target: the inventory surfaces the newest snapshot by mtime, so + the load must resolve the same one instead of an older (possibly + incomplete) main revision.""" + import os + import time + + old_main = _build_cached_repo( + tmp_path, + "org/newer-rev", + {"config.json": "{}"}, + rev = "a" * 40, + ) + stale = time.time() - 1000 + os.utime(old_main, (stale, stale)) + newer = _build_cached_repo( + tmp_path, + "org/newer-rev", + {"config.json": "{}", "model.safetensors": "weights"}, + with_refs = False, + rev = "b" * 40, + ) + resolved = resolve_local_snapshot_path("org/newer-rev", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == newer.resolve() + + +def test_revision_only_snapshot_resolves_without_refs(tmp_path): + """snapshot_download(local_files_only = True) needs refs/main, but the + inventory scanner accepts revision-only layouts (pruned refs), so the + resolver must fall back to the snapshot directory itself.""" + snapshot = _build_cached_repo( + tmp_path, + "org/no-refs", + {"config.json": "{}", "model.safetensors": "weights"}, + with_refs = False, + ) + resolved = resolve_local_snapshot_path("org/no-refs", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == snapshot.resolve() + + +def test_refless_fallback_picks_newest_snapshot_with_config(tmp_path): + """With several revision dirs, the fallback must pick the newest one that + actually holds a config.json, skipping empty or partial revisions.""" + import os + import time + + old = _build_cached_repo( + tmp_path, + "org/multi-rev", + {"config.json": "{}"}, + with_refs = False, + rev = "a" * 40, + ) + stale = time.time() - 1000 + os.utime(old, (stale, stale)) + new = _build_cached_repo( + tmp_path, + "org/multi-rev", + {"config.json": "{}"}, + with_refs = False, + rev = "b" * 40, + ) + configless = _build_cached_repo( + tmp_path, + "org/multi-rev", + {"tokenizer.json": "{}"}, + with_refs = False, + rev = "c" * 40, + ) + assert configless.exists() + resolved = resolve_local_snapshot_path("org/multi-rev", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == new.resolve() + + +def test_gguf_rows_select_gguf_bearing_snapshot_in_mixed_repos(tmp_path): + """A mixed repo caching a newer safetensors revision beside an older GGUF + revision: the GGUF row's snapshot selection must return the GGUF-bearing + revision, not the safetensors one the model row prefers.""" + import os + import sys + import time + + backend_dir = str(Path(__file__).resolve().parent.parent) + if backend_dir not in sys.path: + sys.path.insert(0, backend_dir) + from hub.services.models.cache_inventory import ( + _cached_gguf_repo_snapshot_path, + _cached_model_snapshot_path, + ) + + repo_dir = tmp_path / "models--org--mixed" + gguf_rev = repo_dir / "snapshots" / ("a" * 40) + gguf_rev.mkdir(parents = True) + (gguf_rev / "mixed-Q4_K_M.gguf").write_text("gguf-bytes") + stale = time.time() - 1000 + os.utime(gguf_rev, (stale, stale)) + st_rev = repo_dir / "snapshots" / ("b" * 40) + st_rev.mkdir(parents = True) + (st_rev / "config.json").write_text("{}") + (st_rev / "model.safetensors").write_text("weights") + + gguf_pick = _cached_gguf_repo_snapshot_path(repo_dir) + assert gguf_pick is not None + assert Path(gguf_pick).resolve() == gguf_rev.resolve() + model_pick = _cached_model_snapshot_path(repo_dir) + assert model_pick is not None + assert Path(model_pick).resolve() == st_rev.resolve() + + +def test_weightless_newest_snapshot_does_not_shadow_complete_older_one(tmp_path): + """A newest metadata-only revision (config.json, no weights) must not win + over an older revision holding the inventoried safetensors weights: the + inventory made the row eligible from the weightful revision, so the load + must resolve that one instead of failing on the weightless dir.""" + import os + import time + + complete = _build_cached_repo( + tmp_path, + "org/meta-newest", + {"config.json": "{}", "model.safetensors": "weights"}, + with_refs = False, + rev = "a" * 40, + ) + stale = time.time() - 1000 + os.utime(complete, (stale, stale)) + _build_cached_repo( + tmp_path, + "org/meta-newest", + {"config.json": "{}"}, + with_refs = False, + rev = "b" * 40, + ) + resolved = resolve_local_snapshot_path("org/meta-newest", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == complete.resolve() + + +def test_refless_fallback_without_config_resolves_to_none(tmp_path): + """A snapshots dir with no config.json anywhere is not a loadable text + model cache; resolution must stay None (409 upstream), not guess.""" + _build_cached_repo( + tmp_path, + "org/no-config", + {"tokenizer.json": "{}"}, + with_refs = False, + ) + assert resolve_local_snapshot_path("org/no-config", cache_dir = str(tmp_path)) is None + + +def test_resolution_never_uses_the_network(tmp_path, monkeypatch): + """local_files_only resolution must not open any connection even when the + repo is absent (the tempting fallback would be a Hub metadata call).""" + import socket + + def _no_network(*_args, **_kwargs): + raise AssertionError("network access attempted during local resolution") + + monkeypatch.setattr(socket.socket, "connect", _no_network) + _build_cached_repo(tmp_path, "org/offline-ok", {"config.json": "{}"}) + assert resolve_local_snapshot_path("org/offline-ok", cache_dir = str(tmp_path)) + assert resolve_local_snapshot_path("org/absent", cache_dir = str(tmp_path)) is None diff --git a/studio/backend/tests/test_offline_guard_refcount.py b/studio/backend/tests/test_offline_guard_refcount.py new file mode 100644 index 0000000000..fedadd87d7 --- /dev/null +++ b/studio/backend/tests/test_offline_guard_refcount.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Refcounted HF offline guard semantics. + +Overlapping local-only validations/loads share one process-global env +override. The refcount means a request finishing first cannot restore the +environment while another local-only request still runs (which would let its +remaining metadata checks reach the Hub), and forced mode overrides an +explicitly falsy HF_HUB_OFFLINE=0 then restores it. The guard is extracted +from source and exercised with stubbed logging/DNS so no ML dependencies are +needed. +""" + +from __future__ import annotations + +import contextlib +import os +import threading +from pathlib import Path + +import pytest + +_LLAMA_CPP = Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + + +class _NullLogger: + def info(self, *args, **kwargs): + pass + + def warning(self, *args, **kwargs): + pass + + +def _load_guard(dns_dead: bool = False): + src = _LLAMA_CPP.read_text() + start = src.index("def _hub_offline_env_truthy") + end = src.index("_SLOT_SAVE_MAX_BYTES") + end = src.rindex("try:", start, end) + block = src[start:end] + ns = { + "threading": threading, + "contextlib": contextlib, + "os": os, + "logger": _NullLogger(), + "_hf_env_offline": lambda: os.environ.get("HF_HUB_OFFLINE", "").strip().lower() + in {"1", "true", "yes", "on"}, + "_probe_dns_dead": lambda: dns_dead, + } + exec(block, ns) + return ns["_hf_offline_if_dns_dead"] + + +@pytest.fixture +def clean_env(monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + + +def test_overlapping_guards_restore_only_after_last_exit(clean_env): + guard = _load_guard() + a = guard(force = True) + b = guard(force = True) + assert a.__enter__() is True + assert os.environ.get("HF_HUB_OFFLINE") == "1" + assert b.__enter__() is True + a.__exit__(None, None, None) + assert ( + os.environ.get("HF_HUB_OFFLINE") == "1" + ), "first exit must not restore while another guard is active" + b.__exit__(None, None, None) + assert "HF_HUB_OFFLINE" not in os.environ + + +def test_force_overrides_and_restores_falsy_env(clean_env): + guard = _load_guard() + os.environ["HF_HUB_OFFLINE"] = "0" + g = guard(force = True) + assert g.__enter__() is True + assert os.environ["HF_HUB_OFFLINE"] == "1" + g.__exit__(None, None, None) + assert os.environ["HF_HUB_OFFLINE"] == "0" + + +def test_falsy_env_stays_authoritative_for_ordinary_loads(clean_env): + guard = _load_guard(dns_dead = True) + os.environ["HF_HUB_OFFLINE"] = "0" + g = guard(force = False) + assert g.__enter__() is False + assert os.environ["HF_HUB_OFFLINE"] == "0" + g.__exit__(None, None, None) + + +def test_truthy_user_env_is_a_noop(clean_env): + guard = _load_guard() + os.environ["HF_HUB_OFFLINE"] = "1" + g = guard(force = True) + assert g.__enter__() is False + g.__exit__(None, None, None) + assert os.environ["HF_HUB_OFFLINE"] == "1" + + +def test_transformers_only_env_does_not_satisfy_forced_guard(clean_env): + """TRANSFORMERS_OFFLINE=1 alone is not hub-offline: huggingface_hub + ignores it, so a forced guard must still install HF_HUB_OFFLINE for the + block and restore the prior state after.""" + guard = _load_guard() + os.environ["TRANSFORMERS_OFFLINE"] = "1" + g = guard(force = True) + assert g.__enter__() is True + assert os.environ.get("HF_HUB_OFFLINE") == "1" + g.__exit__(None, None, None) + assert "HF_HUB_OFFLINE" not in os.environ + assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" + + +def test_nonforce_never_extends_forced_window(clean_env): + """An ordinary (non-force) guard entering while a forced override is + active must no-op, not join: its block tolerates either env state, and + joining would only widen the exposure of concurrent online work to the + process-global override.""" + guard = _load_guard() + forced = guard(force = True) + plain = guard(force = False) + assert forced.__enter__() is True + assert plain.__enter__() is False + forced.__exit__(None, None, None) + assert ( + "HF_HUB_OFFLINE" not in os.environ + ), "the forced owner's exit restores; the ordinary no-op guard holds nothing" + plain.__exit__(None, None, None) + assert "HF_HUB_OFFLINE" not in os.environ + + +def test_forced_guards_share_windows_with_dns_dead_owners(clean_env): + """A forced guard joining a DNS-dead override keeps the env until the + forced guard itself exits, even when the owner exits first.""" + guard = _load_guard(dns_dead = True) + plain = guard(force = False) + forced = guard(force = True) + assert plain.__enter__() is True + assert forced.__enter__() is True + plain.__exit__(None, None, None) + assert os.environ.get("HF_HUB_OFFLINE") == "1" + forced.__exit__(None, None, None) + assert "HF_HUB_OFFLINE" not in os.environ diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index fb9331ecc2..bb186cde70 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2,7 +2,21 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { getAuthToken } from "@/features/auth"; +import { + type CachedGgufRepo, + type CachedModelRepo, + type LocalModelInfo, + listCachedGguf, + listCachedModels, + listLocalModels, +} from "@/features/hub/inventory/api"; +import { + ensureHiddenModelMatchers, + isHiddenModelId, +} from "@/features/hub/lib/hidden-models"; import { resolveInitialConfig } from "@/features/model-picker"; +import { isMlxId } from "@/features/model-picker/components/model-selector/recommended-fit"; +import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { projectHasSources } from "@/features/rag/api/rag-api"; import { apiUrl } from "@/lib/api-base"; import { parseParamCountB } from "@/lib/model-size"; @@ -43,6 +57,7 @@ import { type PendingImageEditReference, type RagAutoInject, GPU_LAYERS_AUTO, + isLocalModelPath, loadedGpuMemoryFields, reconcilePersistedGpuIds, resolveLoadedSpeculativeSettings, @@ -79,9 +94,12 @@ import { updateStoredChatThread, } from "../utils/chat-history-storage"; import { + isManagedCacheSource, readLastLocalModelLoad, recordLastLocalModelLoad, type LastLocalModelKind, + type LastLocalModelLoad, + type LastLocalModelSource, } from "../utils/last-local-model-load"; import { getImageInputUnavailableReason } from "../utils/image-input-support"; import { @@ -92,8 +110,6 @@ import { resolveLoadMaxSeqLength } from "../presets/preset-policy"; import { generateAudio, GenerationLengthError, - listCachedGguf, - listCachedModels, listGgufVariants, loadModel, streamChatCompletions, @@ -1427,11 +1443,6 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { }); } -/** - * Auto-load the smallest downloaded model when the user chats without - * selecting one. Prefers GGUF (smallest cached variant), then smallest - * cached safetensors model. - */ // Cap cascade so broken cached repos can't spam /api/inference/load. const MAX_AUTO_LOAD_ATTEMPTS = 3; const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi; @@ -1445,6 +1456,8 @@ type AutoLoadCandidate = { ggufVariant: string | null; maxSeqLength: number; successLabel: string; + inventoryId?: string | null; + source: LastLocalModelSource; }; function autoLoadCandidateKey( @@ -1452,7 +1465,20 @@ function autoLoadCandidateKey( id: string, ggufVariant?: string | null, ): string { - return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`; + // Path-shape-aware case handling: on a case-sensitive filesystem, a skip + // key recorded for /models/Foo must not also skip /models/foo. + return `${kind}:${normalizeLoadTargetKey(id)}:${(ggufVariant ?? "").toLowerCase()}`; +} + +// Skip keys use the backend load target, not the display id: a cached repo +// and an indexed local row aliasing the same files share the key, so a file +// that failed through one row is not retried through the other. +function autoLoadSkipKey(candidate: AutoLoadCandidate): string { + return autoLoadCandidateKey( + candidate.kind, + candidate.loadId ?? candidate.id, + candidate.ggufVariant, + ); } function findCachedRepo( @@ -1463,6 +1489,300 @@ function findCachedRepo( return repos.find((repo) => repo.repo_id.toLowerCase() === normalized); } +/** + * Managed-cache rows eligible for background auto-load: complete, not + * hidden infrastructure, not declared non-chat by the backend, and not an + * adapter (loading an adapter resolves its base model, which for an + * uncached Hub base would start an implicit remote fetch). + */ +function isAutoLoadableCachedRepo(repo: { + repo_id: string; + partial?: boolean; + model_format?: string | null; + capabilities?: { can_chat?: boolean } | null; +}): boolean { + if (repo.partial) return false; + if (repo.model_format === "adapter") return false; + // Cached checkpoint repos (pickle .bin/.pt weights) stay interactive-only, + // like local checkpoint rows: forced-offline validation cannot consult the + // Hub security scan, and deserializing a pickle can execute code. + if (repo.model_format === "checkpoint") { + return false; + } + if (repo.capabilities?.can_chat === false) return false; + return !isHiddenModelId(repo.repo_id); +} + +// Same on-device scan sources the unified picker exposes +// (use-chat-picker-inventory's PICKER_LOCAL_SOURCES). hf_cache rows are +// covered by the cached lists; ollama links are not directly loadable. +const AUTO_LOAD_LOCAL_SOURCES: ReadonlySet = new Set([ + "models_dir", + "lmstudio", + "custom", +]); + +/** The picker's chat-only platform snapshot, read once per auto-load run. */ +type AutoLoadPlatform = { + chatOnly: boolean; + isMac: boolean; +}; + +// Mirrors the picker's localModelIsGguf / localModelIsMlx checks: the backend +// format hint is authoritative for indexed rows, with the same name/path +// fallbacks the picker applies. +function localRowIsGgufLike(row: LocalModelInfo): boolean { + return ( + row.model_format === "gguf" || row.path.toLowerCase().endsWith(".gguf") + ); +} + +function localRowIsMlxNamed(row: LocalModelInfo): boolean { + return ( + isMlxId(row.id) || + isMlxId(row.display_name ?? "") || + isMlxId(row.model_id ?? "") + ); +} + +/** + * Backend-indexed local rows eligible for background auto-load: same policy + * as the on-device picker (complete, chat-capable, not hidden infra), plus + * no variant requirement, since a background load cannot ask for a quant. + */ +function isAutoLoadableLocalRow( + row: LocalModelInfo, + platform: AutoLoadPlatform, +): boolean { + if (!AUTO_LOAD_LOCAL_SOURCES.has(row.source)) return false; + if (row.capabilities?.can_chat !== true) return false; + if (row.partial) return false; + // Chat-only installs run GGUF (any host) and MLX (Mac only); the picker + // hides other local formats there, so the background load must not pick a + // row the user could not have selected (mirrors sortedLocalDir's gate). + if ( + platform.chatOnly && + !localRowIsGgufLike(row) && + !(platform.isMac && localRowIsMlxNamed(row)) + ) { + return false; + } + // Adapters are chat-capable but load by resolving their base model, which + // for a Hub-id base can trigger the implicit remote fetch a background + // auto-load must never start. Adapters stay interactive-only. + if (row.model_format === "adapter") return false; + // Checkpoint rows (pickle .bin/.pt weights) are chat-capable but a local + // scan-folder checkpoint has no Hub security scan, and deserializing a + // pickle can execute code. Loading one must be an explicit user action, + // never a background pick; they stay interactive-only like adapters. + if (row.model_format === "checkpoint") { + return false; + } + if (isHiddenModelId(row.model_id, row.id, row.path)) return false; + // The name-based big-endian marker only applies to direct .gguf files: a + // DIRECTORY named e.g. /models/foo-be says nothing about the files inside + // it, and those are already filtered per-file by isAutoLoadableGgufVariant + // when the folder's quants are resolved. + if ( + row.model_format === "gguf" && + row.path.toLowerCase().endsWith(".gguf") && + hasBigEndianGgufMarker(row.path, row.format_variant) + ) { + return false; + } + return true; +} + +function localRowLoadTarget(row: LocalModelInfo): string { + return row.load_id || row.id; +} + +function localRowToCandidate( + row: LocalModelInfo, + ggufVariant: string | null = null, +): AutoLoadCandidate { + const isGguf = row.model_format === "gguf"; + return { + id: row.id, + loadId: localRowLoadTarget(row), + kind: isGguf ? "gguf" : "model", + // The backend load target identifies the GGUF itself; no Hub quant is + // required (a remembered quant is passed through for multi-quant dirs). + ggufVariant: isGguf ? ggufVariant : null, + maxSeqLength: isGguf ? 0 : 4096, + successLabel: `Loaded ${row.display_name || row.id}`, + inventoryId: row.inventory_id ?? null, + source: AUTO_LOAD_LOCAL_SOURCES.has(row.source) + ? (row.source as LastLocalModelSource) + : "local", + }; +} + +/** + * Build a loadable candidate for a backend-indexed local row. Directory-based + * GGUFs (LM Studio, models dir, custom folders) are flagged requires_variant + * by the backend, so without a remembered quant the folder is scanned through + * the same variants API the picker card uses and the smallest complete, + * auto-loadable quant is chosen. Returns null when no quant can be resolved. + */ +// Unknown sizes (0) sort last so a sizeless row can't shadow a real one. +function sizeOrUnknownBytes(bytes?: number | null): number { + return bytes && bytes > 0 ? bytes : Number.MAX_SAFE_INTEGER; +} + +/** + * Identifier-matching key with path-shape-aware case semantics: Windows-style + * paths (drive letter or UNC) and Hub repo ids compare case-insensitively, + * while POSIX paths keep their case, since Linux filesystems distinguish + * /models/Foo from /models/foo and folding them can match the wrong model. + */ +function normalizeLoadTargetKey(value: string): string { + const looksWindowsPath = /^(?:[A-Za-z]:[\\/]|\\\\)/.test(value); + if (!looksWindowsPath && (value.startsWith("/") || value.startsWith("~"))) { + return value; + } + return value.toLowerCase(); +} + +// Inventory scans hit disk on the backend; an unbounded fan-out over many +// indexed folders can saturate the connection pool and disk. +const AUTO_LOAD_VARIANT_SCAN_CONCURRENCY = 4; + +// How long after the bounded inventory calls resolve the best-effort +// prefetches (hidden-model matchers, platform snapshot) may still be +// waited on. Both use unbounded fetches of their own, so past this bound +// their client-side fallbacks apply (static needles; boot-detected +// platform) rather than letting a stalled request hang Send. +const BEST_EFFORT_PREFETCH_GRACE_MS = 1_000; + +// Settle window before the fallback starts consuming resolved candidates: +// when scans finish quickly (the common case) the pool is complete first and +// keeps the exact smallest-first order; slow scans stop blocking the send +// path after this window and join the pool in size order as they resolve. +const AUTO_LOAD_RESOLVE_GRACE_MS = 2500; + +// Upper bound on ONE GGUF variant scan. Pending scans hold the model-kind +// gate, and the underlying fetch has no timeout of its own, so a single hung +// request would otherwise gate Send forever. Matches the inventory calls' +// own 30s bound; on expiry the scan job settles as failed and a ready +// safetensors candidate can proceed. +const AUTO_LOAD_VARIANT_SCAN_TIMEOUT_MS = 30_000; + +// An HF cache snapshot dir: captures the cache repo ROOT before /snapshots/. +const HF_SNAPSHOT_PATH_RE = /^(.*)[\\/]snapshots[\\/][^\\/]+[\\/]?$/; + +// A custom scan folder registered at an HF cache snapshot aliases the cached +// row for the same repo, but the cached row contributes the cache repo ROOT +// (cache_path) while the custom row contributes its snapshots/ dir. +// Expanding a snapshot path with its cache root lets the two rows collide on +// one seen-set key, so shared files consume one attempt. +function expandSeenValues(value: string): string[] { + const snapshotMatch = HF_SNAPSHOT_PATH_RE.exec(value); + return snapshotMatch ? [value, snapshotMatch[1]] : [value]; +} + +async function listGgufVariantsBounded( + repoId: string, + options: { preferLocalCache?: boolean; localPath?: string | null }, +) { + const controller = new AbortController(); + const timer = setTimeout(() => { + controller.abort(); + }, AUTO_LOAD_VARIANT_SCAN_TIMEOUT_MS); + try { + return await listGgufVariants(repoId, undefined, { + ...options, + signal: controller.signal, + }); + } finally { + clearTimeout(timer); + } +} + +type ResolvedLocalCandidate = { + candidate: AutoLoadCandidate; + /** Size of what would actually load: the resolved quant's own size for a + * multi-quant folder (whose row size_bytes SUMS every quant), else the + * row size. Orders the smallest-first cascade. */ + sizeBytes: number; +}; + +async function resolveLocalRowCandidate( + row: LocalModelInfo, + rememberedVariant: string | null = null, + isSkippedCandidate?: (candidate: AutoLoadCandidate) => boolean, +): Promise { + const isGguf = row.model_format === "gguf"; + if (row.capabilities?.requires_variant === true) { + // Only GGUF folders have an automatic quant resolution path. + if (!isGguf) return null; + if (!rememberedVariant) { + // Scan the folder the row will actually load from: the cache-first + // prefer_local_cache flow could return a quant present in the HF + // cache but missing at row's own path. A local-path repo id routes + // the backend straight to the filesystem scan of that folder. + const variantScanTarget = isLocalModelPath(row.id) ? row.id : row.path; + const variants = await listGgufVariantsBounded(variantScanTarget, { + preferLocalCache: true, + localPath: row.path, + }); + const downloaded = variants.variants + .filter( + (entry) => + entry.downloaded && !entry.partial && isAutoLoadableGgufVariant(entry), + ) + .sort((a, b) => a.size_bytes - b.size_bytes); + // Smallest first, skipping quants that already failed or were + // blocked, so one bad file cannot sink a folder with other quants. + for (const entry of downloaded) { + const candidate = localRowToCandidate(row, entry.quant); + if (isSkippedCandidate?.(candidate)) continue; + return { candidate, sizeBytes: sizeOrUnknownBytes(entry.size_bytes) }; + } + return null; + } + } + const candidate = localRowToCandidate(row, isGguf ? rememberedVariant : null); + // Single-candidate rows resolve to null once their candidate is skipped, + // so the retry loop in the fallback terminates instead of re-attempting. + if (isSkippedCandidate?.(candidate)) return null; + return { + candidate, + sizeBytes: sizeOrUnknownBytes(row.size_bytes), + }; +} + +/** Resolve a remembered local model against current backend inventory. */ +function matchesRememberedLocalRow( + row: LocalModelInfo, + remembered: LastLocalModelLoad, +): boolean { + // A folder holding both GGUF and safetensors weights yields two rows with + // the same path/load target, so the format must agree with the remembered + // kind before identifier matching can accept the row. + if ((row.model_format === "gguf") !== (remembered.kind === "gguf")) { + return false; + } + // Inventory ids come from one backend generator on both sides, so they + // compare exactly; case folding could merge distinct case-sensitive paths + // embedded in the id. + if ( + remembered.inventoryId && + row.inventory_id && + row.inventory_id === remembered.inventoryId + ) { + return true; + } + const targets = new Set( + [remembered.loadId, remembered.id] + .filter((value): value is string => Boolean(value)) + .map((value) => normalizeLoadTargetKey(value)), + ); + return [row.load_id, row.id, row.path, row.model_id].some( + (value) => !!value && targets.has(normalizeLoadTargetKey(value)), + ); +} + function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean { const normalized = filename.replace(/\\/g, "/").toLowerCase(); const separatorIndex = normalized.lastIndexOf("/"); @@ -1499,9 +1819,21 @@ function isAutoLoadableGgufVariant(variant: GgufVariantDetail | null): boolean { return !hasBigEndianGgufMarker(filename, variant.quant); } -async function autoLoadSmallestModel(): Promise<{ +/** + * Auto-load a model already on this device when the user chats without + * selecting one: adopt the server-active model, then the last successfully + * loaded on-device model (managed HF caches, models dir, LM Studio, custom + * scan folders), then the smallest complete chat-capable on-device model + * (GGUF first, then safetensors). Never downloads: with no valid on-device + * candidate the caller shows the actionable "no model" error instead. + * + * Exported for tests. + */ +export async function autoLoadOnDeviceModel(): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean; + /** True when an inventory failure was already surfaced to the user. */ + inventoryErrorSurfaced?: boolean; }> { if (await tryAdoptServerActiveModel()) { return { loaded: true, blockedByTrustRemoteCode: false }; @@ -1515,7 +1847,7 @@ async function autoLoadSmallestModel(): Promise<{ const toastId = toast("Loading a model…", { description: lastLoaded ? "Loading last used model." - : "Auto-selecting the smallest downloaded model.", + : "Auto-selecting the smallest on-device model.", duration: 5000, closeButton: true, }); @@ -1541,6 +1873,10 @@ async function autoLoadSmallestModel(): Promise<{ ...payload, hf_token: hfToken, load_in_4bit: true, + // Same local-only policy the follow-up /load enforces, so ineligible + // candidates (uncached, or audio models whose codecs would fetch) are + // rejected here without consuming a load attempt. + local_files_only: true, trust_remote_code: trustRemoteCode, }); // Background auto-load never runs a repo's custom code or loads Hub-flagged unsafe @@ -1633,9 +1969,7 @@ async function autoLoadSmallestModel(): Promise<{ : {}), })) ) { - skippedAutoLoadCandidates.add( - autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant), - ); + skippedAutoLoadCandidates.add(autoLoadSkipKey(candidate)); return false; } loadAttempts += 1; @@ -1646,6 +1980,11 @@ async function autoLoadSmallestModel(): Promise<{ load_in_4bit: true, is_lora: false, gguf_variant: candidate.ggufVariant, + // A background load never downloads: the backend resolves a cached + // repo id against the local snapshot only, so a cache populated + // outside Studio with missing shards fails over to the next + // candidate instead of silently fetching the gaps. + local_files_only: true, trust_remote_code: trustRemoteCode, chat_template_override: effectiveChatTemplateOverride, cache_type_kv: config.kvCacheDtype, @@ -1772,29 +2111,194 @@ async function autoLoadSmallestModel(): Promise<{ id: candidate.id, kind: candidate.kind, ggufVariant: candidate.ggufVariant, + loadId: candidate.loadId ?? null, + inventoryId: candidate.inventoryId ?? null, + source: candidate.source, }); } toast.success(candidate.successLabel, { id: toastId }); return true; } + // An inventory failure is NOT an empty inventory: surface it and stop the + // automatic selection path instead of concluding nothing is on device. + let allGgufRepos: CachedGgufRepo[]; + let allModelRepos: CachedModelRepo[]; + let allLocalRows: LocalModelInfo[]; try { - const [ggufRepos, modelRepos] = await Promise.all([ - listCachedGguf().catch(() => []), - listCachedModels().catch(() => []), + // Dynamic hidden-model matchers are best-effort; the static needles + // still filter the built-in infra models when the fetch fails. The + // fetch has no timeout of its own, so it runs alongside the + // (30s-bounded) inventory calls and is only awaited through a short + // grace afterwards: a stalled matcher request must not hang Send. + const hiddenMatchersReady = ensureHiddenModelMatchers().catch( + () => undefined, + ); + // The platform fetch behind the picker's format gates is unbounded + // too (raw fetch, no signal); run it alongside as well. Past the + // grace, the store's boot-detected client-side platform applies. + const platformReady: Promise = fetchDeviceType().then( + () => undefined, + () => undefined, + ); + const [cachedGguf, cachedModels, localList] = await Promise.all([ + listCachedGguf(hfToken), + listCachedModels(hfToken), + listLocalModels(), ]); + const bestEffortPrefetches = Promise.all([ + hiddenMatchersReady, + platformReady, + ]); + await new Promise((resolve) => { + const matcherTimer = setTimeout(resolve, BEST_EFFORT_PREFETCH_GRACE_MS); + bestEffortPrefetches.then(() => { + clearTimeout(matcherTimer); + resolve(); + }); + }); + allGgufRepos = cachedGguf; + allModelRepos = cachedModels; + allLocalRows = localList.models; + } catch (error) { + const message = + error instanceof Error + ? error.message + : "Could not read the on-device model inventory."; + toast.error("Couldn't check on-device models", { + id: toastId, + description: message, + }); + return { + loaded: false, + blockedByTrustRemoteCode: false, + inventoryErrorSurfaced: true, + }; + } + // The platform snapshot the picker's format gates key on. Format gates + // only apply once the BACKEND-reported platform has been fetched: the + // boot-detected fallback describes the browser, which may differ from the + // host (a Mac browser against a remote Linux/CUDA backend would wrongly + // gate out every cached non-GGUF candidate). While the platform is + // unknown, candidates flow ungated and an actually chat-only backend + // rejects ineligible ones at validation without consuming an attempt. + const platformState = usePlatformStore.getState(); + const platform: AutoLoadPlatform = { + chatOnly: platformState.fetched ? platformState.isChatOnly() : false, + isMac: platformState.deviceType === "mac", + }; + const ggufRepos = allGgufRepos.filter(isAutoLoadableCachedRepo); + const modelRepos = allModelRepos.filter(isAutoLoadableCachedRepo); + const localRows = allLocalRows.filter((row) => + isAutoLoadableLocalRow(row, platform), + ); + // Dedupe candidates that resolve to the SAME load target (e.g. a custom + // scan folder pointing into an HF cache). Keyed on kind + load target / + // on-disk path: a shared model_id does not mean the same files (a distinct + // local copy stays available when the cached copy fails), and a folder + // emitting both GGUF and safetensors rows shares a path while holding two + // different models, so the format is part of the key. + const seenLoadTargets = new Set(); + const markSeen = ( + kind: LastLocalModelKind, + ...values: (string | null | undefined)[] + ): void => { + for (const value of values) { + if (value) { + for (const alias of expandSeenValues(value)) { + seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(alias)}`); + } + } + } + }; + const isSeen = ( + kind: LastLocalModelKind, + ...values: (string | null | undefined)[] + ): boolean => + values.some( + (value) => + !!value && + expandSeenValues(value).some((alias) => + seenLoadTargets.has(`${kind}:${normalizeLoadTargetKey(alias)}`), + ), + ); + + // One variant scan per cached repo per run: the remembered-model lookup + // and the cascade share the same result, so a stalled repository times + // out ONCE instead of gating Send through a second identical bounded + // request. Rejections are memoized too, on purpose: a repo whose scan + // already failed this run is not rescanned. + const repoVariantScans = new Map< + string, + ReturnType + >(); + const scanRepoVariants = ( + repoId: string, + localPath: string | null | undefined, + ) => { + const key = `${repoId}|${localPath ?? ""}`; + let pending = repoVariantScans.get(key); + if (!pending) { + pending = listGgufVariantsBounded(repoId, { + preferLocalCache: true, + localPath, + }); + repoVariantScans.set(key, pending); + } + return pending; + }; + + try { if (lastLoaded) { - if (lastLoaded.kind === "gguf") { + if (!isManagedCacheSource(lastLoaded.source)) { + const row = localRows.find((candidateRow) => + matchesRememberedLocalRow(candidateRow, lastLoaded), + ); + if (row) { + // Not marked seen here: if this exact quant fails, the fallback + // loop may still pick another complete quant from the same folder + // (only the failed candidate key below is excluded, mirroring the + // managed-cache remembered path). + let rememberedCandidate: AutoLoadCandidate | null = null; + try { + rememberedCandidate = + (await resolveLocalRowCandidate(row, lastLoaded.ggufVariant)) + ?.candidate ?? null; + if (rememberedCandidate) { + toast("Loading last used model…", { + id: toastId, + description: row.display_name || row.id, + duration: 5000, + }); + if (await loadAutoLoadCandidate(rememberedCandidate)) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + rememberedCandidate + ? autoLoadSkipKey(rememberedCandidate) + : autoLoadCandidateKey( + row.model_format === "gguf" ? "gguf" : "model", + localRowLoadTarget(row), + lastLoaded.ggufVariant, + ), + ); + } + } + } else if (lastLoaded.kind === "gguf") { const repo = findCachedRepo(ggufRepos, lastLoaded.id); if (repo && lastLoaded.ggufVariant) { + // Not marked seen: if this exact quant fails, the fallback pool + // may still pick another complete quant from this repo (only the + // failed candidate key below is excluded). try { - const variants = await listGgufVariants(repo.repo_id, undefined, { - preferLocalCache: true, - localPath: repo.cache_path, - }); + const variants = await scanRepoVariants(repo.repo_id, repo.cache_path); const variant = variants.variants.find( (entry) => entry.downloaded && + !entry.partial && entry.quant?.toLowerCase() === lastLoaded.ggufVariant?.toLowerCase() && isAutoLoadableGgufVariant(entry), @@ -1813,6 +2317,8 @@ async function autoLoadSmallestModel(): Promise<{ ggufVariant: variant.quant, maxSeqLength: 0, successLabel: `Loaded ${repo.repo_id} (${variant.quant})`, + inventoryId: repo.inventory_id ?? null, + source: "hf_cache", }) ) { return { loaded: true, blockedByTrustRemoteCode: false }; @@ -1821,7 +2327,11 @@ async function autoLoadSmallestModel(): Promise<{ } catch { hadNonTrustFailure = true; skippedAutoLoadCandidates.add( - autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant), + autoLoadCandidateKey( + "gguf", + repo.load_id || repo.repo_id, + lastLoaded.ggufVariant, + ), ); } } @@ -1842,6 +2352,8 @@ async function autoLoadSmallestModel(): Promise<{ ggufVariant: null, maxSeqLength: store.params.maxSeqLength, successLabel: `Loaded ${repo.repo_id}`, + inventoryId: repo.inventory_id ?? null, + source: "hf_cache", }) ) { return { loaded: true, blockedByTrustRemoteCode: false }; @@ -1849,218 +2361,422 @@ async function autoLoadSmallestModel(): Promise<{ } catch { hadNonTrustFailure = true; skippedAutoLoadCandidates.add( - autoLoadCandidateKey("model", repo.repo_id), + autoLoadCandidateKey("model", repo.load_id || repo.repo_id), ); } } } toast("Loading a model…", { id: toastId, - description: "Auto-selecting the smallest downloaded model.", + description: "Auto-selecting the smallest on-device model.", duration: 5000, }); } - // GGUF first: smallest-total-size repo, then its smallest variant. - if (ggufRepos.length > 0) { - const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes); - for (const repo of sorted) { - if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; - try { - const variants = await listGgufVariants(repo.repo_id, undefined, { - preferLocalCache: true, - localPath: repo.cache_path, - }); - const downloaded = variants.variants - .filter((v) => v.downloaded && isAutoLoadableGgufVariant(v)) - .sort((a, b) => a.size_bytes - b.size_bytes); - if (downloaded.length > 0) { - const variant = downloaded[0]; - if ( - skippedAutoLoadCandidates.has( - autoLoadCandidateKey("gguf", repo.repo_id, variant.quant), - ) - ) { + // On-device fallback: complete/loadable GGUF models first, then + // complete/loadable non-GGUF models, smallest first within each group, + // merging managed-cache repos with backend-indexed local rows. Both + // cached repos and local rows order on the size of the quant that will + // actually load: a multi-quant repo's row size_bytes SUMS every quant, + // which would push a repo holding one small quant behind larger models. + type FallbackCandidate = + | { + type: "cached-gguf"; + repo: CachedGgufRepo; + variant: GgufVariantDetail; + sizeBytes: number; + /** Re-queued quant of an already-visited repo (skips the seen gate). */ + retry?: boolean; + } + | { type: "cached-model"; repo: CachedModelRepo; sizeBytes: number } + | { + type: "local"; + row: LocalModelInfo; + candidate: AutoLoadCandidate; + sizeBytes: number; + /** Re-queued quant of an already-visited row (skips the seen gate). */ + retry?: boolean; + }; + const isModelKindEntry = (entry: FallbackCandidate): boolean => + entry.type === "cached-model" || + (entry.type === "local" && entry.candidate.kind === "model"); + const isSkippedAutoLoadCandidate = (c: AutoLoadCandidate): boolean => + skippedAutoLoadCandidates.has(autoLoadSkipKey(c)); + // Candidates are resolved through a bounded worker pool and consumed + // incrementally from a size-ordered pool: awaiting every folder scan + // before the first attempt let one slow folder stall the send path + // behind the transport timeout. Ordered insertion keeps GGUF entries + // ahead of safetensors entries and each group smallest-first, so + // late-resolving scans and requeued quants land in the same global + // order a full pre-sort would give. + const readyPool: FallbackCandidate[] = []; + const insertReady = (entry: FallbackCandidate): void => { + let at = 0; + if (isModelKindEntry(entry)) { + while (at < readyPool.length && !isModelKindEntry(readyPool[at])) { + at += 1; + } + } + while ( + at < readyPool.length && + isModelKindEntry(readyPool[at]) === isModelKindEntry(entry) && + readyPool[at].sizeBytes <= entry.sizeBytes + ) { + at += 1; + } + readyPool.splice(at, 0, entry); + }; + // Non-GGUF cached repos need no scan: their snapshot loads whole. The + // picker hides cached non-GGUF rows entirely on chat-only installs, so + // the automatic cascade must not pick one there either; the remembered + // path above stays ungated since a recorded load is user precedent that + // the model runs on this install (e.g. an MLX repo loaded on a Mac). + for (const repo of platform.chatOnly ? [] : modelRepos) { + insertReady({ + type: "cached-model", + repo, + // Order by the snapshot the load will actually resolve: the row's + // size_bytes sums weight blobs across EVERY cached revision, so a + // small current revision beside a huge stale one would otherwise be + // ranked as their total and sink behind genuinely larger candidates. + sizeBytes: sizeOrUnknownBytes( + repo.snapshot_size_bytes ?? repo.size_bytes, + ), + }); + } + // Smallest complete, auto-loadable, not-yet-skipped quant of a managed + // cache repo; null when none remains. + const resolveCachedGgufEntry = async ( + repo: CachedGgufRepo, + ): Promise | null> => { + const variants = await scanRepoVariants(repo.repo_id, repo.cache_path); + const downloaded = variants.variants + .filter( + (v) => v.downloaded && !v.partial && isAutoLoadableGgufVariant(v), + ) + .sort((a, b) => a.size_bytes - b.size_bytes); + for (const variant of downloaded) { + if ( + skippedAutoLoadCandidates.has( + autoLoadCandidateKey( + "gguf", + repo.load_id || repo.repo_id, + variant.quant, + ), + ) + ) { + continue; + } + return { + type: "cached-gguf", + repo, + variant, + sizeBytes: sizeOrUnknownBytes(variant.size_bytes), + }; + } + return null; + }; + // Directory-based GGUF rows resolve a quant automatically; only non-GGUF + // variant-requiring rows have no background resolution path. + const cascadeLocalRows = localRows.filter( + (row) => + row.model_format === "gguf" || + row.capabilities?.requires_variant !== true, + ); + // Only GGUF directory rows hit the backend folder scan; everything else + // resolves in-process and is seeded straight into the pool so it can + // never queue behind slow scans. + const needsVariantScan = (row: LocalModelInfo): boolean => + row.model_format === "gguf" && + row.capabilities?.requires_variant === true; + await Promise.all( + cascadeLocalRows + .filter((row) => !needsVariantScan(row)) + .map(async (row) => { + try { + const resolved = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (resolved) { + insertReady({ + type: "local", + row, + candidate: resolved.candidate, + sizeBytes: resolved.sizeBytes, + }); + } + } catch { + hadNonTrustFailure = true; + } + }), + ); + // The remaining jobs all need a backend scan and can all yield GGUF + // candidates. Cached repos and local folders interleave so a run of + // slow scans from one source cannot monopolize every worker. + const cachedScanJobs = ggufRepos.map( + (repo) => () => resolveCachedGgufEntry(repo), + ); + const localScanJobs = cascadeLocalRows.filter(needsVariantScan).map( + (row) => async (): Promise => { + const resolved = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (!resolved) { + return null; + } + return { + type: "local", + row, + candidate: resolved.candidate, + sizeBytes: resolved.sizeBytes, + }; + }, + ); + const resolutionJobs: Array<() => Promise> = []; + for ( + let jobIndex = 0; + jobIndex < Math.max(cachedScanJobs.length, localScanJobs.length); + jobIndex += 1 + ) { + if (jobIndex < cachedScanJobs.length) { + resolutionJobs.push(cachedScanJobs[jobIndex]); + } + if (jobIndex < localScanJobs.length) { + resolutionJobs.push(localScanJobs[jobIndex]); + } + } + let pendingJobs = resolutionJobs.length; + // Once auto-load reaches a terminal result (a model loaded, the attempt + // cap was hit, or the pool drained), the workers stop claiming jobs so a + // successful early load does not leave background scans hammering the + // backend while inference is already running. In-flight scans finish. + let resolutionStopped = false; + let progressWaiters: Array<() => void> = []; + const signalProgress = (): void => { + const waiters = progressWaiters; + progressWaiters = []; + for (const resolve of waiters) { + resolve(); + } + }; + const nextProgress = (): Promise => + new Promise((resolve) => progressWaiters.push(resolve)); + const runResolutionJobs = async (): Promise => { + let nextJob = 0; + await Promise.all( + Array.from( + { + length: Math.max( + 1, + Math.min(AUTO_LOAD_VARIANT_SCAN_CONCURRENCY, resolutionJobs.length), + ), + }, + async () => { + while (!resolutionStopped && nextJob < resolutionJobs.length) { + const job = resolutionJobs[nextJob]; + nextJob += 1; + let entry: FallbackCandidate | null = null; + try { + entry = await job(); + } catch { + hadNonTrustFailure = true; + } + pendingJobs -= 1; + if (entry) { + insertReady(entry); + } + signalProgress(); + } + }, + ), + ); + }; + const resolutionDone = runResolutionJobs(); + if (pendingJobs > 0) { + await new Promise((resolve) => { + const graceTimer = setTimeout(resolve, AUTO_LOAD_RESOLVE_GRACE_MS); + resolutionDone.then(() => { + clearTimeout(graceTimer); + resolve(); + }); + }); + } + try { + while (loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) { + const candidate = readyPool[0]; + if (!candidate) { + if (pendingJobs <= 0) { + break; + } + await nextProgress(); + continue; + } + // Only the first attempt may leapfrog pending scans: it is the + // latency-critical one and it almost always succeeds. Once any + // budget has been spent, consumption waits for resolution to + // settle, so the remaining attempts follow the complete global + // smallest-first order instead of exhausting the cap on larger + // candidates while smaller ones are still resolving. + if (pendingJobs > 0 && loadAttempts > 0) { + await nextProgress(); + continue; + } + // Every pending scan job can still yield a GGUF candidate (no-scan + // rows were seeded upfront), and GGUF outranks safetensors in the + // documented order, so the model-kind group stays gated until the + // scans settle; resolved GGUF entries keep flowing immediately. + if (isModelKindEntry(candidate) && pendingJobs > 0) { + await nextProgress(); + continue; + } + readyPool.shift(); + if (candidate.type === "cached-gguf") { + const repo = candidate.repo; + if (!candidate.retry) { + // A shared load target may already have been visited through an + // indexed local row (e.g. a scan folder aliasing this cache). + if (isSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)) { continue; } + markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path); + } + const skipKey = autoLoadCandidateKey( + "gguf", + repo.load_id || repo.repo_id, + candidate.variant.quant, + ); + if (skippedAutoLoadCandidates.has(skipKey)) { + continue; + } + try { if ( await loadAutoLoadCandidate({ id: repo.repo_id, loadId: repo.load_id, kind: "gguf", - ggufVariant: variant.quant, + ggufVariant: candidate.variant.quant, maxSeqLength: 0, - successLabel: `Loaded ${repo.repo_id} (${variant.quant})`, + successLabel: `Loaded ${repo.repo_id} (${candidate.variant.quant})`, + inventoryId: repo.inventory_id ?? null, + source: "hf_cache", }) ) { return { loaded: true, blockedByTrustRemoteCode: false }; } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add(skipKey); + // A quant that passed validation can still fail /load (corrupt + // file, llama.cpp startup error). Re-enter the repo's next + // complete quant into the global size order, so one repo of + // failing quants cannot starve a smaller model elsewhere. + // Validation blocks are model-scoped, so they get no requeue. + try { + const next = await resolveCachedGgufEntry(repo); + if (next) { + insertReady({ ...next, retry: true }); + } + } catch { + hadNonTrustFailure = true; + } } - } catch { - hadNonTrustFailure = true; continue; } - } - } - - // Fall back to safetensors models. - if (modelRepos.length > 0) { - const sorted = [...modelRepos].sort( - (a, b) => a.size_bytes - b.size_bytes, - ); - for (const repo of sorted) { - if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; - try { + if (candidate.type === "cached-model") { + const repo = candidate.repo; + if (isSeen("model", repo.load_id || repo.repo_id, repo.cache_path)) { + continue; + } + markSeen("model", repo.load_id || repo.repo_id, repo.cache_path); if ( skippedAutoLoadCandidates.has( - autoLoadCandidateKey("model", repo.repo_id), + autoLoadCandidateKey("model", repo.load_id || repo.repo_id), ) ) { continue; } - if ( - await loadAutoLoadCandidate({ - id: repo.repo_id, - loadId: repo.load_id, - kind: "model", - ggufVariant: null, - maxSeqLength: 4096, - successLabel: `Loaded ${repo.repo_id}`, - }) - ) { + try { + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + loadId: repo.load_id, + kind: "model", + ggufVariant: null, + maxSeqLength: 4096, + successLabel: `Loaded ${repo.repo_id}`, + inventoryId: repo.inventory_id ?? null, + source: "hf_cache", + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey("model", repo.load_id || repo.repo_id), + ); + } + continue; + } + const row = candidate.row; + const localCandidate = candidate.candidate; + if (!candidate.retry) { + if (isSeen(localCandidate.kind, row.load_id, row.id, row.path)) { + continue; + } + markSeen(localCandidate.kind, row.load_id, row.id, row.path); + } + if (isSkippedAutoLoadCandidate(localCandidate)) { + continue; + } + try { + if (await loadAutoLoadCandidate(localCandidate)) { return { loaded: true, blockedByTrustRemoteCode: false }; } } catch { hadNonTrustFailure = true; - continue; + skippedAutoLoadCandidates.add(autoLoadSkipKey(localCandidate)); + // Same requeue as the cached-gguf branch: the folder's next complete + // quant re-enters the global size order instead of retrying inline. + try { + const next = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (next) { + insertReady({ + type: "local", + row, + candidate: next.candidate, + sizeBytes: next.sizeBytes, + retry: true, + }); + } + } catch { + hadNonTrustFailure = true; + } } } + } finally { + // Runs on every exit (successful return, cap, drained pool, or a + // thrown error) so no worker keeps scanning after the outcome is set. + resolutionStopped = true; } - // Cap also gates the default download, so total /api/inference/load - // budget across cached + fallback is MAX_AUTO_LOAD_ATTEMPTS, not +1. - if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) { - toast.dismiss(toastId); - return { - loaded: false, - blockedByTrustRemoteCode: - blockedByTrustRemoteCode && !hadNonTrustFailure, - }; - } - - // No cached models — try downloading a small default GGUF. - toast("Downloading a small model…", { - id: toastId, - description: - "No downloaded models found. Fetching Qwen3.5-4B-MTP (UD-Q4_K_XL).", - duration: 30000, - }); - try { - const rt = useChatRuntimeStore.getState(); - if ( - !(await canAutoLoad({ - model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", - max_seq_length: 0, - is_lora: false, - gguf_variant: "UD-Q4_K_XL", - // The same live-store GPU pick the load below sends (a fresh default - // model has no remembered settings to prefer). - gpu_ids: rt.selectedGpuIds ?? undefined, - gpu_memory_mode: rt.gpuMemoryMode, - })) - ) { - toast.dismiss(toastId); - return { loaded: false, blockedByTrustRemoteCode }; - } - loadAttempts += 1; - const loadResp = await loadModel({ - model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", - hf_token: hfToken, - // Model default under both modes: Auto layers + no pin means - // resolveFitMaxSeqLength returns 0 for every mode (the canAutoLoad - // preflight above sends the same). - max_seq_length: 0, - load_in_4bit: true, - is_lora: false, - gguf_variant: "UD-Q4_K_XL", - trust_remote_code: trustRemoteCode, - speculative_type: specSettings.speculativeType, - spec_draft_n_max: specSettings.specDraftNMax, - // GPU Memory mode is a standing preference, so honor it on auto-load. - // The layer/MoE/split knobs and the context pin are per-model: the live - // store may hold edits drafted for a staged pick, and a fresh default - // model has no remembered settings, so those stay at their defaults like - // the cached-candidate path. The GPU pick deliberately differs (it's the - // picker's current on-screen selection, which the canAutoLoad preflight - // above already committed to). - gpu_memory_mode: rt.gpuMemoryMode, - gpu_layers: GPU_LAYERS_AUTO, - n_cpu_moe: 0, - gpu_ids: rt.selectedGpuIds ?? undefined, - }); - saveSpeculativeType(specSettings.speculativeType); - persistGpuMemoryModeOnLoad(loadResp, rt.gpuMemoryMode); - useChatRuntimeStore - .getState() - .setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL"); - const store = useChatRuntimeStore.getState(); - store.setModelRequiresTrustRemoteCode( - loadResp.requires_trust_remote_code ?? false, - ); - store.setParams({ - ...store.params, - maxTokens: loadResp.context_length ?? 131072, - }); - const defaultModel: ChatModelSummary = { - id: "unsloth/Qwen3.5-4B-MTP-GGUF", - name: loadResp.display_name ?? "Qwen3.5-4B-MTP-GGUF", - isVision: loadResp.is_vision ?? false, - isLora: false, - isGguf: true, - }; - if (!store.models.some((m) => m.id === "unsloth/Qwen3.5-4B-MTP-GGUF")) { - store.setModels([...store.models, defaultModel]); - } - useChatRuntimeStore.setState({ - ggufContextLength: loadResp.context_length ?? 131072, - ggufMaxContextLength: - loadResp.max_context_length ?? loadResp.context_length ?? 131072, - supportsReasoning: loadResp.supports_reasoning ?? false, - reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, - reasoningEnabled: loadResp.supports_reasoning ?? false, - ...reasoningCapsFromLoad(loadResp), - supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false, - supportsTools: loadResp.supports_tools ?? false, - ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), - kvCacheDtype: loadResp.cache_type_kv ?? null, - loadedKvCacheDtype: loadResp.cache_type_kv ?? null, - tensorParallel: loadResp.tensor_parallel ?? false, - loadedTensorParallel: loadResp.tensor_parallel ?? false, - ...loadedGpuMemoryFields(loadResp), - // Drives the GPU Memory controls' diffusion gate; set alongside the - // GPU fields on every load path so the gate can't read stale. - loadedIsDiffusion: loadResp.is_diffusion ?? false, - defaultChatTemplate: loadResp.chat_template ?? null, - chatTemplateOverride: null, - loadedIsMultimodal: isMultimodalResponse(loadResp), - ...resolveLoadedSpeculativeSettings(loadResp), - }); - recordLastLocalModelLoad({ - id: "unsloth/Qwen3.5-4B-MTP-GGUF", - kind: "gguf", - ggufVariant: "UD-Q4_K_XL", - }); - toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId }); - return { loaded: true, blockedByTrustRemoteCode: false }; - } catch { - toast.dismiss(toastId); - hadNonTrustFailure = true; - return { - loaded: false, - blockedByTrustRemoteCode: - blockedByTrustRemoteCode && !hadNonTrustFailure, - }; - } + // No auto-loadable on-device model (or the attempt cap was hit). Never + // fall back to a remote download from the send path: the caller shows + // the actionable "no model" error and any download stays an explicit + // user action. + toast.dismiss(toastId); + return { + loaded: false, + blockedByTrustRemoteCode: blockedByTrustRemoteCode && !hadNonTrustFailure, + }; } catch { toast.dismiss(toastId); hadNonTrustFailure = true; @@ -2105,19 +2821,23 @@ export function createOpenAIStreamAdapter( await waitForModelReady(abortSignal); } if (!useChatRuntimeStore.getState().params.checkpoint) { - const { loaded, blockedByTrustRemoteCode } = - await autoLoadSmallestModel(); + // Same on-device-only auto-load as the plain send path: deep + // research must never trigger a background download either. + const { loaded, blockedByTrustRemoteCode, inventoryErrorSurfaced } = + await autoLoadOnDeviceModel(); if (!loaded) { - toast.error( - blockedByTrustRemoteCode - ? "This model needs custom code approval" - : "No model loaded", - { - description: blockedByTrustRemoteCode - ? "Select it from the top bar to review and approve its custom code, or pick another model." - : "Pick a model in the top bar, then retry.", - }, - ); + if (!inventoryErrorSurfaced) { + toast.error( + blockedByTrustRemoteCode + ? "This model needs custom code approval" + : "No model loaded", + { + description: blockedByTrustRemoteCode + ? "Select it from the top bar to review and approve its custom code, or pick another model." + : "Select a model in the top bar, or download one from the Hub, then retry.", + }, + ); + } throw new Error("Load a model first."); } } @@ -2372,24 +3092,27 @@ export function createOpenAIStreamAdapter( // Prefer a model already loaded by the CLI/API before auto-loading. let loaded: boolean; let blockedByTrustRemoteCode: boolean; + let inventoryErrorSurfaced: boolean | undefined; try { - ({ loaded, blockedByTrustRemoteCode } = - await autoLoadSmallestModel()); + ({ loaded, blockedByTrustRemoteCode, inventoryErrorSurfaced } = + await autoLoadOnDeviceModel()); } catch (error) { clearSelectedImageEditReference(); throw error; } if (!loaded) { - toast.error( - blockedByTrustRemoteCode - ? "This model needs custom code approval" - : "No model loaded", - { - description: blockedByTrustRemoteCode - ? "Select it from the top bar to review and approve its custom code, or pick another model." - : "Pick a model in the top bar, then retry.", - }, - ); + if (!inventoryErrorSurfaced) { + toast.error( + blockedByTrustRemoteCode + ? "This model needs custom code approval" + : "No model loaded", + { + description: blockedByTrustRemoteCode + ? "Select it from the top bar to review and approve its custom code, or pick another model." + : "Select a model in the top bar, or download one from the Hub, then retry.", + }, + ); + } clearSelectedImageEditReference(); throw new Error("Load a model first."); } diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 4f558545ca..ee05ad8b0d 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -169,6 +169,9 @@ export async function validateModel( // --fit, while a pinned layer count is owned by the user. Tell validate // so it applies the same training-guard policy as /load. gpu_memory_mode: payload.gpu_memory_mode, + // Background auto-loads: resolve metadata from the local cache only + // and reject candidates whose runtime pulls remote auxiliaries. + local_files_only: payload.local_files_only ?? false, }), }); return parseJsonOrThrow(response); @@ -380,6 +383,9 @@ export interface CachedModelRepo { repo_id: string; load_id?: string | null; size_bytes: number; + /** Weight bytes of the newest cached snapshot only (what a load resolves); + * size_bytes sums selected-format blobs across every cached revision. */ + snapshot_size_bytes?: number | null; /** Epoch seconds of the newest downloaded weight file; sorts Downloaded * newest-first. Optional for older-backend compatibility. */ last_modified?: number; @@ -946,6 +952,9 @@ export async function listGgufVariants( options?: { preferLocalCache?: boolean; localPath?: string | null; + /** Aborts the underlying fetch; background scans pass a timeout signal so + * one hung request cannot gate Send forever. */ + signal?: AbortSignal; }, ): Promise { const params = new URLSearchParams({ repo_id: repoId }); @@ -958,6 +967,7 @@ export async function listGgufVariants( } const response = await authFetch(`/api/models/gguf-variants?${params}`, { headers: hubTokenHeader(hfToken), + signal: options?.signal, }); return parseJsonOrThrow(response); } diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 48a6168555..4906a9cdc0 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1080,21 +1080,39 @@ export function useChatModelRuntime() { } } await refresh({ signal: abortCtrl.signal }); + // Native-picked files are never remembered: their access depends + // on a signed, expiring path lease that a raw remembered path + // would bypass. Backend-indexed local rows (models dir, LM + // Studio, custom scan folders; selection source "local") are + // remembered and re-resolved through inventory on auto-load. + // Other arbitrary paths keep the isLocalModelPath protection. + const indexedLocalSelection = + typeof selection !== "string" && selection.source === "local"; if ( !isLora && !(loadResponse.is_lora ?? false) && !nativePathToken && - !isLocalModelPath(modelId) && - !isExternalModelId(modelId) + !isExternalModelId(modelId) && + (indexedLocalSelection || !isLocalModelPath(modelId)) ) { - if (loadResponse.is_gguf || isGguf || ggufVariant) { + const kind = + loadResponse.is_gguf || isGguf || ggufVariant + ? ("gguf" as const) + : ("model" as const); + if (isLocalModelPath(modelId)) { recordLastLocalModelLoad({ id: modelId, - kind: "gguf", + kind, ggufVariant: ggufVariant ?? null, + loadId: modelId, + source: "local", }); } else { - recordLastLocalModelLoad({ id: modelId, kind: "model" }); + recordLastLocalModelLoad({ + id: modelId, + kind, + ggufVariant: ggufVariant ?? null, + }); } } } catch (error) { diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index d554eb777e..93a49248ca 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -41,6 +41,9 @@ export interface LoadModelRequest { load_in_4bit: boolean; is_lora: boolean; gguf_variant?: string | null; + /** Resolve a cached repo id against the local snapshot only and never + * download missing files (background auto-loads). */ + local_files_only?: boolean; /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trust_remote_code?: boolean; /** sha256 fingerprint pinning user approval of this exact custom-code version. */ diff --git a/studio/frontend/src/features/chat/utils/last-local-model-load.ts b/studio/frontend/src/features/chat/utils/last-local-model-load.ts index 099386fbc7..0fec5dffc7 100644 --- a/studio/frontend/src/features/chat/utils/last-local-model-load.ts +++ b/studio/frontend/src/features/chat/utils/last-local-model-load.ts @@ -3,13 +3,46 @@ export type LastLocalModelKind = "gguf" | "model"; +/** + * Where the remembered model came from: + * - "hf_cache": a managed Hugging Face cache repo (active or inactive cache), + * resolved through the cached-gguf / cached-models inventory. + * - "models_dir" / "lmstudio" / "custom": a backend-indexed local inventory + * row, resolved through the /api/hub/local inventory. + * - "local": an indexed local row whose exact scan source was not known at + * record time (interactive picker loads); resolved like the other local + * sources. + * + * Native file-picker selections are never recorded here: their access depends + * on a signed, expiring native path lease that must not be bypassed by a raw + * remembered path (the caller skips recording when a lease token is present). + */ +export type LastLocalModelSource = + | "hf_cache" + | "models_dir" + | "lmstudio" + | "custom" + | "local"; + export type LastLocalModelLoad = { + /** Display / repository identity (HF repo id, or a local path for indexed rows). */ id: string; kind: LastLocalModelKind; + /** Managed-cache GGUF quant. Null is valid when the load target itself + * identifies the actual GGUF (a local file or directory). */ ggufVariant: string | null; + /** Backend-provided load target when it differs from `id` (e.g. an + * inactive-cache row's load_id or an indexed local path). */ + loadId: string | null; + /** Stable backend inventory row identity, when available. */ + inventoryId: string | null; + source: LastLocalModelSource; loadedAt: number; }; +// Kept at v1: new fields are parsed backward-compatibly, so existing records +// (which predate loadId/inventoryId/source) keep resolving as managed-cache +// entries without a migration pass. const STORAGE_KEY = "unsloth.last-local-model-load.v1"; function storage(): Storage | null { @@ -24,6 +57,26 @@ function isLastLocalModelKind(value: unknown): value is LastLocalModelKind { return value === "gguf" || value === "model"; } +const LOCAL_MODEL_SOURCES: readonly LastLocalModelSource[] = [ + "hf_cache", + "models_dir", + "lmstudio", + "custom", + "local", +]; + +function isLastLocalModelSource(value: unknown): value is LastLocalModelSource { + return LOCAL_MODEL_SOURCES.includes(value as LastLocalModelSource); +} + +export function isManagedCacheSource(source: LastLocalModelSource): boolean { + return source === "hf_cache"; +} + +function normalizeOptionalString(value: unknown): string | null { + return typeof value === "string" && value.trim() ? value : null; +} + export function readLastLocalModelLoad(): LastLocalModelLoad | null { try { const raw = storage()?.getItem(STORAGE_KEY); @@ -39,17 +92,25 @@ export function readLastLocalModelLoad(): LastLocalModelLoad | null { ) { return null; } - if ( - parsed.kind === "gguf" && - (typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim()) - ) { + // Legacy v1 records carry no source; they were only ever written for + // managed-cache repos. + const source = isLastLocalModelSource(parsed.source) + ? parsed.source + : "hf_cache"; + const ggufVariant = normalizeOptionalString(parsed.ggufVariant); + // A managed-cache GGUF loads by repo + quant, so the quant is required. + // An indexed local GGUF's load target identifies the file itself, so a + // null variant stays valid. + if (parsed.kind === "gguf" && source === "hf_cache" && !ggufVariant) { return null; } return { id: parsed.id, kind: parsed.kind, - ggufVariant: - typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null, + ggufVariant, + loadId: normalizeOptionalString(parsed.loadId), + inventoryId: normalizeOptionalString(parsed.inventoryId), + source, loadedAt: parsed.loadedAt, }; } catch { @@ -61,22 +122,31 @@ export function recordLastLocalModelLoad(input: { id: string; kind: LastLocalModelKind; ggufVariant?: string | null; + loadId?: string | null; + inventoryId?: string | null; + source?: LastLocalModelSource; }): void { const id = input.id.trim(); if (!id) { return; } + const source = input.source ?? "hf_cache"; const ggufVariant = input.ggufVariant?.trim() || null; - if (input.kind === "gguf" && !ggufVariant) { + if (input.kind === "gguf" && source === "hf_cache" && !ggufVariant) { return; } try { + // Only inventory identity is stored: never tokens, native path leases, or + // security approvals. storage()?.setItem( STORAGE_KEY, JSON.stringify({ id, kind: input.kind, ggufVariant: input.kind === "gguf" ? ggufVariant : null, + loadId: input.loadId?.trim() || null, + inventoryId: input.inventoryId?.trim() || null, + source, loadedAt: Date.now(), } satisfies LastLocalModelLoad), ); diff --git a/studio/frontend/src/features/hub/inventory/api.ts b/studio/frontend/src/features/hub/inventory/api.ts index d3a1abb540..42f3cb4d64 100644 --- a/studio/frontend/src/features/hub/inventory/api.ts +++ b/studio/frontend/src/features/hub/inventory/api.ts @@ -65,6 +65,9 @@ export interface CachedModelRepo { format_variant?: string | null; capabilities?: BackendModelCapabilities | null; size_bytes: number; + /** Weight bytes of the newest cached snapshot only (what a load resolves); + * size_bytes sums selected-format blobs across every cached revision. */ + snapshot_size_bytes?: number | null; cache_path?: string; last_modified?: number | null; partial?: boolean; diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index e1aba66b1b..3e74f52cfd 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -296,9 +296,14 @@ def test_chat_autoload_scopes_variant_lookup_to_cached_repo_path(): """Autoload must probe the exact cache row it will load, including rows retained from a previously selected Hugging Face cache.""" src = _read("features/chat/api/chat-adapter.ts") - auto_load = src.split("async function autoLoadSmallestModel", 1)[1] - assert auto_load.count("preferLocalCache: true") >= 2 - assert auto_load.count("localPath: repo.cache_path") >= 2 + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + # Both cached-repo lookups (remembered model and cascade) route through + # the memoized scanRepoVariants, which carries the cache-scoped params. + scan_fn = auto_load.split("const scanRepoVariants = (", 1)[1] + scan_fn = scan_fn.split("return pending;", 1)[0] + assert "preferLocalCache: true" in scan_fn + assert "localPath," in scan_fn + assert auto_load.count("await scanRepoVariants(repo.repo_id, repo.cache_path)") == 2 chat_api = _read("features/chat/api/chat-api.ts") variants_fn = chat_api.split("export async function listGgufVariants", 1)[1] @@ -580,6 +585,802 @@ def test_legacy_migration_is_idempotent_and_non_destructive(): assert "if (isDefaultConfig(migrated) || Object.hasOwn(map, key)) {" in src +# --------------------------------------------------------------------------- +# Send-with-no-model auto-load (issue #7374): on-device discovery must cover +# every picker inventory source, the remembered model must survive local +# (non-cache) loads, and the send path must never start a remote download. +# --------------------------------------------------------------------------- + + +def _autoload_section() -> str: + src = _read("features/chat/api/chat-adapter.ts") + return src.split("async function autoLoadOnDeviceModel", 1)[1] + + +def test_send_path_cannot_reach_hardcoded_default_download(): + """Pressing Send with no model loaded must never fetch the hard-coded + default repo from Hugging Face (the unconsented download in the bug + report). Any recommended download must stay an explicit user action.""" + src = _read("features/chat/api/chat-adapter.ts") + assert "Qwen3.5-4B-MTP-GGUF" not in src + assert "Downloading a small model" not in src + assert "No downloaded models found" not in src + # The old entry point must not linger anywhere. + assert "autoLoadSmallestModel" not in src + # The renamed entry point runs exactly once per send, so the submitted + # prompt executes exactly once after a successful load. + assert src.count("await autoLoadOnDeviceModel())") == 1 + + +def test_autoload_no_model_error_is_actionable(): + """With no valid on-device candidate the user is told to select or + explicitly download a model instead of getting a silent remote load.""" + src = _read("features/chat/api/chat-adapter.ts") + assert "Select a model in the top bar, or download one from the Hub, then retry." in src + + +def test_autoload_inventory_failure_is_not_empty_inventory(): + """A failed cached/local inventory request must stop the automatic + selection path, not be swallowed into an empty list that used to fall + through to the remote default download.""" + src = _read("features/chat/api/chat-adapter.ts") + assert ".catch(() => [])" not in src + auto_load = _autoload_section() + assert "inventoryErrorSurfaced: true" in auto_load + # All three inventory sources are queried together and fail closed. + for needle in ( + "listCachedGguf(hfToken)", + "listCachedModels(hfToken)", + "listLocalModels()", + ): + assert needle in auto_load, needle + + +def test_autoload_uses_unified_backend_inventory(): + """Auto-load must consume the same non-React backend inventory the + unified picker uses (no second frontend filesystem scanner), covering + the models dir, LM Studio dirs, and custom scan folders.""" + src = _read("features/chat/api/chat-adapter.ts") + assert re.search( + r'import \{[^}]*listLocalModels[^}]*\} from "@/features/hub/inventory/api"', + src, + re.S, + ) + sources = re.search(r"const AUTO_LOAD_LOCAL_SOURCES[^;]*;", src, re.S) + assert sources, "AUTO_LOAD_LOCAL_SOURCES not found" + for source in ('"models_dir"', '"lmstudio"', '"custom"'): + assert source in sources.group(0), source + + +def test_autoload_filters_match_picker_policy(): + """Only complete, chat-capable, non-hidden rows may auto-load: partial + downloads, weightless/non-chat folders, and infrastructure models are + excluded with the same policy the picker applies.""" + src = _read("features/chat/api/chat-adapter.ts") + local_fn = src.split("function isAutoLoadableLocalRow", 1)[1] + local_fn = local_fn.split("\nfunction ", 1)[0] + assert "row.capabilities?.can_chat !== true" in local_fn + assert "row.partial" in local_fn + # Adapters resolve their base model on load; a Hub-id base would start + # the implicit remote fetch a background auto-load must never trigger. + assert 'row.model_format === "adapter"' in local_fn + assert "isHiddenModelId(row.model_id, row.id, row.path)" in local_fn + # The name-based marker only applies to direct .gguf files; a directory + # named e.g. /models/foo-be says nothing about the files inside it, which + # are filtered per-file during variant resolution. + assert 'row.path.toLowerCase().endsWith(".gguf") &&' in local_fn + assert "hasBigEndianGgufMarker(row.path, row.format_variant)" in local_fn + cached_fn = src.split("function isAutoLoadableCachedRepo", 1)[1] + cached_fn = cached_fn.split("\nconst ", 1)[0] + assert "repo.partial" in cached_fn + # Cached adapter repos are chat-capable too and resolve a base model on + # load, so they must be excluded exactly like local adapter rows. + assert 'repo.model_format === "adapter"' in cached_fn + assert "repo.capabilities?.can_chat === false" in cached_fn + assert "isHiddenModelId(repo.repo_id)" in cached_fn + + +def test_autoload_local_rows_load_via_backend_target(): + """Indexed local rows (models dir, LM Studio, custom scan folders) must + load through the backend-provided target and record their stable + inventory identity, never a reconstructed path or synthetic variant.""" + src = _read("features/chat/api/chat-adapter.ts") + assert "function localRowLoadTarget" in src + assert "row.load_id || row.id" in src + candidate_fn = src.split("function localRowToCandidate", 1)[1] + candidate_fn = candidate_fn.split("\n/**", 1)[0] + assert "loadId: localRowLoadTarget(row)" in candidate_fn + assert "inventoryId: row.inventory_id ?? null" in candidate_fn + # Inactive-cache rows keep loading by their backend load_id. + auto_load = _autoload_section() + assert auto_load.count("loadId: repo.load_id") >= 3 + + +def test_autoload_remembers_last_model_across_all_sources(): + """The remembered model resolves against managed caches AND the indexed + local inventory; a stale entry only falls through to other on-device + candidates (there is no remote branch left to reach).""" + auto_load = _autoload_section() + assert "isManagedCacheSource(lastLoaded.source)" in auto_load + assert "matchesRememberedLocalRow(candidateRow, lastLoaded)" in auto_load + assert "await resolveLocalRowCandidate(" in auto_load + assert "lastLoaded.ggufVariant," in auto_load + # Managed-cache candidates record their provenance for later resolution. + assert auto_load.count('source: "hf_cache"') >= 4 + + +def test_autoload_deduplicates_cached_and_local_candidates(): + """Candidates resolving to the same load target (e.g. a custom scan + folder pointing into an HF cache) must not be tried twice, but the + dedupe must key on actual load targets/paths only: a local copy that + merely shares a repo model_id is a distinct set of files and must stay + available when the cached copy fails or has no usable quant.""" + auto_load = _autoload_section() + assert "const seenLoadTargets = new Set()" in auto_load + # Keys carry the model kind: a folder emitting both GGUF and safetensors + # rows shares a path while holding two different models. + assert "seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(alias)}`)" in auto_load + assert 'markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)' in auto_load + assert 'markSeen("model", repo.load_id || repo.repo_id, repo.cache_path)' in auto_load + assert "isSeen(localCandidate.kind, row.load_id, row.id, row.path)" in auto_load + # The repo-id-based dedupe that shadowed distinct local copies is gone. + assert "isSeen(row.load_id, row.id, row.path, row.model_id)" not in auto_load + assert "markSeen(repo.repo_id," not in auto_load + + +def test_local_quant_resolution_skips_failed_quants(): + """When a folder's smallest quant already failed or was blocked, the + resolver must return the next complete quant instead of abandoning the + whole folder (which made Send falsely report no model).""" + src = _read("features/chat/api/chat-adapter.ts") + resolve_fn = src.split("async function resolveLocalRowCandidate", 1)[1] + resolve_fn = resolve_fn.split("\nfunction ", 1)[0] + assert "for (const entry of downloaded)" in resolve_fn + assert "if (isSkippedCandidate?.(candidate)) continue;" in resolve_fn + # The fallback loop feeds the skip set into resolution. + auto_load = _autoload_section() + assert "isSkippedAutoLoadCandidate," in auto_load + + +def test_autoload_trust_guard_still_blocks_background_loads(): + """A model needing custom-code approval or a security review is never + silently auto-loaded, and a blocked candidate can only cascade to other + on-device candidates.""" + auto_load = _autoload_section() + assert "validation.requires_trust_remote_code" in auto_load + assert "validation.requires_security_review" in auto_load + assert "MAX_AUTO_LOAD_ATTEMPTS" in auto_load + + +def test_remembered_model_record_supports_local_sources(): + """last-local-model-load must represent managed-cache models AND + backend-indexed local models: a local GGUF is valid with a null variant, + legacy v1 records keep resolving as managed-cache entries, and no + secrets (tokens/leases) are ever persisted.""" + src = _read("features/chat/utils/last-local-model-load.ts") + # Same storage key: v1 records parse backward-compatibly, no migration. + assert 'const STORAGE_KEY = "unsloth.last-local-model-load.v1";' in src + # Legacy records carry no source and default to the managed cache. + assert "isLastLocalModelSource(parsed.source)" in src + assert ': "hf_cache";' in src + # The GGUF-variant requirement is scoped to managed-cache records; a + # local GGUF's load target identifies the file, so null stays valid. + assert src.count('source === "hf_cache" && !ggufVariant') == 2 + # Indexed local scan sources are representable. + for source in ('"models_dir"', '"lmstudio"', '"custom"'): + assert source in src, source + # Identity only: never tokens, native path leases, or approvals. + assert "nativePath" not in src + assert "hfToken" not in src and "hf_token" not in src + assert "fingerprint" not in src + + +def test_interactive_local_loads_are_remembered_without_lease_bypass(): + """A successful interactive load of a backend-indexed local model + (picker source "local") must be remembered so auto-load can reuse it, + while native-picked files (signed, expiring path lease) and other + arbitrary paths must never be recorded.""" + src = _read("features/chat/hooks/use-chat-model-runtime.ts") + record_block = src.split("const indexedLocalSelection", 1)[1] + record_block = record_block.split("} catch (error) {", 1)[0] + assert 'selection.source === "local"' in src + assert "!nativePathToken &&" in record_block + assert "(indexedLocalSelection || !isLocalModelPath(modelId))" in record_block + assert 'source: "local",' in record_block + + +def test_remembered_local_row_match_requires_kind_agreement(): + """A folder holding both GGUF and safetensors weights yields two inventory + rows with the same path/load target, so the remembered kind must gate the + identifier match or a remembered safetensors load can resolve to the GGUF + row (and vice versa).""" + src = _read("features/chat/api/chat-adapter.ts") + match_fn = src.split("function matchesRememberedLocalRow", 1)[1] + match_fn = match_fn.split("\nasync function ", 1)[0].split("\nfunction ", 1)[0] + assert '(row.model_format === "gguf") !== (remembered.kind === "gguf")' in match_fn + + +def test_directory_gguf_rows_resolve_variant_like_picker(): + """Directory-based local GGUFs (LM Studio, models dir, custom folders) are + flagged requires_variant by the backend, so the fallback must resolve a + quant through the variants API (as the picker card does) instead of + silently dropping every directory row; non-GGUF variant-requiring rows + have no background resolution and stay excluded.""" + src = _read("features/chat/api/chat-adapter.ts") + resolve_fn = src.split("async function resolveLocalRowCandidate", 1)[1] + resolve_fn = resolve_fn.split("\nfunction ", 1)[0] + assert "row.capabilities?.requires_variant === true" in resolve_fn + assert "if (!isGguf) return null;" in resolve_fn + # Quants must be resolved from the folder the row will load from, not + # from a same-id HF cache repo whose quants may be absent locally. + assert "const variantScanTarget = isLocalModelPath(row.id) ? row.id : row.path;" in resolve_fn + assert "listGgufVariantsBounded(variantScanTarget" in resolve_fn + assert "localPath: row.path" in resolve_fn + assert "entry.downloaded && !entry.partial && isAutoLoadableGgufVariant(entry)" in resolve_fn + # The cascade must keep directory GGUF rows as candidates. + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert 'row.model_format === "gguf" ||' in auto_load + assert "await resolveLocalRowCandidate(" in auto_load + + +def test_remembered_local_failure_does_not_block_folder_fallback(): + """A failed remembered model must exclude only that exact candidate key, + not mark the whole row or repo as seen; otherwise a folder or cache repo + with another complete quant can never fall back and Send falsely reports + no model. Applies to local rows and managed-cache repos alike.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + remembered_block = auto_load.split("if (lastLoaded) {", 1)[1] + remembered_block = remembered_block.split("// On-device fallback", 1)[0] + assert ( + "markSeen(" not in remembered_block + ), "remembered paths must not pre-mark their row/repo as deduped" + assert "autoLoadSkipKey(rememberedCandidate)" in remembered_block + + +def test_fallback_orders_by_resolved_quant_size(): + """A GGUF folder or cache repo row's size_bytes sums every quant in it, so + the smallest-first fallback must order both local rows and cached repos by + the resolved quant's own size; otherwise a repo holding one small quant + loses to a larger single-quant model.""" + src = _read("features/chat/api/chat-adapter.ts") + resolve_fn = src.split("async function resolveLocalRowCandidate", 1)[1] + resolve_fn = resolve_fn.split("\nfunction ", 1)[0] + assert "sizeBytes: sizeOrUnknownBytes(entry.size_bytes)" in resolve_fn + auto_load = _autoload_section() + # Local rows order on the resolved quant size. + assert "sizeBytes: resolved.sizeBytes" in auto_load + # Cached GGUF repos order on the resolved quant size too. + assert "const resolveCachedGgufEntry" in auto_load + assert "sizeBytes: sizeOrUnknownBytes(variant.size_bytes)" in auto_load + # Non-GGUF cached repos order on the SELECTED snapshot's size (falling + # back to the all-revisions row sum for older backends). + seed_block = auto_load.split("for (const repo of platform.chatOnly ? [] : modelRepos)", 1)[1] + seed_block = seed_block.split("const resolveCachedGgufEntry", 1)[0] + assert "repo.snapshot_size_bytes ?? repo.size_bytes" in seed_block + + +def test_cascade_retries_next_quant_after_load_failure(): + """A failed /api/inference/load (not just a blocked validation) must mark + that quant skipped and re-enter the folder's or repo's next complete quant + into the GLOBAL size order (still ahead of the safetensors group) instead + of retrying inline, so one folder of failing quants cannot starve a + smaller model elsewhere; single-candidate rows resolve to null once + skipped, and the attempt cap bounds total loads.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + # No inline retry loop: retries re-enter the shared ordered pool. + assert "while (localCandidate" not in auto_load + assert "retry: true," in auto_load + assert "insertReady({ ...next, retry: true })" in auto_load + # Ordered insertion respects the GGUF-before-safetensors group boundary + # and the ascending size order among the remaining candidates. + assert "!isModelKindEntry(readyPool[at])" in auto_load + assert "readyPool[at].sizeBytes <= entry.sizeBytes" in auto_load + # Requeued entries bypass the seen gate; fresh rows still dedupe. + assert "if (!candidate.retry) {" in auto_load + # The cascade catch records the failed quant before requeueing. + assert "skippedAutoLoadCandidates.add(skipKey)" in auto_load + assert "skippedAutoLoadCandidates.add(autoLoadSkipKey(localCandidate))" in auto_load + # Termination guard: a skipped single candidate resolves to null. + resolve_fn = src.split("async function resolveLocalRowCandidate", 1)[1] + resolve_fn = resolve_fn.split("\nfunction ", 1)[0] + assert "if (isSkippedCandidate?.(candidate)) return null;" in resolve_fn + + +def test_cached_rows_deduped_against_local_aliases(): + """A cached repo and an indexed local row can alias the same files (e.g. + a scan folder pointing into an HF cache). The fallback must not spend a + second load attempt re-trying files already visited or failed through the + other row: cached branches apply the same seen gate local rows use, and + skip keys are scoped to the backend load target both rows share.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert 'if (isSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path))' in auto_load + assert 'if (isSeen("model", repo.load_id || repo.repo_id, repo.cache_path))' in auto_load + key_fn = src.split("function autoLoadSkipKey", 1)[1] + key_fn = key_fn.split("\nfunction ", 1)[0] + assert "candidate.loadId ?? candidate.id" in key_fn + + +def test_send_not_blocked_by_full_inventory_resolution(): + """Pressing Send must not wait for every /gguf-variants folder scan before + the first load attempt: candidates resolve through a bounded worker pool + and are consumed incrementally after a short settle grace, so one slow + folder cannot stall the send path behind the transport timeout.""" + src = _read("features/chat/api/chat-adapter.ts") + assert "const AUTO_LOAD_RESOLVE_GRACE_MS" in src + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + # The consumer never awaits full resolution; it waits for the grace + # window (cut short when resolution finishes) and then per-completion. + assert "await resolutionDone" not in auto_load + assert "clearTimeout(graceTimer)" in auto_load + assert "await nextProgress();" in auto_load + assert "if (pendingJobs <= 0) {" in auto_load + # Rows that need no backend scan seed the pool before the workers start, + # so they can never queue behind slow folder scans. + assert "const needsVariantScan" in auto_load + assert ".filter((row) => !needsVariantScan(row))" in auto_load + assert auto_load.index(".filter((row) => !needsVariantScan(row))") < auto_load.index( + "const cachedScanJobs" + ) + # Cached and local scans interleave so one slow source cannot + # monopolize every worker. + assert "resolutionJobs.push(cachedScanJobs[jobIndex])" in auto_load + assert "resolutionJobs.push(localScanJobs[jobIndex])" in auto_load + + +def test_pending_gguf_scans_gate_safetensors_candidates(): + """GGUF-first is the documented preference order, and incremental + consumption must not let an instantly-resolved safetensors row claim a + load slot while a pending folder scan can still yield a GGUF candidate; + resolved GGUF entries are never gated.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "if (isModelKindEntry(candidate) && pendingJobs > 0) {" in auto_load + + +def test_only_first_attempt_leapfrogs_pending_scans(): + """Fast-resolving candidates whose loads fail must not exhaust the attempt + cap while pending scans can still yield smaller loadable quants: only the + first (latency-critical) attempt may run ahead of pending scans; once any + budget is spent, the remaining attempts wait for the settled global + smallest-first order.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "if (pendingJobs > 0 && loadAttempts > 0) {" in auto_load + + +def test_hidden_matcher_fetch_never_blocks_send_unbounded(): + """ensureHiddenModelMatchers has no timeout of its own (unlike the + 30s-bounded inventory calls), so the send path must not await it serially: + it runs alongside inventory discovery and is only awaited through a short + grace, after which the static needles filter alone.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "await ensureHiddenModelMatchers()" not in auto_load + assert "const hiddenMatchersReady = ensureHiddenModelMatchers().catch(" in auto_load + # The platform probe backing the picker's format gates is unbounded too + # (raw fetch, no signal), so it joins the same bounded prefetch wait. + assert "await fetchDeviceType();" not in auto_load + assert "const platformReady: Promise = fetchDeviceType().then(" in auto_load + assert "const BEST_EFFORT_PREFETCH_GRACE_MS" in src + assert "setTimeout(resolve, BEST_EFFORT_PREFETCH_GRACE_MS)" in auto_load + assert "clearTimeout(matcherTimer)" in auto_load + + +def test_resolution_workers_stop_on_terminal_result(): + """A successful early load (or any other terminal outcome) must stop the + workers from claiming further folder scans, so autoload cannot leave + background scans contending with inference for the backend and disk.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "let resolutionStopped = false;" in auto_load + assert "while (!resolutionStopped && nextJob < resolutionJobs.length)" in auto_load + # The flag is set in a finally so every exit path (return, break, throw) + # stops the workers. + assert "resolutionStopped = true;" in auto_load + assert auto_load.index("} finally {") < auto_load.index("resolutionStopped = true;") + + +def test_local_rows_apply_picker_platform_gate(): + """Chat-only installs run GGUF (any host) and MLX (Mac only); the picker + hides other local formats and all cached non-GGUF rows there, so the + background cascade must not load a row the user could not have picked. + The remembered path stays ungated: a recorded load is user precedent.""" + src = _read("features/chat/api/chat-adapter.ts") + local_fn = src.split("function isAutoLoadableLocalRow", 1)[1] + local_fn = local_fn.split("\nfunction ", 1)[0] + assert "platform.chatOnly" in local_fn + assert "localRowIsGgufLike(row)" in local_fn + assert "platform.isMac && localRowIsMlxNamed(row)" in local_fn + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + # The platform snapshot hydrates through the bounded best-effort prefetch + # (see test_hidden_matcher_fetch_never_blocks_send_unbounded). + assert "fetchDeviceType().then(" in auto_load + # Cascade seeding of cached non-GGUF repos mirrors the picker's + # chat-only exclusion; the remembered lookup above it stays unfiltered. + assert "platform.chatOnly ? [] : modelRepos" in auto_load + + +def test_autoload_keys_preserve_posix_path_case(): + """Linux filesystems distinguish /models/Foo from /models/foo, so seen + keys and remembered-model matching must not fold case on POSIX paths; + Windows-style paths and Hub repo ids keep case-insensitive matching.""" + src = _read("features/chat/api/chat-adapter.ts") + norm_fn = src.split("function normalizeLoadTargetKey", 1)[1] + norm_fn = norm_fn.split("\nfunction ", 1)[0].split("\nconst ", 1)[0] + assert "looksWindowsPath" in norm_fn + assert 'value.startsWith("/") || value.startsWith("~")' in norm_fn + assert "return value;" in norm_fn + assert "return value.toLowerCase();" in norm_fn + match_fn = src.split("function matchesRememberedLocalRow", 1)[1] + match_fn = match_fn.split("\nfunction ", 1)[0] + assert "normalizeLoadTargetKey" in match_fn + # Inventory ids compare exactly (same backend generator on both sides). + assert "row.inventory_id === remembered.inventoryId" in match_fn + assert "row.inventory_id.toLowerCase()" not in match_fn + # Skip keys use the same semantics: a failure recorded for /models/Foo + # must not also skip /models/foo. + key_fn = src.split("function autoLoadCandidateKey", 1)[1] + key_fn = key_fn.split("\nfunction ", 1)[0] + assert "normalizeLoadTargetKey(id)" in key_fn + assert "id.toLowerCase()" not in key_fn + + +def test_local_variant_scans_bounded_concurrency(): + """Each /gguf-variants call triggers a recursive backend directory scan, + so pre-resolution must not fan out unbounded over every indexed folder + at once.""" + src = _read("features/chat/api/chat-adapter.ts") + assert "const AUTO_LOAD_VARIANT_SCAN_CONCURRENCY" in src + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "Math.min(AUTO_LOAD_VARIANT_SCAN_CONCURRENCY, resolutionJobs.length)" in auto_load + assert "await Promise.all(\n cascadeLocalRows.map(" not in auto_load + + +BACKEND = WORKDIR / "studio" / "backend" + + +def _read_backend(rel: str) -> str: + path = BACKEND / rel + assert path.exists(), f"missing backend source file: {path}" + return path.read_text() + + +def test_background_loads_resolve_local_files_only(): + """A cache populated outside Studio can pass the partial check while + missing shard files, and from_pretrained on a repo id downloads the gaps. + Background auto-loads therefore send local_files_only and the load route + rewrites the path to the locally resolved snapshot (identity intact), so + an incomplete cache fails over to the next candidate instead of + downloading on Send.""" + src = _read("features/chat/api/chat-adapter.ts") + load_fn = src.split("async function loadAutoLoadCandidate", 1)[1] + load_fn = load_fn.split("loadAttempts += 1;", 1)[1] + assert "local_files_only: true," in load_fn + types = _read("features/chat/types/api.ts") + assert "local_files_only?: boolean;" in types + + request_model = _read_backend("models/inference.py") + assert "local_files_only: bool = Field(" in request_model + + route = _read_backend("routes/inference.py") + assert "request.local_files_only" in route + assert "resolve_local_snapshot_path" in route + assert "config.path = local_snapshot" in route + # Uncached repos fail closed with a conflict, never a download. + rewrite = route.split("request.local_files_only", 1)[1] + rewrite = rewrite.split("config.path = local_snapshot", 1)[0] + assert "status_code = 409" in rewrite + + helper = _read_backend("hub/utils/local_snapshot.py") + assert "local_files_only = True" in helper + # The rewrite resolves against the LIVE Studio-managed cache location, + # not huggingface_hub's import-time default. + assert "str(get_hf_cache_paths().hub_cache)," in route + + +def test_background_candidate_filters_have_no_side_effects(): + """Round-13 gates. Local checkpoint rows carry pickle weights with no Hub + security scan, so they are never background-picked. The canAutoLoad probe + validates with the same local-only policy /load enforces, and the validate + route runs it offline so the metadata probe cannot reach the Hub. Non-GGUF + audio models are refused under local-only in both routes because their + codec runtimes download auxiliaries at load time. Revision-only caches + (refs pruned) still resolve through the snapshot-dir fallback.""" + adapter = _read("features/chat/api/chat-adapter.ts") + assert 'if (row.model_format === "checkpoint") {' in adapter + + can_fn = adapter.split("async function canAutoLoad", 1)[1] + can_fn = can_fn.split("async function", 1)[0] + assert "local_files_only: true," in can_fn + + chat_api = _read("features/chat/api/chat-api.ts") + assert "local_files_only: payload.local_files_only ?? false," in chat_api + + request_model = _read_backend("models/inference.py") + validate_schema = request_model.split("class ValidateModelRequest", 1)[1] + assert "local_files_only: bool = Field(" in validate_schema + + route = _read_backend("routes/inference.py") + # /load forces offline resolution under local-only; /validate covers its + # whole preflight via the ExitStack wrap (asserted separately below). + assert route.count("with _hf_offline_if_dns_dead(force = request.local_files_only):") == 1 + # Audio gate present in both routes, GGUF exempt (llama.cpp path already + # resolves companions cached-or-skipped under the flag). + assert route.count("needs audio codec downloads") == 2 + for gate in route.split("needs audio codec downloads")[:-1]: + tail = gate.rsplit("if request.local_files_only", 1)[1] + assert "is_gguf" in tail and "is_audio" in tail + + llama = _read_backend("core/inference/llama_cpp.py") + assert "def _hf_offline_if_dns_dead(force: bool = False):" in llama + # force must also override an explicitly falsy HF_HUB_OFFLINE=0 (only a + # TRUTHY env value short-circuits), and overlapping guards are refcounted + # so the env is restored only when the LAST one exits; behavior is + # exercised directly in test_offline_guard_refcount.py. + assert "elif _hf_env_offline() and (not force or _hub_offline_env_truthy()):" in llama + assert "_OFFLINE_GUARD_LOCK" in llama + assert '_OFFLINE_GUARD_STATE["count"] += 1' in llama + + helper = _read_backend("hub/utils/local_snapshot.py") + assert "def _snapshot_dir_fallback(" in helper + assert "config.json" in helper + + +def test_local_only_covers_every_load_and_validate_network_path(): + """Round-14 gates. The Vulkan-ordinal GGUF preflight downloads FIRST, so it + takes the same local-only flag and forced-offline wrap as the Phase 2 + download (whose cache-size check would otherwise call get_paths_info). + The validate route keeps its whole metadata/security preflight offline, not + just the identifier probe. The python load path survives the process + boundary: the orchestrator forwards the flag and the route-resolved + snapshot path into the worker, which runs the entire load under a scoped + offline env and passes the flag to the vision processor fallback.""" + llama = _read_backend("core/inference/llama_cpp.py") + # Both GGUF download blocks in load_model force offline under local-only. + assert llama.count("with _hf_offline_if_dns_dead(force = local_files_only):") == 2 + preflight = llama.split("_preflight_model_path = self._download_gguf(", 1)[1] + preflight = preflight.split(")", 1)[0] + assert "local_files_only = local_files_only," in preflight + + route = _read_backend("routes/inference.py") + validate_src = route.split('operation = "validate-model"', 1)[1] + assert ( + "_local_only_offline.enter_context(_hf_offline_if_dns_dead(force = True))" in validate_src + ) + assert "_local_only_offline.close()" in validate_src + # The non-GGUF load threads the flag into the subprocess backend. + load_call = route.split("backend.load_model,\n config = config,", 1)[1] + load_call = load_call.split(")", 1)[0] + assert "local_files_only = request.local_files_only," in load_call + + orchestrator = _read_backend("core/inference/orchestrator.py") + assert '"local_files_only": bool(local_files_only),' in orchestrator + assert '"local_snapshot_path"' in orchestrator + + worker = _read_backend("core/inference/worker.py") + assert "def _local_only_offline_env(" in worker + assert "_offline_guard.enter_context(" in worker + assert "_offline_guard.close()" in worker + # The route's snapshot rewrite is re-applied after the worker rebuilds + # its ModelConfig from the identifier. + assert 'snapshot_override = config.get("local_snapshot_path")' in worker + assert "mc.path = snapshot_override" in worker + assert '"local_files_only": bool(config.get("local_files_only", False)),' in worker + + inference = _read_backend("core/inference/inference.py") + processor_call = inference.split("processor = AutoProcessor.from_pretrained(", 1)[1] + processor_call = processor_call.split("logger.info", 1)[0] + assert "local_files_only = local_files_only," in processor_call + + +def test_background_picks_mirror_inventory_and_skip_installers(): + """Round-15 gates. Cached checkpoint repos (pickle weights) are excluded + from background picks like local checkpoint rows. The worker keeps its + ENTIRE bootstrap offline under local-only and never pip-installs SSM + kernels (missing fatal kernels fail into candidate failover). Snapshot + resolution prefers the newest snapshot dir, matching the inventory + scanner's latest_snapshot_dir selection, before consulting refs/main. + MLX loads read config.path so the live-cache rewrite is honored. Cached + non-GGUF ordering uses the selected snapshot's size, not the + all-revisions blob total.""" + adapter = _read("features/chat/api/chat-adapter.ts") + cached_filter = adapter.split("function isAutoLoadableCachedRepo", 1)[1] + cached_filter = cached_filter.split("AUTO_LOAD_LOCAL_SOURCES", 1)[0] + assert 'if (repo.model_format === "checkpoint") {' in cached_filter + assert "repo.snapshot_size_bytes ?? repo.size_bytes" in adapter + + worker = _read_backend("core/inference/worker.py") + assert "_bootstrap_offline = contextlib.ExitStack()" in worker + # Entered before base resolution / gates / kernels, closed before BOTH + # command loops (MLX and GPU paths). + bootstrap = worker.split("_bootstrap_offline = contextlib.ExitStack()", 1)[1] + assert bootstrap.count("_bootstrap_offline.close()") == 2 + ssm_sig = worker.split("def _ensure_ssm_kernels(", 1)[1].split(") -> bool:", 1)[0] + assert "local_files_only: bool = False" in ssm_sig + ssm = worker.split("def _ensure_ssm_kernels", 1)[1] + ssm = ssm.split("def _run_security_gates", 1)[0] + assert "if local_files_only:" in ssm + assert 'importlib.util.find_spec("mamba_ssm") is None' in ssm + + helper = _read_backend("hub/utils/local_snapshot.py") + resolve = helper.split("def resolve_local_snapshot_path", 1)[1] + # Newest-snapshot scan runs BEFORE the refs/main-based resolver (compare + # the actual calls, not docstring mentions). + assert resolve.index("resolved = _snapshot_dir_fallback(") < resolve.index( + "return snapshot_download(" + ) + + mlx = _read_backend("core/inference/mlx_inference.py") + assert 'load_source = getattr(config, "path", None) or model_name' in mlx + + inventory = _read_backend("hub/services/models/cache_inventory.py") + assert "snapshot_size_bytes" in inventory + assert "def _snapshot_dir_mtime(" in inventory + schema = _read_backend("hub/schemas/inventory.py") + assert "snapshot_size_bytes" in schema + + +def test_forced_offline_is_hub_specific_and_covers_parent_preflight(): + """Round-16 gates. A forced guard may only no-op when HF_HUB_OFFLINE + itself is truthy (huggingface_hub ignores TRANSFORMERS_OFFLINE), snapshot + selection prefers revisions that hold the inventoried safetensors weights + so a metadata-only newest revision cannot shadow a complete older one, + and the orchestrator's parent-side preflight (transformers tier probe, + GPU sizing via hf model_info) runs under the local-only guard, closed + before the worker spawn so the child does not inherit the env.""" + llama = _read_backend("core/inference/llama_cpp.py") + assert "def _hub_offline_env_truthy(" in llama + assert "not force or _hub_offline_env_truthy()" in llama + + helper = _read_backend("hub/utils/local_snapshot.py") + assert "def _snapshot_has_weights(" in helper + assert "max(weightful or candidates, key = os.path.getmtime)" in helper + + orchestrator = _read_backend("core/inference/orchestrator.py") + preflight = orchestrator.split("def _preflight_offline():", 1)[1] + assert "_hf_offline_if_dns_dead(force = True)" in preflight + # Both metadata call sites are guarded, and the guard closes before spawn. + assert preflight.count("with _preflight_offline():") == 2 + tier_probe = preflight.split("with _preflight_offline():", 1)[1] + assert "needs_transformers_5" in tier_probe.split("with _preflight_offline():", 1)[0] + assert "prepare_gpu_selection(" in tier_probe.split("_spawn_subprocess", 1)[0] + + +def test_generation_and_scan_paths_stay_bounded_and_local(): + """Round-17 gates. Inactive-cache rows emit a weight-bearing snapshot as + their load_id (the load consumes it directly, bypassing the repo-id + resolver). Ordinary offline guards never extend a forced window. The + generation-time native template reload honors the load's local-only flag + and resolved path instead of refetching the repo id online. Every GGUF + variant scan behind the model-kind gate is bounded by an abortable + timeout, and an HF cache snapshot registered as a custom folder dedupes + against its cached row through the shared cache root.""" + inventory = _read_backend("hub/services/models/cache_inventory.py") + assert "def _weightful_snapshot_path(" in inventory + resolver = inventory.split("def _cached_model_snapshot_path(", 1)[1] + resolver = resolver.split("def ", 1)[0] + assert "_weightful_snapshot_path(repo_path)" in resolver + + llama = _read_backend("core/inference/llama_cpp.py") + join_branch = llama.split('if _OFFLINE_GUARD_STATE["count"] > 0:', 1)[1] + join_branch = join_branch.split("elif", 1)[0] + assert "if force:" in join_branch + + helpers = _read_backend("core/inference/chat_template_helpers.py") + reload_block = helpers.split("native_chat_template", 1)[1] + reload_block = reload_block.split("model_info[", 1)[0] + assert 'local_files_only = bool(model_info.get("local_files_only", False))' in reload_block + assert 'template_source = model_info.get("model_path") or template_source' in reload_block + assert "local_files_only = local_files_only," in reload_block + inference = _read_backend("core/inference/inference.py") + assert '"local_files_only": local_files_only,' in inference + mlx = _read_backend("core/inference/mlx_inference.py") + assert '"local_files_only": local_files_only,' in mlx + assert '"model_path": load_source,' in mlx + + adapter = _read("features/chat/api/chat-adapter.ts") + assert "AUTO_LOAD_VARIANT_SCAN_TIMEOUT_MS" in adapter + assert "async function listGgufVariantsBounded(" in adapter + assert "controller.abort()" in adapter + # No unbounded scan calls remain in the adapter: every call site routes + # through the bounded wrapper (the wrapper itself holds the one direct + # call, with the timeout signal attached). + assert adapter.count("await listGgufVariants(") == 1 + chat_api = _read("features/chat/api/chat-api.ts") + assert "signal: options?.signal," in chat_api + + assert "function expandSeenValues(" in adapter + assert adapter.count("expandSeenValues(value)") == 2 + + +def test_local_only_gguf_reuse_and_platform_gates_are_authoritative(): + """Round-18 gates. get_paths_info performs no offline-mode check and + HF_HUB_OFFLINE is baked into hub constants at import, so local-only GGUF + reuse skips the remote size verification per-call instead of relying on + the env guard. Platform format gates only apply once the BACKEND-reported + platform is fetched (the browser fallback may describe a different + machine). snapshot_size_bytes uses the resolvers' complete-revision + predicate (config plus weights in the SAME revision). Cached-repo variant + scans are memoized per run so a stalled repo times out once.""" + llama = _read_backend("core/inference/llama_cpp.py") + assert "verify_sizes = not local_files_only," in llama + reuse = llama.split("_cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards)", 1)[ + 1 + ] + reuse = reuse.split("cached_main is not None", 1)[0] + assert "local_files_only" in reuse and "_cached_candidate_matches_revision_size" in reuse + + adapter = _read("features/chat/api/chat-adapter.ts") + assert "chatOnly: platformState.fetched ? platformState.isChatOnly() : false," in adapter + assert "const repoVariantScans = new Map<" in adapter + assert "const scanRepoVariants = (" in adapter + assert adapter.count("await scanRepoVariants(repo.repo_id, repo.cache_path)") == 2 + + inventory = _read_backend("hub/services/models/cache_inventory.py") + assert "rev_has_config" in inventory + assert 'if selected_category == "safetensors":' in inventory + + +def test_route_preflights_and_gguf_rows_stay_format_true(): + """Round-19 gates. The /load route's sidecar tier probe and training + guard size local-only candidates against the resolved snapshot path (a + repo id would reach _remote_lora_base's raw HTTP request and hf + model_info, neither of which honors offline mode). GGUF cached rows + select a GGUF-bearing snapshot, so a mixed repo's safetensors revision + cannot become the GGUF row's load target.""" + route = _read_backend("routes/inference.py") + assert "_tier_target = config.path if request.local_files_only else config.identifier" in route + guard_block = route.split("_guard_identifier = (", 1)[1] + guard_block = guard_block.split("await asyncio.to_thread(", 1)[0] + assert "config.path" in guard_block + assert "request.local_files_only and not config.is_gguf" in guard_block + assert "model_identifier = _guard_identifier," in route + + inventory = _read_backend("hub/services/models/cache_inventory.py") + assert "def _newest_snapshot_where(" in inventory + assert "def _gguf_snapshot_path(" in inventory + assert "def _cached_gguf_repo_snapshot_path(" in inventory + gguf_scan = inventory.split("def _scan_cached_gguf(", 1)[1] + gguf_scan = gguf_scan.split("def ", 1)[0] + assert "_cached_gguf_repo_snapshot_path(repo_path)" in gguf_scan + assert "_cached_model_snapshot_path(repo_path)" not in gguf_scan + + +def test_gguf_background_loads_never_download_companions(): + """A cached GGUF load can still fetch from the Hub through its optional + companions (mmproj, MTP drafter) or a cache-miss main quant. Background + loads pass local_files_only into the llama.cpp path: companions resolve + cached-or-skipped, and a main-quant cache miss raises instead of + downloading.""" + route = _read_backend("routes/inference.py") + gguf_source = route.split("if config.gguf_hf_repo:", 1)[1] + gguf_source = gguf_source.split("else:", 1)[0] + assert "local_files_only = request.local_files_only," in gguf_source + + llama = _read_backend("core/inference/llama_cpp.py") + # The flag flows to the main quant and both companion helpers, and the + # crash-replay kwargs keep it so a reload stays local-only. + assert llama.count("local_files_only = local_files_only,") >= 3 + assert '"local_files_only": local_files_only,' in llama + # Companions: cached-or-skipped, never fetched. + assert 'logger.info("Skipping %s fetch (local-only load)", label)' in llama + assert "if local_files_only or _hf_env_offline():" in llama + # Main quant: a cache miss fails closed. + download = llama.split("def _download_gguf", 1)[1] + download = download.split("def _download_companion_gguf", 1)[0] + assert "if local_files_only:" in download + assert "select it explicitly to download it." in download + + inference = _read_backend("core/inference/inference.py") + # The vision processor fallback stays on the (possibly rewritten local) + # load path instead of refetching by repo id. + assert "config.base_model if config.is_lora else config.path" in inference + assert "config.base_model if config.is_lora else config.identifier" not in inference + + def test_vulkan_inference_devices_are_the_pickable_set(): """GGUF loads run through llama-server, so on a Vulkan build the picker must offer the inference inventory (ggml ordinals, the space `--device Vulkan`