diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 1efd78c06d..9561280bcc 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -371,6 +371,178 @@ class LlamaCppBackend: # Pipe closed — process is terminating pass + # ── HF download (no lock held) ─────────────────────────────── + + def _download_gguf( + self, + *, + hf_repo: str, + hf_variant: Optional[str] = None, + hf_token: Optional[str] = None, + ) -> str: + """Download GGUF file(s) from HuggingFace. Returns local path. + + Runs WITHOUT self._lock so that unload_model() can set + _cancel_event at any time. Checks _cancel_event between + each shard download. + """ + try: + from huggingface_hub import hf_hub_download + except ImportError: + raise RuntimeError( + "huggingface_hub is required for HF model loading. " + "Install it with: pip install huggingface_hub" + ) + + # Determine the filename from the variant + gguf_filename = None + gguf_extra_shards: list[str] = [] + if hf_variant: + try: + import re + from huggingface_hub import list_repo_files + + files = list_repo_files(hf_repo, token = hf_token) + variant_lower = hf_variant.lower() + boundary = re.compile( + r"(? 0: + cache_dir = os.environ.get( + "HF_HUB_CACHE", + str(Path.home() / ".cache" / "huggingface" / "hub"), + ) + Path(cache_dir).mkdir(parents = True, exist_ok = True) + free_bytes = shutil.disk_usage(cache_dir).free + + total_gb = total_download_bytes / (1024**3) + free_gb = free_bytes / (1024**3) + + logger.info( + f"GGUF download: {total_gb:.1f} GB needed, " + f"{free_gb:.1f} GB free on disk" + ) + + if total_download_bytes > free_bytes: + smaller = self._find_smallest_fitting_variant( + hf_repo, + free_bytes, + hf_token, + ) + if smaller: + fallback_file, fallback_size = smaller + logger.info( + f"Selected variant too large ({total_gb:.1f} GB), " + f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" + ) + gguf_filename = fallback_file + import re as _re + + _shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") + _m = _shard_pat.match(gguf_filename) + _prefix = _m.group(1) if _m else None + if _prefix: + gguf_extra_shards = sorted( + f + for f in all_gguf_files + if f.startswith(_prefix) + and f != gguf_filename + and "mmproj" not in f.lower() + ) + else: + gguf_extra_shards = [] + else: + raise RuntimeError( + f"Not enough disk space to download any variant. " + f"Only {free_gb:.1f} GB free in {cache_dir}" + ) + except RuntimeError: + raise + except Exception as e: + logger.warning(f"Could not check disk space: {e}") + + logger.info( + f"Downloading GGUF: {hf_repo}/{gguf_filename}" + + ( + f" (+{len(gguf_extra_shards)} shards)" + if gguf_extra_shards + else "" + ) + ) + try: + if self._cancel_event.is_set(): + raise RuntimeError("Cancelled") + local_path = hf_hub_download( + repo_id = hf_repo, + filename = gguf_filename, + token = hf_token, + ) + for shard in gguf_extra_shards: + if self._cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Downloading GGUF shard: {shard}") + hf_hub_download( + repo_id = hf_repo, + filename = shard, + token = hf_token, + ) + except RuntimeError as e: + if "Cancelled" in str(e): + raise + raise RuntimeError( + f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" + ) + except Exception as e: + raise RuntimeError( + f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" + ) + + logger.info(f"GGUF downloaded to: {local_path}") + return local_path + # ── Lifecycle ───────────────────────────────────────────────── def load_model( @@ -404,200 +576,47 @@ class LlamaCppBackend: Returns True if server started and health check passed. """ self._cancel_event.clear() + + # ── Phase 1: kill old process (under lock, fast) ────────── with self._lock: self._kill_process() - binary = self._find_llama_server_binary() - if not binary: - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) + binary = self._find_llama_server_binary() + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + + # ── Phase 2: download (NO lock held, so cancel can proceed) ── + if hf_repo: + model_path = self._download_gguf( + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + ) + elif gguf_path: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + model_path = gguf_path + else: + raise ValueError("Either gguf_path or hf_repo must be provided") + + # Check cancel after download + if self._cancel_event.is_set(): + logger.info("Load cancelled after download phase") + return False + + # ── Phase 3: start llama-server (under lock) ────────────── + with self._lock: + # Re-check cancel inside lock + if self._cancel_event.is_set(): + logger.info("Load cancelled before server start") + return False self._port = self._find_free_port() - # Build command based on mode - if hf_repo: - # Download the GGUF file ourselves using huggingface_hub - # (llama-server's -hf flag requires HTTPS/curl which may not - # be available, e.g. Windows builds with -DLLAMA_CURL=OFF) - try: - from huggingface_hub import hf_hub_download - except ImportError: - raise RuntimeError( - "huggingface_hub is required for HF model loading. " - "Install it with: pip install huggingface_hub" - ) - - # Determine the filename from the variant (e.g., "Q4_K_M" -> find matching file) - # For split GGUFs (e.g., *-00001-of-00003.gguf) we must download ALL shards. - gguf_filename = None - gguf_extra_shards: list[str] = [] - if hf_variant: - # Try common naming patterns - try: - import re - from huggingface_hub import list_repo_files - - files = list_repo_files(hf_repo, token = hf_token) - variant_lower = hf_variant.lower() - # Use word-boundary matching so "Q8_0" doesn't also - # match "IQ8_0" or other superset variant names. - boundary = re.compile( - r"(? try model name - repo_name = hf_repo.split("/")[-1].replace("-GGUF", "") - gguf_filename = f"{repo_name}-{hf_variant}.gguf" - - # Check disk space and fall back to a smaller variant if needed - all_gguf_files = [gguf_filename] + gguf_extra_shards - try: - import os - - from huggingface_hub import get_paths_info - - path_infos = list( - get_paths_info(hf_repo, all_gguf_files, token = hf_token) - ) - total_download_bytes = sum((p.size or 0) for p in path_infos) - - if total_download_bytes > 0: - cache_dir = os.environ.get( - "HF_HUB_CACHE", - str(Path.home() / ".cache" / "huggingface" / "hub"), - ) - Path(cache_dir).mkdir(parents = True, exist_ok = True) - free_bytes = shutil.disk_usage(cache_dir).free - - total_gb = total_download_bytes / (1024**3) - free_gb = free_bytes / (1024**3) - - logger.info( - f"GGUF download: {total_gb:.1f} GB needed, " - f"{free_gb:.1f} GB free on disk" - ) - - if total_download_bytes > free_bytes: - # Try to find a smaller variant that fits - smaller = self._find_smallest_fitting_variant( - hf_repo, - free_bytes, - hf_token, - ) - if smaller: - fallback_file, fallback_size = smaller - logger.info( - f"Selected variant too large ({total_gb:.1f} GB), " - f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" - ) - gguf_filename = fallback_file - # Re-discover shards for the fallback variant - import re as _re - - _shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") - _m = _shard_pat.match(gguf_filename) - _prefix = _m.group(1) if _m else None - if _prefix: - gguf_extra_shards = sorted( - f - for f in all_gguf_files - if f.startswith(_prefix) - and f != gguf_filename - and "mmproj" not in f.lower() - ) - else: - gguf_extra_shards = [] - else: - raise RuntimeError( - f"Not enough disk space to download any variant. " - f"Only {free_gb:.1f} GB free in {cache_dir}" - ) - except RuntimeError: - raise - except Exception as e: - logger.warning(f"Could not check disk space: {e}") - - logger.info( - f"Downloading GGUF: {hf_repo}/{gguf_filename}" - + ( - f" (+{len(gguf_extra_shards)} shards)" - if gguf_extra_shards - else "" - ) - ) - try: - if self._cancel_event.is_set(): - raise RuntimeError("Cancelled") - local_path = hf_hub_download( - repo_id = hf_repo, - filename = gguf_filename, - token = hf_token, - ) - # Download remaining shards for split GGUFs — llama-server - # auto-discovers them when they are in the same directory. - for shard in gguf_extra_shards: - if self._cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Downloading GGUF shard: {shard}") - hf_hub_download( - repo_id = hf_repo, - filename = shard, - token = hf_token, - ) - except RuntimeError as e: - if "Cancelled" in str(e): - raise - raise RuntimeError( - f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" - ) - except Exception as e: - raise RuntimeError( - f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" - ) - - logger.info(f"GGUF downloaded to: {local_path}") - model_path = local_path - elif gguf_path: - if not Path(gguf_path).is_file(): - raise FileNotFoundError(f"GGUF file not found: {gguf_path}") - model_path = gguf_path - else: - raise ValueError("Either gguf_path or hf_repo must be provided") - # Select GPU(s) based on model size and free memory try: model_size = self._get_gguf_size_bytes(model_path) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4955d8f210..6c50f2b753 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -74,14 +74,16 @@ class InferenceOrchestrator: self.models: dict = {} self.loading_models: set = set() self.loaded_local_models: list = [] - self.default_models = [ + self._static_models = [ "unsloth/Qwen3-4B-Instruct-2507", - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", + "unsloth/Llama-3.1-8B-Instruct-bnb-4bit", "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", "unsloth/Phi-3.5-mini-instruct", "unsloth/Gemma-3-4B-it", "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit", ] + self._top_gguf_cache: Optional[list[str]] = None + self._top_gguf_fetched = False # Version tracking for subprocess reuse self._current_transformers_major: Optional[str] = None # "4" or "5" @@ -89,6 +91,50 @@ class InferenceOrchestrator: atexit.register(self._cleanup) logger.info("InferenceOrchestrator initialized (subprocess mode)") + # Kick off background fetch of top GGUF models + threading.Thread( + target = self._fetch_top_gguf, daemon = True, name = "top-gguf" + ).start() + + # ------------------------------------------------------------------ + # Default models (top GGUFs fetched dynamically from HF) + # ------------------------------------------------------------------ + + @property + def default_models(self) -> list[str]: + top = self._top_gguf_cache or [] + seen = set(top) + return top + [m for m in self._static_models if m not in seen] + + def _fetch_top_gguf(self) -> None: + """Fetch top 4 GGUF repos from unsloth by downloads (background).""" + try: + import httpx + + resp = httpx.get( + "https://huggingface.co/api/models", + params = { + "author": "unsloth", + "sort": "downloads", + "direction": "-1", + "limit": "40", + }, + timeout = 15, + ) + if resp.status_code == 200: + models = resp.json() + gguf_ids = [ + m["id"] for m in models + if m.get("id", "").upper().endswith("-GGUF") + ][:4] + if gguf_ids: + self._top_gguf_cache = gguf_ids + logger.info("Top GGUF models: %s", gguf_ids) + except Exception as e: + logger.warning("Failed to fetch top GGUF models: %s", e) + finally: + self._top_gguf_fetched = True + # ------------------------------------------------------------------ # Subprocess lifecycle # ------------------------------------------------------------------ diff --git a/studio/backend/main.py b/studio/backend/main.py index a54d5182f3..3c2732eb5a 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -159,12 +159,24 @@ async def get_system_info(): import psutil from utils.hardware import get_device, get_gpu_memory_info, DeviceType - # GPU Info — try nvidia-smi first to get ALL physical GPUs (not filtered - # by CUDA_VISIBLE_DEVICES), since llama-server can use all of them. + # GPU Info — query nvidia-smi for physical GPUs, filtered by + # CUDA_VISIBLE_DEVICES when set (the frontend uses this for GGUF + # fit estimation and llama-server respects CVD too). + import os + gpu_info: dict = {"available": False, "devices": []} device = get_device() if device == DeviceType.CUDA: + # Parse CUDA_VISIBLE_DEVICES allowlist + allowed_indices = None + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None and cvd.strip(): + try: + allowed_indices = set(int(x.strip()) for x in cvd.split(",")) + except ValueError: + pass # Non-numeric (e.g. GPU-uuid), show all + try: result = subprocess.run( [ @@ -180,9 +192,12 @@ async def get_system_info(): for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] if len(parts) == 3: + idx = int(parts[0]) + if allowed_indices is not None and idx not in allowed_indices: + continue gpu_info["devices"].append( { - "index": int(parts[0]), + "index": idx, "name": parts[1], "memory_total_gb": round(int(parts[2]) / 1024, 2), } diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5299d168f5..c7fa9ed720 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -122,9 +122,13 @@ async def load_model( unsloth_backend.unload_model(unsloth_backend.active_model_name) # Route to HF mode or local mode based on config + # Run in a thread so the event loop stays free for progress + # polling and other requests during the (potentially long) + # GGUF download + llama-server startup. if config.gguf_hf_repo: - # HF mode: llama-server downloads via -hf "repo:quant" - success = llama_backend.load_model( + # HF mode: download via huggingface_hub then start llama-server + success = await asyncio.to_thread( + llama_backend.load_model, hf_repo = config.gguf_hf_repo, hf_variant = config.gguf_variant, hf_token = request.hf_token, @@ -134,7 +138,8 @@ async def load_model( ) else: # Local mode: llama-server loads via -m - success = llama_backend.load_model( + success = await asyncio.to_thread( + llama_backend.load_model, gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, model_identifier = config.identifier, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 6b54faa0a9..2b1fb09b79 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -13,6 +13,12 @@ from typing import List, Optional import structlog from loggers import get_logger +import re as _re +_VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") + +def _is_valid_repo_id(repo_id: str) -> bool: + return bool(_VALID_REPO_ID.fullmatch(repo_id)) + # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -555,17 +561,19 @@ async def get_gguf_variants( best = _pick_best_gguf(filenames) default_variant = _extract_quant_label(best) if best else None - # Check which variants are already downloaded in the HF cache + # Check which variants are fully downloaded in the HF cache. + # For split GGUFs, ALL shards must be present -- sum cached bytes + # per variant and compare against the expected total. # HF cache dir uses the exact case from the repo_id at download time, # which may differ from the canonical HF repo_id, so do a # case-insensitive match. - cached_files: set = set() + cached_bytes_by_quant: dict[str, int] = {} try: import re as _re from huggingface_hub import constants as hf_constants # Sanitize repo_id: must be "owner/name" with safe chars only - if not _re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", repo_id): + if not _is_valid_repo_id(repo_id): raise ValueError(f"Invalid repo_id format: {repo_id}") cache_dir = Path(hf_constants.HF_HUB_CACHE) @@ -576,11 +584,21 @@ async def get_gguf_variants( if snapshots.is_dir(): for snap in snapshots.iterdir(): for f in snap.rglob("*.gguf"): - cached_files.add(f.name) + q = _extract_quant_label(f.name) + cached_bytes_by_quant[q] = ( + cached_bytes_by_quant.get(q, 0) + f.stat().st_size + ) break except Exception: pass + def _is_fully_downloaded(variant) -> bool: + cached = cached_bytes_by_quant.get(variant.quant, 0) + if cached == 0 or variant.size_bytes == 0: + return False + # Allow small rounding tolerance (symlinks vs real sizes) + return cached >= variant.size_bytes * 0.99 + return GgufVariantsResponse( repo_id = repo_id, variants = [ @@ -588,7 +606,7 @@ async def get_gguf_variants( filename = v.filename, quant = v.quant, size_bytes = v.size_bytes, - downloaded = Path(v.filename).name in cached_files, + downloaded = _is_fully_downloaded(v), ) for v in variants ], @@ -616,10 +634,8 @@ async def get_gguf_download_progress( Tracks completed shard downloads in snapshots and in-progress downloads in the blobs directory (incomplete files). """ - import re as _re - try: - if not _re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", repo_id): + if not _is_valid_repo_id(repo_id): return { "downloaded_bytes": 0, "expected_bytes": expected_bytes, @@ -670,41 +686,43 @@ async def get_gguf_download_progress( async def list_cached_gguf( current_subject: str = Depends(get_current_subject), ): - """List GGUF repos that have already been downloaded to the HF cache.""" - try: - from huggingface_hub import constants as hf_constants + """List GGUF repos that have already been downloaded to the HF cache. - cache_dir = Path(hf_constants.HF_HUB_CACHE) - cached = [] - if cache_dir.is_dir(): - for entry in sorted(cache_dir.iterdir()): - if not entry.name.startswith("models--"): - continue - # models--unsloth--Qwen3-8B-GGUF -> unsloth/Qwen3-8B-GGUF - parts = entry.name.split("--", 1) - if len(parts) < 2: - continue - repo_id = parts[1].replace("--", "/") - if not repo_id.lower().endswith("-gguf"): - continue - # Check if there are actual .gguf files in snapshots - snapshots = entry / "snapshots" - if not snapshots.is_dir(): - continue - total_size = 0 - has_gguf = False - for snap in snapshots.iterdir(): - for f in snap.rglob("*.gguf"): + Uses scan_cache_dir() for proper repo IDs, then deduplicates by + lowercased key (HF cache dirs are lowercased but the canonical repo + ID preserves casing). + """ + try: + from huggingface_hub import scan_cache_dir + + hf_cache = scan_cache_dir() + seen_lower: dict[str, dict] = {} + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if not repo_id.upper().endswith("-GGUF"): + continue + # Check for actual .gguf files and sum sizes + total_size = 0 + has_gguf = False + for revision in repo_info.revisions: + for f in revision.files: + if f.file_name.endswith(".gguf"): has_gguf = True - total_size += f.stat().st_size - if has_gguf: - cached.append( - { - "repo_id": repo_id, - "size_bytes": total_size, - "cache_path": str(entry), - } - ) + total_size += f.size_on_disk + if not has_gguf: + continue + # Deduplicate: keep the entry with the most data + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + "cache_path": str(repo_info.repo_path), + } + cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) return {"cached": cached} except Exception as e: logger.error(f"Error listing cached GGUF repos: {e}", exc_info = True) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 1c52d53915..76a32b6129 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -250,8 +250,10 @@ function GgufVariantExpander({ const bIsRec = b.quant === effectiveRecommended; if (aIsRec !== bIsRec) return aIsRec ? -1 : 1; - // fits/tight: largest first (best quality); OOM: smallest first - return aTier === 4 ? a.size_bytes - b.size_bytes : b.size_bytes - a.size_bytes; + // fits: largest first (best quality that fits in GPU) + // tight/OOM: smallest first (closest to fitting, fastest to run) + const fitsInGpu = aTier === 0 || aTier === 2; + return fitsInGpu ? b.size_bytes - a.size_bytes : a.size_bytes - b.size_bytes; }); }, [variants, effectiveRecommended, getGgufFit]); diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index c875c5ae84..2a1e309782 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -62,7 +62,6 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ }} > {!hideWelcome && ( @@ -140,10 +139,10 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { className="size-20" />

