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

Merge pull request #22 from unslothai/feature/auth

Browse source 
Added jwt authentication
This commit is contained in:
Roland Tannous 2026-02-07 14:27:35 +04:00 • committed by GitHub
commit d54671b1a6
13 changed files with 674 additions and 220 deletions
Download patch file Download diff file Expand all files Collapse all files

24
studio/backend/auth/__init__.py
View file

@ -0,0 +1,24 @@
"""
Authentication module for JWT-based auth with SQLite storage.
"""
from .authentication import create_access_token, get_current_subject, reload_secret
from .storage import (
is_initialized,
create_initial_user,
get_user_and_secret,
load_jwt_secret,
)
from .hashing import hash_password, verify_password
__all__ = [
"create_access_token",
"get_current_subject",
"reload_secret",
"is_initialized",
"create_initial_user",
"get_user_and_secret",
"load_jwt_secret",
"hash_password",
"verify_password",
]

BIN
studio/backend/auth/auth.db Normal file
View file

Binary file not shown.

83
studio/backend/auth/authentication.py Normal file
View file

@ -0,0 +1,83 @@
import secrets
from datetime import UTC, datetime, timedelta
from typing import Optional
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
import jwt
from .storage import load_jwt_secret
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60
# Load stable secret from SQLite (set during first-time setup)
# This will raise RuntimeError if auth hasn't been initialized yet
try:
SECRET_KEY = load_jwt_secret()
except RuntimeError:
# Fallback: use a temporary secret until setup is complete
# This allows the app to start, but protected routes will fail until setup
SECRET_KEY = secrets.token_urlsafe(64)
security = HTTPBearer() # Reads Authorization: Bearer <token>
def create_access_token(
subject: str,
expires_delta: Optional[timedelta] = None,
) -> str:
"""
Create a signed JWT for the given subject (e.g. username).
Tokens are valid across restarts because SECRET_KEY is stored in SQLite.
"""
to_encode = {"sub": subject}
expire = datetime.now(UTC) + (
expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
)
to_encode.update({"exp": expire})
return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
def reload_secret() -> None:
"""
Reload the JWT secret from SQLite.
Call this after setup to ensure new tokens use the persistent secret.
"""
global SECRET_KEY
SECRET_KEY = load_jwt_secret()
async def get_current_subject(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:
"""
FastAPI dependency to validate the JWT and return the subject.
Use this as a dependency on routes that should be protected, e.g.:
@router.get("/secure")
async def secure_endpoint(current_subject: str = Depends(get_current_subject)):
...
"""
token = credentials.credentials
try:
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
subject: Optional[str] = payload.get("sub")
if subject is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid token payload",
)
return subject
except jwt.InvalidTokenError:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid or expired token",
)
# token = create_access_token("local-user")
# print(token)

40
studio/backend/auth/hashing.py Normal file
View file

@ -0,0 +1,40 @@
"""
Password hashing utilities using PBKDF2.
"""
import hashlib
import hmac
import secrets
from typing import Tuple
def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]:
"""
Hash a password using PBKDF2-HMAC-SHA256.
Returns (salt, hex_hash) tuple.
"""
if salt is None:
salt = secrets.token_hex(16)
dk = hashlib.pbkdf2_hmac(
"sha256",
password.encode("utf-8"),
salt.encode("utf-8"),
100_000, # 100k iterations
)
return salt, dk.hex()
def verify_password(password: str, salt: str, hashed: str) -> bool:
"""
Verify a password against a stored salt and hash.
Uses constant-time comparison to prevent timing attacks.
"""
dk = hashlib.pbkdf2_hmac(
"sha256",
password.encode("utf-8"),
salt.encode("utf-8"),
100_000,
)
return hmac.compare_digest(dk.hex(), hashed)

101
studio/backend/auth/storage.py Normal file
View file

