unsloth/studio/backend/models/export.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

132 lines
4 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Pydantic schemas for Export API.
"""
from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Dict, Any
class LoadCheckpointRequest(BaseModel):
"""Request for loading a checkpoint into the export backend."""
checkpoint_path: str = Field(..., description = "Path to the checkpoint directory")
max_seq_length: int = Field(
2048,
ge = 128,
le = 32768,
description = "Maximum sequence length for loading the model",
)
load_in_4bit: bool = Field(
True,
description = "Whether to load the model in 4-bit quantization",
)
trust_remote_code: bool = Field(
False,
description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
)
class ExportStatusResponse(BaseModel):
"""Current export backend status."""
current_checkpoint: Optional[str] = Field(
None,
description = "Path to the currently loaded checkpoint, if any",
)
is_vision: bool = Field(
False,
description = "True if the loaded checkpoint is a vision model",
)
is_peft: bool = Field(
False,
description = "True if the loaded checkpoint is a PEFT (LoRA) model",
)
class ExportOperationResponse(BaseModel):
"""Generic response for export operations."""
success: bool = Field(..., description = "True if the operation succeeded")
message: str = Field(..., description = "Human-readable status or error message")
details: Optional[Dict[str, Any]] = Field(
default = None,
description = "Optional extra details about the operation",
)
class ExportCommonOptions(BaseModel):
"""Common options for export operations that save locally and/or push to Hub."""
save_directory: str = Field(
...,
description = "Local directory where the exported artifacts will be written",
)
push_to_hub: bool = Field(
False,
description = "If True, also push the exported model to the Hugging Face Hub",
)
repo_id: Optional[str] = Field(
None,
description = "Hugging Face Hub repository ID (username/model-name)",
)
hf_token: Optional[str] = Field(
None,
description = "Hugging Face access token used for Hub operations",
)
private: bool = Field(
False,
description = "If True, create a private repository on the Hub (where applicable)",
)
base_model_id: Optional[str] = Field(
None,
description = "HuggingFace model ID of the base model (for model card metadata)",
)
class ExportMergedModelRequest(ExportCommonOptions):
"""Request for exporting a merged PEFT model."""
format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field(
"16-bit (FP16)",
description = "Export precision / format for the merged model",
)
class ExportBaseModelRequest(ExportCommonOptions):
"""Request for exporting a non-PEFT (base) model."""
# Uses fields from ExportCommonOptions only
class ExportGGUFRequest(BaseModel):
"""Request for exporting the current model to GGUF format."""
save_directory: str = Field(
...,
description = "Directory where GGUF files will be saved",
)
quantization_method: str = Field(
"Q4_K_M",
description = 'GGUF quantization method (e.g. "Q4_K_M")',
)
push_to_hub: bool = Field(
False,
description = "If True, also push GGUF artifacts to the Hugging Face Hub",
)
repo_id: Optional[str] = Field(
None,
description = "Hugging Face Hub repository ID for GGUF upload",
)
hf_token: Optional[str] = Field(
None,
description = "Hugging Face token for GGUF upload",
)
class ExportLoRAAdapterRequest(ExportCommonOptions):
"""Request for exporting only the LoRA adapter (not merged)."""
# Uses fields from ExportCommonOptions only