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

Merge remote-tracking branch 'origin/nightly' into feature/training-frontend-integration

Browse source
This commit is contained in:
Shine1i 2026-02-13 12:41:09 +01:00
commit b0cb7a7305
9 changed files with 139 additions and 220 deletions
Download patch file Download diff file Expand all files Collapse all files

4
cli/commands/ui.py
View file

@ -15,7 +15,9 @@ def ui(
from studio.backend.run import run_server
if not silent:
typer.echo(f"Starting Unsloth UI on http://{host}:{port}")
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,

42
setup.sh
View file

@ -3,6 +3,23 @@ set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# ── Helper: run command quietly, show output only on failure ──
run_quiet() {
local label="$1"
shift
local tmplog
tmplog=$(mktemp)
if "$@" > "$tmplog" 2>&1; then
rm -f "$tmplog"
else
local exit_code=$?
echo "❌ $label failed (exit code $exit_code):"
cat "$tmplog"
rm -f "$tmplog"
exit $exit_code
fi
}
echo "╔══════════════════════════════════════╗"
echo "║ Unsloth Studio Setup Script ║"
echo "╚══════════════════════════════════════╝"
@ -25,7 +42,7 @@ fi
if [ "$NEED_NODE" = true ]; then
# ── 2. Install nvm ──
echo "Installing nvm..."
curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash > /dev/null 2>&1
# Load nvm (source ~/.bashrc won't work inside a script)
export NVM_DIR="$HOME/.nvm"
@ -33,8 +50,8 @@ if [ "$NEED_NODE" = true ]; then
# ── 3. Install Node LTS ──
echo "Installing Node LTS..."
nvm install --lts
nvm use --lts
run_quiet "nvm install" nvm install --lts
nvm use --lts > /dev/null 2>&1
# ── 4. Verify versions ──
NODE_MAJOR=$(node -v | sed 's/v//' | cut -d. -f1)
@ -46,7 +63,7 @@ if [ "$NEED_NODE" = true ]; then
fi
if [ "$NPM_MAJOR" -lt 11 ]; then
echo "⚠️ npm version is $(npm -v), expected >= 11. Updating..."
npm install -g npm@latest
run_quiet "npm update" npm install -g npm@latest
fi
fi
@ -56,8 +73,8 @@ echo "✅ Node $(node -v) | npm $(npm -v)"
echo ""
echo "Building frontend..."
cd "$SCRIPT_DIR/studio/frontend"
npm install
npm run build
run_quiet "npm install" npm install
run_quiet "npm run build" npm run build
cd "$SCRIPT_DIR"
echo "✅ Frontend built to studio/frontend/dist"
@ -66,9 +83,11 @@ echo ""
echo "Setting up Python environment..."
python3 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install unsloth-zoo unsloth
pip install typer fastapi uvicorn pydantic matplotlib pandas "datasets==4.3.0"
run_quiet "pip upgrade" pip install --upgrade pip
echo " Installing unsloth-zoo + unsloth..."
run_quiet "pip install unsloth" pip install unsloth-zoo unsloth
echo " Installing studio dependencies..."
run_quiet "pip install extras" pip install typer fastapi uvicorn pydantic matplotlib pandas nest_asyncio "datasets==4.3.0" pyjwt
echo "✅ Python dependencies installed"
# ── 7. Add shell alias ──
@ -76,7 +95,7 @@ echo "✅ Python dependencies installed"
# This alias hardcodes the venv python path so users don't need to activate.
echo ""
REPO_DIR="$SCRIPT_DIR"
ALIAS_LINE="alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui'"
ALIAS_LINE="unsloth-ui() { ${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist \"\$@\"; }"
if ! grep -qF "unsloth-ui" ~/.bashrc 2>/dev/null; then
echo "" >> ~/.bashrc
@ -94,6 +113,5 @@ echo "╠═══════════════════════
echo "║ Run 'source ~/.bashrc' or open a ║"
echo "║ new terminal, then launch with: ║"
echo "║ ║"
echo "║ unsloth-ui -H 0.0.0.0 -p 8000 \ ║"
echo "║ -f studio/frontend/dist ║"
echo "║ unsloth-ui -H 0.0.0.0 -p 8000 ║"
echo "╚══════════════════════════════════════╝"

14
studio/backend/auth/storage.py
View file

@ -84,6 +84,20 @@ def create_initial_user(username: str, password: str, jwt_secret: str) -> None:
conn.close()
def delete_user(username: str) -> None:
"""
Delete a user from the database.
Used for rollback when setup fails after user creation.
"""
conn = get_connection()
try:
conn.execute("DELETE FROM auth_user WHERE username = ?", (username,))
conn.commit()
finally:
conn.close()
def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str]]:
"""
Get user's password salt, hash, and JWT secret.

3
studio/backend/core/__init__.py
View file

@ -6,7 +6,7 @@ Unified core module for Unsloth backend
from .inference import InferenceBackend, get_inference_backend
# Training
from .training import UnslothTrainer, get_trainer, TrainingBackend, get_training_backend, create_training_handlers, TrainingProgress
from .training import UnslothTrainer, get_trainer, TrainingBackend, get_training_backend, TrainingProgress
# Configuration (from utils)
from utils.models import is_vision_model, ModelConfig, scan_trained_loras, load_model_defaults, get_base_model_from_lora
@ -27,7 +27,6 @@ __all__ = [
'get_trainer',
'get_training_backend',
'TrainingBackend',
'create_training_handlers',
'TrainingProgress',
# Config

3
studio/backend/core/training/__init__.py
View file

@ -2,7 +2,7 @@
Training submodule - Training backends and trainer classes
"""
from .trainer import UnslothTrainer, get_trainer, TrainingProgress
from .training import TrainingBackend, get_training_backend, create_training_handlers
from .training import TrainingBackend, get_training_backend
__all__ = [
'UnslothTrainer',
@ -10,5 +10,4 @@ __all__ = [
'TrainingProgress',
'TrainingBackend',
'get_training_backend',
'create_training_handlers',
]

4
studio/backend/core/training/trainer.py
View file

@ -1,6 +1,6 @@
"""
Unsloth Training Backend
Integrates Unsloth training capabilities with the Gradio UI
Integrates Unsloth training capabilities with the FastAPI backend
"""
import torch
from utils.hardware import clear_gpu_cache
@ -46,7 +46,7 @@ class TrainingProgress:
class UnslothTrainer:
"""
Unsloth Training Backend for Gradio UI Integration
Unsloth Training Backend
"""
def __init__(self):

214
studio/backend/core/training/training.py
View file

@ -1,9 +1,8 @@
"""
Training backend and UI integration
Training backend for FastAPI integration
"""
import gradio as gr
import matplotlib.pyplot as plt
from typing import Dict, Any, Generator, Tuple
from typing import Any, Generator, Tuple
import logging
from .trainer import get_trainer, TrainingProgress
@ -17,7 +16,7 @@ PLOT_HEIGHT = 3.5 # Inches
class TrainingBackend:
"""
Training orchestration and UI integration.
Training orchestration backend.
Handles both text and vision models, LoRA and full finetuning.
"""
@ -91,12 +90,12 @@ class TrainingBackend:
wandb_token: str,
wandb_project: str,
enable_tensorboard: bool,
tensorboard_dir: str) -> Generator[Tuple, None, None]:
tensorboard_dir: str) -> bool:
"""
Start training - yields UI updates as generator.
Start training.
Yields:
Tuple of (start_btn_update, stop_btn_update, progress_visible, config_visible)
Returns:
True if training started successfully, False otherwise.
"""
try:
# Reset stop flag and clear history
@ -107,20 +106,12 @@ class TrainingBackend:
import time
output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}"
# NEW: Derive use_lora from training_type
# Derive use_lora from training_type
use_lora_actual = (training_type == "LoRA/QLoRA")
if use_lora_actual: print("using Lora")
else: print("using full finetuning")
logger.info(f"Starting training - Type: {training_type}, Model: {model_name}")
# Yield initial status - buttons toggle immediately
yield (
gr.update(interactive=False), # Start button disabled
gr.update(interactive=True), # Stop button enabled
gr.update(visible=True), # Training progress visible
#gr.update(visible=False) # Config selection hidden
)
# ========== LOAD MODEL ==========
logger.info("Loading model...")
success = self.trainer.load_model(
@ -132,17 +123,7 @@ class TrainingBackend:
if not success or self.trainer.should_stop:
logger.error("Failed to load model or stopped by user")
return
# Capture if this is a vision model
#self.current_training_session['is_vlm'] = self.trainer.is_vlm
yield (
gr.update(interactive=False),
gr.update(interactive=True),
gr.update(visible=True),
#gr.update(visible=False)
)
return False
# ========== PREPARE MODEL FOR TRAINING ==========
if use_lora_actual:
@ -171,14 +152,7 @@ class TrainingBackend:
if not success or self.trainer.should_stop:
logger.error("Failed to prepare model or stopped by user")
return
yield (
gr.update(interactive=False),
gr.update(interactive=True),
gr.update(visible=True),
#gr.update(visible=False)
)
return False
# ========== LOAD DATASET ==========
logger.info("Loading dataset...")
@ -191,14 +165,7 @@ class TrainingBackend:
if dataset is None or self.trainer.should_stop:
logger.error("Failed to load dataset or stopped by user")
return
yield (
gr.update(interactive=False),
gr.update(interactive=True),
gr.update(visible=True),
#gr.update(visible=False)
)
return False
# ========== START TRAINING ==========
# Convert learning rate string to float
@ -241,12 +208,9 @@ class TrainingBackend:
if not success:
logger.error("Failed to start training")
yield (
gr.update(interactive=True),
gr.update(interactive=False),
gr.update(visible=False),
#gr.update(visible=True)
)
return False
return True
except Exception as e:
logger.error(f"Error in start_training: {e}", exc_info=True)
@ -254,40 +218,24 @@ class TrainingBackend:
error=str(e),
is_training=False
)
yield (
gr.update(interactive=True),
gr.update(interactive=False),
gr.update(visible=False),
#gr.update(visible=True)
)
return False
def stop_training(self) -> Tuple:
def stop_training(self) -> bool:
"""
Stop ongoing training.
Returns:
Tuple of (start_btn_update, stop_btn_update, progress_visible, config_visible)
True if training was successfully stopped.
"""
try:
logger.info("Stopping training...")
self.trainer.stop_training()
return (
gr.update(interactive=True), # Start button enabled
gr.update(interactive=False), # Stop button disabled
gr.update(visible=False), # Training progress hidden
#gr.update(visible=True) # Config selection visible
)
return True
except Exception as e:
logger.error(f"Error stopping training: {e}")
return (
gr.update(interactive=True),
gr.update(interactive=False),
gr.update(visible=False),
#gr.update(visible=True)
)
return False
def get_training_status(self, theme: str = "light") -> Tuple[plt.Figure, gr.update, gr.update, gr.update]:
def get_training_status(self, theme: str = "light") -> Tuple:
"""
Get current training status and loss plot.
@ -295,7 +243,7 @@ class TrainingBackend:
theme: "light" or "dark" for plot styling
Returns:
Tuple of (plot, start_btn, stop_btn, progress_visible)
Tuple of (plot, progress)
"""
try:
@ -303,26 +251,15 @@ class TrainingBackend:
# If not training and not completed, return no updates
if not (progress.is_training or progress.is_completed or progress.error):
return (None, gr.update(), gr.update(), gr.update())
return (None, progress)
# Generate plot
plot = self._create_loss_plot(progress, theme)
# If completed or error, enable start button
if progress.is_completed or progress.error:
return (
plot,
gr.update(interactive=True), # Start button enabled
gr.update(interactive=False), # Stop button disabled
gr.update(visible=True), # Training progress visible
)
# Still training - no button updates
return (plot, gr.update(), gr.update(), gr.update())
return (plot, progress)
except Exception as e:
logger.error(f"Error getting training status: {e}")
return (None, gr.update(), gr.update(), gr.update())
return (None, None)
def refresh_plot_for_theme(self, theme: str) -> plt.Figure:
"""
@ -578,106 +515,3 @@ def get_training_backend() -> TrainingBackend:
if _training_backend is None:
_training_backend = TrainingBackend()
return _training_backend
# ========== UI HANDLER CREATION ==========
def create_training_handlers(train_components: Dict[str, Any]) -> Dict[str, Any]:
"""
Create training event handlers for Gradio UI components.
Args:
train_components: Dictionary of Gradio components from train page
Returns:
Dictionary of handler functions
"""
backend = get_training_backend()
def start_training_handler(*args):
"""Handler for start training button - yields status updates"""
try:
# Extract parameters in the order they're passed from the UI
(model_name, training_type, hf_token, load_4bit, max_seq_length,
hf_dataset, local_datasets, format_type,
num_epochs, learning_rate, batch_size, gradient_accumulation_steps,
warmup_steps, warmup_ratio, max_steps, save_steps, weight_decay, random_seed, packing,
optim, lr_scheduler_type,
use_lora, lora_r, lora_alpha, lora_dropout, target_modules,
gradient_checkpointing, use_rslora, use_loftq, train_on_completions,
finetune_vision_layers, finetune_language_layers,
finetune_attention_modules, finetune_mlp_modules,
enable_wandb, wandb_token, wandb_project,
enable_tensorboard, tensorboard_dir) = args
# Start training with correctly named parameters - this is a generator
for update_tuple in backend.start_training(
model_name=model_name,
training_type=training_type,
hf_token=hf_token,
load_in_4bit=load_4bit,
max_seq_length=max_seq_length,
hf_dataset=hf_dataset,
local_datasets=local_datasets,
format_type=format_type,
num_epochs=num_epochs,
learning_rate=learning_rate,
batch_size=batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
warmup_steps=warmup_steps,
warmup_ratio=warmup_ratio,
max_steps=max_steps,
save_steps=save_steps,
weight_decay=weight_decay,
random_seed=random_seed,
packing=packing,
optim=optim,
lr_scheduler_type=lr_scheduler_type,
use_lora=use_lora,
lora_r=lora_r,
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
target_modules=target_modules,
gradient_checkpointing=gradient_checkpointing,
use_rslora=use_rslora,
use_loftq=use_loftq,
train_on_completions=train_on_completions,
finetune_vision_layers=finetune_vision_layers,
finetune_language_layers=finetune_language_layers,
finetune_attention_modules=finetune_attention_modules,
finetune_mlp_modules=finetune_mlp_modules,
enable_wandb=enable_wandb,
wandb_token=wandb_token,
wandb_project=wandb_project,
enable_tensorboard=enable_tensorboard,
tensorboard_dir=tensorboard_dir
):
# Yield each status update to Gradio
yield update_tuple
except Exception as e:
logger.error(f"Error in start_training_handler: {e}", exc_info=True)
yield (
gr.update(interactive=True), # Start button
gr.update(interactive=False), # Stop button
gr.update(visible=False), # Training progress
#gr.update(visible=True) # Config selection
)
def stop_training_handler():
"""Handler for stop training button"""
return backend.stop_training()
def update_training_status():
"""Periodic update of training status and plot"""
return backend.get_training_status(backend.current_theme)
def refresh_plot_for_theme(theme):
"""Refresh plot with new theme"""
return backend.refresh_plot_for_theme(theme)
return {
'start_training': start_training_handler,
'stop_training': stop_training_handler,
'update_status': update_training_status,
'refresh_plot': refresh_plot_for_theme
}

20
studio/backend/routes/auth.py
View file

@ -57,25 +57,29 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
# Generate a strong random JWT secret for this installation
jwt_secret = secrets.token_urlsafe(64)
# Save username/password hash and secret in SQLite
# Create user + generate tokens atomically — rollback if anything fails
try:
storage.create_initial_user(
username=payload.username,
password=payload.password,
jwt_secret=jwt_secret,
)
# Reload JWT secret from DB (so authentication.py picks it up)
reload_secret()
# Issue access + refresh tokens for the new user
access_token = create_access_token(subject=payload.username)
refresh_token = create_refresh_token(subject=payload.username)
except Exception as e:
# Rollback: remove the user row so setup can be retried
storage.delete_user(payload.username)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to create user: {str(e)}",
detail=f"Setup failed (rolled back): {str(e)}",
)
# Reload JWT secret from DB (so authentication.py picks it up)
reload_secret()
# Issue access + refresh tokens for the new user
access_token = create_access_token(subject=payload.username)
refresh_token = create_refresh_token(subject=payload.username)
return Token(
access_token=access_token,
refresh_token=refresh_token,

55
studio/backend/run.py
View file

@ -11,6 +11,51 @@ if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
def _resolve_external_ip() -> str:
"""
Resolve the machine's external IP address.
Tries (in order):
1. GCE metadata server (instant, works on Google Cloud VMs)
2. ifconfig.me (works anywhere with internet)
3. LAN IP via UDP socket trick (fallback)
"""
import urllib.request
import socket
# 1. Try GCE metadata server (responds in <10ms on GCE, times out fast elsewhere)
try:
req = urllib.request.Request(
"http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip",
headers={"Metadata-Flavor": "Google"},
)
with urllib.request.urlopen(req, timeout=1) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
except Exception:
pass
# 2. Try public IP service
try:
with urllib.request.urlopen("https://ifconfig.me", timeout=3) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
except Exception:
pass
# 3. Fallback: LAN IP via UDP socket trick
try:
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.connect(("8.8.8.8", 80))
ip = s.getsockname()[0]
s.close()
return ip
except Exception:
return "0.0.0.0"
def run_server(
host: str = "0.0.0.0",
port: int = 8000,
@ -57,11 +102,15 @@ def run_server(
time.sleep(3)
if not silent:
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
print("")
print("=" * 50)
print(f"🦥 Unsloth UI Backend is running on port {port}")
print(f" API: http://{host}:{port}/api")
print(f" Health: http://{host}:{port}/api/health")
print(f"🦥 Unsloth Studio is running on port {port}")
print(f" Local: http://localhost:{port}")
print(f" External: http://{display_host}:{port}")
print(f" API: http://{display_host}:{port}/api")
print(f" Health: http://{display_host}:{port}/api/health")
print("=" * 50)
return app