@ -0,0 +1,101 @@
"""
SQLite storage for authentication data (user credentials + JWT secret).
"""
import sqlite3
from pathlib import Path
from typing import Optional, Tuple
DB_PATH = Path(__file__).parent / "auth.db"
def get_connection() -> sqlite3.Connection:
"""Get a connection to the auth database, creating tables if needed."""
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
conn.execute(
"""
CREATE TABLE IF NOT EXISTS auth_user (
id INTEGER PRIMARY KEY,
username TEXT UNIQUE NOT NULL,
password_salt TEXT NOT NULL,
password_hash TEXT NOT NULL,
jwt_secret TEXT NOT NULL
);
"""
)
conn.commit()
return conn
def is_initialized() -> bool:
"""Check if auth has been set up (user exists in DB)."""
conn = get_connection()
cur = conn.execute("SELECT COUNT(*) AS c FROM auth_user")
row = cur.fetchone()
conn.close()
return bool(row["c"])
def create_initial_user(username: str, password: str, jwt_secret: str) -> None:
"""
Create the initial admin user in the database.
Raises sqlite3.IntegrityError if username already exists.
"""
from .hashing import hash_password
salt, pwd_hash = hash_password(password)
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO auth_user (username, password_salt, password_hash, jwt_secret)
VALUES (?, ?, ?, ?)
""",
(username, salt, pwd_hash, jwt_secret),
)
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.
Returns (password_salt, password_hash, jwt_secret) or None if user not found.
"""
conn = get_connection()
try:
cur = conn.execute(
"""
SELECT password_salt, password_hash, jwt_secret
FROM auth_user
WHERE username = ?
""",
(username,),
)
row = cur.fetchone()
if not row:
return None
return row["password_salt"], row["password_hash"], row["jwt_secret"]
finally:
conn.close()
def load_jwt_secret() -> str:
"""
Load the JWT secret from the database.
Raises RuntimeError if auth is not initialized.
"""
conn = get_connection()
try:
cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1")
row = cur.fetchone()
if not row:
raise RuntimeError("Auth is not initialized. Please set up a password first.")
return row["jwt_secret"]
finally:
conn.close()

3
studio/backend/main.py
View file

