diff --git a/install.ps1 b/install.ps1 index a4ed2658c9..bb0acf1237 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1,10 +1,37 @@ # Unsloth Studio Installer for Windows PowerShell # Usage: irm https://raw.githubusercontent.com/unslothai/unsloth/main/install.ps1 | iex -# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 +# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local +# Test: .\install.ps1 --package roland-sloth function Install-UnslothStudio { $ErrorActionPreference = "Stop" + # ── Parse flags ── + $StudioLocalInstall = $false + $PackageName = "unsloth" + $RepoRoot = "" + $argList = $args + for ($i = 0; $i -lt $argList.Count; $i++) { + switch ($argList[$i]) { + "--local" { $StudioLocalInstall = $true } + "--package" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red + return + } + $PackageName = $argList[$i] + } + } + } + if ($StudioLocalInstall) { + $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path + if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) { + Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red + return + } + } + $PythonVersion = "3.13" $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" $VenvDir = Join-Path $StudioHome "unsloth_studio" @@ -581,6 +608,10 @@ shell.Run cmd, 0, False # in the new venv location, while preserving existing torch/CUDA Write-Host "==> Upgrading unsloth in migrated environment..." uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.11" unsloth-zoo + if ($StudioLocalInstall) { + Write-Host "==> Overlaying local repo (editable)..." + uv pip install --python $VenvPython -e $RepoRoot --no-deps + } } elseif ($TorchIndexUrl) { Write-Host "==> Installing PyTorch ($TorchIndexUrl)..." uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl @@ -590,11 +621,23 @@ shell.Run cmd, 0, False } Write-Host "==> Installing unsloth (this may take a few minutes)..." - uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11" + if ($StudioLocalInstall) { + uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11" unsloth-zoo + Write-Host "==> Overlaying local repo (editable)..." + uv pip install --python $VenvPython -e $RepoRoot --no-deps + } else { + uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" + } } else { # Fallback: GPU detection failed to produce a URL -- let uv resolve torch Write-Host "==> Installing unsloth (this may take a few minutes)..." - uv pip install --python $VenvPython "unsloth>=2026.3.11" --torch-backend=auto + if ($StudioLocalInstall) { + uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.11" --torch-backend=auto + Write-Host "==> Overlaying local repo (editable)..." + uv pip install --python $VenvPython -e $RepoRoot --no-deps + } else { + uv pip install --python $VenvPython "$PackageName" --torch-backend=auto + } } if ($LASTEXITCODE -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red @@ -615,6 +658,11 @@ shell.Run cmd, 0, False } # Tell setup.ps1 to skip base package installation (install.ps1 already did it) $env:SKIP_STUDIO_BASE = "1" + $env:STUDIO_PACKAGE_NAME = $PackageName + if ($StudioLocalInstall) { + $env:STUDIO_LOCAL_INSTALL = "1" + $env:STUDIO_LOCAL_REPO = $RepoRoot + } & $UnslothExe studio setup if ($LASTEXITCODE -ne 0) { Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red @@ -623,6 +671,19 @@ shell.Run cmd, 0, False New-StudioShortcuts -UnslothExePath $UnslothExe + # ── Add venv Scripts dir to User PATH so `unsloth studio` works from any terminal ── + $ScriptsDir = Join-Path $VenvDir "Scripts" + $UserPath = [System.Environment]::GetEnvironmentVariable("Path", "User") + if (-not $UserPath -or $UserPath -notlike "*$ScriptsDir*") { + if ($UserPath) { + [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir;$UserPath", "User") + } else { + [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User") + } + Refresh-SessionPath + Write-Host "[OK] Added unsloth to PATH" -ForegroundColor Green + } + Write-Host "" Write-Host "=========================================" Write-Host " Unsloth Studio installed!" @@ -645,4 +706,4 @@ shell.Run cmd, 0, False } } -Install-UnslothStudio +Install-UnslothStudio @args diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 1d5643ac09..3f89cd3e5d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -848,7 +848,7 @@ class LlamaCppBackend: "--port", str(self._port), "-c", - "0", # 0 = use model's native context size + str(n_ctx) if n_ctx > 0 else "0", # 0 = model's native context size "--parallel", "1", # Single-user studio, saves VRAM "--flash-attn", diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b0498319ca..b4e496b051 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -22,7 +22,10 @@ class LoadRequest(BaseModel): None, description = "HuggingFace token for gated models" ) max_seq_length: int = Field( - 4096, ge = 128, le = 32768, description = "Maximum sequence length" + 0, + ge = 0, + le = 1048576, + description = "Maximum sequence length (0 = model default for GGUF)", ) load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") is_lora: bool = Field(False, description = "Whether this is a LoRA adapter") diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index daa8eec907..046e36137d 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -165,7 +165,7 @@ class LocalModelInfo(BaseModel): id: str = Field(..., description = "Identifier to use for loading/training") display_name: str = Field(..., description = "Display label") path: str = Field(..., description = "Local path where model data was discovered") - source: Literal["models_dir", "hf_cache"] = Field( + source: Literal["models_dir", "hf_cache", "lmstudio"] = Field( ..., description = "Discovery source", ) @@ -189,6 +189,10 @@ class LocalModelListResponse(BaseModel): None, description = "HF cache root that was scanned", ) + lmstudio_dirs: List[str] = Field( + default_factory = list, + description = "LM Studio model directories that were scanned", + ) models: List[LocalModelInfo] = Field( default_factory = list, description = "Discovered local/cached models", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index aa8c34a3c5..78d95fedbd 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,27 @@ import asyncio import threading +import re as _re + + +def _friendly_error(exc: Exception) -> str: + """Extract a user-friendly message from known llama-server errors.""" + msg = str(exc) + m = _re.search( + r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)", + msg, + ) + if m: + return ( + f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token " + f"context window. Try increasing the Context Length in Model settings, " + f"or shorten the conversation." + ) + if "Lost connection to llama-server" in msg: + return "Lost connection to the model server. It may have crashed -- try reloading the model." + return "An internal error occurred" + + # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -550,7 +571,7 @@ async def generate_stream( except Exception as e: backend.reset_generation_state() logger.error(f"Error during generation: {e}", exc_info = True) - yield f"data: {json.dumps({'error': 'An internal error occurred'})}\n\n" + yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" return StreamingResponse( stream(), @@ -944,7 +965,7 @@ async def openai_chat_completions( logger.error( f"Error during audio input streaming: {e}", exc_info = True ) - yield f"data: {json.dumps({'error': {'message': 'An internal error occurred', 'type': 'server_error'}})}\n\n" + yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" return StreamingResponse( audio_input_stream(), @@ -1176,7 +1197,7 @@ async def openai_chat_completions( logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") error_chunk = { "error": { - "message": "An internal error occurred", + "message": _friendly_error(e), "type": "server_error", }, } @@ -1314,7 +1335,7 @@ async def openai_chat_completions( logger.error(f"Error during GGUF streaming: {e}", exc_info = True) error_chunk = { "error": { - "message": "An internal error occurred", + "message": _friendly_error(e), "type": "server_error", }, } @@ -1495,7 +1516,7 @@ async def openai_chat_completions( logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) error_chunk = { "error": { - "message": "An internal error occurred", + "message": _friendly_error(e), "type": "server_error", }, } diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index e705762447..f76034c95b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -210,6 +210,76 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: return found +def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: + """Scan an LM Studio models directory for model files. + + LM Studio uses a ``publisher/model-name`` folder structure containing + GGUF files, or standalone GGUF files at the top level. + """ + if not lm_dir.exists() or not lm_dir.is_dir(): + return [] + + found: List[LocalModelInfo] = [] + for child in lm_dir.iterdir(): + if not child.is_dir(): + if child.suffix == ".gguf" and child.is_file(): + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(child), + display_name = child.stem, + path = str(child), + source = "lmstudio", + updated_at = updated_at, + ), + ) + continue + + # child is a publisher directory — scan its sub-directories + for model_dir in child.iterdir(): + if model_dir.is_dir(): + has_model = ( + any(model_dir.glob("*.gguf")) + or (model_dir / "config.json").exists() + or any(model_dir.glob("*.safetensors")) + ) + if not has_model: + continue + model_id = f"{child.name}/{model_dir.name}" + try: + updated_at = model_dir.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(model_dir), + model_id = model_id, + display_name = model_dir.name, + path = str(model_dir), + source = "lmstudio", + updated_at = updated_at, + ), + ) + elif model_dir.suffix == ".gguf" and model_dir.is_file(): + try: + updated_at = model_dir.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(model_dir), + display_name = model_dir.stem, + path = str(model_dir), + source = "lmstudio", + updated_at = updated_at, + ), + ) + return found + + @router.get("/local", response_model = LocalModelListResponse) async def list_local_models( models_dir: str = Query( @@ -218,13 +288,29 @@ async def list_local_models( current_subject: str = Depends(get_current_subject), ): """ - List local model candidates from custom models dir and HF cache. + List local model candidates from custom models dir, HF cache, + legacy Unsloth HF cache, and LM Studio directories. """ + from utils.paths import ( + legacy_hf_cache_dir, + hf_default_cache_dir, + lmstudio_model_dirs, + ) + + # Resolve all scan directories up front. + hf_cache_dir = _resolve_hf_cache_dir() + legacy_hf = legacy_hf_cache_dir() + hf_default = hf_default_cache_dir() + lm_dirs = lmstudio_model_dirs() + # Validate models_dir against an allowlist of trusted directories. # Only the trusted Path objects are used for filesystem access -- the # user-supplied string is only used for matching, never for path construction. - hf_cache_dir = _resolve_hf_cache_dir() - allowed_roots = [Path("./models").resolve(), hf_cache_dir] + allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir] + if legacy_hf.is_dir(): + allowed_roots.append(legacy_hf) + if hf_default.is_dir(): + allowed_roots.append(hf_default) try: from utils.paths import studio_root, outputs_root @@ -248,6 +334,22 @@ async def list_local_models( try: local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + # Scan legacy Unsloth HF cache for backward compatibility + if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve(): + local_models += _scan_hf_cache(legacy_hf) + + # Scan HF system default cache (may differ when env vars are overridden) + if ( + hf_default.is_dir() + and hf_default.resolve() != hf_cache_dir.resolve() + and hf_default.resolve() != legacy_hf.resolve() + ): + local_models += _scan_hf_cache(hf_default) + + # Scan LM Studio directories + for lm_dir in lm_dirs: + local_models += _scan_lmstudio_dir(lm_dir) + deduped: dict[str, LocalModelInfo] = {} for model in local_models: if model.id not in deduped: @@ -262,6 +364,7 @@ async def list_local_models( return LocalModelListResponse( models_dir = str(models_root), hf_cache_dir = str(hf_cache_dir), + lmstudio_dirs = [str(d) for d in lm_dirs], models = models, ) except Exception as e: @@ -846,46 +949,65 @@ def _get_repo_size_cached(repo_id: str) -> int: return 0 +def _all_hf_cache_scans(): + """Return scan_cache_dir results for the active, legacy, and default HF caches.""" + from huggingface_hub import scan_cache_dir + from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir + + scans = [scan_cache_dir()] + seen: set[str] = set() + try: + # Resolve the active cache dir so we can dedup + from huggingface_hub.constants import HF_HUB_CACHE + + seen.add(str(Path(HF_HUB_CACHE).resolve())) + except Exception: + pass + + for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): + extra = extra_fn() + if extra.is_dir() and str(extra.resolve()) not in seen: + seen.add(str(extra.resolve())) + try: + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra, exc) + return scans + + @router.get("/cached-gguf") async def list_cached_gguf( current_subject: str = Depends(get_current_subject), ): - """List GGUF repos that have already been downloaded to the HF cache. - - 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). - """ + """List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" try: - from huggingface_hub import scan_cache_dir + cache_scans = _all_hf_cache_scans() - 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.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), - } + for hf_cache in cache_scans: + 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 + 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.size_on_disk + if not has_gguf: + continue + 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: @@ -897,44 +1019,39 @@ async def list_cached_gguf( async def list_cached_models( current_subject: str = Depends(get_current_subject), ): - """List non-GGUF model repos that have been downloaded to the HF cache. - - Only includes repos that actually contain model weight files - (.safetensors, .bin), not repos with only config/metadata. - """ + """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" _WEIGHT_EXTENSIONS = (".safetensors", ".bin") try: - from huggingface_hub import scan_cache_dir + cache_scans = _all_hf_cache_scans() - 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 repo_id.upper().endswith("-GGUF"): - continue - total_size = sum( - f.size_on_disk for rev in repo_info.revisions for f in rev.files - ) - if total_size == 0: - continue - # Skip repos that only have config/metadata files (no weights) - has_weights = any( - f.file_name.endswith(_WEIGHT_EXTENSIONS) - for rev in repo_info.revisions - for f in rev.files - ) - if not has_weights: - continue - 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, - } + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if repo_id.upper().endswith("-GGUF"): + continue + total_size = sum( + f.size_on_disk for rev in repo_info.revisions for f in rev.files + ) + if total_size == 0: + continue + has_weights = any( + f.file_name.endswith(_WEIGHT_EXTENSIONS) + for rev in repo_info.revisions + for f in rev.files + ) + if not has_weights: + continue + 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, + } cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) return {"cached": cached} except Exception as e: @@ -989,15 +1106,17 @@ async def delete_cached_model( pass try: - from huggingface_hub import scan_cache_dir + cache_scans = _all_hf_cache_scans() - hf_cache = scan_cache_dir() target_repo = None - for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": - continue - if repo_info.repo_id.lower() == repo_id.lower(): - target_repo = repo_info + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + target_repo = repo_info + break + if target_repo is not None: break if target_repo is None: diff --git a/studio/backend/run.py b/studio/backend/run.py index e32b912c37..b892037565 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -158,6 +158,29 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int: ) +_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid" + + +def _write_pid_file(): + """Write the current process PID to the studio PID file.""" + try: + _PID_FILE.parent.mkdir(parents = True, exist_ok = True) + _PID_FILE.write_text(str(os.getpid())) + except OSError: + pass + + +def _remove_pid_file(): + """Remove the PID file if it belongs to this process.""" + try: + if _PID_FILE.is_file(): + stored = _PID_FILE.read_text().strip() + if stored == str(os.getpid()): + _PID_FILE.unlink(missing_ok = True) + except OSError: + pass + + def _graceful_shutdown(server = None): """Explicitly shut down all subprocess backends and the uvicorn server. @@ -165,6 +188,7 @@ def _graceful_shutdown(server = None): before the parent exits. This is critical on Windows where atexit handlers are unreliable after Ctrl+C. """ + _remove_pid_file() logger.info("Graceful shutdown initiated — cleaning up subprocesses...") # 1. Shut down uvicorn server (releases the listening socket) @@ -307,6 +331,11 @@ def run_server( thread.start() time.sleep(3) + _write_pid_file() + import atexit + + atexit.register(_remove_pid_file) + if not silent: display_host = _resolve_external_ip() if host == "0.0.0.0" else host diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 789052f372..44a7c8e287 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -23,6 +23,9 @@ from .storage_roots import ( unstructured_uploads_root, oxc_validator_tmp_root, tensorboard_root, + legacy_hf_cache_dir, + hf_default_cache_dir, + lmstudio_model_dirs, ensure_dir, ensure_studio_directories, resolve_under_root, @@ -53,6 +56,9 @@ __all__ = [ "unstructured_uploads_root", "oxc_validator_tmp_root", "tensorboard_root", + "legacy_hf_cache_dir", + "hf_default_cache_dir", + "lmstudio_model_dirs", "ensure_dir", "ensure_studio_directories", "resolve_under_root", diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 626e868275..9bcf3758ad 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -3,6 +3,7 @@ from __future__ import annotations +import json import os from pathlib import Path import tempfile @@ -82,22 +83,77 @@ def ensure_dir(path: Path) -> Path: return path +def legacy_hf_cache_dir() -> Path: + """Old Unsloth-specific HF hub cache, kept for backward-compat scanning.""" + return cache_root() / "huggingface" / "hub" + + +def hf_default_cache_dir() -> Path: + """Return the platform default HuggingFace hub cache (ignoring env overrides). + + This is the location HF uses when no ``HF_HUB_CACHE`` / ``HF_HOME`` + env var is set. We scan it so that models a user downloaded *before* + installing Unsloth Studio are still discovered. + """ + return Path.home() / ".cache" / "huggingface" / "hub" + + +def lmstudio_model_dirs() -> list[Path]: + """Return LM Studio model directories that exist on disk.""" + dirs: list[Path] = [] + seen: set[Path] = set() + + def _add(p: Path) -> None: + resolved = p.resolve() + if resolved not in seen and p.is_dir(): + seen.add(resolved) + dirs.append(p) + + # 1. Check LM Studio settings.json for custom downloads folder + settings_path = Path.home() / ".lmstudio" / "settings.json" + if settings_path.is_file(): + try: + with open(settings_path) as f: + settings = json.load(f) + downloads = settings.get("downloadsFolder", "") + if downloads: + _add(Path(downloads).expanduser()) + except Exception: + pass + + # 2. LM Studio current default models directory (all platforms) + _add(Path.home() / ".lmstudio" / "models") + + # 3. Legacy LM Studio cache location + _add(Path.home() / ".cache" / "lm-studio" / "models") + + return dirs + + def _setup_cache_env() -> None: """Set cache environment variables for HuggingFace, uv, and vLLM. + HuggingFace cache variables are only set when the legacy Unsloth HF + cache already exists, preserving existing model locations. New + installations leave HF at its own defaults. + Only sets variables that are not already set by the user, so explicit overrides (e.g. HF_HOME=/data/hf) are respected. Works on Linux, macOS, and Windows. """ root = cache_root() hf_dir = root / "huggingface" - defaults = { - "HF_HOME": str(hf_dir), - "HF_HUB_CACHE": str(hf_dir / "hub"), - "HF_XET_CACHE": str(hf_dir / "xet"), + defaults: dict[str, str] = { "UV_CACHE_DIR": str(root / "uv"), "VLLM_CACHE_ROOT": str(root / "vllm"), } + # Preserve legacy HF cache for existing installations + legacy_hub = hf_dir / "hub" + if legacy_hub.is_dir() and any(legacy_hub.iterdir()): + defaults["HF_HOME"] = str(hf_dir) + defaults["HF_HUB_CACHE"] = str(legacy_hub) + defaults["HF_XET_CACHE"] = str(hf_dir / "xet") + for key, value in defaults.items(): if key not in os.environ: os.environ[key] = value 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 328cba3acd..3ca9aadb1f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -454,7 +454,8 @@ export function HubModelPicker({ const recommendedIds = useMemo(() => { const all = dedupe([...models.map((model) => model.id), value ?? ""]) .filter((id) => !downloadedSet.has(id.toLowerCase())) - .filter((id) => !chatOnly || isGgufRepo(id)); + .filter((id) => !chatOnly || isGgufRepo(id)) + .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); // Sort: GGUFs first, then hub models const gguf: string[] = []; const hub: string[] = []; @@ -498,7 +499,8 @@ export function HubModelPicker({ return results .map((result) => result.id) .filter((id) => !recommendedSet.has(id)) - .filter((id) => !chatOnly || isGgufRepo(id)); + .filter((id) => !chatOnly || isGgufRepo(id)) + .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); }, [recommendedSet, results, showHfSection, chatOnly]); const metricsById = useMemo( diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 15ac416b1f..95af560305 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -253,6 +253,7 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { * falls back to smallest cached safetensors model. */ async function autoLoadSmallestModel(): Promise { + const hfToken = useChatRuntimeStore.getState().hfToken || null; const toastId = toast("Loading a model…", { description: "Auto-selecting the smallest downloaded model.", duration: 5000, @@ -278,8 +279,8 @@ async function autoLoadSmallestModel(): Promise { const variant = downloaded[0]; const loadResp = await loadModel({ model_path: repo.repo_id, - hf_token: null, - max_seq_length: 4096, + hf_token: hfToken, + max_seq_length: 0, load_in_4bit: true, is_lora: false, gguf_variant: variant.quant, @@ -308,8 +309,10 @@ async function autoLoadSmallestModel(): Promise { supportsReasoning: loadResp.supports_reasoning ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, supportsTools: loadResp.supports_tools ?? false, - toolsEnabled: false, - codeToolsEnabled: false, + toolsEnabled: loadResp.supports_tools ?? false, + codeToolsEnabled: loadResp.supports_tools ?? false, + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, }); @@ -329,7 +332,7 @@ async function autoLoadSmallestModel(): Promise { try { const sfLoadResp = await loadModel({ model_path: repo.repo_id, - hf_token: null, + hf_token: hfToken, max_seq_length: 4096, load_in_4bit: true, is_lora: false, @@ -366,8 +369,8 @@ async function autoLoadSmallestModel(): Promise { try { const loadResp = await loadModel({ model_path: "unsloth/Qwen3.5-4B-GGUF", - hf_token: null, - max_seq_length: 4096, + hf_token: hfToken, + max_seq_length: 0, load_in_4bit: true, is_lora: false, gguf_variant: "UD-Q4_K_XL", @@ -391,7 +394,10 @@ async function autoLoadSmallestModel(): Promise { supportsReasoning: loadResp.supports_reasoning ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, supportsTools: loadResp.supports_tools ?? false, - toolsEnabled: false, + toolsEnabled: loadResp.supports_tools ?? false, + codeToolsEnabled: loadResp.supports_tools ?? false, + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, }); @@ -410,8 +416,7 @@ async function autoLoadSmallestModel(): Promise { export function createOpenAIStreamAdapter(): ChatModelAdapter { return { async *run({ messages, abortSignal, unstable_threadId }) { - const runtime = useChatRuntimeStore.getState(); - const { params } = runtime; + let runtime = useChatRuntimeStore.getState(); // Wait for in-progress model load to finish before inferring if (runtime.modelLoading) { @@ -430,6 +435,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } } + // Re-read store after potential auto-load / model ready wait + runtime = useChatRuntimeStore.getState(); + const { params } = runtime; const { supportsTools, toolsEnabled, diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 081e3efc26..a315aeb005 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -279,6 +279,14 @@ export function ChatSettingsPanel({ const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype); + const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype); + const customContextLength = useChatRuntimeStore((s) => s.customContextLength); + const setCustomContextLength = useChatRuntimeStore((s) => s.setCustomContextLength); + + const ctxDisplayValue = customContextLength ?? ggufContextLength ?? ""; + const kvDirty = kvCacheDtype !== loadedKvCacheDtype; + const ctxDirty = customContextLength !== null; + const modelSettingsDirty = kvDirty || ctxDirty; const [customPresets, setCustomPresets] = useState(() => loadSavedCustomPresets(), ); @@ -467,32 +475,53 @@ export function ChatSettingsPanel({
{isGguf && ( <> -
-
-
Context Length
-
- Reported by the loaded GGUF model. -
+
+
+ Context Length + { + const raw = e.target.value; + if (raw === "") { + setCustomContextLength(null); + return; + } + const v = parseInt(raw, 10); + if (!Number.isNaN(v) && v >= 0) { + const maxCtx = ggufContextLength ?? Infinity; + const clamped = Math.min(v, maxCtx); + setCustomContextLength(clamped === (ggufContextLength ?? 0) ? null : clamped); + } + }} + />
- { + setCustomContextLength(v === (ggufContextLength ?? 0) ? null : v); + }} />
KV Cache Dtype
- Quantize KV cache to reduce VRAM. Reload to apply. + Quantize KV cache to reduce VRAM.
+ {modelSettingsDirty && ( +
+ + +
+ )} )} - {!isGguf && ( + {!isGguf && params.checkpoint && (
-
Trust remote code
+
Enable custom code
- Allow models with custom code (e.g. Nemotron). Only enable for repos you trust. + Allow models with custom code (e.g. Nemotron). Only enable if sure.
+
@@ -775,6 +826,29 @@ function AutoHealToolCallsToggle() { ); } +function HfTokenField() { + const hfToken = useChatRuntimeStore((s) => s.hfToken); + const setHfToken = useChatRuntimeStore((s) => s.setHfToken); + + return ( +
+
+
Hugging Face Token
+
+ For downloading gated or private models. +
+
+ setHfToken(e.target.value)} + /> +
+ ); +} + function ChatTemplateSection({ onReloadModel, }: { 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 cfdd4e774a..25c776948f 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 @@ -354,12 +354,13 @@ export function useChatModelRuntime() { useChatRuntimeStore.getState().params.checkpoint; const paramsBeforeLoad = useChatRuntimeStore.getState().params; const maxSeqLength = paramsBeforeLoad.maxSeqLength; + const hfToken = useChatRuntimeStore.getState().hfToken || null; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). await validateModel({ model_path: modelId, - hf_token: null, + hf_token: hfToken, max_seq_length: maxSeqLength, load_in_4bit: true, is_lora: isLora, @@ -371,11 +372,16 @@ export function useChatModelRuntime() { previousWasUnloaded = true; } - const { chatTemplateOverride, kvCacheDtype } = useChatRuntimeStore.getState(); + const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength } = useChatRuntimeStore.getState(); + // GGUF: use custom context length, or 0 = model's native context + // Non-GGUF: use the Max Seq Length slider value + const effectiveMaxSeqLength = customContextLength != null + ? customContextLength + : ggufVariant != null ? (ggufContextLength ?? 0) : maxSeqLength; const loadResponse = await loadModel({ model_path: modelId, - hf_token: null, - max_seq_length: maxSeqLength, + hf_token: hfToken, + max_seq_length: effectiveMaxSeqLength, load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, @@ -403,15 +409,27 @@ export function useChatModelRuntime() { } } } + const loadedKv = loadResponse.cache_type_kv ?? null; + const nativeCtx = loadResponse.is_gguf + ? (loadResponse.context_length ?? 131072) + : null; + // Keep customContextLength if the user set one and it differs + // from the model's native context; otherwise clear it so the + // display shows the native value without a dirty marker. + const keepCustomCtx = customContextLength != null + && customContextLength !== nativeCtx + ? customContextLength + : null; useChatRuntimeStore.setState({ - ggufContextLength: loadResponse.is_gguf - ? (loadResponse.context_length ?? 131072) - : null, + ggufContextLength: nativeCtx, supportsReasoning: loadResponse.supports_reasoning ?? false, reasoningEnabled: reasoningDefault, supportsTools: loadResponse.supports_tools ?? false, - toolsEnabled: false, - kvCacheDtype: loadResponse.cache_type_kv ?? null, + toolsEnabled: loadResponse.supports_tools ?? false, + codeToolsEnabled: loadResponse.supports_tools ?? false, + kvCacheDtype: loadedKv, + loadedKvCacheDtype: loadedKv, + customContextLength: keepCustomCtx, defaultChatTemplate: loadResponse.chat_template ?? null, chatTemplateOverride: null, }); @@ -432,7 +450,7 @@ export function useChatModelRuntime() { try { await loadModel({ model_path: previousCheckpoint, - hf_token: null, + hf_token: hfToken, max_seq_length: maxSeqLength, load_in_4bit: true, is_lora: previousIsLora, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 78cfcc66d2..5ac8c79160 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -337,7 +337,7 @@ export function SharedComposer({ async function ensureModelLoaded(sel: CompareModelSelection): Promise { const resp = await loadModel({ model_path: sel.id, - hf_token: null, + hf_token: useChatRuntimeStore.getState().hfToken || null, max_seq_length: maxSeqLength, load_in_4bit: true, is_lora: sel.isLora, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 920737a279..2d60d52043 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -14,6 +14,7 @@ const AUTO_TITLE_KEY = "unsloth_chat_auto_title"; const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls"; const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message"; const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout"; +const HF_TOKEN_KEY = "unsloth_hf_token"; const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params"; let hasShownInferencePersistenceWarning = false; @@ -62,6 +63,24 @@ function saveInt(key: string, value: number): void { } } +function loadString(key: string, fallback: string): string { + if (!canUseStorage()) return fallback; + try { + return localStorage.getItem(key) ?? fallback; + } catch { + return fallback; + } +} + +function saveString(key: string, value: string): void { + if (!canUseStorage()) return; + try { + localStorage.setItem(key, value); + } catch { + // ignore + } +} + function asFiniteNumber(value: unknown, fallback: number): number { return typeof value === "number" && Number.isFinite(value) ? value : fallback; } @@ -127,6 +146,7 @@ type ChatRuntimeStore = { loras: ChatLoraSummary[]; runningByThreadId: Record; autoTitle: boolean; + hfToken: string; modelsError: string | null; activeGgufVariant: string | null; ggufContextLength: number | null; @@ -141,6 +161,8 @@ type ChatRuntimeStore = { maxToolCallsPerMessage: number; toolCallTimeout: number; kvCacheDtype: string | null; + loadedKvCacheDtype: string | null; + customContextLength: number | null; defaultChatTemplate: string | null; chatTemplateOverride: string | null; activeThreadId: string | null; @@ -159,6 +181,7 @@ type ChatRuntimeStore = { setLoras: (loras: ChatLoraSummary[]) => void; setThreadRunning: (threadId: string, running: boolean) => void; setAutoTitle: (enabled: boolean) => void; + setHfToken: (token: string) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string, ggufVariant?: string | null) => void; setActiveThreadId: (threadId: string | null) => void; @@ -172,6 +195,7 @@ type ChatRuntimeStore = { setMaxToolCallsPerMessage: (value: number) => void; setToolCallTimeout: (value: number) => void; setKvCacheDtype: (dtype: string | null) => void; + setCustomContextLength: (v: number | null) => void; setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; @@ -184,6 +208,7 @@ export const useChatRuntimeStore = create((set) => ({ loras: [], runningByThreadId: {}, autoTitle: loadBool(AUTO_TITLE_KEY, false), + hfToken: loadString(HF_TOKEN_KEY, ""), modelsError: null, activeGgufVariant: null, ggufContextLength: null, @@ -198,6 +223,8 @@ export const useChatRuntimeStore = create((set) => ({ maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 10), toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5), kvCacheDtype: null, + loadedKvCacheDtype: null, + customContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, activeThreadId: null, @@ -235,6 +262,11 @@ export const useChatRuntimeStore = create((set) => ({ saveBool(AUTO_TITLE_KEY, autoTitle); return { autoTitle }; }), + setHfToken: (hfToken) => + set(() => { + saveString(HF_TOKEN_KEY, hfToken); + return { hfToken }; + }), setModelsError: (modelsError) => set({ modelsError }), setCheckpoint: (modelId, ggufVariant) => set((state) => ({ @@ -261,6 +293,8 @@ export const useChatRuntimeStore = create((set) => ({ codeToolsEnabled: false, toolStatus: null, kvCacheDtype: null, + loadedKvCacheDtype: null, + customContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, })), @@ -285,6 +319,7 @@ export const useChatRuntimeStore = create((set) => ({ return { toolCallTimeout }; }), setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }), + setCustomContextLength: (customContextLength) => set({ customContextLength }), setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }), setPendingAudio: (base64, name) => set({ pendingAudioBase64: base64, pendingAudioName: name }), diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx index 53c5521dc7..ba97d2ee6e 100644 --- a/studio/frontend/src/features/chat/thread-sidebar.tsx +++ b/studio/frontend/src/features/chat/thread-sidebar.tsx @@ -172,7 +172,7 @@ export function ThreadSidebar({ Learn more in docs localModels.filter((m) => { + if (m.source === "lmstudio") return false; if (m.path.endsWith(".gguf")) return false; if (m.id.toLowerCase().includes("-gguf")) return false; return true; @@ -334,7 +335,11 @@ export function ModelSection() { {(id: string) => { const model = localMetaById.get(id); const source = - model?.source === "hf_cache" ? "HF cache" : "Local dir"; + model?.source === "hf_cache" + ? "HF cache" + : model?.source === "lmstudio" + ? "LM Studio" + : "Local dir"; return ( diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index 84051e3e1d..2a9ad7c0d6 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -79,7 +79,7 @@ export interface LocalModelInfo { id: string; display_name: string; path: string; - source: "models_dir" | "hf_cache"; + source: "models_dir" | "hf_cache" | "lmstudio"; model_id?: string | null; updated_at?: number | null; } @@ -87,6 +87,7 @@ export interface LocalModelInfo { interface LocalModelListResponse { models_dir: string; hf_cache_dir?: string | null; + lmstudio_dirs: string[]; models: LocalModelInfo[]; } diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index c6d398eebd..a2f0873e22 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -166,6 +166,74 @@ def studio_default( typer.echo("\nShutting down...") +# ── unsloth studio stop ─────────────────────────────────────────────── + +_PID_FILE = STUDIO_HOME / "studio.pid" + + +@studio_app.command() +def stop(): + """Stop a running Unsloth Studio server. + + Reads the PID from ~/.unsloth/studio/studio.pid and sends SIGTERM + (or TerminateProcess on Windows) to shut it down gracefully. + """ + import signal as _signal + + if not _PID_FILE.is_file(): + typer.echo("No running Studio server found (no PID file).") + raise typer.Exit(0) + + pid_text = _PID_FILE.read_text().strip() + if not pid_text.isdigit(): + typer.echo(f"Invalid PID file contents: {pid_text}") + _PID_FILE.unlink(missing_ok = True) + raise typer.Exit(1) + + pid = int(pid_text) + + # Check if the process is still alive + try: + os.kill(pid, 0) + except ProcessLookupError: + typer.echo( + f"Studio server (PID {pid}) is not running. Cleaning up stale PID file." + ) + _PID_FILE.unlink(missing_ok = True) + raise typer.Exit(0) + except PermissionError: + pass # process exists but we may not own it; try to signal anyway + + # Send SIGTERM (graceful shutdown) or TerminateProcess on Windows + try: + if sys.platform == "win32": + subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True) + else: + os.kill(pid, _signal.SIGTERM) + typer.echo(f"Sent shutdown signal to Studio server (PID {pid}).") + except ProcessLookupError: + typer.echo(f"Studio server (PID {pid}) already exited.") + _PID_FILE.unlink(missing_ok = True) + raise typer.Exit(0) + except Exception as e: + typer.echo(f"Failed to stop Studio server (PID {pid}): {e}", err = True) + raise typer.Exit(1) + + # Wait briefly for the process to exit and clean up + for _ in range(10): + time.sleep(0.5) + try: + os.kill(pid, 0) + except ProcessLookupError: + _PID_FILE.unlink(missing_ok = True) + typer.echo("Studio server stopped.") + raise typer.Exit(0) + except PermissionError: + break + + typer.echo("Studio server is shutting down (may take a few seconds).") + + # ── unsloth studio setup / update ─────────────────────────────────────