ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/routes/models.py
Daniel Han f08aef1804 Studio (#4237) 
* Rebuild Studio branch on top of main

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix security and code quality issues for Studio PR #4237

- Validate models_dir query param against allowed directory roots
  to prevent path traversal in /api/models/local endpoint
- Replace string startswith() with Path.is_relative_to() for
  frontend path traversal check in serve_frontend
- Sanitize SSE error messages to not leak exception details to
  clients (4 locations in inference.py)
- Bind port-discovery socket to 127.0.0.1 instead of all interfaces
  in llama_cpp backend
- Import datasets_root and resolve_output_dir in embedding training
  function to fix NameError and use managed output directory
- Remove stale .gitignore entries for package-lock.json and test
  directories so tests can be tracked in version control
- Add venv-reexecution logic to ui CLI command matching the studio
  command behavior

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move models_dir path validation before try/except block

The HTTPException(403) was inside the try/except Exception handler,
so it would be caught and re-raised as a 500. Moving the validation
before the try block ensures the 403 is returned directly and also
makes the control flow clearer for static analysis (path is validated
before any filesystem operations).

* Use os.path.realpath + startswith for models_dir validation

CodeQL py/path-injection does not recognize Path.is_relative_to() as
a sanitizer. Switched to os.path.realpath + str.startswith which is
a recognized sanitizer pattern in CodeQL's taint analysis. The
startswith check uses root_str + os.sep to prevent prefix collisions
(e.g. /app/models_evil matching /app/models).

* Never pass user input to Path constructor in models_dir validation

CodeQL traces taint through Path(resolved) even after a startswith
barrier guard. Fix: the user-supplied models_dir is only used as a
string for comparison against allowed roots. The Path object passed
to _scan_models_dir comes from the trusted allowed_roots list, not
from user input. This fully breaks the taint chain.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-12 03:36:19 -07:00

620 lines
20 KiB
Python
Raw Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Model Management API routes
"""
import os
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Query
from typing import List, Optional
import structlog
from loggers import get_logger
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
from auth.authentication import get_current_subject
# Import backend functions
try:
from utils.models import (
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
is_embedding_model,
scan_checkpoints,
list_gguf_variants,
ModelConfig,
)
from utils.models.model_config import (
_pick_best_gguf,
_extract_quant_label,
is_audio_input_type,
)
from core.inference import get_inference_backend
from utils.paths import (
outputs_root,
exports_root,
resolve_output_dir,
resolve_export_dir,
)
except ImportError:
# Fallback: try to import from parent directory
parent_backend = backend_path.parent / "backend"
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
from utils.models import (
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
is_embedding_model,
scan_checkpoints,
list_gguf_variants,
ModelConfig,
)
from utils.models.model_config import (
_pick_best_gguf,
_extract_quant_label,
is_audio_input_type,
)
from core.inference import get_inference_backend
from utils.paths import (
outputs_root,
exports_root,
resolve_output_dir,
resolve_export_dir,
)
from models import (
CheckpointInfo,
CheckpointListResponse,
LocalModelInfo,
LocalModelListResponse,
ModelCheckpoints,
ModelDetails,
LoRAScanResponse,
LoRAInfo,
ModelListResponse,
)
from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType
from models.responses import (
LoRABaseModelResponse,
VisionCheckResponse,
EmbeddingCheckResponse,
)
router = APIRouter()
logger = get_logger(__name__)
def derive_model_type(
is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
) -> ModelType:
"""Collapse individual capability flags into a single model modality string."""
if is_embedding:
return "embeddings"
if audio_type is not None:
return "audio"
if is_vision:
return "vision"
return "text"
def _resolve_hf_cache_dir() -> Path:
"""Resolve local HF cache root used by hub downloads."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
if not models_dir.exists() or not models_dir.is_dir():
return []
found: List[LocalModelInfo] = []
for child in models_dir.iterdir():
if not child.is_dir():
continue
has_model_files = (
(child / "config.json").exists()
or (child / "adapter_config.json").exists()
or any(child.glob("*.safetensors"))
or any(child.glob("*.bin"))
or any(child.glob("*.gguf"))
)
if not has_model_files:
continue
try:
updated_at = child.stat().st_mtime
except OSError:
updated_at = None
found.append(
LocalModelInfo(
id = str(child),
display_name = child.name,
path = str(child),
source = "models_dir",
updated_at = updated_at,
),
)
# Also scan for standalone .gguf files directly in the models directory
for gguf_file in models_dir.glob("*.gguf"):
if gguf_file.is_file():
try:
updated_at = gguf_file.stat().st_mtime
except OSError:
updated_at = None
found.append(
LocalModelInfo(
id = str(gguf_file),
display_name = gguf_file.stem,
path = str(gguf_file),
source = "models_dir",
updated_at = updated_at,
),
)
return found
def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
if not cache_dir.exists() or not cache_dir.is_dir():
return []
found: List[LocalModelInfo] = []
for repo_dir in cache_dir.glob("models--*"):
if not repo_dir.is_dir():
continue
repo_name = repo_dir.name[len("models--") :]
if not repo_name:
continue
model_id = repo_name.replace("--", "/")
try:
updated_at = repo_dir.stat().st_mtime
except OSError:
updated_at = None
found.append(
LocalModelInfo(
id = model_id,
model_id = model_id,
display_name = model_id.split("/")[-1],
path = str(repo_dir),
source = "hf_cache",
updated_at = updated_at,
),
)
return found
@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(
default = "./models", description = "Directory to scan for local model folders"
),
current_subject: str = Depends(get_current_subject),
):
"""
List local model candidates from custom models dir and HF cache.
"""
# 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]
try:
from utils.paths import studio_root, outputs_root
allowed_roots.extend([studio_root(), outputs_root()])
except Exception:
pass
requested = os.path.realpath(os.path.expanduser(models_dir))
models_root = None
for root in allowed_roots:
root_str = os.path.realpath(str(root))
if requested == root_str or requested.startswith(root_str + os.sep):
models_root = root # Use the trusted root, not the user-supplied path
break
if models_root is None:
raise HTTPException(
status_code = 403,
detail = "Directory not allowed",
)
try:
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
if model.id not in deduped:
deduped[model.id] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
return LocalModelListResponse(
models_dir = str(models_root),
hf_cache_dir = str(hf_cache_dir),
models = models,
)
except Exception as e:
logger.error(f"Error listing local models: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = f"Failed to list local models: {str(e)}",
)
@router.get("/list")
async def list_models(
current_subject: str = Depends(get_current_subject),
):
"""
List available models (default models and loaded models).
This endpoint returns the default models and any currently loaded models.
"""
try:
inference_backend = get_inference_backend()
# Get default models
default_models = inference_backend.default_models
# Get loaded models
loaded_models = []
for model_name, model_data in inference_backend.models.items():
_is_vision = model_data.get("is_vision", False)
_audio_type = model_data.get("audio_type")
model_info = ModelDetails(
id = model_name,
name = model_name.split("/")[-1] if "/" in model_name else model_name,
is_vision = _is_vision,
is_lora = model_data.get("is_lora", False),
is_audio = model_data.get("is_audio", False),
audio_type = _audio_type,
has_audio_input = model_data.get("has_audio_input", False),
model_type = derive_model_type(_is_vision, _audio_type),
)
loaded_models.append(model_info)
# Combine default and loaded models
all_models = []
seen_ids = set()
# Add default models
for model_id in default_models:
if model_id not in seen_ids:
model_info = ModelDetails(
id = model_id,
name = model_id.split("/")[-1] if "/" in model_id else model_id,
)
all_models.append(model_info)
seen_ids.add(model_id)
# Add loaded models
for model_info in loaded_models:
if model_info.id not in seen_ids:
all_models.append(model_info)
seen_ids.add(model_info.id)
return ModelListResponse(models = all_models, default_models = default_models)
except Exception as e:
logger.error(f"Error listing models: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = f"Failed to list models: {str(e)}")
@router.get("/config/{model_name:path}")
async def get_model_config(
model_name: str,
hf_token: Optional[str] = Query(None),
current_subject: str = Depends(get_current_subject),
):
"""
Get configuration for a specific model.
This endpoint wraps the backend load_model_defaults function.
"""
try:
from utils.models.model_config import is_local_path
if not is_local_path(model_name):
model_name = model_name.lower()
logger.info(f"Getting model config for: {model_name}")
from utils.models.model_config import detect_audio_type
# Load model defaults from backend
config_dict = load_model_defaults(model_name)
# Detect model capabilities (pass HF token for gated models)
is_vision = is_vision_model(model_name)
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
audio_type = detect_audio_type(model_name, hf_token = hf_token)
# Check if it's a LoRA adapter
is_lora = False
base_model = None
try:
model_config = ModelConfig.from_identifier(model_name)
is_lora = model_config.is_lora
base_model = model_config.base_model if is_lora else None
except Exception:
pass
logger.info(
f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}"
)
return ModelDetails(
id = model_name,
model_name = model_name,
config = config_dict,
is_vision = is_vision,
is_embedding = is_embedding,
is_lora = is_lora,
is_audio = audio_type is not None,
audio_type = audio_type,
has_audio_input = is_audio_input_type(audio_type),
model_type = derive_model_type(is_vision, audio_type, is_embedding),
base_model = base_model,
)
except Exception as e:
logger.error(f"Error getting model config: {e}", exc_info = True)
raise HTTPException(
status_code = 500, detail = f"Failed to get model config: {str(e)}"
)
@router.get("/loras")
async def scan_loras(
outputs_dir: str = Query(
default = str(outputs_root()), description = "Directory to scan for LoRA adapters"
),
exports_dir: str = Query(
default = str(exports_root()), description = "Directory to scan for exported models"
),
current_subject: str = Depends(get_current_subject),
):
"""
Scan for trained LoRA adapters and exported models.
Returns both training outputs (from outputs_dir) and exported models
(from exports_dir) in a single list, distinguished by source field.
"""
try:
resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
resolved_exports_dir = str(resolve_export_dir(exports_dir))
lora_list = []
# Scan training outputs
trained_loras = scan_trained_loras(outputs_dir = resolved_outputs_dir)
for display_name, adapter_path in trained_loras:
base_model = get_base_model_from_lora(adapter_path)
lora_list.append(
LoRAInfo(
display_name = display_name,
adapter_path = adapter_path,
base_model = base_model,
source = "training",
)
)
# Scan exported models (merged, LoRA, base — skips GGUF)
exported = scan_exported_models(exports_dir = resolved_exports_dir)
for display_name, model_path, export_type, base_model in exported:
lora_list.append(
LoRAInfo(
display_name = display_name,
adapter_path = model_path,
base_model = base_model,
source = "exported",
export_type = export_type,
)
)
return LoRAScanResponse(loras = lora_list, outputs_dir = resolved_outputs_dir)
except Exception as e:
logger.error(f"Error scanning LoRAs: {e}", exc_info = True)
raise HTTPException(
status_code = 500, detail = f"Failed to scan LoRA adapters: {str(e)}"
)
@router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse)
async def get_lora_base_model(
lora_path: str,
current_subject: str = Depends(get_current_subject),
):
"""
Get the base model for a LoRA adapter.
This endpoint wraps the backend get_base_model_from_lora function.
"""
try:
base_model = get_base_model_from_lora(lora_path)
if base_model is None:
raise HTTPException(
status_code = 404,
detail = f"Could not determine base model for LoRA: {lora_path}",
)
return LoRABaseModelResponse(
lora_path = lora_path,
base_model = base_model,
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting LoRA base model: {e}", exc_info = True)
raise HTTPException(
status_code = 500, detail = f"Failed to get base model: {str(e)}"
)
@router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
async def check_vision_model(
model_name: str,
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is a vision model.
This endpoint wraps the backend is_vision_model function.
"""
try:
logger.info(f"Checking if vision model: {model_name}")
is_vision = is_vision_model(model_name)
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
return VisionCheckResponse(
model_name = model_name,
is_vision = is_vision,
)
except Exception as e:
logger.error(f"Error checking vision model: {e}", exc_info = True)
raise HTTPException(
status_code = 500, detail = f"Failed to check vision model: {str(e)}"
)
@router.get("/check-embedding/{model_name:path}", response_model = EmbeddingCheckResponse)
async def check_embedding_model(
model_name: str,
hf_token: Optional[str] = Query(None),
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is an embedding model.
This endpoint wraps the backend is_embedding_model function.
"""
try:
logger.info(f"Checking if embedding model: {model_name}")
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
logger.info(
f"Embedding check result for {model_name}: is_embedding={is_embedding}"
)
return EmbeddingCheckResponse(
model_name = model_name,
is_embedding = is_embedding,
)
except Exception as e:
logger.error(f"Error checking embedding model: {e}", exc_info = True)
raise HTTPException(
status_code = 500, detail = f"Failed to check embedding model: {str(e)}"
)
@router.get("/gguf-variants", response_model = GgufVariantsResponse)
async def get_gguf_variants(
repo_id: str = Query(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
hf_token: Optional[str] = Query(
None, description = "HuggingFace token for private repos"
),
current_subject: str = Depends(get_current_subject),
):
"""
List available GGUF quantization variants for a HuggingFace repo.
Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.)
with file sizes, whether the model supports vision, and the recommended
default variant.
"""
try:
variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
# Determine default variant
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id = repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
size_bytes = v.size_bytes,
)
for v in variants
],
has_vision = has_vision,
default_variant = default_variant,
)
except Exception as e:
logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = f"Failed to list GGUF variants: {str(e)}",
)
@router.get("/checkpoints", response_model = CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(
default = str(outputs_root()),
description = "Directory to scan for checkpoints",
),
current_subject: str = Depends(get_current_subject),
):
"""
List available checkpoints in the outputs directory.
Scans the outputs folder for training runs and their checkpoints.
"""
try:
resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
raw_models = scan_checkpoints(outputs_dir = resolved_outputs_dir)
models = [
ModelCheckpoints(
name = model_name,
checkpoints = [
CheckpointInfo(display_name = display_name, path = path, loss = loss)
for display_name, path, loss in checkpoints
],
base_model = metadata.get("base_model"),
peft_type = metadata.get("peft_type"),
lora_rank = metadata.get("lora_rank"),
)
for model_name, checkpoints, metadata in raw_models
]
return CheckpointListResponse(
outputs_dir = resolved_outputs_dir,
models = models,
)
except Exception as e:
logger.error(f"Error listing checkpoints: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = f"Failed to list checkpoints: {str(e)}",
)
Reference in a new issue View git blame Copy permalink