ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-03-25 08:34:48 -07:00
commit 93a70fbe4e
19 changed files with 644 additions and 134 deletions
Download patch file Download diff file Expand all files Collapse all files

69
install.ps1
View file

@ -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

2
studio/backend/core/inference/llama_cpp.py
View file

@ -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",

5
studio/backend/models/inference.py
View file

@ -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")

6
studio/backend/models/models.py
View file

@ -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",

31
studio/backend/routes/inference.py
View file

@ -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",
},
}

271
studio/backend/routes/models.py
View file

@ -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:

29
studio/backend/run.py
View file

@ -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

6
studio/backend/utils/paths/__init__.py
View file

@ -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",

64
studio/backend/utils/paths/storage_roots.py
View file

@ -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

6
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

@ -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(

28
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -253,6 +253,7 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise<void> {
* falls back to smallest cached safetensors model.
*/
async function autoLoadSmallestModel(): Promise<boolean> {
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<boolean> {
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<boolean> {
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<boolean> {
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<boolean> {
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<boolean> {
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<boolean> {
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,

106
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -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<Preset[]>(() =>
loadSavedCustomPresets(),
);
@ -467,32 +475,53 @@ export function ChatSettingsPanel({
<div className="flex flex-col gap-3 py-1">
{isGguf && (
<>
<div className="flex items-center justify-between gap-3">
<div className="min-w-0">
<div className="text-xs font-medium">Context Length</div>
<div className="text-[11px] text-muted-foreground">
Reported by the loaded GGUF model.
</div>
<div className="space-y-2">
<div className="flex items-center justify-between">
<span className="text-xs font-medium">Context Length</span>
<Input
type="number"
value={typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? "")}
placeholder="..."
min={128}
max={ggufContextLength ?? undefined}
step={1024}
className="h-6 w-[100px] text-right text-xs tabular-nums"
onChange={(e) => {
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);
}
}}
/>
</div>
<Input
value={ggufContextLength ?? ""}
placeholder="Loading..."
disabled={true}
className="h-7 w-[90px] text-xs"
<Slider
min={1024}
max={ggufContextLength ?? 4096}
step={1024}
value={[Math.min(typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? 4096), ggufContextLength ?? 4096)]}
onValueChange={([v]) => {
setCustomContextLength(v === (ggufContextLength ?? 0) ? null : v);
}}
/>
</div>
<div className="flex items-center justify-between gap-3">
<div className="min-w-0">
<div className="text-xs font-medium">KV Cache Dtype</div>
<div className="text-[11px] text-muted-foreground">
Quantize KV cache to reduce VRAM. Reload to apply.
Quantize KV cache to reduce VRAM.
</div>
</div>
<Select
value={kvCacheDtype ?? "f16"}
onValueChange={(v) => {
setKvCacheDtype(v === "f16" ? null : v);
onReloadModel?.();
}}
>
<SelectTrigger className="h-7 w-[90px] text-xs">
@ -507,14 +536,35 @@ export function ChatSettingsPanel({
</SelectContent>
</Select>
</div>
{modelSettingsDirty && (
<div className="flex flex-wrap gap-1.5 pt-1">
<button
type="button"
onClick={() => onReloadModel?.()}
className="rounded-md bg-primary px-2.5 py-1 text-[11px] font-medium text-primary-foreground transition-colors hover:bg-primary/90"
>
Apply
</button>
<button
type="button"
onClick={() => {
setCustomContextLength(null);
setKvCacheDtype(loadedKvCacheDtype);
}}
className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
>
Reset
</button>
</div>
)}
</>
)}
{!isGguf && (
{!isGguf && params.checkpoint && (
<div className="flex items-center justify-between gap-3">
<div className="min-w-0">
<div className="text-xs font-medium">Trust remote code</div>
<div className="text-xs font-medium">Enable custom code</div>
<div className="text-[11px] text-muted-foreground">
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.
</div>
</div>
<Switch
@ -632,6 +682,7 @@ export function ChatSettingsPanel({
onCheckedChange={onAutoTitleChange}
/>
</div>
<HfTokenField />
</div>
</CollapsibleSection>
@ -775,6 +826,29 @@ function AutoHealToolCallsToggle() {
);
}
function HfTokenField() {
const hfToken = useChatRuntimeStore((s) => s.hfToken);
const setHfToken = useChatRuntimeStore((s) => s.setHfToken);
return (
<div className="flex flex-col gap-1.5">
<div className="min-w-0">
<div className="text-xs font-medium">Hugging Face Token</div>
<div className="text-[11px] text-muted-foreground">
For downloading gated or private models.
</div>
</div>
<Input
type="password"
value={hfToken}
placeholder="hf_..."
className="h-7 text-xs font-mono"
onChange={(e) => setHfToken(e.target.value)}
/>
</div>
);
}
function ChatTemplateSection({
onReloadModel,
}: {

38
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -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,

2
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -337,7 +337,7 @@ export function SharedComposer({
async function ensureModelLoaded(sel: CompareModelSelection): Promise<string> {
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,

35
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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<string, boolean>;
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<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((set) => ({
codeToolsEnabled: false,
toolStatus: null,
kvCacheDtype: null,
loadedKvCacheDtype: null,
customContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,
})),
@ -285,6 +319,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
return { toolCallTimeout };
}),
setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
setCustomContextLength: (customContextLength) => set({ customContextLength }),
setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }),
setPendingAudio: (base64, name) =>
set({ pendingAudioBase64: base64, pendingAudioName: name }),

2
studio/frontend/src/features/chat/thread-sidebar.tsx
View file

@ -172,7 +172,7 @@ export function ThreadSidebar({
<span>Learn more in docs</span>
</a>
<a
href="https://unsloth.ai/blog"
href="https://unsloth.ai/docs/new/changelog"
target="_blank"
rel="noopener noreferrer"
className="flex items-center gap-2 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"

7
studio/frontend/src/features/studio/sections/model-section.tsx
View file

@ -181,6 +181,7 @@ export function ModelSection() {
const trainableLocalModels = useMemo(
() =>
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 (
<ComboboxItem key={id} value={id} className="gap-2">
<Tooltip>

3
studio/frontend/src/features/training/api/models-api.ts
View file

@ -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[];
}

68
unsloth_cli/commands/studio.py
View file

@ -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 ─────────────────────────────────────