@ -9,7 +9,7 @@ from pathlib import Path
from datetime import datetime
# Import routers
from routes import training_router, models_router, inference_router, datasets_router
from routes import training_router, models_router, inference_router, datasets_router, auth_router
# Create FastAPI app
app = FastAPI(
@ -30,6 +30,7 @@ app.add_middleware(
# ============ Register API Routes ============
# Register routers
app.include_router(auth_router, prefix="/api/auth", tags=["auth"])
app.include_router(training_router, prefix="/api/train", tags=["training"])
app.include_router(models_router, prefix="/api/models", tags=["models"])
app.include_router(inference_router, prefix="/api/inference", tags=["inference"])

41
studio/backend/models/__init__.py
View file

@ -3,35 +3,36 @@ Pydantic models for API request/response schemas
"""
from .training import (
TrainingStartRequest,
TrainingStartResponse,
TrainingStatusResponse,
TrainingMetricsResponse,
TrainingProgressResponse,
TrainingJobResponse,
TrainingStatus,
TrainingProgress,
)
from .models import (
ModelSearchRequest,
ModelSearchResponse,
ModelListResponse,
ModelConfigResponse,
LoRAScanResponse,
ModelDetails,
LoRAInfo,
ModelInfo,
LoRAScanResponse,
ModelListResponse,
)
from .auth import (
AuthSetupRequest,
AuthLoginRequest,
AuthStatusResponse,
)
__all__ = [
# Training schemas
"TrainingStartRequest",
"TrainingStartResponse",
"TrainingStatusResponse",
"TrainingMetricsResponse",
"TrainingProgressResponse",
"TrainingJobResponse",
"TrainingStatus",
"TrainingProgress",
# Model management schemas
"ModelSearchRequest",
"ModelSearchResponse",
"ModelListResponse",
"ModelConfigResponse",
"LoRAScanResponse",
"ModelDetails",
"LoRAInfo",
"ModelInfo",
"LoRAScanResponse",
"ModelListResponse",
# Auth schemas
"AuthSetupRequest",
"AuthLoginRequest",
"AuthStatusResponse",
]

22
studio/backend/models/auth.py Normal file
View file

@ -0,0 +1,22 @@
"""
Pydantic schemas for Authentication API
"""
from pydantic import BaseModel, Field
class AuthSetupRequest(BaseModel):
"""First-time setup: create the initial admin user + password."""
username: str = Field(..., description="Admin username")
password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)")
class AuthLoginRequest(BaseModel):
"""Login payload: username/password to obtain a JWT."""
username: str = Field(..., description="Username")
password: str = Field(..., description="Password")
class AuthStatusResponse(BaseModel):
"""Indicate whether auth has been initialized."""
initialized: bool = Field(..., description="True if auth setup has been completed")

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

@ -6,9 +6,11 @@ from typing import Optional, List, Dict, Any
class ModelDetails(BaseModel):
"""Detailed model configuration and metadata"""
model_name: str = Field(..., description="Model identifier")
config: Dict[str, Any] = Field(..., description="Model configuration dictionary")
"""Detailed model configuration and metadata - can be used for both list and detail views"""
id: str = Field(..., description="Model identifier")
model_name: Optional[str] = Field(None, description="Model identifier (alias for id, for backward compatibility)")
name: Optional[str] = Field(None, description="Display name for the model")
config: Optional[Dict[str, Any]] = Field(None, description="Model configuration dictionary")
is_vision: bool = Field(False, description="Whether model is a vision model")
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter")
@ -26,3 +28,9 @@ class LoRAScanResponse(BaseModel):
loras: List[LoRAInfo] = Field(default_factory=list, description="List of found LoRA adapters")
outputs_dir: str = Field(..., description="Directory that was scanned")
class ModelListResponse(BaseModel):
"""Response schema for listing models"""
models: List[ModelDetails] = Field(default_factory=list, description="List of models")
default_models: List[str] = Field(default_factory=list, description="List of default model IDs")

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

@ -6,5 +6,6 @@ from routes.training import router as training_router
from routes.models import router as models_router
from routes.inference import router as inference_router
from routes.datasets import router as datasets_router
from routes.auth import router as auth_router
__all__ = ["training_router", "models_router", "inference_router", "datasets_router"]
__all__ = ["training_router", "models_router", "inference_router", "datasets_router", "auth_router"]

88
studio/backend/routes/auth.py Normal file
View file

@ -0,0 +1,88 @@
"""
Authentication API routes
"""
from fastapi import APIRouter, HTTPException, status
import secrets
from models.auth import (
AuthSetupRequest,
AuthLoginRequest,
AuthStatusResponse,
)
from models.users import Token
from auth import storage, hashing
from auth.authentication import create_access_token, reload_secret
router = APIRouter()
@router.get("/status", response_model=AuthStatusResponse)
async def auth_status() -> AuthStatusResponse:
"""
Check whether auth has already been initialized.
- initialized = False -> frontend should show "Set admin password" screen.
- initialized = True -> frontend should show normal login.
"""
return AuthStatusResponse(initialized=storage.is_initialized())
@router.post("/setup", response_model=Token, status_code=status.HTTP_201_CREATED)
async def setup_auth(payload: AuthSetupRequest) -> Token:
"""
First-time setup: create the admin user and a JWT secret.
Can only be called once. Subsequent calls will return 400.
"""
if storage.is_initialized():
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Auth is already initialized.",
)
# Generate a strong random JWT secret for this installation
jwt_secret = secrets.token_urlsafe(64)
# Save username/password hash and secret in SQLite
try:
storage.create_initial_user(
username=payload.username,
password=payload.password,
jwt_secret=jwt_secret,
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to create user: {str(e)}",
)
# Reload JWT secret from DB (so authentication.py picks it up)
reload_secret()
# Issue a token for the new user
access_token = create_access_token(subject=payload.username)
return Token(access_token=access_token, token_type="bearer")
@router.post("/login", response_model=Token)
async def login(payload: AuthLoginRequest) -> Token:
"""
Login with username/password and receive a JWT.
"""
record = storage.get_user_and_secret(payload.username)
if record is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
)
salt, pwd_hash, _jwt_secret = record
if not hashing.verify_password(payload.password, salt, pwd_hash):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
)
access_token = create_access_token(subject=payload.username)
return Token(access_token=access_token, token_type="bearer")

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

@ -3,18 +3,21 @@ Model Management API routes
"""
import sys
from pathlib import Path
from fastapi import APIRouter, HTTPException, Query
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from typing import List, Optional
import logging
from pydantic import BaseModel
# 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.utils import search_hf_models
from utils.models import (
scan_trained_loras,
load_model_defaults,
@ -28,7 +31,6 @@ except ImportError:
parent_backend = backend_path.parent / "backend"
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
from utils.utils import search_hf_models
from utils.models import (
scan_trained_loras,
load_model_defaults,
@ -38,16 +40,12 @@ except ImportError:
)
from core.inference import get_inference_backend
from models.models import (
ModelSearchRequest,
ModelSearchResponse,
ModelInfo,
ModelListResponse,
ModelConfigResponse,
from models import (
ModelDetails,
LoRAScanResponse,
LoRAInfo,
ModelListResponse,
)
router = APIRouter()
logger = logging.getLogger(__name__)
@ -61,63 +59,12 @@ if not logger.handlers:
logger.setLevel(logging.INFO)
@router.post("/search")
async def search_models(request: ModelSearchRequest):
"""
Search for models on HuggingFace Hub.
This endpoint wraps the backend search_hf_models function.
"""
try:
# Call backend search function
gradio_update = search_hf_models(
search_query=request.query,
hf_token=request.hf_token
)
# Convert Gradio update to list of model IDs
model_list = []
if gradio_update and hasattr(gradio_update, 'choices'):
choices = gradio_update.choices
elif isinstance(gradio_update, dict) and 'choices' in gradio_update:
choices = gradio_update['choices']
elif isinstance(gradio_update, list):
choices = gradio_update
else:
choices = []
# Process choices - they may be tuples (display_name, model_id) or just strings
for choice in choices:
if isinstance(choice, tuple) and len(choice) >= 2:
# Format: (display_name, model_id)
model_id = choice[1] if len(choice) > 1 else choice[0]
display_name = choice[0]
model_info = ModelInfo(
id=model_id,
name=display_name
)
elif isinstance(choice, str):
# Just a model ID string
model_info = ModelInfo(id=choice)
else:
continue
model_list.append(model_info)
return ModelSearchResponse(
models=model_list,
total=len(model_list)
)
except Exception as e:
logger.error(f"Error searching models: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to search models: {str(e)}"
)
@router.get("/list")
async def list_models():
async def list_models(
current_subject: str = Depends(get_current_subject),
):
"""
List available models (default models and loaded models).
@ -132,7 +79,7 @@ async def list_models():
# Get loaded models
loaded_models = []
for model_name, model_data in inference_backend.models.items():
model_info = ModelInfo(
model_info = ModelDetails(
id=model_name,
name=model_name.split("/")[-1] if "/" in model_name else model_name,
is_vision=model_data.get("is_vision", False),
@ -147,7 +94,7 @@ async def list_models():
# Add default models
for model_id in default_models:
if model_id not in seen_ids:
model_info = ModelInfo(
model_info = ModelDetails(
id=model_id,
name=model_id.split("/")[-1] if "/" in model_id else model_id
)
@ -174,7 +121,10 @@ async def list_models():
@router.get("/config/{model_name:path}")
async def get_model_config(model_name: str):
async def get_model_config(
model_name: str,
current_subject: str = Depends(get_current_subject),
):
"""
Get configuration for a specific model.
@ -200,12 +150,13 @@ async def get_model_config(model_name: str):
# If ModelConfig creation fails, use defaults
pass
return ModelConfigResponse(
return ModelDetails(
id=model_name,
model_name=model_name,
config=config_dict,
is_vision=is_vision,
is_lora=is_lora,
base_model=base_model
base_model=base_model,
)
except Exception as e:
@ -218,7 +169,8 @@ async def get_model_config(model_name: str):
@router.get("/loras")
async def scan_loras(
outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters")
outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters"),
current_subject: str = Depends(get_current_subject),
):
"""
Scan for trained LoRA adapters in the outputs directory.
@ -256,7 +208,10 @@ async def scan_loras(
@router.get("/loras/{lora_path:path}/base-model")
async def get_lora_base_model(lora_path: str):
async def get_lora_base_model(
lora_path: str,
current_subject: str = Depends(get_current_subject),
):
"""
Get the base model for a LoRA adapter.
@ -287,7 +242,10 @@ async def get_lora_base_model(lora_path: str):
@router.get("/check-vision/{model_name:path}")
async def check_vision_model(model_name: str):
async def check_vision_model(
model_name: str,
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is a vision model.

371
studio/backend/routes/training.py
View file

@ -3,9 +3,9 @@ Training API routes
"""
import sys
from pathlib import Path
from fastapi import APIRouter, HTTPException
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from typing import Dict
from typing import Dict, Optional
import logging
import asyncio
from datetime import datetime
@ -27,12 +27,14 @@ except ImportError:
sys.path.insert(0, str(parent_backend))
from core.training import get_training_backend
from models.training import (
# Auth
from auth.authentication import get_current_subject
from models import (
TrainingStartRequest,
TrainingStartResponse,
TrainingStatusResponse,
TrainingMetricsResponse,
TrainingProgressResponse,
TrainingJobResponse,
TrainingStatus,
TrainingProgress,
)
router = APIRouter()
@ -49,35 +51,47 @@ if not logger.handlers:
@router.post("/start")
async def start_training(request: TrainingStartRequest):
async def start_training(
request: TrainingStartRequest,
current_subject: str = Depends(get_current_subject),
):
"""
Start a training job.
This endpoint initiates training in the background and returns immediately.
Use the /status endpoint to check training progress.
"""
try:
logger.info(f"Starting training job with model: {request.model_name}")
backend = get_training_backend()
# Generate job ID and attach to backend for later status/progress calls
job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
backend.current_job_id = job_id
# Check if training is already active
if backend.is_training_active():
return TrainingStartResponse(
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
return TrainingJobResponse(
job_id=existing_job_id or job_id,
status="error",
message="Training is already in progress. Stop current training before starting a new one.",
error="Training already active"
message=(
"Training is already in progress. "
"Stop current training before starting a new one."
),
error="Training already active",
)
# Validate dataset paths if provided
if request.local_datasets:
validated_datasets = []
# Get the backend directory (where this file is located)
backend_dir = Path(__file__).parent.parent
assets_datasets_dir = backend_dir / "assets" / "datasets"
for dataset_path in request.local_datasets:
dataset_file = Path(dataset_path)
# If not absolute, try multiple locations
if not dataset_file.is_absolute():
# First try: relative to current working directory
@ -89,14 +103,16 @@ async def start_training(request: TrainingStartRequest):
# Third try: just the filename in assets/datasets
candidate = assets_datasets_dir / dataset_file.name
dataset_file = candidate
if not dataset_file.exists():
logger.warning(f"Dataset file not found: {dataset_path} (resolved: {dataset_file})")
logger.warning(
f"Dataset file not found: {dataset_path} (resolved: {dataset_file})"
)
else:
logger.info(f"Found dataset file: {dataset_file}")
validated_datasets.append(str(dataset_file))
request.local_datasets = validated_datasets
# Convert request to kwargs for backend
training_kwargs = {
"model_name": request.model_name,
@ -125,7 +141,9 @@ async def start_training(request: TrainingStartRequest):
"lora_alpha": request.lora_alpha,
"lora_dropout": request.lora_dropout,
"target_modules": request.target_modules if request.target_modules else None,
"gradient_checkpointing": request.gradient_checkpointing.strip() if request.gradient_checkpointing and request.gradient_checkpointing.strip() else "unsloth",
"gradient_checkpointing": request.gradient_checkpointing.strip()
if request.gradient_checkpointing and request.gradient_checkpointing.strip()
else "unsloth",
"use_rslora": request.use_rslora,
"use_loftq": request.use_loftq,
"train_on_completions": request.train_on_completions,
@ -139,84 +157,95 @@ async def start_training(request: TrainingStartRequest):
"enable_tensorboard": request.enable_tensorboard,
"tensorboard_dir": request.tensorboard_dir or "",
}
# Generate job ID
job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
# Set initial "preparing" state
try:
backend.trainer._update_progress(
status_message="Initializing training...",
is_training=False
is_training=False,
)
except:
except Exception:
pass
def run_training():
try:
logger.info(f"Starting training job {job_id} with model {request.model_name}")
logger.info(
f"Starting training job {job_id} with model {request.model_name}"
)
# Update status to show we're loading model
try:
backend.trainer._update_progress(status_message="Loading model...")
except Exception as e:
logger.error(f"Error updating progress: {e}")
# Consume the generator - this actually runs the training
update_count = 0
for update_tuple in backend.start_training(**training_kwargs):
for _update_tuple in backend.start_training(**training_kwargs):
update_count += 1
if update_count % 10 == 0:
logger.info(f"Training progress update #{update_count}")
logger.info(f"Training job {job_id} completed successfully")
except Exception as e:
logger.error(f"Training error in job {job_id}: {e}", exc_info=True)
try:
backend.trainer._update_progress(
error=str(e),
is_training=False
is_training=False,
)
except Exception as update_error:
logger.error(f"Failed to update progress: {update_error}")
# Start training in a daemon thread
training_thread = threading.Thread(target=run_training, daemon=True, name=f"Training-{job_id}")
training_thread = threading.Thread(
target=run_training,
daemon=True,
name=f"Training-{job_id}",
)
training_thread.start()
# Store thread reference for status checking
backend._training_thread = training_thread
# Give it a moment to start
import time
time.sleep(0.5)
# Verify training thread is alive
if not training_thread.is_alive():
logger.warning(f"Training thread died immediately for job {job_id}")
return TrainingStartResponse(
return TrainingJobResponse(
job_id=job_id,
status="error",
message="Training thread failed to start. Check server logs for details.",
error="Thread not alive"
message=(
"Training thread failed to start. "
"Check server logs for details."
),
error="Thread not alive",
)
return TrainingStartResponse(
status="started",
return TrainingJobResponse(
job_id=job_id,
message="Training job started successfully"
status="queued",
message="Training job queued and starting in background",
error=None,
)
except Exception as e:
logger.error(f"Error starting training: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to start training: {str(e)}"
detail=f"Failed to start training: {str(e)}",
)
@router.post("/stop")
async def stop_training():
async def stop_training(
current_subject: str = Depends(get_current_subject),
):
"""
Stop the currently running training job.
"""
@ -246,55 +275,75 @@ async def stop_training():
@router.get("/status")
async def get_training_status():
async def get_training_status(
current_subject: str = Depends(get_current_subject),
):
"""
Get the current training status.
"""
try:
backend = get_training_backend()
job_id: str = getattr(backend, "current_job_id", "")
# Check if training is active
is_active = backend.is_training_active()
# Check if there's a training thread running (preparation phase)
has_thread = hasattr(backend, '_training_thread') and backend._training_thread and backend._training_thread.is_alive()
# Get progress info
has_thread = (
hasattr(backend, "_training_thread")
and backend._training_thread
and backend._training_thread.is_alive()
)
# Get progress info from trainer
try:
progress = backend.trainer.get_training_progress()
status_message = progress.status_message or "Ready to train"
except:
except Exception:
progress = None
status_message = "Unknown"
if is_active:
# Actual training is running
trainer = backend.trainer
current_step = getattr(trainer.training_progress, 'step', None) or (progress.step if progress else None)
total_steps = getattr(trainer.training_progress, 'total_steps', None) or (progress.total_steps if progress else None)
return TrainingStatusResponse(
status="training",
is_active=True,
message=status_message or "Training is in progress",
current_step=current_step,
total_steps=total_steps
)
elif has_thread or (progress and status_message and any(keyword in status_message.lower() for keyword in ["loading", "preparing", "initializing"])):
# Training thread is running but not yet in active training phase
return TrainingStatusResponse(
status="preparing",
is_active=False,
message=status_message or "Preparing training...",
current_step=None,
total_steps=None
)
status_message = (
getattr(progress, "status_message", None) if progress else None
) or "Ready to train"
error_message = getattr(progress, "error", None) if progress else None
# Derive high-level phase
if error_message:
phase = "error"
elif is_active:
msg_lower = status_message.lower()
if "loading" in msg_lower:
phase = "loading_model"
elif any(
k in msg_lower for k in ["preparing", "initializing", "configuring"]
):
phase = "configuring"
else:
phase = "training"
elif progress and getattr(progress, "is_completed", False):
phase = "completed"
elif has_thread:
phase = "loading_model"
else:
return TrainingStatusResponse(
status="idle",
is_active=False,
message="No training job is currently running"
)
phase = "idle"
details = None
if progress:
details = {
"epoch": getattr(progress, "epoch", 0),
"step": getattr(progress, "step", 0),
"total_steps": getattr(progress, "total_steps", 0),
"loss": getattr(progress, "loss", 0.0),
"learning_rate": getattr(progress, "learning_rate", 0.0),
}
return TrainingStatus(
job_id=job_id,
phase=phase,
is_training_running=is_active,
message=status_message,
error=error_message,
details=details,
)
except Exception as e:
logger.error(f"Error getting training status: {e}", exc_info=True)
@ -305,7 +354,9 @@ async def get_training_status():
@router.get("/metrics")
async def get_training_metrics():
async def get_training_metrics(
current_subject: str = Depends(get_current_subject),
):
"""
Get training metrics (loss, learning rate, steps).
"""
@ -316,20 +367,21 @@ async def get_training_metrics():
loss_history = backend.loss_history
lr_history = backend.lr_history
step_history = backend.step_history
# Get current values
current_loss = loss_history[-1] if loss_history else None
current_lr = lr_history[-1] if lr_history else None
current_step = step_history[-1] if step_history else None
return TrainingMetricsResponse(
loss_history=loss_history,
lr_history=lr_history,
step_history=step_history,
current_loss=current_loss,
current_lr=current_lr,
current_step=current_step
)
# Keep metrics as a simple JSON payload instead of a Pydantic model
return {
"loss_history": loss_history,
"lr_history": lr_history,
"step_history": step_history,
"current_loss": current_loss,
"current_lr": current_lr,
"current_step": current_step,
}
except Exception as e:
logger.error(f"Error getting training metrics: {e}", exc_info=True)
@ -340,7 +392,9 @@ async def get_training_metrics():
@router.get("/progress")
async def stream_training_progress():
async def stream_training_progress(
current_subject: str = Depends(get_current_subject),
):
"""
Stream training progress updates using Server-Sent Events (SSE).
@ -348,12 +402,53 @@ async def stream_training_progress():
"""
async def event_generator():
backend = get_training_backend()
job_id: str = getattr(backend, "current_job_id", "")
# Helper to build a TrainingProgress payload from raw values
def build_progress(
step: int,
loss: float,
learning_rate: float,
total_steps: int,
epoch: Optional[int] = None,
) -> TrainingProgress:
total = max(total_steps, 0)
if step < 0 or total == 0:
progress_percent = 0.0
else:
progress_percent = (
float(step) / float(total) * 100.0 if total > 0 else 0.0
)
return TrainingProgress(
job_id=job_id,
step=step,
total_steps=total,
loss=loss,
learning_rate=learning_rate,
progress_percent=progress_percent,
epoch=epoch,
elapsed_seconds=None,
eta_seconds=None,
grad_norm=None,
num_tokens=None,
)
# Send initial status
is_active = backend.is_training_active()
initial_message = 'Connecting...' if is_active else 'No training in progress'
yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message=initial_message).model_dump_json()}\n\n"
tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
initial_total_steps = getattr(tp, "total_steps", 0) if tp else 0
initial_epoch = getattr(tp, "epoch", None) if tp else None
initial_progress = build_progress(
step=0,
loss=0.0,
learning_rate=0.0,
total_steps=initial_total_steps,
epoch=initial_epoch,
)
yield f"data: {initial_progress.model_dump_json()}\n\n"
# If not active, check if there's any history
if not is_active:
if backend.step_history:
@ -361,9 +456,13 @@ async def stream_training_progress():
final_step = backend.step_history[-1]
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
yield f"data: {TrainingProgressResponse(step=final_step, loss=final_loss, learning_rate=final_lr, status_message='Training completed').model_dump_json()}\n\n"
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
yield f"data: {build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch).model_dump_json()}\n\n"
else:
yield f"data: {TrainingProgressResponse(step=-1, loss=0.0, learning_rate=0.0, status_message='No training in progress').model_dump_json()}\n\n"
yield f"data: {build_progress(-1, 0.0, 0.0, 0).model_dump_json()}\n\n"
return
# Poll for updates while training is active
@ -378,53 +477,81 @@ async def stream_training_progress():
current_step = backend.step_history[-1]
current_loss = backend.loss_history[-1] if backend.loss_history else 0.0
current_lr = backend.lr_history[-1] if backend.lr_history else 0.0
tp_inner = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
current_total_steps = (
getattr(tp_inner, "total_steps", current_step)
if tp_inner
else current_step
)
current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None
# Only send if step changed
if current_step != last_step:
progress = TrainingProgressResponse(
step=current_step,
loss=current_loss,
learning_rate=current_lr,
status_message=f"Training step {current_step}"
progress_payload = build_progress(
current_step,
current_loss,
current_lr,
current_total_steps,
current_epoch,
)
yield f"data: {progress.model_dump_json()}\n\n"
yield f"data: {progress_payload.model_dump_json()}\n\n"
last_step = current_step
no_update_count = 0
else:
no_update_count += 1
# Send heartbeat every 10 seconds
if no_update_count % 10 == 0:
progress = TrainingProgressResponse(
step=current_step,
loss=current_loss,
learning_rate=current_lr,
status_message=f"Training step {current_step} (waiting for next update...)"
heartbeat_payload = build_progress(
current_step,
current_loss,
current_lr,
current_total_steps,
current_epoch,
)
yield f"data: {progress.model_dump_json()}\n\n"
yield f"data: {heartbeat_payload.model_dump_json()}\n\n"
else:
# No steps yet, but training is active
no_update_count += 1
if no_update_count % 5 == 0:
yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message='Preparing training...').model_dump_json()}\n\n"
preparing_payload = build_progress(0, 0.0, 0.0, 0)
yield f"data: {preparing_payload.model_dump_json()}\n\n"
# Timeout check
if no_update_count > max_no_updates:
logger.warning("Progress stream timeout - no updates received")
yield f"data: {TrainingProgressResponse(step=last_step, loss=0.0, learning_rate=0.0, status_message='Progress timeout - training may have stopped').model_dump_json()}\n\n"
timeout_payload = build_progress(last_step, 0.0, 0.0, 0)
yield f"data: {timeout_payload.model_dump_json()}\n\n"
break
await asyncio.sleep(1) # Poll every second
except Exception as e:
logger.error(f"Error in progress stream: {e}", exc_info=True)
yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message=f'Error: {str(e)}').model_dump_json()}\n\n"
error_payload = build_progress(0, 0.0, 0.0, 0)
yield f"data: {error_payload.model_dump_json()}\n\n"
break
# Send final status
final_step = backend.step_history[-1] if backend.step_history else last_step
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
yield f"data: {TrainingProgressResponse(step=final_step, loss=final_loss, learning_rate=final_lr, status_message='Training completed').model_dump_json()}\n\n"
final_tp = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
final_total_steps = (
getattr(final_tp, "total_steps", final_step) if final_tp else final_step
)
final_epoch = getattr(final_tp, "epoch", None) if final_tp else None
final_payload = build_progress(
final_step,
final_loss,
final_lr,
final_total_steps,
final_epoch,
)
yield f"data: {final_payload.model_dump_json()}\n\n"
return StreamingResponse(
event_generator(),