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>
This commit is contained in:
Daniel Han
GitHub
parent
bced78373f
commit
f08aef1804
664 changed files with 103567 additions and 4 deletions
2
cli/commands/__init__.py
Normal file
2
cli/commands/__init__.py
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
132
cli/commands/export.py
Normal file
132
cli/commands/export.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
|
||||
|
||||
EXPORT_FORMATS = ["merged-16bit", "merged-4bit", "gguf", "lora"]
|
||||
GGUF_QUANTS = ["q4_k_m", "q5_k_m", "q8_0", "f16"]
|
||||
|
||||
|
||||
def list_checkpoints(
|
||||
outputs_dir: Path = typer.Option(
|
||||
Path("./outputs"), "--outputs-dir", help = "Directory that holds training runs."
|
||||
),
|
||||
):
|
||||
"""List checkpoints detected in the outputs directory."""
|
||||
from studio.backend.core.export import ExportBackend
|
||||
|
||||
backend = ExportBackend()
|
||||
checkpoints = backend.scan_checkpoints(outputs_dir = str(outputs_dir))
|
||||
if not checkpoints:
|
||||
typer.echo("No checkpoints found.")
|
||||
raise typer.Exit()
|
||||
|
||||
for model_name, ckpt_list, metadata in checkpoints:
|
||||
typer.echo(f"\n{model_name}:")
|
||||
for display, path, loss in ckpt_list:
|
||||
loss_str = f" (loss: {loss:.4f})" if loss is not None else ""
|
||||
typer.echo(f" {display}{loss_str}: {path}")
|
||||
|
||||
|
||||
def export(
|
||||
checkpoint: Path = typer.Argument(..., help = "Path to checkpoint directory."),
|
||||
output_dir: Path = typer.Argument(..., help = "Directory to save exported model."),
|
||||
format: str = typer.Option(
|
||||
"merged-16bit",
|
||||
"--format",
|
||||
"-f",
|
||||
help = f"Export format: {', '.join(EXPORT_FORMATS)}",
|
||||
),
|
||||
quantization: str = typer.Option(
|
||||
"q4_k_m",
|
||||
"--quantization",
|
||||
"-q",
|
||||
help = f"GGUF quantization method: {', '.join(GGUF_QUANTS)}",
|
||||
),
|
||||
push_to_hub: bool = typer.Option(
|
||||
False, "--push-to-hub", help = "Push exported model to HuggingFace Hub."
|
||||
),
|
||||
repo_id: Optional[str] = typer.Option(
|
||||
None, "--repo-id", help = "HuggingFace repo ID (username/model-name)."
|
||||
),
|
||||
hf_token: Optional[str] = typer.Option(
|
||||
None, "--hf-token", envvar = "HF_TOKEN", help = "HuggingFace token."
|
||||
),
|
||||
private: bool = typer.Option(
|
||||
False, "--private", help = "Make the HuggingFace repo private."
|
||||
),
|
||||
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
|
||||
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
|
||||
):
|
||||
"""Export a checkpoint to various formats (merged, GGUF, LoRA adapter)."""
|
||||
if format not in EXPORT_FORMATS:
|
||||
typer.echo(
|
||||
f"Error: Invalid format '{format}'. Choose from: {', '.join(EXPORT_FORMATS)}",
|
||||
err = True,
|
||||
)
|
||||
raise typer.Exit(code = 2)
|
||||
|
||||
if push_to_hub and not repo_id:
|
||||
typer.echo("Error: --repo-id required when using --push-to-hub", err = True)
|
||||
raise typer.Exit(code = 2)
|
||||
|
||||
from studio.backend.core.export import ExportBackend
|
||||
|
||||
backend = ExportBackend()
|
||||
|
||||
typer.echo(f"Loading checkpoint: {checkpoint}")
|
||||
success, message = backend.load_checkpoint(
|
||||
checkpoint_path = str(checkpoint),
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = load_in_4bit,
|
||||
)
|
||||
if not success:
|
||||
typer.echo(f"Error: {message}", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
typer.echo(message)
|
||||
|
||||
typer.echo(f"Exporting as {format}...")
|
||||
if format == "merged-16bit":
|
||||
success, message = backend.export_merged_model(
|
||||
save_directory = str(output_dir),
|
||||
format_type = "16-bit (FP16)",
|
||||
push_to_hub = push_to_hub,
|
||||
repo_id = repo_id,
|
||||
hf_token = hf_token,
|
||||
private = private,
|
||||
)
|
||||
elif format == "merged-4bit":
|
||||
success, message = backend.export_merged_model(
|
||||
save_directory = str(output_dir),
|
||||
format_type = "4-bit (FP4)",
|
||||
push_to_hub = push_to_hub,
|
||||
repo_id = repo_id,
|
||||
hf_token = hf_token,
|
||||
private = private,
|
||||
)
|
||||
elif format == "gguf":
|
||||
success, message = backend.export_gguf(
|
||||
save_directory = str(output_dir),
|
||||
quantization_method = quantization.upper(),
|
||||
push_to_hub = push_to_hub,
|
||||
repo_id = repo_id,
|
||||
hf_token = hf_token,
|
||||
)
|
||||
elif format == "lora":
|
||||
success, message = backend.export_lora_adapter(
|
||||
save_directory = str(output_dir),
|
||||
push_to_hub = push_to_hub,
|
||||
repo_id = repo_id,
|
||||
hf_token = hf_token,
|
||||
private = private,
|
||||
)
|
||||
|
||||
if not success:
|
||||
typer.echo(f"Error: {message}", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
typer.echo(message)
|
||||
69
cli/commands/inference.py
Normal file
69
cli/commands/inference.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
|
||||
|
||||
def inference(
|
||||
model: str = typer.Argument(..., help = "HF model id or local path."),
|
||||
prompt: str = typer.Argument(..., help = "Prompt to send to the model."),
|
||||
hf_token: Optional[str] = typer.Option(
|
||||
None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed."
|
||||
),
|
||||
temperature: float = typer.Option(0.7, "--temperature"),
|
||||
top_p: float = typer.Option(0.9, "--top-p"),
|
||||
top_k: int = typer.Option(40, "--top-k"),
|
||||
max_new_tokens: int = typer.Option(256, "--max-new-tokens"),
|
||||
repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"),
|
||||
system_prompt: str = typer.Option(
|
||||
"",
|
||||
"--system-prompt",
|
||||
help = "Optional system prompt to prepend.",
|
||||
),
|
||||
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
|
||||
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
|
||||
):
|
||||
"""Run a single inference using the specified model."""
|
||||
from studio.backend.core import ModelConfig, get_inference_backend
|
||||
|
||||
inference_backend = get_inference_backend()
|
||||
model_config = ModelConfig.from_ui_selection(
|
||||
dropdown_value = model, search_value = None, hf_token = hf_token, is_lora = False
|
||||
)
|
||||
if not model_config:
|
||||
typer.echo("Could not resolve model config", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
if not inference_backend.load_model(
|
||||
config = model_config,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
):
|
||||
typer.echo("Model load failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
stream = inference_backend.generate_chat_response(
|
||||
messages = messages,
|
||||
system_prompt = system_prompt,
|
||||
temperature = temperature,
|
||||
top_p = top_p,
|
||||
top_k = top_k,
|
||||
max_new_tokens = max_new_tokens,
|
||||
repetition_penalty = repetition_penalty,
|
||||
)
|
||||
|
||||
typer.echo("Assistant:", nl = True)
|
||||
previous = ""
|
||||
for chunk in stream:
|
||||
delta = chunk[len(previous) :]
|
||||
if delta:
|
||||
sys.stdout.write(delta)
|
||||
sys.stdout.flush()
|
||||
previous = chunk
|
||||
sys.stdout.write("\n")
|
||||
sys.stdout.flush()
|
||||
374
cli/commands/studio.py
Normal file
374
cli/commands/studio.py
Normal file
|
|
@ -0,0 +1,374 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
import typer
|
||||
|
||||
studio_app = typer.Typer(help = "Unsloth Studio commands.")
|
||||
|
||||
STUDIO_HOME = Path.home() / ".unsloth" / "studio"
|
||||
|
||||
# __file__ is cli/commands/studio.py — two parents up is the package root
|
||||
# (either site-packages or the repo root for editable installs).
|
||||
_PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
|
||||
|
||||
def _is_repo_root(path: Path) -> bool:
|
||||
"""Check if a directory looks like the repo root (actual git clone, not site-packages)."""
|
||||
return (
|
||||
(path / ".git").exists()
|
||||
and (path / "pyproject.toml").is_file()
|
||||
and (
|
||||
(path / "studio" / "setup.sh").is_file()
|
||||
or (path / "studio" / "setup.ps1").is_file()
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _get_repo_root() -> Optional[Path]:
|
||||
"""Find the git clone repo root.
|
||||
|
||||
Used only by setup() — checks __file__ first (editable install),
|
||||
then walks CWD parents (wheel install, user is inside the clone).
|
||||
"""
|
||||
# Check 1: __file__ is in the repo (editable install)
|
||||
if _is_repo_root(_PACKAGE_ROOT):
|
||||
return _PACKAGE_ROOT
|
||||
# Check 2: CWD or any parent is the repo
|
||||
cwd = Path.cwd().resolve()
|
||||
for parent in (cwd, *cwd.parents):
|
||||
if _is_repo_root(parent):
|
||||
return parent
|
||||
return None
|
||||
|
||||
|
||||
def _studio_venv_python() -> Optional[Path]:
|
||||
"""Return the studio venv Python binary, or None if not set up."""
|
||||
if platform.system() == "Windows":
|
||||
p = STUDIO_HOME / ".venv" / "Scripts" / "python.exe"
|
||||
else:
|
||||
p = STUDIO_HOME / ".venv" / "bin" / "python"
|
||||
return p if p.is_file() else None
|
||||
|
||||
|
||||
def _find_run_py() -> Optional[Path]:
|
||||
"""Find studio/backend/run.py.
|
||||
|
||||
No CWD dependency — works from any directory.
|
||||
Since studio/ is now a proper package (has __init__.py), it lives in
|
||||
site-packages after pip install, right next to cli/.
|
||||
"""
|
||||
# 1. Relative to __file__ (site-packages or editable repo root)
|
||||
run_py = _PACKAGE_ROOT / "studio" / "backend" / "run.py"
|
||||
if run_py.is_file():
|
||||
return run_py
|
||||
# 2. Studio venv's site-packages (Linux + Windows layouts)
|
||||
for pattern in (
|
||||
"lib/python*/site-packages/studio/backend/run.py",
|
||||
"Lib/site-packages/studio/backend/run.py",
|
||||
):
|
||||
for match in (STUDIO_HOME / ".venv").glob(pattern):
|
||||
return match
|
||||
return None
|
||||
|
||||
|
||||
def _find_install_script() -> Optional[Path]:
|
||||
"""Find studio/install_python_stack.py.
|
||||
|
||||
No CWD dependency — works from any directory.
|
||||
"""
|
||||
# 1. Relative to __file__ (site-packages or editable repo root)
|
||||
s = _PACKAGE_ROOT / "studio" / "install_python_stack.py"
|
||||
if s.is_file():
|
||||
return s
|
||||
# 2. Studio venv's site-packages
|
||||
for pattern in (
|
||||
"lib/python*/site-packages/studio/install_python_stack.py",
|
||||
"Lib/site-packages/studio/install_python_stack.py",
|
||||
):
|
||||
for match in (STUDIO_HOME / ".venv").glob(pattern):
|
||||
return match
|
||||
return None
|
||||
|
||||
|
||||
def _find_setup_script() -> Optional[Path]:
|
||||
"""Find studio/setup.sh or studio/setup.ps1.
|
||||
|
||||
No CWD dependency — works from any directory.
|
||||
"""
|
||||
name = "setup.ps1" if platform.system() == "Windows" else "setup.sh"
|
||||
# 1. Relative to __file__ (site-packages or editable repo root)
|
||||
s = _PACKAGE_ROOT / "studio" / name
|
||||
if s.is_file():
|
||||
return s
|
||||
# 2. Studio venv's site-packages
|
||||
for pattern in (
|
||||
f"lib/python*/site-packages/studio/{name}",
|
||||
f"Lib/site-packages/studio/{name}",
|
||||
):
|
||||
for match in (STUDIO_HOME / ".venv").glob(pattern):
|
||||
return match
|
||||
return None
|
||||
|
||||
|
||||
# ── unsloth studio (server) ──────────────────────────────────────────
|
||||
|
||||
|
||||
@studio_app.callback(invoke_without_command = True)
|
||||
def studio_default(
|
||||
ctx: typer.Context,
|
||||
port: int = typer.Option(8000, "--port", "-p"),
|
||||
host: str = typer.Option("0.0.0.0", "--host", "-H"),
|
||||
frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"),
|
||||
silent: bool = typer.Option(False, "--silent", "-q"),
|
||||
):
|
||||
"""Launch the Unsloth Studio server."""
|
||||
if ctx.invoked_subcommand is not None:
|
||||
return
|
||||
|
||||
# Always use the studio venv if it exists and we're not already in it
|
||||
studio_venv_dir = STUDIO_HOME / ".venv"
|
||||
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
|
||||
|
||||
if not in_studio_venv:
|
||||
studio_python = _studio_venv_python()
|
||||
run_py = _find_run_py()
|
||||
if studio_python and run_py:
|
||||
if not silent:
|
||||
typer.echo("Launching with studio venv...")
|
||||
args = [
|
||||
str(studio_python),
|
||||
str(run_py),
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
]
|
||||
if frontend:
|
||||
args.extend(["--frontend", str(frontend)])
|
||||
if silent:
|
||||
args.append("--silent")
|
||||
os.execvp(str(studio_python), args)
|
||||
else:
|
||||
typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
from studio.backend.run import run_server
|
||||
|
||||
if not silent:
|
||||
from studio.backend.run import _resolve_external_ip
|
||||
|
||||
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
|
||||
|
||||
run_server(
|
||||
host = host,
|
||||
port = port,
|
||||
frontend_path = frontend,
|
||||
silent = silent,
|
||||
)
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
typer.echo("\nShutting down...")
|
||||
|
||||
|
||||
# ── unsloth studio setup ─────────────────────────────────────────────
|
||||
|
||||
|
||||
@studio_app.command()
|
||||
def setup():
|
||||
"""Run one-time Studio environment setup."""
|
||||
# If we're inside a git clone, use the full setup script (builds frontend, etc.)
|
||||
repo = _get_repo_root()
|
||||
if repo:
|
||||
_dev_setup(repo)
|
||||
else:
|
||||
_pip_setup()
|
||||
|
||||
|
||||
def _dev_setup(repo_root: Path):
|
||||
"""Git-clone: run setup.sh / setup.ps1."""
|
||||
studio_dir = repo_root / "studio"
|
||||
if platform.system() == "Windows":
|
||||
script = studio_dir / "setup.ps1"
|
||||
subprocess.run(
|
||||
["powershell", "-ExecutionPolicy", "Bypass", "-File", str(script)],
|
||||
check = True,
|
||||
)
|
||||
else:
|
||||
script = studio_dir / "setup.sh"
|
||||
subprocess.run(["bash", str(script)], check = True)
|
||||
|
||||
|
||||
def _pip_setup():
|
||||
"""Pip-install: create studio venv, install all deps, build extras."""
|
||||
import venv as _venv
|
||||
|
||||
venv_dir = STUDIO_HOME / ".venv"
|
||||
venv_t5_dir = STUDIO_HOME / ".venv_t5"
|
||||
|
||||
if platform.system() == "Windows":
|
||||
venv_python = venv_dir / "Scripts" / "python.exe"
|
||||
venv_pip = venv_dir / "Scripts" / "pip.exe"
|
||||
else:
|
||||
venv_python = venv_dir / "bin" / "python"
|
||||
venv_pip = venv_dir / "bin" / "pip"
|
||||
|
||||
typer.echo("Setting up Unsloth Studio...")
|
||||
|
||||
# 1. Create venv
|
||||
if not venv_python.is_file():
|
||||
typer.echo(f" Creating venv at {venv_dir}...")
|
||||
STUDIO_HOME.mkdir(parents = True, exist_ok = True)
|
||||
_venv.create(str(venv_dir), with_pip = True)
|
||||
|
||||
# 2. Install all Python deps via install_python_stack.py
|
||||
install_script = _find_install_script()
|
||||
if install_script:
|
||||
typer.echo(" Installing Python dependencies...")
|
||||
subprocess.run([str(venv_python), str(install_script)], check = True)
|
||||
else:
|
||||
typer.echo("Error: Could not find install_python_stack.py")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# 3. Pre-install transformers 5.x overlay
|
||||
if venv_t5_dir.is_dir() and any(venv_t5_dir.iterdir()):
|
||||
typer.echo(f" Transformers 5.x overlay already at {venv_t5_dir}")
|
||||
else:
|
||||
typer.echo(" Installing transformers 5.x overlay...")
|
||||
venv_t5_dir.mkdir(parents = True, exist_ok = True)
|
||||
subprocess.run(
|
||||
[
|
||||
str(venv_pip),
|
||||
"install",
|
||||
"--target",
|
||||
str(venv_t5_dir),
|
||||
"--no-deps",
|
||||
"transformers==5.2.0",
|
||||
],
|
||||
check = True,
|
||||
)
|
||||
subprocess.run(
|
||||
[
|
||||
str(venv_pip),
|
||||
"install",
|
||||
"--target",
|
||||
str(venv_t5_dir),
|
||||
"--no-deps",
|
||||
"huggingface_hub==1.3.0",
|
||||
],
|
||||
check = True,
|
||||
)
|
||||
typer.echo(f" Installed to {venv_t5_dir}")
|
||||
|
||||
# 4. Build llama.cpp
|
||||
_build_llama_cpp()
|
||||
|
||||
typer.echo("")
|
||||
typer.echo("Setup complete! Run 'unsloth studio' to start.")
|
||||
|
||||
|
||||
def _build_llama_cpp():
|
||||
"""Build llama.cpp at ~/.unsloth/llama.cpp/."""
|
||||
import shutil
|
||||
|
||||
unsloth_home = Path.home() / ".unsloth"
|
||||
llama_dir = unsloth_home / "llama.cpp"
|
||||
|
||||
if not shutil.which("cmake"):
|
||||
typer.echo(" cmake not found — skipping llama.cpp build")
|
||||
return
|
||||
if not shutil.which("git"):
|
||||
typer.echo(" git not found — skipping llama.cpp build")
|
||||
return
|
||||
|
||||
typer.echo(" Building llama.cpp for GGUF inference...")
|
||||
|
||||
if llama_dir.exists():
|
||||
shutil.rmtree(llama_dir)
|
||||
unsloth_home.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
"git",
|
||||
"clone",
|
||||
"--depth",
|
||||
"1",
|
||||
"https://github.com/ggml-org/llama.cpp.git",
|
||||
str(llama_dir),
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
typer.echo(" Failed to clone llama.cpp")
|
||||
return
|
||||
|
||||
cmake_args = []
|
||||
nvcc_path = shutil.which("nvcc")
|
||||
if not nvcc_path and Path("/usr/local/cuda/bin/nvcc").is_file():
|
||||
nvcc_path = "/usr/local/cuda/bin/nvcc"
|
||||
if nvcc_path:
|
||||
typer.echo(f" Building with CUDA (nvcc: {nvcc_path})...")
|
||||
cmake_args.append("-DGGML_CUDA=ON")
|
||||
else:
|
||||
typer.echo(" Building CPU-only...")
|
||||
|
||||
build_dir = llama_dir / "build"
|
||||
result = subprocess.run(
|
||||
["cmake", "-S", str(llama_dir), "-B", str(build_dir)] + cmake_args,
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
typer.echo(" cmake configure failed")
|
||||
return
|
||||
|
||||
ncpu = str(os.cpu_count() or 4)
|
||||
result = subprocess.run(
|
||||
[
|
||||
"cmake",
|
||||
"--build",
|
||||
str(build_dir),
|
||||
"--config",
|
||||
"Release",
|
||||
"--target",
|
||||
"llama-server",
|
||||
f"-j{ncpu}",
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
typer.echo(" llama-server build failed")
|
||||
return
|
||||
|
||||
subprocess.run(
|
||||
[
|
||||
"cmake",
|
||||
"--build",
|
||||
str(build_dir),
|
||||
"--config",
|
||||
"Release",
|
||||
"--target",
|
||||
"llama-quantize",
|
||||
f"-j{ncpu}",
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
)
|
||||
|
||||
server_bin = build_dir / "bin" / "llama-server"
|
||||
if server_bin.is_file():
|
||||
typer.echo(f" llama-server built at {server_bin}")
|
||||
else:
|
||||
typer.echo(" llama-server binary not found after build")
|
||||
144
cli/commands/train.py
Normal file
144
cli/commands/train.py
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
|
||||
from cli.config import Config, load_config
|
||||
from cli.options import add_options_from_config
|
||||
|
||||
|
||||
@add_options_from_config(Config)
|
||||
def train(
|
||||
config: Optional[Path] = typer.Option(
|
||||
None,
|
||||
"--config",
|
||||
"-c",
|
||||
help = "Path to YAML/JSON config file. CLI flags override config values.",
|
||||
),
|
||||
hf_token: Optional[str] = typer.Option(
|
||||
None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed."
|
||||
),
|
||||
wandb_token: Optional[str] = typer.Option(
|
||||
None, "--wandb-token", envvar = "WANDB_API_KEY", help = "Weights & Biases API key."
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
False,
|
||||
"--dry-run",
|
||||
help = "Show resolved config and exit without training.",
|
||||
),
|
||||
config_overrides: dict = None,
|
||||
):
|
||||
"""Launch training using the existing Unsloth training backend."""
|
||||
try:
|
||||
cfg = load_config(config)
|
||||
except FileNotFoundError as e:
|
||||
typer.echo(f"Error: {e}", err = True)
|
||||
raise typer.Exit(code = 2)
|
||||
|
||||
cfg.apply_overrides(**config_overrides)
|
||||
|
||||
# CLI/env tokens take precedence over config
|
||||
# Handle case where typer.Option isn't resolved (decorator interaction)
|
||||
from typer.models import OptionInfo
|
||||
|
||||
if isinstance(hf_token, OptionInfo):
|
||||
hf_token = None
|
||||
if isinstance(wandb_token, OptionInfo):
|
||||
wandb_token = None
|
||||
hf_token = hf_token or cfg.logging.hf_token
|
||||
wandb_token = wandb_token or cfg.logging.wandb_token
|
||||
|
||||
if dry_run:
|
||||
import yaml
|
||||
|
||||
data = cfg.model_dump()
|
||||
data["training"]["output_dir"] = str(data["training"]["output_dir"])
|
||||
typer.echo(yaml.dump(data, default_flow_style = False, sort_keys = False))
|
||||
raise typer.Exit(code = 0)
|
||||
|
||||
if not cfg.model:
|
||||
typer.echo("Error: provide --model or set model in --config", err = True)
|
||||
raise typer.Exit(code = 2)
|
||||
|
||||
if not cfg.data.dataset and not cfg.data.local_dataset:
|
||||
typer.echo(
|
||||
"Error: provide --dataset or --local-dataset (or via --config)", err = True
|
||||
)
|
||||
raise typer.Exit(code = 2)
|
||||
|
||||
# Check if the model path is a LoRA adapter (has adapter_config.json)
|
||||
model_path = Path(cfg.model) if cfg.model else None
|
||||
model_is_lora = (
|
||||
model_path
|
||||
and model_path.is_dir()
|
||||
and (model_path / "adapter_config.json").exists()
|
||||
)
|
||||
use_lora = cfg.training.training_type.lower() == "lora"
|
||||
|
||||
if model_is_lora and not use_lora:
|
||||
typer.echo(
|
||||
"Error: Cannot do full finetuning on a LoRA adapter. "
|
||||
"Use --training-type lora or provide a base model.",
|
||||
err = True,
|
||||
)
|
||||
raise typer.Exit(code = 2)
|
||||
|
||||
from studio.backend.core.training.trainer import UnslothTrainer
|
||||
|
||||
trainer = UnslothTrainer()
|
||||
|
||||
# Load model (trainer.is_vlm is set after this)
|
||||
if not trainer.load_model(
|
||||
model_name = cfg.model,
|
||||
max_seq_length = cfg.training.max_seq_length,
|
||||
load_in_4bit = cfg.training.load_in_4bit if use_lora else False,
|
||||
hf_token = hf_token,
|
||||
):
|
||||
typer.echo("Model load failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
is_vision = trainer.is_vlm
|
||||
|
||||
if not trainer.prepare_model_for_training(**cfg.model_kwargs(use_lora, is_vision)):
|
||||
typer.echo("Model preparation failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
result = trainer.load_and_format_dataset(
|
||||
dataset_source = cfg.data.dataset or "",
|
||||
format_type = cfg.data.format_type,
|
||||
local_datasets = cfg.data.local_dataset,
|
||||
)
|
||||
if result is None:
|
||||
typer.echo("Dataset load failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
ds, eval_ds = result
|
||||
|
||||
training_kwargs = cfg.training_kwargs()
|
||||
training_kwargs["wandb_token"] = wandb_token # CLI/env takes precedence
|
||||
started = trainer.start_training(
|
||||
dataset = ds, eval_dataset = eval_ds, **training_kwargs
|
||||
)
|
||||
|
||||
if not started:
|
||||
typer.echo("Training failed to start", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
try:
|
||||
while trainer.training_thread and trainer.training_thread.is_alive():
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
typer.echo("Stopping training (Ctrl+C detected)...")
|
||||
trainer.stop_training()
|
||||
finally:
|
||||
if trainer.training_thread:
|
||||
trainer.training_thread.join()
|
||||
|
||||
final = trainer.get_training_progress()
|
||||
if getattr(final, "error", None):
|
||||
typer.echo(f"Training error: {final.error}", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
76
cli/commands/ui.py
Normal file
76
cli/commands/ui.py
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
|
||||
|
||||
def ui(
|
||||
port: int = typer.Option(
|
||||
8000, "--port", "-p", help = "Port to run the UI server on."
|
||||
),
|
||||
host: str = typer.Option(
|
||||
"0.0.0.0", "--host", "-H", help = "Host address to bind to."
|
||||
),
|
||||
frontend: Optional[Path] = typer.Option(
|
||||
None, "--frontend", "-f", help = "Path to frontend build directory."
|
||||
),
|
||||
silent: bool = typer.Option(
|
||||
False, "--silent", "-q", help = "Suppress startup messages."
|
||||
),
|
||||
):
|
||||
"""Launch the Unsloth web UI backend server (alias for 'unsloth studio')."""
|
||||
from cli.commands.studio import _studio_venv_python, _find_run_py, STUDIO_HOME
|
||||
|
||||
# Re-execute in studio venv if available and not already inside it
|
||||
studio_venv_dir = STUDIO_HOME / ".venv"
|
||||
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
|
||||
|
||||
if not in_studio_venv:
|
||||
studio_python = _studio_venv_python()
|
||||
run_py = _find_run_py()
|
||||
if studio_python and run_py:
|
||||
if not silent:
|
||||
typer.echo("Launching with studio venv...")
|
||||
args = [
|
||||
str(studio_python),
|
||||
str(run_py),
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
]
|
||||
if frontend:
|
||||
args.extend(["--frontend", str(frontend)])
|
||||
if silent:
|
||||
args.append("--silent")
|
||||
os.execvp(str(studio_python), args)
|
||||
else:
|
||||
typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
from studio.backend.run import run_server
|
||||
|
||||
if not silent:
|
||||
from studio.backend.run import _resolve_external_ip
|
||||
|
||||
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
|
||||
|
||||
run_server(
|
||||
host = host,
|
||||
port = port,
|
||||
frontend_path = frontend,
|
||||
silent = silent,
|
||||
)
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
typer.echo("\nShutting down...")
|
||||
Loading…
Add table
Add a link
Reference in a new issue