- Test Your Fine-tuned Model + Run LLMs or test your fine-tune

- Start a conversation to see how your model performs. + Run GGUFs, safetensors, vision and audio models!

diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx index e35e11cb5a..a20f106d42 100644 --- a/studio/frontend/src/components/ui/sonner.tsx +++ b/studio/frontend/src/components/ui/sonner.tsx @@ -1,6 +1,6 @@ -// SPDX-License-Identifier: AGPL-3.0-only -// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + import { Alert02Icon, CheckmarkCircle02Icon, @@ -19,6 +19,7 @@ const Toaster = ({ ...props }: ToasterProps) => { (null); - const [loadAbortController, setLoadAbortController] = + const [_loadAbortController, setLoadAbortController] = useState(null); + const loadAbortRef = useRef(null); + const loadingModelRef = useRef(null); + const loadToastIdRef = useRef(null); const refresh = useCallback(async () => { setModelsError(null); @@ -224,9 +227,12 @@ export function useChatModelRuntime() { .join(" "); setModelsError(null); - setLoadingModel({ id: modelId, displayName, isDownloaded }); + const loadInfo = { id: modelId, displayName, isDownloaded }; + setLoadingModel(loadInfo); + loadingModelRef.current = loadInfo; const abortCtrl = new AbortController(); setLoadAbortController(abortCtrl); + loadAbortRef.current = abortCtrl; try { async function performLoad(): Promise { if (abortCtrl.signal.aborted) throw new Error("Cancelled"); @@ -262,12 +268,18 @@ export function useChatModelRuntime() { trust_remote_code: paramsBeforeLoad.trustRemoteCode ?? false, }); + // If cancelled while loading, don't update UI to show + // the model as active -- it's being unloaded. + if (abortCtrl.signal.aborted) throw new Error("Cancelled"); + const currentParams = useChatRuntimeStore.getState().params; setParams( mergeRecommendedInference(currentParams, loadResponse, modelId), ); await refresh(); } catch (error) { + // Skip rollback if user cancelled -- model is already being unloaded. + if (abortCtrl.signal.aborted) throw error; // If we unloaded a previous model and the new load failed, attempt a rollback. if (previousWasUnloaded && previousCheckpoint) { try { @@ -292,13 +304,16 @@ export function useChatModelRuntime() { isDownloaded ? "Loading model…" : "Downloading model…", { description: loadingDescription, - duration: Infinity, + duration: 10000, action: { label: "Cancel", onClick: () => { abortCtrl.abort(); setLoadingModel(null); setLoadAbortController(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; unloadModel({ model_path: modelId }).catch(() => {}); clearCheckpoint(); toast.dismiss(toastId); @@ -307,6 +322,7 @@ export function useChatModelRuntime() { }, }, ); + loadToastIdRef.current = toastId; // Poll download progress for non-cached models let progressInterval: ReturnType | null = null; @@ -330,13 +346,16 @@ export function useChatModelRuntime() { { id: toastId, description: `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB`, - duration: Infinity, + duration: 10000, action: { label: "Cancel", onClick: () => { abortCtrl.abort(); setLoadingModel(null); setLoadAbortController(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; unloadModel({ model_path: modelId }).catch(() => {}); clearCheckpoint(); toast.dismiss(toastId); @@ -349,7 +368,7 @@ export function useChatModelRuntime() { toast.loading("Loading model…", { id: toastId, description: "Download complete. Starting inference server…", - duration: Infinity, + duration: 10000, }); if (progressInterval) clearInterval(progressInterval); } @@ -375,9 +394,14 @@ export function useChatModelRuntime() { if (progressInterval) clearInterval(progressInterval); setLoadingModel(null); setLoadAbortController(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; } } catch (error) { + if (abortCtrl.signal.aborted) return; // User cancelled, nothing to report setLoadingModel(null); + loadingModelRef.current = null; const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); @@ -412,19 +436,22 @@ export function useChatModelRuntime() { } }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]); - const cancelLoading = useCallback(async () => { - if (!loadingModel) return; - loadAbortController?.abort(); + const cancelLoading = useCallback(() => { + const model = loadingModelRef.current; + if (!model) return; + loadAbortRef.current?.abort(); + loadAbortRef.current = null; + loadingModelRef.current = null; + const tid = loadToastIdRef.current; + loadToastIdRef.current = null; setLoadingModel(null); setLoadAbortController(null); - try { - await unloadModel({ model_path: loadingModel.id }); - } catch { - // Best-effort cleanup - } clearCheckpoint(); + if (tid != null) toast.dismiss(tid); toast.info("Model loading cancelled"); - }, [loadingModel, loadAbortController, clearCheckpoint]); + // Fire-and-forget: tell backend to stop, don't block UI + unloadModel({ model_path: model.id }).catch(() => {}); + }, [clearCheckpoint]); return { refresh, diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index d040ba1d80..dbc631a1d5 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -32,7 +32,7 @@ import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import type { MessageRecord, ModelType } from "./types"; const DEFAULT_SUGGESTIONS = [ - "Draw a simple flowchart of a login system using Mermaid", + "Draw an ASCII art of a cute sloth", "Solve the integral of x²·sin(x) step by step", "Write a Python function that finds the longest palindrome in a string", "Format a comparison of 3 databases as a markdown table with pros and cons",