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

merge nightly

Browse source
This commit is contained in:
Shine1i 2026-02-14 16:42:20 +01:00
commit ca30e9f004
83 changed files with 5382 additions and 1118 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -36,6 +36,7 @@ Thumbs.db
resources/
tmp/
auth.db
studio/frontend/package-lock.json
# Local working docs
**/CLAUDE.md

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,

118
setup.sh Executable file
View file

@ -0,0 +1,118 @@
#!/usr/bin/env bash
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 "╚══════════════════════════════════════╝"
# ── 1. Check existing Node/npm versions ──
NEED_NODE=true
if command -v node &>/dev/null && command -v npm &>/dev/null; then
NODE_MAJOR=$(node -v | sed 's/v//' | cut -d. -f1)
NPM_MAJOR=$(npm -v | cut -d. -f1)
if [ "$NODE_MAJOR" -ge 20 ] && [ "$NPM_MAJOR" -ge 11 ]; then
echo "✅ Node $(node -v) and npm $(npm -v) already meet requirements. Skipping nvm install."
NEED_NODE=false
else
echo "⚠️ Node $(node -v) / npm $(npm -v) too old. Installing via nvm..."
fi
else
echo "⚠️ Node/npm not found. Installing via nvm..."
fi
if [ "$NEED_NODE" = true ]; then
# ── 2. Install nvm ──
echo "Installing nvm..."
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"
[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"
# ── 3. Install Node LTS ──
echo "Installing Node 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)
NPM_MAJOR=$(npm -v | cut -d. -f1)
if [ "$NODE_MAJOR" -lt 20 ]; then
echo "❌ ERROR: Node version must be >= 20 (got $(node -v))"
exit 1
fi
if [ "$NPM_MAJOR" -lt 11 ]; then
echo "⚠️ npm version is $(npm -v), expected >= 11. Updating..."
run_quiet "npm update" npm install -g npm@latest
fi
fi
echo "✅ Node $(node -v) | npm $(npm -v)"
# ── 5. Build frontend ──
echo ""
echo "Building frontend..."
cd "$SCRIPT_DIR/studio/frontend"
run_quiet "npm install" npm install
run_quiet "npm run build" npm run build
cd "$SCRIPT_DIR"
echo "✅ Frontend built to studio/frontend/dist"
# ── 6. Python venv + deps ──
echo ""
echo "Setting up Python environment..."
python3 -m venv .venv
source .venv/bin/activate
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 easydict addict
echo "✅ Python dependencies installed"
# ── 7. Add shell alias ──
# Note: venv activation does NOT persist across terminal sessions.
# This alias hardcodes the venv python path so users don't need to activate.
echo ""
REPO_DIR="$SCRIPT_DIR"
if ! grep -qF "unsloth-ui" ~/.bashrc 2>/dev/null; then
cat >> ~/.bashrc <<UNSLOTH_EOF
# Unsloth Studio launcher
alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist'
UNSLOTH_EOF
echo "✅ Alias 'unsloth-ui' added to ~/.bashrc"
else
echo "✅ Alias 'unsloth-ui' already exists in ~/.bashrc"
fi
echo ""
echo "╔══════════════════════════════════════╗"
echo "║ Setup Complete! ║"
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 "╚══════════════════════════════════════╝"

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

43
studio/backend/core/inference/inference.py
View file

@ -526,7 +526,7 @@ class InferenceBackend:
def _generate_vision_response(self, messages, system_prompt, image,
temperature, top_p, top_k, max_new_tokens,
repetition_penalty) -> Generator[str, None, None]:
"""Handle vision model generation."""
"""Handle vision model generation with true token-by-token streaming."""
model_info = self.models[self.active_model_name]
model = model_info["model"]
processor = model_info["processor"]
@ -565,31 +565,44 @@ class InferenceBackend:
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
inputs = processor.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
# Generate with streaming
captured_output = StringIO()
original_stdout = sys.stdout
# Stream with TextIteratorStreamer + background thread
try:
sys.stdout = captured_output
from transformers import TextIteratorStreamer
import threading
text_streamer = TextStreamer(processor.tokenizer, skip_prompt=True)
model.generate(
streamer = TextIteratorStreamer(
processor.tokenizer, skip_prompt=True, skip_special_tokens=True
)
generation_kwargs = dict(
**inputs,
streamer=text_streamer,
streamer=streamer,
max_new_tokens=max_new_tokens,
use_cache=True,
temperature=temperature,
top_p=top_p,
top_k=top_k
top_k=top_k,
)
sys.stdout = original_stdout
generated_text = captured_output.getvalue()
cleaned = self._clean_generated_text(generated_text)
yield cleaned
def generate_fn():
try:
model.generate(**generation_kwargs)
except Exception as e:
logger.error(f"Vision generation error in thread: {e}")
thread = threading.Thread(target=generate_fn)
thread.start()
output = ""
for new_token in streamer:
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
yield cleaned
thread.join()
except Exception as e:
sys.stdout = original_stdout
logger.error(f"Vision generation error: {e}")
yield f"Error: {str(e)}"
pass

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',
]

10
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
@ -34,7 +34,7 @@ logger = logging.getLogger(__name__)
@dataclass
class TrainingProgress:
"""Training progress tracking"""
epoch: int = 0
epoch: float = 0
step: int = 0
total_steps: int = 0
loss: float = 0.0
@ -46,7 +46,7 @@ class TrainingProgress:
class UnslothTrainer:
"""
Unsloth Training Backend for Gradio UI Integration
Unsloth Training Backend
"""
def __init__(self):
@ -301,7 +301,8 @@ class UnslothTrainer:
def load_and_format_dataset(self,
dataset_source: str,
format_type: str = "auto",
local_datasets: list = None) -> Optional[Dataset]:
local_datasets: list = None,
custom_format_mapping: dict = None) -> Optional[Dataset]:
"""
Load and prepare dataset for training
"""
@ -374,6 +375,7 @@ class UnslothTrainer:
is_vlm=self.is_vlm,
format_type=format_type, # "auto", "alpaca", "chatml", "sharegpt"
dataset_name=dataset_source,
custom_format_mapping=custom_format_mapping,
)
# Check if stopped during formatting

224
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.
"""
@ -36,7 +35,7 @@ class TrainingBackend:
def _on_progress_update(self, progress: TrainingProgress):
"""Callback for progress updates"""
if progress.step > 0 and progress.loss > 0:
if progress.step >= 0 and progress.loss > 0:
self.loss_history.append(progress.loss)
self.lr_history.append(progress.learning_rate)
self.step_history.append(progress.step)
@ -91,12 +90,15 @@ class TrainingBackend:
wandb_token: str,
wandb_project: str,
enable_tensorboard: bool,
tensorboard_dir: str) -> Generator[Tuple, None, None]:
"""
Start training - yields UI updates as generator.
tensorboard_dir: str,
Yields:
Tuple of (start_btn_update, stop_btn_update, progress_visible, config_visible)
# Optional: user-provided column mapping
custom_format_mapping: dict = None) -> bool:
"""
Start training.
Returns:
True if training started successfully, False otherwise.
"""
try:
# Reset stop flag and clear history
@ -107,20 +109,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 +126,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 +155,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...")
@ -186,19 +163,13 @@ class TrainingBackend:
dataset = self.trainer.load_and_format_dataset(
dataset_source=hf_dataset if hf_dataset.strip() else None,
format_type=format_type,
local_datasets=local_datasets if local_datasets else None
local_datasets=local_datasets if local_datasets else None,
custom_format_mapping=custom_format_mapping,
)
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 +212,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 +222,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 +247,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 +255,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 +519,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
}

4
studio/backend/models/datasets.py
View file

@ -18,6 +18,10 @@ class CheckFormatResponse(BaseModel):
requires_manual_mapping: bool
detected_format: str
columns: List[str]
is_multimodal: bool = False
multimodal_columns: Optional[List[str]] = None
suggested_mapping: Optional[Dict[str, str]] = None
detected_image_column: Optional[str] = None
detected_text_column: Optional[str] = None
preview_samples: Optional[List[Dict]] = None
total_rows: Optional[int] = None

142
studio/backend/models/inference.py
View file

@ -1,8 +1,13 @@
"""
Pydantic schemas for Inference API
"""
from pydantic import BaseModel, Field
from typing import Optional, List
from __future__ import annotations
import time
import uuid
from typing import Annotated, Literal, Optional, List, Union
from pydantic import BaseModel, Discriminator, Field, Tag
class LoadRequest(BaseModel):
@ -20,7 +25,7 @@ class UnloadRequest(BaseModel):
class GenerateRequest(BaseModel):
"""Request for text generation"""
"""Request for text generation (legacy /generate/stream endpoint)"""
messages: List[dict] = Field(..., description="Chat messages in OpenAI format")
system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt")
temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature")
@ -52,3 +57,134 @@ class InferenceStatusResponse(BaseModel):
is_vision: bool = Field(False, description="Whether the active model is a vision model")
loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
loaded: List[str] = Field(default_factory=list, description="Models currently loaded")
# =====================================================================
# OpenAI-Compatible Chat Completions Models
# =====================================================================
# ── Multimodal content parts (OpenAI vision format) ──────────────
class TextContentPart(BaseModel):
"""Text content part in a multimodal message."""
type: Literal["text"]
text: str
class ImageUrl(BaseModel):
"""Image URL object — supports data URIs and remote URLs."""
url: str = Field(..., description="data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
class ImageContentPart(BaseModel):
"""Image content part in a multimodal message."""
type: Literal["image_url"]
image_url: ImageUrl
def _content_part_discriminator(v):
if isinstance(v, dict):
return v.get("type")
return getattr(v, "type", None)
ContentPart = Annotated[
Union[
Annotated[TextContentPart, Tag("text")],
Annotated[ImageContentPart, Tag("image_url")],
],
Discriminator(_content_part_discriminator),
]
"""Union type for multimodal content parts, discriminated by the 'type' field."""
# ── Messages ─────────────────────────────────────────────────────
class ChatMessage(BaseModel):
"""
A single message in the conversation.
``content`` may be a plain string (text-only) or a list of
content parts for multimodal messages (OpenAI vision format).
"""
role: Literal["system", "user", "assistant"] = Field(..., description="Message role")
content: Union[str, list[ContentPart]] = Field(..., description="Message content (string or multimodal parts)")
class ChatCompletionRequest(BaseModel):
"""
OpenAI-compatible chat completion request.
Extensions (non-OpenAI fields) are marked with 'x-unsloth'.
"""
model: str = Field("default", description="Model identifier (informational; the active model is used)")
messages: list[ChatMessage] = Field(..., description="Conversation messages")
stream: bool = Field(True, description="Whether to stream the response via SSE")
temperature: float = Field(0.7, ge=0.0, le=2.0)
top_p: float = Field(0.9, ge=0.0, le=1.0)
max_tokens: Optional[int] = Field(512, ge=1, le=4096, description="Maximum tokens to generate")
# ── Unsloth extensions (ignored by standard OpenAI clients) ──
top_k: int = Field(40, ge=1, le=100, description="[x-unsloth] Top-k sampling")
repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="[x-unsloth] Repetition penalty")
image_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded image for vision models")
# ── Streaming response chunks ────────────────────────────────────
class ChoiceDelta(BaseModel):
"""Delta content for a streaming chunk."""
role: Optional[str] = None
content: Optional[str] = None
class ChunkChoice(BaseModel):
"""A single choice in a streaming chunk."""
index: int = 0
delta: ChoiceDelta
finish_reason: Optional[Literal["stop", "length"]] = None
class ChatCompletionChunk(BaseModel):
"""A single SSE chunk in OpenAI streaming format."""
id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
created: int = Field(default_factory=lambda: int(time.time()))
model: str = "default"
choices: list[ChunkChoice]
# ── Non-streaming response ───────────────────────────────────────
class CompletionMessage(BaseModel):
"""The assistant's complete response message."""
role: Literal["assistant"] = "assistant"
content: str
class CompletionChoice(BaseModel):
"""A single choice in a non-streaming response."""
index: int = 0
message: CompletionMessage
finish_reason: Literal["stop", "length"] = "stop"
class CompletionUsage(BaseModel):
"""Token usage statistics (approximate)."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
class ChatCompletion(BaseModel):
"""Non-streaming chat completion response."""
id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion"] = "chat.completion"
created: int = Field(default_factory=lambda: int(time.time()))
model: str = "default"
choices: list[CompletionChoice]
usage: CompletionUsage = Field(default_factory=CompletionUsage)

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

@ -2,7 +2,7 @@
Pydantic schemas for Training API
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Literal
from typing import Optional, List, Dict, Literal
class TrainingStartRequest(BaseModel):
@ -18,7 +18,10 @@ class TrainingStartRequest(BaseModel):
hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier")
local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths")
format_type: str = Field(..., description="Dataset format type")
custom_format_mapping: Optional[Dict[str, str]] = Field(
None,
description="User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM"
)
# Training parameters
num_epochs: int = Field(1, description="Number of training epochs")
learning_rate: str = Field("2e-4", description="Learning rate")
@ -84,6 +87,11 @@ class TrainingStatus(BaseModel):
message: str = Field(..., description="Human-readable status message")
error: Optional[str] = Field(None, description="Error details if phase is 'error'")
details: Optional[dict] = Field(None, description="Phase-specific info, e.g. {'model_size': '8B'}")
metric_history: Optional[dict] = Field(
None,
description="Full metric history arrays for chart recovery after SSE reconnection. "
"Keys: 'steps', 'loss', 'lr' — each a list of numeric values.",
)
class TrainingProgress(BaseModel):
@ -94,7 +102,7 @@ class TrainingProgress(BaseModel):
loss: float = Field(..., description="Current loss value")
learning_rate: float = Field(..., description="Current learning rate")
progress_percent: float = Field(..., description="Progress percentage (0.0 to 100.0)")
epoch: Optional[int] = Field(None, description="Current epoch")
epoch: Optional[float] = Field(None, description="Current epoch")
elapsed_seconds: Optional[float] = Field(None, description="Time elapsed since training started")
eta_seconds: Optional[float] = Field(None, description="Estimated time remaining")
grad_norm: Optional[float] = Field(None, description="L2 norm of gradients, computed before gradient clipping")

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,

80
studio/backend/routes/datasets.py
View file

@ -1,6 +1,8 @@
"""
Datasets API routes
"""
import base64
import io
import sys
from pathlib import Path
from fastapi import APIRouter, HTTPException
@ -30,6 +32,42 @@ if not logger.handlers:
from models.datasets import CheckFormatRequest, CheckFormatResponse
def _serialize_preview_value(value):
"""make it json safe for client preview ⊂(◉‿◉)つ"""
if value is None or isinstance(value, (str, int, float, bool)):
return value
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
buffer = io.BytesIO()
value.convert("RGB").save(buffer, format="JPEG", quality=85)
return {
"type": "image",
"mime": "image/jpeg",
"width": value.width,
"height": value.height,
"data": base64.b64encode(buffer.getvalue()).decode("ascii"),
}
except Exception:
pass
if isinstance(value, dict):
return {str(key): _serialize_preview_value(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [_serialize_preview_value(item) for item in value]
return str(value)
def _serialize_preview_rows(rows):
return [
{str(key): _serialize_preview_value(value) for key, value in dict(row).items()}
for row in rows
]
# --- Endpoints ---
@router.post("/check-format", response_model=CheckFormatResponse)
@ -37,12 +75,15 @@ async def check_format(request: CheckFormatRequest):
"""
Check if a dataset requires manual column mapping.
This is a lightweight check that only runs format detection,
not full processing. Use before starting training to determine
if the user needs to manually map columns.
This is a lightweight check that loads only the first 10 rows,
runs format detection, and (if processable) returns processed
preview samples. The full dataset is re-processed at training time.
"""
try:
from datasets import load_dataset
from utils.datasets import format_dataset
PREVIEW_SIZE = 10
logger.info(f"Checking format for dataset: {request.dataset_name}")
@ -69,18 +110,47 @@ async def check_format(request: CheckFormatRequest):
load_kwargs["token"] = request.hf_token
dataset = load_dataset(**load_kwargs)
# Run lightweight format check
result = check_dataset_format(dataset, is_vlm=request.is_vlm)
# Slice to top N rows — all detection and preview runs on this subset
total_rows = len(dataset)
preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
# Run lightweight format check on the preview slice
result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)
logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}")
# Generate preview samples
preview_samples = None
if not result["requires_manual_mapping"]:
# Format detected — return processed preview
try:
format_result = format_dataset(
preview_slice,
format_type="auto",
custom_format_mapping=result.get("suggested_mapping"),
)
processed = format_result["dataset"]
preview_samples = _serialize_preview_rows(processed)
except Exception as e:
logger.warning(f"Processed preview generation failed (non-fatal): {e}")
# Fall back to raw samples so frontend still has something
preview_samples = _serialize_preview_rows(preview_slice)
else:
# Format detection failed — return raw samples so user can
# see actual data and map columns in the frontend
preview_samples = _serialize_preview_rows(preview_slice)
return CheckFormatResponse(
requires_manual_mapping=result["requires_manual_mapping"],
detected_format=result["detected_format"],
columns=result["columns"],
is_multimodal=result.get("is_multimodal", False),
multimodal_columns=result.get("multimodal_columns"),
suggested_mapping=result.get("suggested_mapping"),
detected_image_column=result.get("detected_image_column"),
detected_text_column=result.get("detected_text_column"),
preview_samples=preview_samples,
total_rows=total_rows,
)
except HTTPException:

238
studio/backend/routes/inference.py
View file

@ -2,13 +2,17 @@
Inference API routes for model loading and text generation.
"""
import sys
import time
import uuid
from pathlib import Path
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.responses import StreamingResponse, JSONResponse
from typing import Optional
import json
import logging
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
@ -32,6 +36,13 @@ from models.inference import (
LoadResponse,
UnloadResponse,
InferenceStatusResponse,
ChatCompletionRequest,
ChatCompletionChunk,
ChatCompletion,
ChunkChoice,
ChoiceDelta,
CompletionChoice,
CompletionMessage,
)
router = APIRouter()
@ -222,3 +233,228 @@ async def get_status():
status_code=500,
detail=f"Failed to get status: {str(e)}"
)
# =====================================================================
# OpenAI-Compatible Chat Completions (/chat/completions)
# =====================================================================
def _extract_content_parts(
messages: list,
) -> tuple[str, list[dict], "Optional[str]"]:
"""
Parse OpenAI-format messages into components the inference backend expects.
Handles both plain-string ``content`` and multimodal content-part arrays
(``[{type: "text", ...}, {type: "image_url", ...}]``).
Returns:
system_prompt: The system message text (or a default).
chat_messages: Non-system messages with content flattened to strings.
image_base64: Base64 data of the *first* image found, or ``None``.
"""
system_prompt = "You are a helpful AI assistant."
chat_messages: list[dict] = []
first_image_b64: Optional[str] = None
for msg in messages:
# ── System messages → extract as system_prompt ────────
if msg.role == "system":
if isinstance(msg.content, str):
system_prompt = msg.content
elif isinstance(msg.content, list):
# Unlikely but handle: join text parts
system_prompt = "\n".join(
p.text for p in msg.content if p.type == "text"
)
continue
# ── User / assistant messages ─────────────────────────
if isinstance(msg.content, str):
# Plain string content — pass through
chat_messages.append({"role": msg.role, "content": msg.content})
elif isinstance(msg.content, list):
# Multimodal content parts
text_parts: list[str] = []
for part in msg.content:
if part.type == "text":
text_parts.append(part.text)
elif part.type == "image_url" and first_image_b64 is None:
url = part.image_url.url
if url.startswith("data:"):
# data:image/png;base64,<DATA> → extract <DATA>
first_image_b64 = url.split(",", 1)[1] if "," in url else None
else:
logger.warning(
f"Remote image URLs not yet supported: {url[:80]}..."
)
combined_text = "\n".join(text_parts) if text_parts else ""
chat_messages.append({"role": msg.role, "content": combined_text})
return system_prompt, chat_messages, first_image_b64
@router.post("/chat/completions")
async def openai_chat_completions(request: ChatCompletionRequest):
"""
OpenAI-compatible chat completions endpoint.
Supports multimodal messages: ``content`` may be a plain string or a
list of content parts (``text`` / ``image_url``).
Streaming (default): returns SSE chunks matching OpenAI's format.
Non-streaming: returns a single ChatCompletion JSON object.
"""
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(
status_code=400,
detail="No model loaded. Call POST /inference/load first.",
)
# ── Parse messages (handles multimodal content parts) ─────
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
request.messages
)
# If no non-system messages were provided, error out
if not chat_messages:
raise HTTPException(
status_code=400,
detail="At least one non-system message is required.",
)
# ── Decode image (from content parts OR legacy field) ─────
# Content-part images take priority; fall back to legacy field
image_b64 = extracted_image_b64 or request.image_base64
image = None
if image_b64:
try:
import base64
from PIL import Image
from io import BytesIO
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_vision"):
raise HTTPException(
status_code=400,
detail="Image provided but current model is text-only. Load a vision model.",
)
image_data = base64.b64decode(image_b64)
image = Image.open(BytesIO(image_data))
image = backend.resize_image(image)
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}")
# ── Shared generation kwargs ──────────────────────────────
gen_kwargs = dict(
messages=chat_messages,
system_prompt=system_prompt,
image=image,
temperature=request.temperature,
top_p=request.top_p,
top_k=request.top_k,
max_new_tokens=request.max_tokens or 512,
repetition_penalty=request.repetition_penalty,
)
model_name = backend.active_model_name or request.model
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
# ── Streaming response ────────────────────────────────────────
if request.stream:
async def stream_chunks():
try:
# First chunk: send the role
first_chunk = ChatCompletionChunk(
id=completion_id,
created=created,
model=model_name,
choices=[ChunkChoice(
delta=ChoiceDelta(role="assistant"),
finish_reason=None,
)],
)
yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
# Content chunks — generate_chat_response yields cumulative
# text, so we diff to get incremental deltas.
prev_text = ""
for cumulative in backend.generate_chat_response(**gen_kwargs):
new_text = cumulative[len(prev_text):]
prev_text = cumulative
if not new_text:
continue
chunk = ChatCompletionChunk(
id=completion_id,
created=created,
model=model_name,
choices=[ChunkChoice(
delta=ChoiceDelta(content=new_text),
finish_reason=None,
)],
)
yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
# Final chunk: finish_reason = stop
final_chunk = ChatCompletionChunk(
id=completion_id,
created=created,
model=model_name,
choices=[ChunkChoice(
delta=ChoiceDelta(),
finish_reason="stop",
)],
)
yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
yield "data: [DONE]\n\n"
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI streaming: {e}", exc_info=True)
error_chunk = {
"error": {"message": str(e), "type": "server_error"},
}
yield f"data: {json.dumps(error_chunk)}\n\n"
return StreamingResponse(
stream_chunks(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
# ── Non-streaming response ────────────────────────────────────
else:
try:
full_text = ""
for token in backend.generate_chat_response(**gen_kwargs):
full_text = token # generate_stream yields cumulative text
response = ChatCompletion(
id=completion_id,
created=created,
model=model_name,
choices=[CompletionChoice(
message=CompletionMessage(content=full_text),
finish_reason="stop",
)],
)
return JSONResponse(content=response.model_dump())
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI completion: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))

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

@ -3,7 +3,7 @@ Training API routes
"""
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from typing import Dict, Optional
import logging
@ -124,6 +124,7 @@ async def start_training(
"hf_dataset": request.hf_dataset or "",
"local_datasets": request.local_datasets,
"format_type": request.format_type,
"custom_format_mapping": request.custom_format_mapping,
"num_epochs": request.num_epochs,
"learning_rate": request.learning_rate,
"batch_size": request.batch_size,
@ -180,14 +181,19 @@ async def start_training(
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):
update_count += 1
if update_count % 10 == 0:
logger.info(f"Training progress update #{update_count}")
# start_training returns bool (not generator)
run_result = backend.start_training(**training_kwargs)
logger.info(
"Training job %s backend.start_training returned type=%s value=%r",
job_id,
type(run_result).__name__,
run_result,
)
if not run_result:
progress_error = backend.trainer.training_progress.error
raise RuntimeError(progress_error or "Training failed to start")
logger.info(f"Training job {job_id} completed successfully")
logger.info(f"Training job {job_id} started successfully")
except Exception as e:
logger.error(f"Training error in job {job_id}: {e}", exc_info=True)
@ -337,6 +343,15 @@ async def get_training_status(
"learning_rate": getattr(progress, "learning_rate", 0.0),
}
# Build metric history for chart recovery after SSE reconnection
metric_history = None
if backend.step_history:
metric_history = {
"steps": list(backend.step_history),
"loss": list(backend.loss_history),
"lr": list(backend.lr_history),
}
return TrainingStatus(
job_id=job_id,
phase=phase,
@ -344,6 +359,7 @@ async def get_training_status(
message=status_message,
error=error_message,
details=details,
metric_history=metric_history,
)
except Exception as e:
@ -393,24 +409,40 @@ async def get_training_metrics(
@router.get("/progress")
async def stream_training_progress(
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Stream training progress updates using Server-Sent Events (SSE).
This endpoint provides real-time updates on training progress.
Supports reconnection via the SSE spec:
- Sends `id:` with each event so the browser tracks position.
- Sends `retry:` to control reconnection interval.
- Sends named `event:` types (progress, heartbeat, complete, error).
- Reads `Last-Event-ID` header on reconnect to replay missed steps.
"""
# Read Last-Event-ID header for reconnection resume
last_event_id = request.headers.get("last-event-id")
resume_from_step: Optional[int] = None
if last_event_id is not None:
try:
resume_from_step = int(last_event_id)
logger.info(f"SSE reconnect: resuming from step {resume_from_step}")
except ValueError:
logger.warning(f"Invalid Last-Event-ID: {last_event_id}")
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
# ── Helpers ──────────────────────────────────────────────
def build_progress(
step: int,
loss: float,
learning_rate: float,
total_steps: int,
epoch: Optional[int] = None,
epoch: Optional[float] = None,
) -> TrainingProgress:
total = max(total_steps, 0)
if step < 0 or total == 0:
@ -434,45 +466,86 @@ async def stream_training_progress(
num_tokens=None,
)
# Send initial status
is_active = backend.is_training_active()
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
def format_sse(
data: str,
event: str = "progress",
event_id: Optional[int] = None,
) -> str:
"""Format a single SSE message with id/event/data fields."""
lines = []
if event_id is not None:
lines.append(f"id: {event_id}")
lines.append(f"event: {event}")
lines.append(f"data: {data}")
lines.append("") # trailing blank line
lines.append("") # double newline terminates the event
return "\n".join(lines)
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"
# ── Retry directive ──────────────────────────────────────
# Tell the browser to reconnect after 3 seconds if the connection drops
yield "retry: 3000\n\n"
# If not active, check if there's any history
if not is_active:
if backend.step_history:
# Training completed - send final metrics
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
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: {build_progress(-1, 0.0, 0.0, 0).model_dump_json()}\n\n"
return
# Poll for updates while training is active
last_step = -1
# ── Replay missed steps on reconnect ─────────────────────
if resume_from_step is not None and backend.step_history:
replayed = 0
for i, step_val in enumerate(backend.step_history):
if step_val > resume_from_step:
loss_val = backend.loss_history[i] if i < len(backend.loss_history) else 0.0
lr_val = backend.lr_history[i] if i < len(backend.lr_history) else 0.0
tp_replay = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
total_replay = getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
payload = build_progress(step_val, loss_val, lr_val, total_replay, epoch_replay)
yield format_sse(payload.model_dump_json(), event="progress", event_id=step_val)
replayed += 1
if replayed:
logger.info(f"SSE reconnect: replayed {replayed} missed steps")
# ── Initial status (only on fresh connections) ───────────
if resume_from_step is None:
is_active = backend.is_training_active()
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 format_sse(initial_progress.model_dump_json(), event="progress", event_id=0)
# If not active, send final state and exit
if not is_active:
if backend.step_history:
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
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch)
yield format_sse(payload.model_dump_json(), event="complete", event_id=final_step)
else:
yield format_sse(
build_progress(-1, 0.0, 0.0, 0).model_dump_json(),
event="complete",
event_id=0,
)
return
# ── Live polling loop ────────────────────────────────────
last_step = resume_from_step if resume_from_step is not None else -1
no_update_count = 0
max_no_updates = 300 # Timeout after 5 minutes
max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation)
while backend.is_training_active():
try:
# Get current metrics
if backend.step_history:
current_step = backend.step_history[-1]
current_loss = backend.loss_history[-1] if backend.loss_history else 0.0
@ -496,7 +569,11 @@ async def stream_training_progress(
current_total_steps,
current_epoch,
)
yield f"data: {progress_payload.model_dump_json()}\n\n"
yield format_sse(
progress_payload.model_dump_json(),
event="progress",
event_id=current_step,
)
last_step = current_step
no_update_count = 0
else:
@ -510,30 +587,58 @@ async def stream_training_progress(
current_total_steps,
current_epoch,
)
yield f"data: {heartbeat_payload.model_dump_json()}\n\n"
yield format_sse(
heartbeat_payload.model_dump_json(),
event="heartbeat",
event_id=current_step,
)
else:
# No steps yet, but training is active
# No steps yet, but training is active (model loading, etc.)
no_update_count += 1
if no_update_count % 5 == 0:
preparing_payload = build_progress(0, 0.0, 0.0, 0)
yield f"data: {preparing_payload.model_dump_json()}\n\n"
# Pull total_steps and status from trainer so
# the frontend can show "Tokenizing…" etc.
tp_prep = getattr(
getattr(backend, "trainer", None),
"training_progress", None,
)
prep_total = (
getattr(tp_prep, "total_steps", 0)
if tp_prep else 0
)
preparing_payload = build_progress(
0, 0.0, 0.0, prep_total,
)
yield format_sse(
preparing_payload.model_dump_json(),
event="heartbeat",
event_id=0,
)
# Timeout check
if no_update_count > max_no_updates:
logger.warning("Progress stream timeout - no updates received")
timeout_payload = build_progress(last_step, 0.0, 0.0, 0)
yield f"data: {timeout_payload.model_dump_json()}\n\n"
yield format_sse(
timeout_payload.model_dump_json(),
event="error",
event_id=last_step if last_step >= 0 else 0,
)
break
await asyncio.sleep(1) # Poll every second
except Exception as e:
logger.error(f"Error in progress stream: {e}", exc_info=True)
error_payload = build_progress(0, 0.0, 0.0, 0)
yield f"data: {error_payload.model_dump_json()}\n\n"
yield format_sse(
error_payload.model_dump_json(),
event="error",
event_id=last_step if last_step >= 0 else 0,
)
break
# Send final status
# ── Final "complete" event ───────────────────────────────
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
@ -551,14 +656,18 @@ async def stream_training_progress(
final_total_steps,
final_epoch,
)
yield f"data: {final_payload.model_dump_json()}\n\n"
yield format_sse(
final_payload.model_dump_json(),
event="complete",
event_id=final_step if final_step >= 0 else 0,
)
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
}
)

59
studio/backend/run.py
View file

@ -11,10 +11,55 @@ 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,
frontend_path: Path = None,
frontend_path: Path = "studio/frontend/dist",
silent: bool = False,
):
"""
@ -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
@ -75,7 +124,7 @@ if __name__ == "__main__":
parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
parser.add_argument(
"--frontend", type=str, default=None, help="Path to frontend build"
"--frontend", type=str, default="studio/frontend/dist", help="Path to frontend build"
)
parser.add_argument("--silent", action="store_true", help="Suppress output")

13
studio/backend/utils/datasets/dataset_utils.py
View file

@ -63,6 +63,11 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"""
columns = list(dataset.column_names) if hasattr(dataset, 'column_names') else list(next(iter(dataset)).keys())
# Auto-detect multimodal data regardless of is_vlm flag
multimodal_info = detect_multimodal_dataset(dataset)
if multimodal_info["is_multimodal"]:
is_vlm = True # Route to VLM detection automatically
if is_vlm:
vlm_structure = detect_vlm_dataset_structure(dataset)
requires_mapping = vlm_structure["format"] == "unknown"
@ -74,6 +79,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"suggested_mapping": None,
"detected_image_column": vlm_structure.get("image_column"),
"detected_text_column": vlm_structure.get("text_column"),
"is_multimodal": multimodal_info["is_multimodal"],
"multimodal_columns": multimodal_info.get("multimodal_columns"),
}
else:
# LLM flow
@ -91,6 +98,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"suggested_mapping": heuristic_mapping,
"detected_image_column": None,
"detected_text_column": None,
"is_multimodal": False,
"multimodal_columns": None,
}
else:
# Both detection and heuristic failed
@ -101,6 +110,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"suggested_mapping": None,
"detected_image_column": None,
"detected_text_column": None,
"is_multimodal": False,
"multimodal_columns": None,
}
# Known format detected
@ -111,6 +122,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"suggested_mapping": None,
"detected_image_column": None,
"detected_text_column": None,
"is_multimodal": False,
"multimodal_columns": None,
}
def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):

40
studio/backend/utils/models/model_config.py
View file

@ -526,8 +526,8 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
"""
try:
# Get the script directory to locate configs
script_dir = Path(__file__).parent.parent
defaults_dir = script_dir / "configs" / "model_defaults"
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
# First, check if model is in the mapping
if model_name in _REVERSE_MODEL_MAPPING:
@ -661,10 +661,44 @@ class ModelConfig:
identifier = f"unsloth/{identifier}"
path = identifier
# Auto-detect LoRA for local paths (check adapter_config.json on disk)
if not is_lora and is_local:
detected_base = get_base_model_from_lora(path)
if detected_base:
is_lora = True
logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})")
# Auto-detect LoRA for remote HF models (check repo file listing)
if not is_lora and not is_local:
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(identifier, token=hf_token)
repo_files = [s.rfilename for s in info.siblings]
if "adapter_config.json" in repo_files:
is_lora = True
logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'")
except Exception as e:
logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
# Handle LoRA adapters
base_model = None
if is_lora:
base_model = get_base_model_from_lora(path)
if is_local:
# Local LoRA: read adapter_config.json from disk
base_model = get_base_model_from_lora(path)
else:
# Remote LoRA: download adapter_config.json from HF
try:
from huggingface_hub import hf_hub_download
config_path = hf_hub_download(identifier, "adapter_config.json", token=hf_token)
with open(config_path, 'r') as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
if base_model:
logger.info(f"Resolved remote LoRA base model: '{base_model}'")
except Exception as e:
logger.warning(f"Could not download adapter_config.json for '{identifier}': {e}")
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'")
return None

6
studio/frontend/bun.lock
View file

@ -10,6 +10,7 @@
"@assistant-ui/react-streamdown": "^0.1.2",
"@base-ui/react": "^1.2.0",
"@dagrejs/dagre": "^2.0.4",
"@dagrejs/graphlib": "^3.0.4",
"@fontsource-variable/figtree": "^5.2.10",
"@fontsource-variable/inter": "^5.2.8",
"@fontsource-variable/space-grotesk": "^5.2.10",
@ -27,6 +28,7 @@
"@streamdown/mermaid": "^1.0.2",
"@tailwindcss/vite": "^4.1.18",
"@tanstack/react-router": "^1.159.10",
"@tanstack/react-table": "^8.21.3",
"@toolwind/corner-shape": "^0.0.8-3",
"@types/canvas-confetti": "^1.9.0",
"@xyflow/react": "^12.10.0",
@ -663,10 +665,14 @@
"@tanstack/react-store": ["@tanstack/react-store@0.8.0", "", { "dependencies": { "@tanstack/store": "0.8.0", "use-sync-external-store": "^1.6.0" }, "peerDependencies": { "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-1vG9beLIuB7q69skxK9r5xiLN3ztzIPfSQSs0GfeqWGO2tGIyInZx0x1COhpx97RKaONSoAb8C3dxacWksm1ow=="],
"@tanstack/react-table": ["@tanstack/react-table@8.21.3", "", { "dependencies": { "@tanstack/table-core": "8.21.3" }, "peerDependencies": { "react": ">=16.8", "react-dom": ">=16.8" } }, "sha512-5nNMTSETP4ykGegmVkhjcS8tTLW6Vl4axfEGQN3v0zdHYbK4UfoqfPChclTrJ4EoK9QynqAu9oUf8VEmrpZ5Ww=="],
"@tanstack/router-core": ["@tanstack/router-core@1.159.9", "", { "dependencies": { "@tanstack/history": "1.154.14", "@tanstack/store": "^0.8.0", "cookie-es": "^2.0.0", "seroval": "^1.4.2", "seroval-plugins": "^1.4.2", "tiny-invariant": "^1.3.3", "tiny-warning": "^1.0.3" } }, "sha512-A9B8gvklvMCjSAFG8nDAhfmROI8kjcij8wzznQaw4RfGIOrYXyNe5fCAcbHXGpgNeTE2JnK75b6AjidDPQfrmw=="],
"@tanstack/store": ["@tanstack/store@0.8.0", "", {}, "sha512-Om+BO0YfMZe//X2z0uLF2j+75nQga6TpTJgLJQBiq85aOyZNIhkCgleNcud2KQg4k4v9Y9l+Uhru3qWMPGTOzQ=="],
"@tanstack/table-core": ["@tanstack/table-core@8.21.3", "", {}, "sha512-ldZXEhOBb8Is7xLs01fR3YEc3DERiz5silj8tnGkFZytt1abEvl/GhUmCE0PMLaMPTa3Jk4HbKmRlHmu+gCftg=="],
"@toolwind/corner-shape": ["@toolwind/corner-shape@0.0.8-3", "", { "dependencies": { "@types/node": "^20.4.1" } }, "sha512-MPIF81F2bhtXbzEeXF0vnL+PKpnopCHOzBspOkK8osMzWQvPUujZn2XZOMdsu4DF6wsVbbRYQtdsJr486HmIPQ=="],
"@ts-morph/common": ["@ts-morph/common@0.27.0", "", { "dependencies": { "fast-glob": "^3.3.3", "minimatch": "^10.0.1", "path-browserify": "^1.0.1" } }, "sha512-Wf29UqxWDpc+i61k3oIOzcUfQt79PIT9y/MWfAGlrkjg6lBC1hwDECLXPVJAhWjiGbfBCxZd65F/LIZF3+jeJQ=="],

2
studio/frontend/package.json
View file

@ -18,6 +18,7 @@
"@assistant-ui/react-streamdown": "^0.1.2",
"@base-ui/react": "^1.2.0",
"@dagrejs/dagre": "^2.0.4",
"@dagrejs/graphlib": "^3.0.4",
"@fontsource-variable/figtree": "^5.2.10",
"@fontsource-variable/inter": "^5.2.8",
"@fontsource-variable/space-grotesk": "^5.2.10",
@ -35,6 +36,7 @@
"@streamdown/mermaid": "^1.0.2",
"@tailwindcss/vite": "^4.1.18",
"@tanstack/react-router": "^1.159.10",
"@tanstack/react-table": "^8.21.3",
"@toolwind/corner-shape": "^0.0.8-3",
"@types/canvas-confetti": "^1.9.0",
"@xyflow/react": "^12.10.0",

2
studio/frontend/src/app/provider.tsx
View file

@ -10,7 +10,7 @@ export function AppProvider({ children }: AppProviderProps) {
return (
<ThemeProvider attribute="class" defaultTheme="light">
{children}
<Toaster />
<Toaster position="top-right" />
</ThemeProvider>
);
}

436
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -1,14 +1,22 @@
"use client";
import { Input } from "@/components/ui/input";
import {
Popover,
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
import { cn } from "@/lib/utils";
import { ArrowDown01Icon, Logout01Icon } from "@hugeicons/core-free-icons";
import { Spinner } from "@/components/ui/spinner";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks";
import { cn, formatCompact } from "@/lib/utils";
import {
ArrowDown01Icon,
Logout01Icon,
Search01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type ReactNode, useState } from "react";
import { type ReactNode, useMemo, useState } from "react";
export interface ModelOption {
id: string;
@ -17,11 +25,22 @@ export interface ModelOption {
icon?: ReactNode;
}
export interface LoraModelOption extends ModelOption {
baseModel?: string;
updatedAt?: number;
}
export interface ModelSelectorChangeMeta {
source: "hub" | "lora";
isLora: boolean;
}
interface ModelSelectorProps {
models: ModelOption[];
loraModels?: LoraModelOption[];
value?: string;
defaultValue?: string;
onValueChange?: (value: string) => void;
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
onEject?: () => void;
variant?: "outline" | "ghost" | "muted";
size?: "sm" | "default" | "lg";
@ -29,7 +48,9 @@ interface ModelSelectorProps {
contentClassName?: string;
}
// --- Composable sub-components ---
function dedupe(values: string[]): string[] {
return [...new Set(values.filter(Boolean))];
}
function ModelSelectorTrigger({
currentModel,
@ -63,15 +84,11 @@ function ModelSelectorTrigger({
{isLoaded && (
<span className="size-2 shrink-0 rounded-full bg-emerald-500" />
)}
<span
className={isLoaded ? "text-foreground" : "text-muted-foreground"}
>
{currentModel?.name ?? "Select a model\u2026"}
<span className={isLoaded ? "text-foreground" : "text-muted-foreground"}>
{currentModel?.name ?? "Select model..."}
</span>
{currentModel?.description && (
<span className="text-muted-foreground text-xs">
{currentModel.description}
</span>
<span className="text-muted-foreground text-xs">{currentModel.description}</span>
)}
<HugeiconsIcon
icon={ArrowDown01Icon}
@ -82,95 +99,329 @@ function ModelSelectorTrigger({
);
}
function ListLabel({ children }: { children: ReactNode }) {
return (
<div className="px-2.5 py-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
{children}
</div>
);
}
function ModelRow({
label,
meta,
selected,
onClick,
}: {
label: string;
meta?: string;
selected?: boolean;
onClick: () => void;
}) {
return (
<button
type="button"
onClick={onClick}
className={cn(
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-accent",
selected && "bg-accent/60",
)}
>
<span className="min-w-0 flex-1 truncate">{label}</span>
{meta ? (
<span className="shrink-0 text-[10px] text-muted-foreground">{meta}</span>
) : null}
</button>
);
}
function HubModelPicker({
models,
value,
onSelect,
}: {
models: ModelOption[];
value?: string;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
}) {
const [query, setQuery] = useState("");
const debouncedQuery = useDebouncedValue(query);
const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch(
debouncedQuery,
);
const recommendedIds = useMemo(
() => dedupe([...models.map((model) => model.id), value ?? ""]),
[models, value],
);
const showHfSection = debouncedQuery.trim().length > 0;
const recommendedSet = useMemo(
() => new Set(recommendedIds),
[recommendedIds],
);
const hfIds = useMemo(() => {
if (!showHfSection) {
return [];
}
return results
.map((result) => result.id)
.filter((id) => !recommendedSet.has(id));
}, [recommendedSet, results, showHfSection]);
const metricsById = useMemo(
() =>
new Map(
results.map((result) => [
result.id,
result.totalParams
? formatCompact(result.totalParams)
: `↓${formatCompact(result.downloads)}`,
]),
),
[results],
);
const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
return (
<div className="space-y-2">
<div className="relative">
<HugeiconsIcon
icon={Search01Icon}
className="pointer-events-none absolute left-2.5 top-2.5 size-4 text-muted-foreground"
/>
<Input
value={query}
onChange={(event) => setQuery(event.target.value)}
placeholder="Search Hugging Face models"
className="h-9 pl-8 pr-8"
/>
{isLoading && (
<Spinner className="pointer-events-none absolute right-2.5 top-2.5 size-4 text-muted-foreground" />
)}
</div>
<div
ref={scrollRef}
className="max-h-64 overflow-y-auto"
>
<div className="p-1">
{!showHfSection ? (
<>
<ListLabel>Recommended</ListLabel>
{recommendedIds.length === 0 ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">
No default models.
</div>
) : (
recommendedIds.map((id) => (
<ModelRow
key={id}
label={id}
selected={value === id}
onClick={() => onSelect(id, { source: "hub", isLora: false })}
/>
))
)}
</>
) : null}
{showHfSection ? (
<>
<ListLabel>Hugging Face</ListLabel>
{hfIds.length === 0 && !isLoading ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">
No matching models.
</div>
) : (
hfIds.map((id) => (
<ModelRow
key={id}
label={id}
meta={metricsById.get(id)}
selected={value === id}
onClick={() => onSelect(id, { source: "hub", isLora: false })}
/>
))
)}
<div ref={sentinelRef} className="h-px" />
{isLoadingMore ? (
<div className="flex items-center justify-center py-2">
<Spinner className="size-3.5 text-muted-foreground" />
</div>
) : null}
</>
) : null}
</div>
</div>
</div>
);
}
function LoraModelPicker({
loraModels,
value,
onSelect,
}: {
loraModels: LoraModelOption[];
value?: string;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
}) {
const [query, setQuery] = useState("");
const normalized = useMemo(
() =>
loraModels
.map((model) => ({
...model,
baseModel: model.baseModel || model.description || "Unknown base model",
}))
.sort((a, b) => {
const aTime = a.updatedAt ?? -1;
const bTime = b.updatedAt ?? -1;
if (aTime !== bTime) {
return bTime - aTime;
}
const baseCmp = a.baseModel.localeCompare(b.baseModel);
if (baseCmp !== 0) {
return baseCmp;
}
return a.name.localeCompare(b.name);
}),
[loraModels],
);
const grouped = useMemo(() => {
const needle = query.trim().toLowerCase();
const out = new Map<string, LoraModelOption[]>();
for (const model of normalized) {
const searchText = `${model.name} ${model.baseModel} ${model.id}`.toLowerCase();
if (needle && !searchText.includes(needle)) {
continue;
}
const key = model.baseModel || "Unknown base model";
const prev = out.get(key) ?? [];
prev.push(model);
out.set(key, prev);
}
return [...out.entries()].sort((a, b) => {
const aLatest = Math.max(...a[1].map((model) => model.updatedAt ?? -1));
const bLatest = Math.max(...b[1].map((model) => model.updatedAt ?? -1));
if (aLatest !== bLatest) {
return bLatest - aLatest;
}
return a[0].localeCompare(b[0]);
});
}, [normalized, query]);
return (
<div className="space-y-2">
<div className="relative">
<HugeiconsIcon
icon={Search01Icon}
className="pointer-events-none absolute left-2.5 top-2.5 size-4 text-muted-foreground"
/>
<Input
value={query}
onChange={(event) => setQuery(event.target.value)}
placeholder="Search local adapters"
className="h-9 pl-8"
/>
</div>
<div className="max-h-64 overflow-y-auto">
<div className="p-1">
{grouped.length === 0 ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">No adapters found.</div>
) : (
grouped.map(([baseModel, adapters], index) => (
<div key={baseModel}>
{index > 0 ? <div className="my-1" /> : null}
<ListLabel>{baseModel}</ListLabel>
{adapters.map((adapter) => (
<ModelRow
key={adapter.id}
label={adapter.name}
meta="LoRA"
selected={value === adapter.id}
onClick={() => onSelect(adapter.id, { source: "lora", isLora: true })}
/>
))}
</div>
))
)}
</div>
</div>
</div>
);
}
function ModelSelectorContent({
models,
loraModels,
value,
onSelect,
onEject,
className,
}: {
models: ModelOption[];
loraModels: LoraModelOption[];
value?: string;
onSelect: (id: string) => void;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
onEject?: () => void;
className?: string;
}) {
const hasSelection = Boolean(value);
return (
<PopoverContent
align="start"
className={cn("w-auto min-w-[280px] gap-0 p-1", className)}
className={cn("w-[440px] min-w-[440px] gap-0 p-2", className)}
>
{models.map((model) => (
<ModelSelectorItem
key={model.id}
model={model}
isActive={value === model.id}
onSelect={onSelect}
onEject={onEject}
/>
))}
<Tabs defaultValue="hub" className="w-full">
<TabsList className="mb-2 w-full">
<TabsTrigger value="hub">Hub models</TabsTrigger>
<TabsTrigger value="lora">Fine-tuned</TabsTrigger>
</TabsList>
<TabsContent value="hub" className="m-0">
<HubModelPicker models={models} value={value} onSelect={onSelect} />
</TabsContent>
<TabsContent value="lora" className="m-0">
<LoraModelPicker
loraModels={loraModels}
value={value}
onSelect={onSelect}
/>
</TabsContent>
</Tabs>
{hasSelection && onEject ? (
<div className="mt-2 border-t border-border/70 pt-2">
<button
type="button"
onClick={onEject}
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
title="Eject model"
>
<HugeiconsIcon icon={Logout01Icon} className="size-3.5" />
Eject loaded model
</button>
</div>
) : null}
</PopoverContent>
);
}
function ModelSelectorItem({
model,
isActive,
onSelect,
onEject,
}: {
model: ModelOption;
isActive: boolean;
onSelect: (id: string) => void;
onEject?: () => void;
}) {
return (
<button
type="button"
aria-pressed={isActive}
onClick={() => onSelect(model.id)}
className={cn(
"group flex w-full cursor-pointer items-center gap-2.5 rounded-md px-3 py-2 text-left text-sm transition-colors hover:bg-accent",
isActive && "bg-accent/50",
)}
>
<span
className={cn(
"size-2 shrink-0 rounded-full",
isActive ? "bg-emerald-500" : "bg-transparent",
)}
/>
{model.icon && <span className="shrink-0">{model.icon}</span>}
<div className="min-w-0 flex-1">
<div className="truncate text-sm">{model.name}</div>
{model.description && (
<div className="truncate text-xs text-muted-foreground">
{model.description}
</div>
)}
</div>
{isActive && onEject && (
<button
type="button"
onClick={(e) => {
e.stopPropagation();
onEject();
}}
className="flex items-center gap-1 rounded px-1.5 py-0.5 text-[10px] text-muted-foreground opacity-0 transition-opacity group-hover:opacity-100 hover:bg-muted hover:text-foreground"
title="Eject model"
>
<HugeiconsIcon icon={Logout01Icon} className="size-3" />
Eject
</button>
)}
</button>
);
}
// --- Main component ---
export function ModelSelector({
models,
loraModels = [],
value,
defaultValue,
onValueChange,
@ -182,13 +433,31 @@ export function ModelSelector({
}: ModelSelectorProps) {
const [open, setOpen] = useState(false);
const [uncontrolled, setUncontrolled] = useState(defaultValue ?? "");
const selected = value ?? uncontrolled;
const isLoaded = selected !== "";
const currentModel = models.find((m) => m.id === selected);
function handleSelect(id: string) {
const optionById = useMemo(() => {
const all = new Map<string, ModelOption>();
for (const model of models) {
all.set(model.id, model);
}
for (const lora of loraModels) {
all.set(lora.id, {
...lora,
description: lora.baseModel || lora.description,
});
}
return all;
}, [loraModels, models]);
const currentModel = selected
? optionById.get(selected) ?? { id: selected, name: selected }
: undefined;
function handleSelect(id: string, meta: ModelSelectorChangeMeta) {
if (onValueChange) {
onValueChange(id);
onValueChange(id, meta);
} else {
setUncontrolled(id);
}
@ -211,6 +480,7 @@ export function ModelSelector({
/>
<ModelSelectorContent
models={models}
loraModels={loraModels}
value={selected}
onSelect={handleSelect}
onEject={onEject ? handleEject : undefined}
@ -220,7 +490,5 @@ export function ModelSelector({
);
}
// Composable exports
ModelSelector.Trigger = ModelSelectorTrigger;
ModelSelector.Content = ModelSelectorContent;
ModelSelector.Item = ModelSelectorItem;

59
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -7,7 +7,9 @@ import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { AnimatedShinyText } from "@/components/ui/animated-shiny-text";
import { Button } from "@/components/ui/button";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { cn } from "@/lib/utils";
import {
ActionBarMorePrimitive,
@ -20,6 +22,8 @@ import {
SuggestionPrimitive,
ThreadPrimitive,
useAui,
useAuiEvent,
useAuiState,
} from "@assistant-ui/react";
import { motion } from "framer-motion";
import {
@ -36,7 +40,7 @@ import {
RefreshCwIcon,
SquareIcon,
} from "lucide-react";
import type { FC } from "react";
import { type FC, useRef } from "react";
export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
hideComposer,
@ -69,6 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
<ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4 before:pointer-events-none before:absolute before:inset-x-0 before:bottom-full before:h-20 before:bg-gradient-to-t before:from-background before:to-transparent">
<ThreadScrollToBottom />
<WarmupIndicator />
<AuiIf condition={({ thread }) => !thread.isEmpty}>
{!hideComposer && <ComposerAnimated />}
</AuiIf>
@ -78,6 +83,28 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
);
};
const WarmupIndicator: FC = () => {
const threadId = useAuiState(({ threads }) => threads.mainThreadId);
const isRunning = useAuiState(({ thread }) => thread.isRunning);
const isWarmingUp = useChatRuntimeStore((state) =>
Boolean(state.warmingByThreadId[threadId ?? "__default"]),
);
if (!isRunning || !isWarmingUp) {
return null;
}
return (
<div className="mx-auto -mb-2 w-full max-w-(--thread-max-width) px-2">
<div className="inline-flex items-center rounded-full border border-border/60 bg-background/90 px-3 py-1.5 text-xs text-muted-foreground shadow-sm">
<AnimatedShinyText className="text-xs">
Warming up model...
</AnimatedShinyText>
</div>
</div>
);
};
const ThreadScrollToBottom: FC = () => {
return (
<ThreadPrimitive.ScrollToBottom asChild={true}>
@ -93,13 +120,26 @@ const ThreadScrollToBottom: FC = () => {
};
const SuggestionItem: FC = () => {
const aui = useAui();
const prompt = useAuiState(({ suggestion }) => suggestion.prompt);
const isDisabled = useAuiState(({ thread }) => thread.isDisabled);
const isRunning = useAuiState(({ thread }) => thread.isRunning);
return (
<SuggestionPrimitive.Trigger
send={true}
<button
type="button"
onClick={() => {
if (!isDisabled && !isRunning) {
aui.thread().append(prompt);
aui.composer().setText("");
return;
}
aui.composer().setText(prompt);
}}
className="fade-in slide-in-from-bottom-1 animate-in cursor-pointer corner-squircle rounded-xl border bg-background px-4 py-2.5 text-left text-sm text-foreground shadow-sm transition-colors duration-150 hover:bg-accent"
>
<SuggestionPrimitive.Title />
</SuggestionPrimitive.Trigger>
</button>
);
};
@ -358,6 +398,15 @@ const UserActionBar: FC = () => {
const EditComposer: FC = () => {
const aui = useAui();
const resendAfterCancelRef = useRef(false);
useAuiEvent("thread.runEnd", () => {
if (!resendAfterCancelRef.current) {
return;
}
resendAfterCancelRef.current = false;
aui.composer().send();
});
return (
<MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
@ -384,7 +433,9 @@ const EditComposer: FC = () => {
}
if (aui.thread().getState().isRunning) {
resendAfterCancelRef.current = true;
aui.thread().cancelRun();
return;
}
aui.composer().send();
}}

110
studio/frontend/src/components/ui/data-table.tsx Normal file
View file

@ -0,0 +1,110 @@
import {
type ColumnDef,
type SortingState,
flexRender,
getCoreRowModel,
getSortedRowModel,
useReactTable,
} from "@tanstack/react-table";
import { useState } from "react";
import {
Table,
TableBody,
TableCell,
TableHead,
TableHeader,
TableRow,
} from "@/components/ui/table";
import { cn } from "@/lib/utils";
interface DataTableProps<TData, TValue> {
columns: ColumnDef<TData, TValue>[];
data: TData[];
className?: string;
}
export function DataTable<TData, TValue>({
columns,
data,
className,
}: DataTableProps<TData, TValue>) {
const [sorting, setSorting] = useState<SortingState>([]);
const table = useReactTable({
data,
columns,
getCoreRowModel: getCoreRowModel(),
getSortedRowModel: getSortedRowModel(),
onSortingChange: setSorting,
state: { sorting },
});
return (
<div className={cn("w-full", className)}>
<Table>
<TableHeader className="sticky top-0 z-10">
{table.getHeaderGroups().map((headerGroup) => (
<TableRow
key={headerGroup.id}
className="bg-muted/60 hover:bg-muted/60 border-b border-border/60"
>
{headerGroup.headers.map((header) => (
<TableHead
key={header.id}
className="border-r border-border/40 last:border-r-0 h-11 px-4 text-xs"
style={{
width:
header.getSize() !== 150 ? header.getSize() : undefined,
}}
>
{header.isPlaceholder
? null
: flexRender(
header.column.columnDef.header,
header.getContext(),
)}
</TableHead>
))}
</TableRow>
))}
</TableHeader>
<TableBody>
{table.getRowModel().rows.length ? (
table.getRowModel().rows.map((row, idx) => (
<TableRow
key={row.id}
data-state={row.getIsSelected() ? "selected" : undefined}
className={cn(
"transition-colors border-b border-border/30",
idx % 2 === 0
? "bg-background"
: "bg-muted/20",
"hover:bg-primary/[0.03]",
)}
>
{row.getVisibleCells().map((cell) => (
<TableCell
key={cell.id}
className="border-r border-border/20 last:border-r-0 text-[13px] py-3 px-4 align-top whitespace-normal"
>
{flexRender(cell.column.columnDef.cell, cell.getContext())}
</TableCell>
))}
</TableRow>
))
) : (
<TableRow>
<TableCell
colSpan={columns.length}
className="h-32 text-center text-muted-foreground text-sm"
>
No results.
</TableCell>
</TableRow>
)}
</TableBody>
</Table>
</div>
);
}

227
studio/frontend/src/components/ui/terminal.tsx Normal file
View file

@ -0,0 +1,227 @@
import { cn } from "@/lib/utils"
import {
Children,
cloneElement,
isValidElement,
useEffect,
useRef,
useState,
} from "react"
import type { ElementType, ReactElement, ReactNode } from "react"
type TerminalProps = {
children: ReactNode
className?: string
sequence?: boolean
startOnView?: boolean
}
type InternalLineProps = {
__isActive?: boolean
__onDone?: () => void
__sequence?: boolean
}
function useStartOnView(enabled: boolean): {
ref: React.RefObject<HTMLDivElement | null>
started: boolean
} {
const ref = useRef<HTMLDivElement | null>(null)
const [isInView, setIsInView] = useState(false)
const started = !enabled || isInView
useEffect(() => {
if (!enabled) {
return
}
const node = ref.current
if (!node) {
return
}
const observer = new IntersectionObserver(
([entry]) => {
if (entry?.isIntersecting) {
setIsInView(true)
observer.disconnect()
}
},
{ threshold: 0.2 }
)
observer.observe(node)
return () => observer.disconnect()
}, [enabled])
return { ref, started }
}
export function Terminal({
children,
className,
sequence = true,
startOnView = true,
}: TerminalProps): ReactElement {
const { ref, started } = useStartOnView(startOnView)
const childElements = Children.toArray(children).filter(isValidElement)
const [activeIndex, setActiveIndex] = useState(0)
const visibleIndex = sequence
? started
? activeIndex
: -1
: Number.MAX_SAFE_INTEGER
function handleLineDone(index: number): void {
if (!sequence) {
return
}
setActiveIndex((prev) => {
if (prev !== index) {
return prev
}
return Math.min(index + 1, childElements.length)
})
}
return (
<div
ref={ref}
className={cn(
"w-full rounded-2xl border border-border bg-card px-6 py-5 font-mono text-sm text-foreground shadow-2xl",
className
)}
>
{childElements.map((child, index) =>
cloneElement(child, {
__sequence: sequence,
__isActive: !sequence || visibleIndex >= index,
__onDone: () => handleLineDone(index),
key: child.key ?? index,
} as InternalLineProps)
)}
</div>
)
}
type AnimatedSpanProps = InternalLineProps & {
children: ReactNode
className?: string
delay?: number
startOnView?: boolean
}
export function AnimatedSpan({
children,
className,
delay = 0,
startOnView = false,
__isActive,
__sequence,
__onDone,
}: AnimatedSpanProps): ReactElement {
const { ref, started } = useStartOnView(startOnView)
const [visible, setVisible] = useState(false)
const doneRef = useRef(false)
const onDoneRef = useRef(__onDone)
const shouldStart = __sequence ? __isActive : started
useEffect(() => {
onDoneRef.current = __onDone
}, [__onDone])
useEffect(() => {
if (!shouldStart || doneRef.current) {
return
}
const timeout = window.setTimeout(() => {
setVisible(true)
doneRef.current = true
onDoneRef.current?.()
}, delay)
return () => window.clearTimeout(timeout)
}, [delay, shouldStart])
return (
<div
ref={ref}
className={cn(
"min-h-5 transition-opacity duration-300",
visible ? "opacity-100" : "opacity-0",
className
)}
>
{children}
</div>
)
}
type TypingAnimationProps = InternalLineProps & {
children: string
className?: string
duration?: number
delay?: number
as?: ElementType
startOnView?: boolean
}
export function TypingAnimation({
children,
className,
duration = 60,
delay = 0,
as: Component = "span",
startOnView = true,
__isActive,
__sequence,
__onDone,
}: TypingAnimationProps): ReactElement {
const { ref, started } = useStartOnView(startOnView)
const [typed, setTyped] = useState("")
const doneRef = useRef(false)
const onDoneRef = useRef(__onDone)
const shouldStart = __sequence ? __isActive : started
useEffect(() => {
onDoneRef.current = __onDone
}, [__onDone])
useEffect(() => {
if (!shouldStart || doneRef.current) {
return
}
let index = 0
let intervalId: number | null = null
const startTimer = window.setTimeout(() => {
intervalId = window.setInterval(() => {
index += 1
setTyped(children.slice(0, index))
if (index >= children.length) {
if (intervalId) {
window.clearInterval(intervalId)
}
doneRef.current = true
onDoneRef.current?.()
}
}, duration)
}, delay)
return () => {
window.clearTimeout(startTimer)
if (intervalId) {
window.clearInterval(intervalId)
}
}
}, [children, delay, duration, shouldStart])
return (
<div ref={ref} className="min-h-5">
<Component className={cn("whitespace-pre-wrap", className)}>{typed}</Component>
</div>
)
}

2
studio/frontend/src/features/auth/index.ts
View file

@ -1,6 +1,6 @@
export { LoginPage } from "./login-page";
export { SignupPage } from "./signup-page";
export { refreshSession } from "./api";
export { authFetch, refreshSession } from "./api";
export {
getPostAuthRoute,
hasAuthToken,

121
studio/frontend/src/features/chat/adapter.ts
View file

@ -1,121 +0,0 @@
import type { ChatModelAdapter, ChatModelRunResult } from "@assistant-ui/react";
const API = import.meta.env.VITE_INFERENCE_URL || "/api/chat/generate";
type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
type RunMessage = RunMessages[number];
function collectTextParts(message: RunMessage): string[] {
const textParts = message.content
.filter((c) => c.type === "text")
.map((c) => c.text);
if ("attachments" in message) {
for (const att of message.attachments ?? []) {
for (const part of att.content ?? []) {
if (part.type === "text") {
textParts.push(part.text);
}
}
}
}
return textParts;
}
function messageToPayload(message: RunMessage): {
role: string;
content: string;
} {
return {
role: message.role,
content: collectTextParts(message).join("\n"),
};
}
function makeBody(messages: RunMessages): string {
const payloadMessages: Array<{ role: string; content: string }> = [];
for (const message of messages) {
payloadMessages.push(messageToPayload(message));
}
return JSON.stringify({ messages: payloadMessages });
}
export function parseThinkTags(raw: string): ChatModelRunResult["content"] {
const parts: ContentPart[] = [];
const thinkStart = raw.indexOf("<think>");
if (thinkStart === -1) {
if (raw) {
parts.push({ type: "text", text: raw });
}
return parts;
}
const before = raw.slice(0, thinkStart);
if (before.trim()) {
parts.push({ type: "text", text: before });
}
const thinkEnd = raw.indexOf("</think>");
if (thinkEnd === -1) {
const reasoning = raw.slice(thinkStart + 7);
if (reasoning) {
parts.push({ type: "reasoning", text: reasoning });
}
return parts;
}
const reasoning = raw.slice(thinkStart + 7, thinkEnd);
if (reasoning) {
parts.push({ type: "reasoning", text: reasoning });
}
const after = raw.slice(thinkEnd + 8);
if (after) {
parts.push({ type: "text", text: after });
}
return parts;
}
export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter {
return {
// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: stream loop ok
async *run({ messages, abortSignal }) {
const res = await fetch(apiUrl, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: makeBody(messages),
signal: abortSignal,
});
const reader = res.body?.getReader();
if (!reader) {
throw new Error("Response body is empty");
}
const decoder = new TextDecoder();
let text = "";
let reasoningStart: number | null = null;
let reasoningDuration = 0;
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
text += decoder.decode(value, { stream: true });
const parts = parseThinkTags(text) ?? [];
if (parts.some((p) => p.type === "reasoning") && !reasoningStart) {
reasoningStart = Date.now();
}
if (text.includes("</think>") && reasoningStart && !reasoningDuration) {
reasoningDuration = Math.round((Date.now() - reasoningStart) / 1000);
}
if (parts.length > 0) {
yield {
content: parts,
metadata: { custom: { reasoningDuration } },
};
}
}
},
};
}

165
studio/frontend/src/features/chat/api/chat-adapter.ts Normal file
View file

@ -0,0 +1,165 @@
import type { ChatModelAdapter } from "@assistant-ui/react";
import { streamChatCompletions } from "./chat-api";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
import {
hasClosedThinkTag,
parseAssistantContent,
} from "../utils/parse-assistant-content";
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
type RunMessage = RunMessages[number];
function collectTextParts(message: RunMessage): string[] {
const textParts = message.content
.filter((part) => part.type === "text")
.map((part) => part.text);
if ("attachments" in message && (message.attachments?.length ?? 0) > 0) {
for (const attachment of message.attachments ?? []) {
for (const part of attachment.content ?? []) {
if (part.type === "text") {
textParts.push(part.text);
}
}
}
}
return textParts;
}
function toOpenAIMessage(message: RunMessage): {
role: "system" | "user" | "assistant";
content: string;
} | null {
if (
message.role !== "system" &&
message.role !== "user" &&
message.role !== "assistant"
) {
return null;
}
return {
role: message.role,
content: collectTextParts(message).join("\n"),
};
}
function extractImageBase64(input: string): string | undefined {
if (!input) {
return undefined;
}
if (input.startsWith("data:")) {
const commaIndex = input.indexOf(",");
return commaIndex >= 0 ? input.slice(commaIndex + 1) : undefined;
}
return input;
}
function findLatestUserImageBase64(messages: RunMessages): string | undefined {
for (let i = messages.length - 1; i >= 0; i -= 1) {
const message = messages[i];
if (!message || message.role !== "user") {
continue;
}
if ("attachments" in message && (message.attachments?.length ?? 0) > 0) {
for (const attachment of message.attachments ?? []) {
for (const part of attachment.content ?? []) {
if (part.type !== "image") {
continue;
}
const encoded = extractImageBase64(part.image);
if (encoded) {
return encoded;
}
}
}
}
}
return undefined;
}
export function createOpenAIStreamAdapter(): ChatModelAdapter {
return {
async *run({ messages, abortSignal, unstable_threadId }) {
const state = useChatRuntimeStore.getState();
const { params } = state;
if (!params.checkpoint) {
throw new Error("Load a model first.");
}
const outboundMessages = messages
.map(toOpenAIMessage)
.filter((message): message is NonNullable<typeof message> =>
Boolean(message),
);
if (params.systemPrompt.trim()) {
outboundMessages.unshift({
role: "system",
content: params.systemPrompt.trim(),
});
}
const imageBase64 = findLatestUserImageBase64(messages);
const threadKey = unstable_threadId || "__default";
let waitingFirstChunk = true;
useChatRuntimeStore.getState().setThreadWarming(threadKey, true);
let cumulativeText = "";
let reasoningStartAt: number | null = null;
let reasoningDuration = 0;
try {
const stream = streamChatCompletions(
{
model: params.checkpoint,
messages: outboundMessages,
stream: true,
temperature: params.temperature,
top_p: params.topP,
max_tokens: params.maxTokens,
top_k: params.topK,
repetition_penalty: params.repetitionPenalty,
image_base64: imageBase64,
},
abortSignal,
);
for await (const chunk of stream) {
const delta = chunk.choices?.[0]?.delta?.content;
if (!delta) {
continue;
}
if (waitingFirstChunk) {
waitingFirstChunk = false;
useChatRuntimeStore.getState().setThreadWarming(threadKey, false);
}
cumulativeText += delta;
const parts = parseAssistantContent(cumulativeText);
if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) {
reasoningStartAt = Date.now();
}
if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) {
reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000);
}
if (parts.length > 0) {
yield {
content: parts,
metadata: { custom: { reasoningDuration } },
};
}
}
} finally {
if (waitingFirstChunk) {
useChatRuntimeStore.getState().setThreadWarming(threadKey, false);
}
}
},
};
}

146
studio/frontend/src/features/chat/api/chat-api.ts Normal file
View file

@ -0,0 +1,146 @@
import { authFetch } from "@/features/auth";
import type {
InferenceStatusResponse,
ListLorasResponse,
ListModelsResponse,
LoadModelRequest,
LoadModelResponse,
OpenAIChatChunk,
OpenAIChatCompletionsRequest,
UnloadModelRequest,
} from "../types/api";
function parseErrorText(status: number, body: unknown): string {
if (
body &&
typeof body === "object" &&
"detail" in body &&
typeof body.detail === "string"
) {
return body.detail;
}
if (
body &&
typeof body === "object" &&
"message" in body &&
typeof body.message === "string"
) {
return body.message;
}
return `Request failed (${status})`;
}
async function parseJsonOrThrow<T>(response: Response): Promise<T> {
const body = await response.json().catch(() => null);
if (!response.ok) {
throw new Error(parseErrorText(response.status, body));
}
return body as T;
}
export async function listModels(): Promise<ListModelsResponse> {
const response = await authFetch("/api/models/list");
return parseJsonOrThrow<ListModelsResponse>(response);
}
export async function listLoras(outputsDir = "./outputs"): Promise<ListLorasResponse> {
const query = new URLSearchParams({ outputs_dir: outputsDir }).toString();
const response = await authFetch(`/api/models/loras?${query}`);
return parseJsonOrThrow<ListLorasResponse>(response);
}
export async function getInferenceStatus(): Promise<InferenceStatusResponse> {
const response = await authFetch("/api/inference/status");
return parseJsonOrThrow<InferenceStatusResponse>(response);
}
export async function loadModel(
payload: LoadModelRequest,
): Promise<LoadModelResponse> {
const response = await authFetch("/api/inference/load", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
return parseJsonOrThrow<LoadModelResponse>(response);
}
export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
const response = await authFetch("/api/inference/unload", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
await parseJsonOrThrow<unknown>(response);
}
function parseSseEvent(rawEvent: string): string[] {
const dataLines: string[] = [];
for (const line of rawEvent.split(/\r?\n/)) {
if (line.startsWith("data:")) {
dataLines.push(line.slice(5).trimStart());
}
}
return dataLines;
}
export async function* streamChatCompletions(
payload: OpenAIChatCompletionsRequest,
signal: AbortSignal,
): AsyncGenerator<OpenAIChatChunk> {
const response = await authFetch("/api/inference/chat/completions", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
signal,
});
if (!response.ok) {
const body = await response.json().catch(() => null);
throw new Error(parseErrorText(response.status, body));
}
if (!response.body) {
throw new Error("Stream response missing body");
}
const reader = response.body.getReader();
const decoder = new TextDecoder();
let buffer = "";
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
buffer += decoder.decode(value, { stream: true });
let separatorIndex = buffer.search(/\r?\n\r?\n/);
while (separatorIndex >= 0) {
const rawEvent = buffer.slice(0, separatorIndex);
const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2;
buffer = buffer.slice(separatorIndex + separatorLength);
const dataLines = parseSseEvent(rawEvent);
if (dataLines.length === 0) {
separatorIndex = buffer.search(/\r?\n\r?\n/);
continue;
}
const dataText = dataLines.join("\n");
if (dataText === "[DONE]") {
return;
}
const parsed = JSON.parse(dataText) as
| OpenAIChatChunk
| { error?: { message?: string } };
if ("error" in parsed && parsed.error) {
throw new Error(parsed.error.message || "Stream error");
}
yield parsed as OpenAIChatChunk;
separatorIndex = buffer.search(/\r?\n\r?\n/);
}
}
}

124
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -1,4 +1,5 @@
import {
type LoraModelOption,
type ModelOption,
ModelSelector,
} from "@/components/assistant-ui/model-selector";
@ -28,16 +29,15 @@ import {
memo,
useCallback,
useEffect,
useMemo,
useRef,
useState,
} from "react";
import {
ChatSettingsPanel,
type InferenceParams,
defaultInferenceParams,
} from "./chat-settings-sheet";
import { ChatSettingsPanel } from "./chat-settings-sheet";
import { db } from "./db";
import { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
import { ChatRuntimeProvider } from "./runtime-provider";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import {
type CompareHandle,
CompareHandlesProvider,
@ -47,47 +47,16 @@ import {
import { ThreadSidebar } from "./thread-sidebar";
import type { ChatView } from "./types";
const LORA_MODELS: ModelOption[] = [
{
id: "outputs/llama-3.1-8b-instruct-lora",
name: "meta-llama/Llama-3.1-8B-Instruct",
description: "LoRA v1",
},
{
id: "outputs/qwen2.5-7b-lora",
name: "Qwen/Qwen2.5-7B-Instruct",
description: "LoRA v2",
},
{
id: "outputs/mistral-7b-v0.3-lora",
name: "mistralai/Mistral-7B-Instruct-v0.3",
description: "LoRA v1",
},
];
const GGUF_MODELS: ModelOption[] = [
{
id: "models/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
name: "Meta-Llama-3.1-8B-Instruct",
description: "Q4_K_M",
},
{
id: "models/Qwen2.5-7B-Instruct-Q5_K_M.gguf",
name: "Qwen2.5-7B-Instruct",
description: "Q5_K_M",
},
{
id: "models/Mistral-7B-Instruct-v0.3-Q4_K_M.gguf",
name: "Mistral-7B-Instruct-v0.3",
description: "Q4_K_M",
},
];
const SingleContent = memo(function SingleContent({
threadId,
}: { threadId?: string }): ReactElement {
newThreadNonce,
}: { threadId?: string; newThreadNonce?: string }): ReactElement {
return (
<ChatRuntimeProvider modelType="base" initialThreadId={threadId}>
<ChatRuntimeProvider
modelType="base"
initialThreadId={threadId}
newThreadNonce={newThreadNonce}
>
<div className="min-h-0 flex-1">
<Thread />
</div>
@ -233,33 +202,66 @@ function TopBarActions({
}
export function ChatPage(): ReactElement {
const [view, setView] = useState<ChatView>({ mode: "single" });
const [view, setView] = useState<ChatView>({
mode: "single",
newThreadNonce: crypto.randomUUID(),
});
const [settingsOpen, setSettingsOpen] = useState(false);
const [inferenceParams, setInferenceParams] = useState<InferenceParams>(
defaultInferenceParams,
);
const inferenceParams = useChatRuntimeStore((state) => state.params);
const setInferenceParams = useChatRuntimeStore((state) => state.setParams);
const modelsFromStore = useChatRuntimeStore((state) => state.models);
const lorasFromStore = useChatRuntimeStore((state) => state.loras);
const modelsError = useChatRuntimeStore((state) => state.modelsError);
const { refresh, selectModel, ejectModel } = useChatModelRuntime();
const handleCheckpointChange = useCallback(
(v: string) => setInferenceParams((p) => ({ ...p, checkpoint: v })),
(value: string, meta?: { isLora: boolean }) => {
void selectModel({ id: value, isLora: meta?.isLora });
},
[selectModel],
);
const handleEject = useCallback(() => {
void ejectModel();
}, [ejectModel]);
const handleNewThread = useCallback(
() => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }),
[],
);
const handleEject = useCallback(
() => setInferenceParams((p) => ({ ...p, checkpoint: "" })),
[],
);
const handleNewThread = useCallback(() => setView({ mode: "single" }), []);
const handleNewCompare = useCallback(
() => setView({ mode: "compare", pairId: crypto.randomUUID() }),
[],
);
const models =
inferenceParams.inferenceEngine === "llama-cpp" ? GGUF_MODELS : LORA_MODELS;
const models = useMemo<ModelOption[]>(
() =>
modelsFromStore.map((model) => ({
id: model.id,
name: model.name,
description: model.description,
})),
[modelsFromStore],
);
const loraModels = useMemo<LoraModelOption[]>(
() =>
lorasFromStore.map((lora) => ({
id: lora.id,
name: lora.name,
baseModel: lora.baseModel,
updatedAt: lora.updatedAt,
})),
[lorasFromStore],
);
useEffect(() => {
void refresh();
}, [refresh]);
return (
<div className="h-[calc(100vh-4rem)] bg-background overflow-hidden">
<SidebarProvider
defaultOpen={true}
className="!min-h-0 h-[calc(100vh-4rem)] max-w-7xl mx-auto px-4"
className="!min-h-0 h-full max-w-7xl mx-auto px-4"
style={
{
"--sidebar-width": "14rem",
@ -286,12 +288,18 @@ export function ChatPage(): ReactElement {
/>
<ModelSelector
models={models}
loraModels={loraModels}
value={inferenceParams.checkpoint}
onValueChange={handleCheckpointChange}
onEject={handleEject}
variant="ghost"
/>
</div>
{modelsError && (
<div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">
{modelsError}
</div>
)}
<div className="flex-1" />
<button
type="button"
@ -305,8 +313,9 @@ export function ChatPage(): ReactElement {
{view.mode === "single" ? (
<SingleContent
key={view.threadId ?? "new"}
key={view.threadId ?? view.newThreadNonce ?? "new"}
threadId={view.threadId}
newThreadNonce={view.newThreadNonce}
/>
) : (
<CompareContent key={view.pairId} pairId={view.pairId} />
@ -319,5 +328,6 @@ export function ChatPage(): ReactElement {
onParamsChange={setInferenceParams}
/>
</SidebarProvider>
</div>
);
}

71
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -10,7 +10,6 @@ import { Textarea } from "@/components/ui/textarea";
import {
ArrowDown01Icon,
Delete02Icon,
EngineIcon,
FloppyDiskIcon,
PencilEdit01Icon,
Settings02Icon,
@ -20,28 +19,13 @@ import { HugeiconsIcon } from "@hugeicons/react";
import { AnimatePresence, motion } from "motion/react";
import type { ReactNode } from "react";
import { useState } from "react";
import {
DEFAULT_INFERENCE_PARAMS,
type InferenceParams,
} from "./types/runtime";
export interface InferenceParams {
temperature: number;
topP: number;
topK: number;
repetitionPenalty: number;
maxTokens: number;
systemPrompt: string;
inferenceEngine: string;
checkpoint: string;
}
export const defaultInferenceParams: InferenceParams = {
temperature: 0.7,
topP: 0.9,
topK: 50,
repetitionPenalty: 1.1,
maxTokens: 512,
systemPrompt: "",
inferenceEngine: "unsloth",
checkpoint: "outputs/llama-3.1-8b-instruct-lora",
};
export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS;
export type { InferenceParams } from "./types/runtime";
export interface Preset {
name: string;
@ -72,11 +56,6 @@ const BUILTIN_PRESETS: Preset[] = [
},
];
const ENGINE_OPTIONS = [
{ value: "unsloth", label: "Unsloth" },
{ value: "llama-cpp", label: "llama.cpp (GGUF)" },
];
function ParamSlider({
label,
value,
@ -214,18 +193,19 @@ export function ChatSettingsPanel({
className={`shrink-0 h-full overflow-hidden bg-sidebar rounded-2xl corner-squircle transition-[width] duration-200 ease-linear ${open ? "w-[17rem] border-sidebar-border" : "w-0"}`}
>
<div className="flex h-full w-[17rem] flex-col">
<div className="flex items-center gap-2 px-3 py-2">
<div className="flex items-center gap-2 px-4 py-3">
<HugeiconsIcon
icon={PencilEdit01Icon}
className="size-3.5 text-muted-foreground"
className="size-4 text-muted-foreground/70"
/>
<span className="flex-1 text-xs font-semibold uppercase tracking-wider text-muted-foreground">
Advanced Configuration
<span className="flex-1 text-base font-semibold tracking-tight">
Configuration
</span>
</div>
<div className="flex-1 overflow-y-auto px-1.5">
<div className="px-2 pb-3">
{/* mt-4 matches the Playground sidebar gap (SidebarHeader py-3 + SidebarGroup pt-1) */}
<div className="mt-4 px-2 pb-3">
<div className="flex items-center gap-2">
<Select value={activePreset} onValueChange={applyPreset}>
<SelectTrigger className="h-8 flex-1 corner-squircle text-xs">
@ -285,33 +265,6 @@ export function ChatSettingsPanel({
/>
</div>
<CollapsibleSection
icon={EngineIcon}
label="Inference Engine"
defaultOpen={true}
>
<div>
<span className="mb-1 block text-[11px] text-muted-foreground">
Backend
</span>
<Select
value={params.inferenceEngine}
onValueChange={set("inferenceEngine")}
>
<SelectTrigger className="h-8 w-full text-xs corner-squircle">
<SelectValue />
</SelectTrigger>
<SelectContent>
{ENGINE_OPTIONS.map((o) => (
<SelectItem key={o.value} value={o.value}>
{o.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
</CollapsibleSection>
<CollapsibleSection
icon={SlidersHorizontalIcon}
label="Sampling"

175
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts Normal file
View file

@ -0,0 +1,175 @@
import { useCallback } from "react";
import { toast } from "sonner";
import {
getInferenceStatus,
listLoras,
listModels,
loadModel,
unloadModel,
} from "../api/chat-api";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
import type { ChatLoraSummary, ChatModelSummary } from "../types/runtime";
const DEFAULT_MODEL_MAX_SEQ_LENGTH = 2048;
type SelectedModelInput = {
id: string;
isLora?: boolean;
};
const LORA_SUFFIX_RE = /_(\d{9,})$/;
function parseTrailingEpoch(input: string): number | undefined {
const match = input.match(LORA_SUFFIX_RE);
if (!match) {
return undefined;
}
const parsed = Number.parseInt(match[1], 10);
return Number.isFinite(parsed) ? parsed : undefined;
}
function stripTrailingEpoch(input: string): string {
const cleaned = input.replace(LORA_SUFFIX_RE, "").replace(/[_-]+$/, "").trim();
return cleaned || input;
}
function describeModel(model: {
is_lora?: boolean;
is_vision?: boolean;
}): string | undefined {
const tags: string[] = [];
if (model.is_lora) tags.push("LoRA");
if (model.is_vision) tags.push("Vision");
if (!model.is_lora && !model.is_vision) tags.push("Base");
return tags.join(" · ");
}
function toChatModelSummary(model: {
id: string;
name?: string | null;
is_lora?: boolean;
is_vision?: boolean;
}): ChatModelSummary {
return {
id: model.id,
name: model.name || model.id,
description: describeModel(model),
isLora: Boolean(model.is_lora),
isVision: Boolean(model.is_vision),
};
}
function toLoraSummary(lora: {
display_name: string;
adapter_path: string;
base_model?: string | null;
}): ChatLoraSummary {
const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? "";
const updatedAt =
parseTrailingEpoch(lora.display_name) ?? parseTrailingEpoch(idTail);
return {
id: lora.adapter_path,
name: stripTrailingEpoch(lora.display_name),
baseModel: lora.base_model || "Unknown base model",
updatedAt,
};
}
export function useChatModelRuntime() {
const params = useChatRuntimeStore((state) => state.params);
const models = useChatRuntimeStore((state) => state.models);
const loras = useChatRuntimeStore((state) => state.loras);
const setModels = useChatRuntimeStore((state) => state.setModels);
const setLoras = useChatRuntimeStore((state) => state.setLoras);
const setModelsError = useChatRuntimeStore((state) => state.setModelsError);
const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint);
const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint);
const refresh = useCallback(async () => {
setModelsError(null);
try {
const [listRes, statusRes, lorasRes] = await Promise.all([
listModels(),
getInferenceStatus(),
listLoras(),
]);
setModels(listRes.models.map(toChatModelSummary));
setLoras(lorasRes.loras.map(toLoraSummary));
if (statusRes.active_model) {
setCheckpoint(statusRes.active_model);
}
} catch (error) {
const message =
error instanceof Error ? error.message : "Failed to load models";
setModelsError(message);
}
}, [setCheckpoint, setLoras, setModels, setModelsError]);
const selectModel = useCallback(
async (selection: string | SelectedModelInput) => {
const modelId = typeof selection === "string" ? selection : selection.id;
if (!modelId || params.checkpoint === modelId) {
return;
}
const explicitIsLora =
typeof selection === "string" ? undefined : selection.isLora;
const model = models.find((entry) => entry.id === modelId);
const lora = loras.find((entry) => entry.id === modelId);
const isLora =
explicitIsLora ?? model?.isLora ?? (lora ? true : false);
const displayName = model?.name || lora?.name || modelId;
const loadingToastId = toast.loading(`Loading ${displayName}...`);
setModelsError(null);
try {
if (params.checkpoint) {
await unloadModel({ model_path: params.checkpoint });
}
await loadModel({
model_path: modelId,
hf_token: null,
max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH,
load_in_4bit: true,
is_lora: isLora,
});
setCheckpoint(modelId);
await refresh();
toast.success(`${displayName} loaded`, { id: loadingToastId });
} catch (error) {
const message =
error instanceof Error ? error.message : "Failed to load model";
setModelsError(message);
toast.error(message, { id: loadingToastId });
}
},
[loras, models, params.checkpoint, refresh, setCheckpoint, setModelsError],
);
const ejectModel = useCallback(async () => {
if (!params.checkpoint) {
return;
}
setModelsError(null);
try {
await unloadModel({ model_path: params.checkpoint });
clearCheckpoint();
await refresh();
} catch (error) {
const message =
error instanceof Error ? error.message : "Failed to unload model";
setModelsError(message);
}
}, [clearCheckpoint, params.checkpoint, refresh, setModelsError]);
return {
refresh,
selectModel,
ejectModel,
};
}

2
studio/frontend/src/features/chat/index.ts
View file

@ -5,3 +5,5 @@ export {
type InferenceParams,
type Preset,
} from "./chat-settings-sheet";
export { useChatRuntimeStore } from "./stores/chat-runtime-store";
export { useChatModelRuntime } from "./hooks/use-chat-model-runtime";

98
studio/frontend/src/features/chat/runtime-provider.tsx
View file

@ -7,9 +7,8 @@ import {
type ExportedMessageRepositoryItem,
type PendingAttachment,
RuntimeAdapterProvider,
SimpleImageAttachmentAdapter,
SimpleTextAttachmentAdapter,
Suggestions,
SimpleTextAttachmentAdapter,
type ThreadHistoryAdapter,
type ThreadMessage,
type ThreadUserMessagePart,
@ -24,10 +23,66 @@ import { createAssistantStream } from "assistant-stream";
import mammoth from "mammoth";
import { type ReactElement, type ReactNode, useEffect, useMemo } from "react";
import { extractText, getDocumentProxy } from "unpdf";
import { createStreamAdapter } from "./adapter";
import { createOpenAIStreamAdapter } from "./api/chat-adapter";
import { db } from "./db";
import type { MessageRecord, ModelType } from "./types";
const DEFAULT_SUGGESTIONS = [
"Draw a simple flowchart of a login system using Mermaid",
"Solve the integral of x²·sin(x) step by step",
"Write a Python function that finds the longest palindrome in a string",
"Format a comparison of 3 databases as a markdown table with pros and cons",
];
class VisionImageAdapter implements AttachmentAdapter {
accept = "image/jpeg,image/png,image/webp,image/gif";
async add({ file }: { file: File }): Promise<PendingAttachment> {
const maxSize = 20 * 1024 * 1024;
if (file.size > maxSize) {
throw new Error("Image size exceeds 20MB limit");
}
return {
id: crypto.randomUUID(),
type: "image",
name: file.name,
contentType: file.type,
file,
status: { type: "requires-action", reason: "composer-send" },
};
}
async send(attachment: PendingAttachment): Promise<CompleteAttachment> {
return {
id: attachment.id,
type: "image",
name: attachment.name,
contentType: attachment.contentType,
content: [
{
type: "image",
image: await this.fileToBase64DataURL(attachment.file),
},
],
status: { type: "complete" },
};
}
async remove(): Promise<void> {
return Promise.resolve();
}
private async fileToBase64DataURL(file: File): Promise<string> {
return new Promise((resolve, reject) => {
const reader = new FileReader();
reader.onload = () => resolve(reader.result as string);
reader.onerror = () => reject(new Error("Failed to read image file"));
reader.readAsDataURL(file);
});
}
}
class PDFAttachmentAdapter implements AttachmentAdapter {
accept = "application/pdf";
@ -269,7 +324,7 @@ function ThreadHistoryProvider({
const attachments = useMemo(
() =>
new CompositeAttachmentAdapter([
new SimpleImageAttachmentAdapter(),
new VisionImageAdapter(),
new SimpleTextAttachmentAdapter(),
new PDFAttachmentAdapter(),
new DocxAttachmentAdapter(),
@ -288,9 +343,11 @@ function ThreadHistoryProvider({
);
}
const chatAdapter = createStreamAdapter();
const useRuntimeHook = (): ReturnType<typeof useLocalRuntime> =>
useLocalRuntime(chatAdapter);
const chatAdapter = createOpenAIStreamAdapter();
function useRuntimeHook(): ReturnType<typeof useLocalRuntime> {
return useLocalRuntime(chatAdapter);
}
function ThreadAutoSwitch({
threadId,
@ -308,16 +365,33 @@ function ThreadAutoSwitch({
return null;
}
function ThreadNewChatSwitch({
nonce,
}: { nonce: string }): ReactElement | null {
const aui = useAui();
const isLoading = useAuiState(({ threads }) => threads.isLoading);
useEffect(() => {
if (!isLoading) {
aui.threads().switchToNewThread();
}
}, [aui, isLoading, nonce]);
return null;
}
export function ChatRuntimeProvider({
children,
modelType = "base",
pairId,
initialThreadId,
newThreadNonce,
}: {
children: ReactNode;
modelType?: ModelType;
pairId?: string;
initialThreadId?: string;
newThreadNonce?: string;
}): ReactElement {
const runtime = useRemoteThreadListRuntime({
runtimeHook: useRuntimeHook,
@ -328,17 +402,15 @@ export function ChatRuntimeProvider({
});
const aui = useAui({
suggestions: Suggestions([
"Draw a simple flowchart of a login system using Mermaid",
"Solve the integral of x\u00B2\u00B7sin(x) step by step",
"Write a Python function that finds the longest palindrome in a string",
"Format a comparison of 3 databases as a markdown table with pros and cons",
]),
suggestions: Suggestions(DEFAULT_SUGGESTIONS),
});
return (
<AssistantRuntimeProvider runtime={runtime} aui={aui}>
{initialThreadId && <ThreadAutoSwitch threadId={initialThreadId} />}
{!initialThreadId && newThreadNonce && (
<ThreadNewChatSwitch nonce={newThreadNonce} />
)}
{children}
</AssistantRuntimeProvider>
);

58
studio/frontend/src/features/chat/stores/chat-runtime-store.ts Normal file
View file

@ -0,0 +1,58 @@
import { create } from "zustand";
import {
DEFAULT_INFERENCE_PARAMS,
type ChatLoraSummary,
type ChatModelSummary,
type InferenceParams,
} from "../types/runtime";
type ChatRuntimeStore = {
params: InferenceParams;
models: ChatModelSummary[];
loras: ChatLoraSummary[];
warmingByThreadId: Record<string, boolean>;
modelsError: string | null;
setParams: (params: InferenceParams) => void;
setModels: (models: ChatModelSummary[]) => void;
setLoras: (loras: ChatLoraSummary[]) => void;
setThreadWarming: (threadId: string, warming: boolean) => void;
setModelsError: (error: string | null) => void;
setCheckpoint: (modelId: string) => void;
clearCheckpoint: () => void;
};
export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
params: DEFAULT_INFERENCE_PARAMS,
models: [],
loras: [],
warmingByThreadId: {},
modelsError: null,
setParams: (params) => set({ params }),
setModels: (models) => set({ models }),
setLoras: (loras) => set({ loras }),
setThreadWarming: (threadId, warming) =>
set((state) => {
const next = { ...state.warmingByThreadId };
if (warming) {
next[threadId] = true;
} else {
delete next[threadId];
}
return { warmingByThreadId: next };
}),
setModelsError: (modelsError) => set({ modelsError }),
setCheckpoint: (modelId) =>
set((state) => ({
params: {
...state.params,
checkpoint: modelId,
},
})),
clearCheckpoint: () =>
set((state) => ({
params: {
...state.params,
checkpoint: "",
},
})),
}));

12
studio/frontend/src/features/chat/thread-sidebar.tsx
View file

@ -99,11 +99,11 @@ export function ThreadSidebar({
return (
<>
<SidebarHeader>
<span className="text-sm font-semibold">Playground</span>
<SidebarHeader className="px-4 py-3">
<span className="text-base font-semibold tracking-tight">Playground</span>
</SidebarHeader>
<SidebarContent>
<SidebarGroup>
<SidebarGroup className="px-4 pt-1">
<SidebarGroupContent>
<SidebarMenu>
<SidebarMenuItem>
@ -121,8 +121,8 @@ export function ThreadSidebar({
</SidebarMenu>
</SidebarGroupContent>
</SidebarGroup>
<SidebarGroup className="flex-1">
<SidebarGroupLabel>Your Chats</SidebarGroupLabel>
<SidebarGroup className="flex-1 px-4">
<SidebarGroupLabel className="text-xs font-medium text-muted-foreground/80">Your Chats</SidebarGroupLabel>
<SidebarGroupContent>
<SidebarMenu>
{items.map((item) => (
@ -144,7 +144,7 @@ export function ThreadSidebar({
))}
</SidebarMenu>
{items.length === 0 && (
<p className="px-2 py-4 text-center text-xs text-muted-foreground">
<p className="px-2 py-6 text-center text-xs text-muted-foreground">
No threads yet
</p>
)}

2
studio/frontend/src/features/chat/types.ts
View file

@ -1,7 +1,7 @@
export type ModelType = "base" | "lora";
export type ChatView =
| { mode: "single"; threadId?: string }
| { mode: "single"; threadId?: string; newThreadNonce?: string }
| { mode: "compare"; pairId: string };
export interface ThreadRecord {

80
studio/frontend/src/features/chat/types/api.ts Normal file
View file

@ -0,0 +1,80 @@
export interface BackendModelDetails {
id: string;
name?: string | null;
is_vision?: boolean;
is_lora?: boolean;
}
export interface ListModelsResponse {
models: BackendModelDetails[];
default_models: string[];
}
export interface BackendLoraInfo {
display_name: string;
adapter_path: string;
base_model?: string | null;
}
export interface ListLorasResponse {
loras: BackendLoraInfo[];
outputs_dir: string;
}
export interface LoadModelRequest {
model_path: string;
hf_token: string | null;
max_seq_length: number;
load_in_4bit: boolean;
is_lora: boolean;
}
export interface LoadModelResponse {
status: string;
model: string;
display_name: string;
is_vision: boolean;
is_lora: boolean;
}
export interface UnloadModelRequest {
model_path: string;
}
export interface InferenceStatusResponse {
active_model: string | null;
is_vision: boolean;
loading: string[];
loaded: string[];
}
export interface OpenAIChatMessage {
role: "system" | "user" | "assistant";
content: string;
}
export interface OpenAIChatCompletionsRequest {
model: string;
messages: OpenAIChatMessage[];
stream: boolean;
temperature: number;
top_p: number;
max_tokens: number;
top_k: number;
repetition_penalty: number;
image_base64?: string;
}
export interface OpenAIChatDelta {
role?: string;
content?: string;
}
export interface OpenAIChatChunkChoice {
delta?: OpenAIChatDelta;
finish_reason?: string | null;
}
export interface OpenAIChatChunk {
choices?: OpenAIChatChunkChoice[];
}

34
studio/frontend/src/features/chat/types/runtime.ts Normal file
View file

@ -0,0 +1,34 @@
export interface InferenceParams {
temperature: number;
topP: number;
topK: number;
repetitionPenalty: number;
maxTokens: number;
systemPrompt: string;
checkpoint: string;
}
export const DEFAULT_INFERENCE_PARAMS: InferenceParams = {
temperature: 0.7,
topP: 0.9,
topK: 50,
repetitionPenalty: 1.1,
maxTokens: 512,
systemPrompt: "",
checkpoint: "",
};
export interface ChatModelSummary {
id: string;
name: string;
description?: string;
isVision: boolean;
isLora: boolean;
}
export interface ChatLoraSummary {
id: string;
name: string;
baseModel: string;
updatedAt?: number;
}

54
studio/frontend/src/features/chat/utils/parse-assistant-content.ts Normal file
View file

@ -0,0 +1,54 @@
import type { ChatModelRunResult } from "@assistant-ui/react";
type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
const THINK_OPEN_TAG = "<think>";
const THINK_CLOSE_TAG = "</think>";
function appendTextPart(parts: ContentPart[], text: string): void {
if (text) {
parts.push({ type: "text", text });
}
}
function appendReasoningPart(parts: ContentPart[], text: string): void {
if (text) {
parts.push({ type: "reasoning", text });
}
}
export function parseAssistantContent(
raw: string,
): ContentPart[] {
const parts: ContentPart[] = [];
if (!raw) {
return parts;
}
let cursor = 0;
while (cursor < raw.length) {
const openIndex = raw.indexOf(THINK_OPEN_TAG, cursor);
if (openIndex === -1) {
appendTextPart(parts, raw.slice(cursor));
break;
}
appendTextPart(parts, raw.slice(cursor, openIndex));
const reasoningStart = openIndex + THINK_OPEN_TAG.length;
const closeIndex = raw.indexOf(THINK_CLOSE_TAG, reasoningStart);
if (closeIndex === -1) {
appendReasoningPart(parts, raw.slice(reasoningStart));
break;
}
appendReasoningPart(parts, raw.slice(reasoningStart, closeIndex));
cursor = closeIndex + THINK_CLOSE_TAG.length;
}
return parts;
}
export function hasClosedThinkTag(raw: string): boolean {
return raw.includes(THINK_CLOSE_TAG);
}

4
studio/frontend/src/features/export/components/quant-picker.tsx
View file

@ -71,14 +71,14 @@ export function QuantPicker({ value, onChange }: QuantPickerProps) {
className={cn(
"inline-flex items-center gap-1.5 rounded-full px-3 py-1.5 text-xs font-medium ring-1 transition-all",
active
? "ring-primary bg-primary/10 text-primary"
? "ring-primary bg-primary/10 text-foreground"
: "ring-border text-muted-foreground hover:text-foreground hover:ring-foreground/20",
)}
>
{active && (
<HugeiconsIcon
icon={CheckmarkCircle01Icon}
className="size-3"
className="size-3 text-primary"
/>
)}
{q.label}

22
studio/frontend/src/features/export/export-page.tsx
View file

@ -13,7 +13,8 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
import { useWizardStore } from "@/stores/training";
import { useTrainingRuntimeStore } from "@/features/training";
import { useTrainingConfigStore } from "@/features/training";
import { isAdapterMethod } from "@/types/training";
import { InformationCircleIcon, PackageIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@ -36,36 +37,31 @@ export function ExportPage() {
trainingMethod,
selectedModel,
saveSteps,
trainingMetrics,
epochs,
loraRank,
hfToken,
setHfToken,
} = useWizardStore(
} = useTrainingConfigStore(
useShallow((s) => ({
trainingMethod: s.trainingMethod,
selectedModel: s.selectedModel,
saveSteps: s.saveSteps,
trainingMetrics: s.trainingMetrics,
epochs: s.epochs,
loraRank: s.loraRank,
hfToken: s.hfToken,
setHfToken: s.setHfToken,
})),
);
const totalSteps = useTrainingRuntimeStore((state) => state.totalSteps);
const isAdapter = isAdapterMethod(trainingMethod);
const checkpoints = useMemo(() => {
if (isAdapter) {
const interval = saveSteps > 0 ? saveSteps : 100;
const total = trainingMetrics?.totalSteps ?? 500;
const total = totalSteps > 0 ? totalSteps : 500;
const entries: { value: string; label: string; detail: string }[] = [];
for (let step = interval; step <= total; step += interval) {
const loss = (
1.5 -
(step / total) * 0.7 +
Math.random() * 0.05
).toFixed(2);
const loss = (1.5 - (step / total) * 0.7).toFixed(2);
entries.push({
value: `checkpoint-${step}`,
label: `checkpoint-${step}`,
@ -81,7 +77,7 @@ export function ExportPage() {
detail: "Full fine-tuned weights",
},
];
}, [isAdapter, saveSteps, trainingMetrics?.totalSteps]);
}, [isAdapter, saveSteps, totalSteps]);
const [checkpoint, setCheckpoint] = useState<string | null>(null);
const [exportMethod, setExportMethod] = useState<ExportMethod | null>(null);
@ -109,8 +105,8 @@ export function ExportPage() {
return (
<div className="min-h-screen bg-background">
<main className="mx-auto max-w-7xl px-6 py-8">
<div className="mb-8 flex flex-col gap-1">
<main className="mx-auto max-w-7xl px-6 py-4">
<div className="mb-8 flex flex-col gap-0.5">
<h1 className="text-2xl font-semibold tracking-tight">
Export Model
</h1>

4
studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx
View file

@ -38,7 +38,7 @@ import {
useInfiniteScroll,
} from "@/hooks";
import { cn, formatCompact } from "@/lib/utils";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { DatasetFormat } from "@/types/training";
import {
InformationCircleIcon,
@ -69,7 +69,7 @@ export function DatasetStep() {
setDataset,
uploadedFile,
setUploadedFile,
} = useWizardStore(
} = useTrainingConfigStore(
useShallow((s) => ({
hfToken: s.hfToken,
setHfToken: s.setHfToken,

4
studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
View file

@ -20,7 +20,7 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip";
import { CONTEXT_LENGTHS } from "@/config/training";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import { InformationCircleIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useShallow } from "zustand/react/shallow";
@ -40,7 +40,7 @@ export function HyperparametersStep() {
setLoraAlpha,
loraDropout,
setLoraDropout,
} = useWizardStore(
} = useTrainingConfigStore(
useShallow((s) => ({
trainingMethod: s.trainingMethod,
epochs: s.epochs,

4
studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
View file

@ -37,7 +37,7 @@ import {
useInfiniteScroll,
} from "@/hooks";
import { formatCompact } from "@/lib/utils";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { TrainingMethod } from "@/types/training";
import {
InformationCircleIcon,
@ -57,7 +57,7 @@ export function ModelSelectionStep() {
setTrainingMethod,
hfToken,
setHfToken,
} = useWizardStore(
} = useTrainingConfigStore(
useShallow((s) => ({
modelType: s.modelType,
selectedModel: s.selectedModel,

4
studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx
View file

@ -8,7 +8,7 @@ import {
} from "@/components/ui/tooltip";
import { MODEL_TYPES } from "@/config/training";
import { cn } from "@/lib/utils";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { ModelType } from "@/types/training";
import {
Database02Icon,
@ -38,7 +38,7 @@ const TYPE_TOOLTIPS: Record<ModelType, string> = {
const COMING_SOON: ModelType[] = ["tts", "embeddings"];
export function ModelTypeStep(): ReactElement {
const { modelType, setModelType } = useWizardStore(
const { modelType, setModelType } = useTrainingConfigStore(
useShallow((s) => ({
modelType: s.modelType,
setModelType: s.setModelType,

4
studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
View file

@ -1,7 +1,7 @@
import { Badge } from "@/components/ui/badge";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Separator } from "@/components/ui/separator";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import { isAdapterMethod } from "@/types/training";
import { GpuIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@ -30,7 +30,7 @@ export function SummaryStep() {
loraRank,
loraAlpha,
loraDropout,
} = useWizardStore(
} = useTrainingConfigStore(
useShallow(
({
modelType,

4
studio/frontend/src/features/onboarding/components/wizard-content.tsx
View file

@ -1,5 +1,5 @@
import { STEPS } from "@/config/training";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { StepNumber } from "@/types/training";
import { DatasetStep } from "./steps/dataset-step";
import { HyperparametersStep } from "./steps/hyperparameters-step";
@ -24,7 +24,7 @@ const STEP_MASCOTS: Record<StepNumber, string> = {
};
export function WizardContent() {
const currentStep = useWizardStore((s) => s.currentStep);
const currentStep = useTrainingConfigStore((s) => s.currentStep);
const stepConfig = STEPS[currentStep - 1];
const StepComponent = STEP_COMPONENTS[currentStep];
const mascotSrc = STEP_MASCOTS[currentStep];

4
studio/frontend/src/features/onboarding/components/wizard-footer.tsx
View file

@ -1,14 +1,14 @@
import { Button } from "@/components/ui/button";
import { STEPS } from "@/config/training";
import { markOnboardingDone } from "@/features/auth";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import { ArrowLeft02Icon, ArrowRight02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useNavigate } from "@tanstack/react-router";
import { useShallow } from "zustand/react/shallow";
export function WizardFooter() {
const { currentStep, prevStep, nextStep, canProceed } = useWizardStore(
const { currentStep, prevStep, nextStep, canProceed } = useTrainingConfigStore(
useShallow((s) => ({
currentStep: s.currentStep,
prevStep: s.prevStep,

4
studio/frontend/src/features/onboarding/components/wizard-layout.tsx
View file

@ -6,7 +6,7 @@ import { Suspense, lazy, useEffect, useRef, useState } from "react";
import type { ConfettiRef } from "@/components/ui/confetti";
import { STEPS } from "@/config/training";
import { isOnboardingDone, markOnboardingDone } from "@/features/auth";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import { SplashScreen } from "./splash-screen";
import { WizardContent } from "./wizard-content";
import { WizardFooter } from "./wizard-footer";
@ -19,7 +19,7 @@ const Confetti = lazy(() =>
export function WizardLayout() {
const navigate = useNavigate();
const [showSplash, setShowSplash] = useState(true);
const currentStep = useWizardStore((s) => s.currentStep);
const currentStep = useTrainingConfigStore((s) => s.currentStep);
const confettiRef = useRef<ConfettiRef>(null);
const hasFiredRef = useRef(false);
const isFinalStep = currentStep === STEPS.length;

4
studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx
View file

@ -1,10 +1,10 @@
import { Progress } from "@/components/ui/progress";
import { STEPS } from "@/config/training";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import { WizardStepItem } from "./wizard-step-item";
export function WizardSidebar() {
const currentStep = useWizardStore((s) => s.currentStep);
const currentStep = useTrainingConfigStore((s) => s.currentStep);
const progress = ((currentStep - 1) / (STEPS.length - 1)) * 100;
return (

4
studio/frontend/src/features/onboarding/components/wizard-step-item.tsx
View file

@ -1,5 +1,5 @@
import { cn } from "@/lib/utils";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { StepConfig, StepNumber } from "@/types/training";
import { useShallow } from "zustand/react/shallow";
@ -8,7 +8,7 @@ interface WizardStepItemProps {
}
export function WizardStepItem({ step }: WizardStepItemProps) {
const { currentStep, setStep } = useWizardStore(
const { currentStep, setStep } = useTrainingConfigStore(
useShallow((s) => ({ currentStep: s.currentStep, setStep: s.setStep })),
);
const isActive = currentStep === step.number;

286
studio/frontend/src/features/studio/sections/charts-content.tsx
View file

@ -23,7 +23,6 @@ import {
} from "@/components/ui/dropdown-menu";
import { Label } from "@/components/ui/label";
import { Slider } from "@/components/ui/slider";
import type { TrainingMetrics } from "@/types/training";
import { ChartAverageIcon, Settings02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type ReactElement, useMemo, useState } from "react";
@ -44,9 +43,11 @@ const lossConfig = {
const lrConfig = {
lr: { label: "LR", color: "#8b5cf6" },
} satisfies ChartConfig;
const gradNormConfig = {
gradNorm: { label: "Grad Norm", color: "#f97316" },
} satisfies ChartConfig;
const evalLossConfig = {
loss: { label: "Eval Loss", color: "#ef4444" },
} satisfies ChartConfig;
@ -62,10 +63,81 @@ const placeholderEvalData = [
type LossHistoryItem = { step: number; loss: number };
type SmoothedLossItem = LossHistoryItem & { smoothed: number };
interface TrainingChartSeries {
lossHistory: LossHistoryItem[];
lrHistory: { step: number; lr: number }[];
gradNormHistory: { step: number; gradNorm: number }[];
}
const CHART_SYNC_ID = "train-metrics-sync";
const MAX_RENDER_POINTS = 800;
const DEFAULT_VISIBLE_POINTS = 160;
function formatStepTick(value: number): string {
if (value >= 1_000_000) {
return `${(value / 1_000_000).toFixed(1)}M`;
}
if (value >= 1_000) {
return `${(value / 1_000).toFixed(1)}k`;
}
return String(Math.round(value));
}
function compressSeries<T>(data: T[], maxPoints: number): T[] {
if (data.length <= maxPoints) {
return data;
}
const stride = Math.ceil(data.length / maxPoints);
return data.filter(
(_item, index) => index % stride === 0 || index === data.length - 1,
);
}
function buildStepTicks(min: number, max: number, targetCount = 6): number[] {
if (!Number.isFinite(min) || !Number.isFinite(max)) {
return [0, 1];
}
if (max <= min) {
return [min, max];
}
const stepSize = Math.max(1, Math.ceil((max - min) / (targetCount - 1)));
const ticks: number[] = [];
let current = min;
while (current < max) {
ticks.push(current);
current += stepSize;
}
ticks.push(max);
return Array.from(new Set(ticks));
}
function buildYDomain(values: number[]): [number, number] {
if (values.length === 0) {
return [0, 1];
}
const min = Math.min(...values);
const max = Math.max(...values);
if (min === max) {
const base = Math.abs(min);
const pad = base > 0 ? base * 0.08 : 0.1;
return [min - pad, max + pad];
}
const pad = (max - min) * 0.12;
return [min - pad, max + pad];
}
function ema(data: LossHistoryItem[], alpha: number): SmoothedLossItem[] {
if (data.length === 0) {
return [];
}
let s = data[0].loss;
return data.map((d) => {
s = alpha * d.loss + (1 - alpha) * s;
@ -75,18 +147,113 @@ function ema(data: LossHistoryItem[], alpha: number): SmoothedLossItem[] {
export function ChartsContent({
metrics,
}: { metrics: TrainingMetrics }): ReactElement {
const [smoothing, setSmoothing] = useState(0.6);
}: { metrics: TrainingChartSeries }): ReactElement {
const [smoothing, setSmoothing] = useState(0.75);
const [showRaw, setShowRaw] = useState(true);
const [showSmoothed, setShowSmoothed] = useState(true);
const [showAvgLine, setShowAvgLine] = useState(true);
const lossHistory = metrics.lossHistory;
const smoothedData = useMemo(
() => (lossHistory ? ema(lossHistory, 1 - smoothing) : []),
() => (lossHistory.length > 0 ? ema(lossHistory, 1 - smoothing) : []),
[lossHistory, smoothing],
);
const reducedLossData = useMemo(
() => compressSeries(smoothedData, MAX_RENDER_POINTS),
[smoothedData],
);
const reducedGradNormData = useMemo(
() => compressSeries(metrics.gradNormHistory, MAX_RENDER_POINTS),
[metrics.gradNormHistory],
);
const reducedLrData = useMemo(
() => compressSeries(metrics.lrHistory, MAX_RENDER_POINTS),
[metrics.lrHistory],
);
const visibleStepDomain = useMemo<[number, number]>(() => {
const allSteps = [
...reducedLossData.map((point) => point.step),
...reducedGradNormData.map((point) => point.step),
...reducedLrData.map((point) => point.step),
].sort((a, b) => a - b);
if (allSteps.length === 0) {
return [0, 1];
}
const minStep = allSteps[0] ?? 0;
const endStep = allSteps[allSteps.length - 1] ?? 1;
const startIndex = Math.max(0, allSteps.length - DEFAULT_VISIBLE_POINTS);
const startStep = allSteps[startIndex] ?? minStep;
if (startStep === endStep) {
return [startStep, startStep + 4];
}
if (endStep - startStep < 6) {
return [Math.max(minStep, endStep - 6), endStep];
}
return [startStep, endStep];
}, [reducedGradNormData, reducedLossData, reducedLrData]);
const xAxisTicks = useMemo(
() => buildStepTicks(visibleStepDomain[0], visibleStepDomain[1]),
[visibleStepDomain],
);
const visibleLossValues = useMemo(
() =>
reducedLossData
.filter(
(point) =>
point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1],
)
.map((point) => point.loss),
[reducedLossData, visibleStepDomain],
);
const visibleSmoothValues = useMemo(
() =>
reducedLossData
.filter(
(point) =>
point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1],
)
.map((point) => point.smoothed),
[reducedLossData, visibleStepDomain],
);
const visibleGradValues = useMemo(
() =>
reducedGradNormData
.filter(
(point) =>
point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1],
)
.map((point) => point.gradNorm),
[reducedGradNormData, visibleStepDomain],
);
const visibleLrValues = useMemo(
() =>
reducedLrData
.filter(
(point) =>
point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1],
)
.map((point) => point.lr),
[reducedLrData, visibleStepDomain],
);
const lossDomain = useMemo(
() => buildYDomain([...visibleLossValues, ...visibleSmoothValues]),
[visibleLossValues, visibleSmoothValues],
);
const gradDomain = useMemo(() => buildYDomain(visibleGradValues), [visibleGradValues]);
const lrDomain = useMemo(() => buildYDomain(visibleLrValues), [visibleLrValues]);
const avg =
metrics.lossHistory.length > 0
? +(
@ -96,17 +263,16 @@ export function ChartsContent({
: 0;
return (
<div className="grid grid-cols-1 lg:grid-cols-2 gap-6">
{/* Training Loss */}
<div className="grid grid-cols-1 gap-6 lg:grid-cols-2">
<Card size="sm">
<CardHeader>
<CardTitle className="text-sm">Training Loss</CardTitle>
<CardTitle className="text-sm pl-2">Training Loss</CardTitle>
<CardAction>
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<button
type="button"
className="rounded-md p-1 text-muted-foreground hover:bg-muted hover:text-foreground transition-colors cursor-pointer"
className="cursor-pointer rounded-md p-1 text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
<HugeiconsIcon icon={Settings02Icon} className="size-3.5" />
</button>
@ -155,12 +321,10 @@ export function ChartsContent({
</CardAction>
</CardHeader>
<CardContent>
<ChartContainer
config={lossConfig}
className="h-[200px] w-full -ml-3"
>
<ChartContainer config={lossConfig} className="-ml-3 h-[220px] w-full">
<LineChart
data={smoothedData}
data={reducedLossData}
syncId={CHART_SYNC_ID}
accessibilityLayer={true}
margin={{ left: 0, right: 8 }}
>
@ -168,19 +332,27 @@ export function ChartsContent({
<XAxis
dataKey="step"
type="number"
domain={["dataMin", "dataMax"]}
domain={visibleStepDomain}
ticks={xAxisTicks}
allowDataOverflow={true}
allowDecimals={false}
minTickGap={28}
tickLine={false}
axisLine={false}
tickMargin={8}
fontSize={10}
tickFormatter={(value) => formatStepTick(Number(value))}
interval="preserveStartEnd"
/>
<YAxis
domain={lossDomain}
allowDataOverflow={true}
tickLine={false}
axisLine={false}
tickMargin={4}
fontSize={10}
width={40}
width={52}
tickFormatter={(value) => Number(value).toFixed(2)}
/>
<ChartTooltip
content={
@ -207,22 +379,30 @@ export function ChartsContent({
)}
{showRaw && (
<Line
type="monotone"
type="monotoneX"
dataKey="loss"
stroke="var(--color-loss)"
strokeWidth={1.5}
strokeOpacity={showSmoothed ? 0.3 : 1}
strokeWidth={1.2}
strokeOpacity={showSmoothed ? 0.35 : 1}
dot={false}
activeDot={{ r: 3, strokeWidth: 0 }}
connectNulls={true}
strokeLinecap="round"
strokeLinejoin="round"
isAnimationActive={false}
/>
)}
{showSmoothed && (
<Line
type="monotone"
type="monotoneX"
dataKey="smoothed"
stroke="var(--color-smoothed)"
strokeWidth={2}
strokeWidth={2.2}
dot={false}
activeDot={{ r: 3, strokeWidth: 0 }}
connectNulls={true}
strokeLinecap="round"
strokeLinejoin="round"
isAnimationActive={false}
/>
)}
@ -232,18 +412,18 @@ export function ChartsContent({
</CardContent>
</Card>
{/* Grad Norm */}
<Card size="sm">
<CardHeader>
<CardTitle className="text-sm">Gradient Norm</CardTitle>
<CardTitle className="text-sm pl-2">Gradient Norm</CardTitle>
</CardHeader>
<CardContent>
<ChartContainer
config={gradNormConfig}
className="h-[200px] w-full -ml-3"
className="-ml-3 h-[220px] w-full"
>
<LineChart
data={metrics.gradNormHistory}
data={reducedGradNormData}
syncId={CHART_SYNC_ID}
accessibilityLayer={true}
margin={{ left: 0, right: 8 }}
>
@ -251,19 +431,27 @@ export function ChartsContent({
<XAxis
dataKey="step"
type="number"
domain={["dataMin", "dataMax"]}
domain={visibleStepDomain}
ticks={xAxisTicks}
allowDataOverflow={true}
allowDecimals={false}
minTickGap={28}
tickLine={false}
axisLine={false}
tickMargin={8}
fontSize={10}
tickFormatter={(value) => formatStepTick(Number(value))}
interval="preserveStartEnd"
/>
<YAxis
domain={gradDomain}
allowDataOverflow={true}
tickLine={false}
axisLine={false}
tickMargin={4}
fontSize={10}
width={40}
width={52}
tickFormatter={(value) => Number(value).toFixed(2)}
/>
<ChartTooltip
content={
@ -275,11 +463,15 @@ export function ChartsContent({
}
/>
<Line
type="monotone"
type="monotoneX"
dataKey="gradNorm"
stroke="var(--color-gradNorm)"
strokeWidth={2}
dot={false}
activeDot={{ r: 3, strokeWidth: 0 }}
connectNulls={true}
strokeLinecap="round"
strokeLinejoin="round"
isAnimationActive={false}
/>
<ChartLegend content={<ChartLegendContent />} />
@ -288,18 +480,15 @@ export function ChartsContent({
</CardContent>
</Card>
{/* Learning Rate */}
<Card size="sm">
<CardHeader>
<CardTitle className="text-sm">Learning Rate</CardTitle>
<CardTitle className="text-sm pl-2">Learning Rate</CardTitle>
</CardHeader>
<CardContent>
<ChartContainer
config={lrConfig}
className="h-[200px] w-full -ml-1.5"
>
<ChartContainer config={lrConfig} className="-ml-1.5 h-[220px] w-full">
<LineChart
data={metrics.lrHistory}
data={reducedLrData}
syncId={CHART_SYNC_ID}
accessibilityLayer={true}
margin={{ left: 0, right: 8 }}
>
@ -307,20 +496,27 @@ export function ChartsContent({
<XAxis
dataKey="step"
type="number"
domain={["dataMin", "dataMax"]}
domain={visibleStepDomain}
ticks={xAxisTicks}
allowDataOverflow={true}
allowDecimals={false}
minTickGap={28}
tickLine={false}
axisLine={false}
tickMargin={8}
fontSize={10}
tickFormatter={(value) => formatStepTick(Number(value))}
interval="preserveStartEnd"
/>
<YAxis
domain={lrDomain}
allowDataOverflow={true}
tickLine={false}
axisLine={false}
tickMargin={4}
fontSize={10}
width={40}
tickFormatter={(v) => v.toExponential(0)}
width={52}
tickFormatter={(value) => Number(value).toExponential(0)}
/>
<ChartTooltip
content={
@ -328,19 +524,20 @@ export function ChartsContent({
labelFormatter={(_value, payload) =>
`Step ${payload?.[0]?.payload?.step ?? ""}`
}
formatter={(value) => [
Number(value).toExponential(3),
"LR",
]}
formatter={(value) => [Number(value).toExponential(3), "LR"]}
/>
}
/>
<Line
type="monotone"
type="monotoneX"
dataKey="lr"
stroke="var(--color-lr)"
strokeWidth={2}
dot={false}
activeDot={{ r: 3, strokeWidth: 0 }}
connectNulls={true}
strokeLinecap="round"
strokeLinejoin="round"
isAnimationActive={false}
/>
<ChartLegend content={<ChartLegendContent />} />
@ -349,10 +546,9 @@ export function ChartsContent({
</CardContent>
</Card>
{/* Eval Loss (disabled/blurred) */}
<Card size="sm">
<CardHeader>
<CardTitle className="text-sm text-muted-foreground">
<CardTitle className="text-sm text-muted-foreground pl-2">
Eval Loss
</CardTitle>
</CardHeader>
@ -360,7 +556,7 @@ export function ChartsContent({
<div className="relative">
<ChartContainer
config={evalLossConfig}
className="h-[200px] w-full -ml-3 blur"
className="-ml-3 h-[220px] w-full blur"
>
<LineChart
data={placeholderEvalData}

40
studio/frontend/src/features/studio/sections/charts-section.tsx
View file

@ -1,5 +1,5 @@
import { useWizardStore } from "@/stores/training";
import { type ReactElement, Suspense, lazy } from "react";
import { useTrainingRuntimeStore } from "@/features/training";
import { type ReactElement, Suspense, lazy, useMemo } from "react";
const ChartsContent = lazy(() =>
import("./charts-content").then((module) => ({
@ -14,9 +14,39 @@ const SKELETON_KEYS = [
];
export function ChartsSection(): ReactElement | null {
const metrics = useWizardStore((s) => s.trainingMetrics);
const currentStep = useTrainingRuntimeStore((state) => state.currentStep);
const totalSteps = useTrainingRuntimeStore((state) => state.totalSteps);
const lossHistoryRaw = useTrainingRuntimeStore((state) => state.lossHistory);
const lrHistoryRaw = useTrainingRuntimeStore((state) => state.lrHistory);
const gradNormHistoryRaw = useTrainingRuntimeStore(
(state) => state.gradNormHistory,
);
if (!metrics) {
const series = useMemo(
() => ({
currentStep,
totalSteps,
lossHistory: lossHistoryRaw.map((point) => ({
step: point.step,
loss: point.value,
})),
lrHistory: lrHistoryRaw.map((point) => ({
step: point.step,
lr: point.value,
})),
gradNormHistory: gradNormHistoryRaw.map((point) => ({
step: point.step,
gradNorm: point.value,
})),
}),
[currentStep, gradNormHistoryRaw, lossHistoryRaw, lrHistoryRaw, totalSteps],
);
if (
series.lossHistory.length === 0 &&
series.lrHistory.length === 0 &&
series.gradNormHistory.length === 0
) {
return null;
}
@ -33,7 +63,7 @@ export function ChartsSection(): ReactElement | null {
</div>
}
>
<ChartsContent metrics={metrics} />
<ChartsContent metrics={series} />
</Suspense>
);
}

363
studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx Normal file
View file

@ -0,0 +1,363 @@
import type { ColumnDef } from "@tanstack/react-table";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { DataTable } from "@/components/ui/data-table";
import { Badge } from "@/components/ui/badge";
import { Spinner } from "@/components/ui/spinner";
import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type ReactNode, useEffect, useMemo, useState } from "react";
// ---------------------------------------------------------------------------
// Types (matches CheckFormatResponse from backend)
// ---------------------------------------------------------------------------
type CheckFormatResponse = {
requires_manual_mapping: boolean;
detected_format: string;
columns: string[];
suggested_mapping?: Record<string, string> | null;
detected_image_column?: string | null;
detected_text_column?: string | null;
preview_samples?: Record<string, unknown>[] | null;
total_rows?: number | null;
};
type PreviewImagePayload = {
type: "image";
mime?: string;
width?: number;
height?: number;
data?: string;
};
type DatasetPreviewDialogProps = {
open: boolean;
onOpenChange: (open: boolean) => void;
datasetName: string | null;
hfToken: string | null;
};
// ---------------------------------------------------------------------------
// API -- uses existing /check-format endpoint
// ---------------------------------------------------------------------------
async function fetchCheckFormat(
datasetName: string,
hfToken: string | null,
): Promise<CheckFormatResponse> {
const res = await fetch("/api/datasets/check-format", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
dataset_name: datasetName,
hf_token: hfToken || undefined,
split: "train",
}),
});
if (!res.ok) {
const body = await res.json().catch(() => null);
throw new Error(body?.detail || `Request failed (${res.status})`);
}
return res.json();
}
// ---------------------------------------------------------------------------
// Component
// ---------------------------------------------------------------------------
export function DatasetPreviewDialog({
open,
onOpenChange,
datasetName,
hfToken,
}: DatasetPreviewDialogProps) {
const [data, setData] = useState<CheckFormatResponse | null>(null);
const [loading, setLoading] = useState(false);
const [error, setError] = useState<string | null>(null);
useEffect(() => {
if (!open || !datasetName) {
setData(null);
setError(null);
return;
}
let cancelled = false;
setLoading(true);
setError(null);
fetchCheckFormat(datasetName, hfToken)
.then((res) => {
if (!cancelled) {
setData(res);
setError(null);
}
})
.catch((err) => {
if (!cancelled) setError(err.message || "Failed to load preview");
})
.finally(() => {
if (!cancelled) setLoading(false);
});
return () => {
cancelled = true;
};
}, [open, datasetName, hfToken]);
const rows = data?.preview_samples ?? [];
const columns = data?.columns ?? [];
// Determine source label
const sourceLabel = useMemo(() => {
if (!datasetName) return "";
if (datasetName.includes("/")) return `Hugging Face (${datasetName})`;
return `Local Files (${datasetName})`;
}, [datasetName]);
// Build TanStack Table columns from the column names
const tableColumns = useMemo<ColumnDef<Record<string, unknown>>[]>(() => {
if (!columns.length) return [];
return columns.map((colName) => ({
accessorKey: colName,
header: () => (
<span className="font-heading text-[13px] font-semibold tracking-tight text-foreground">
{colName}
</span>
),
cell: ({ getValue }: { getValue: () => unknown }) => {
const value = getValue();
const images = collectPreviewImages(value);
if (images.length > 0) {
return (
<div className="flex flex-wrap gap-2">
{images.slice(0, 4).map((image, index) => {
const mime = image.mime || "image/jpeg";
const src = image.data ? `data:${mime};base64,${image.data}` : "";
const width = image.width ?? 128;
const height = image.height ?? 128;
return (
<img
key={`${colName}-img-${index}`}
src={src}
alt={`preview-${index}`}
className="h-16 w-auto max-w-40 rounded-md border object-contain bg-muted"
width={width}
height={height}
loading="lazy"
/>
);
})}
{images.length > 4 && (
<span className="text-xs text-muted-foreground self-end">
+{images.length - 4} more
</span>
)}
</div>
);
}
const text = formatCell(value);
if (!text) {
return (
<span className="text-muted-foreground/40 italic text-[13px]">
--
</span>
);
}
const full =
typeof value === "string" ? value : JSON.stringify(value);
return (
<p
className="text-[13px] leading-relaxed line-clamp-6"
title={full}
>
{text}
</p>
);
},
}));
}, [columns]);
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent
className="sm:max-w-5xl w-[90vw] max-h-[88vh] flex flex-col gap-0 p-0 overflow-hidden rounded-3xl corner-squircle"
showCloseButton={true}
>
{/* Header */}
<DialogHeader className="px-6 pt-5 pb-4 shrink-0">
<div className="flex items-center gap-3 pr-10">
<div className="rounded-xl corner-squircle p-2 ring-1 ring-indigo-200 bg-indigo-50 text-indigo-600 dark:ring-indigo-800 dark:bg-indigo-950 dark:text-indigo-400 shrink-0">
<HugeiconsIcon icon={Database02Icon} className="size-4" />
</div>
<DialogTitle className="font-heading text-lg font-semibold tracking-tight">
Dataset Preview
</DialogTitle>
</div>
</DialogHeader>
{/* Body */}
<div className="flex flex-col min-h-0 flex-1 overflow-hidden px-6 pb-6">
{/* Loading */}
{loading && (
<div className="py-24 flex flex-col items-center justify-center gap-3">
<div className="rounded-2xl corner-squircle bg-primary/5 p-4">
<Spinner className="size-5 text-primary" />
</div>
<p className="text-sm text-muted-foreground font-medium">
Loading preview...
</p>
</div>
)}
{/* Error */}
{error && (
<div className="py-20 flex flex-col items-center justify-center gap-3">
<div className="rounded-2xl corner-squircle bg-destructive/10 p-3">
<HugeiconsIcon
icon={AlertCircleIcon}
className="size-5 text-destructive"
/>
</div>
<div className="text-center space-y-1">
<p className="text-sm font-medium text-destructive">{error}</p>
<p className="text-xs text-muted-foreground">
Make sure the backend is running on port 8000.
</p>
</div>
</div>
)}
{/* Content */}
{!loading && !error && data && (
<>
{/* Metadata card */}
<div className="rounded-xl corner-squircle ring-1 ring-border/60 bg-muted/30 px-5 py-4 mb-4 space-y-2">
<MetaRow label="Source" value={sourceLabel} />
<MetaRow
label="Format"
value={data.detected_format || "--"}
/>
<MetaRow
label="Total Rows"
value={
data.total_rows != null
? data.total_rows.toLocaleString()
: "--"
}
/>
<MetaRow
label="Columns"
value={
<span className="flex items-center gap-1.5 flex-wrap">
{columns.map((col) => (
<Badge
key={col}
variant="outline"
className="text-[11px] font-mono h-5"
>
{col}
</Badge>
))}
</span>
}
/>
</div>
{/* Data table */}
<div className="flex-1 min-h-0 rounded-xl corner-squircle ring-1 ring-border/60 overflow-auto">
<DataTable columns={tableColumns} data={rows} />
</div>
{/* Footer */}
<p className="text-[11px] text-muted-foreground/60 mt-3 text-center tabular-nums">
Showing {rows.length}
{data.total_rows != null &&
` of ${data.total_rows.toLocaleString()}`}{" "}
rows
</p>
</>
)}
</div>
</DialogContent>
</Dialog>
);
}
// ---------------------------------------------------------------------------
// Metadata row
// ---------------------------------------------------------------------------
function MetaRow({
label,
value,
}: {
label: string;
value: ReactNode;
}) {
return (
<div className="flex items-baseline gap-3 text-sm">
<span className="text-muted-foreground font-medium text-xs w-24 shrink-0">
{label}:
</span>
<span className="text-foreground text-[13px] min-w-0">{value}</span>
</div>
);
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
function formatCell(value: unknown): string {
if (value == null) return "";
if (typeof value === "string") return value;
if (typeof value === "number" || typeof value === "boolean")
return String(value);
if (Array.isArray(value) || typeof value === "object")
return JSON.stringify(value).slice(0, 500);
return String(value);
}
function isPreviewImagePayload(value: unknown): value is PreviewImagePayload {
if (!value || typeof value !== "object") return false;
const record = value as Record<string, unknown>;
return (
record.type === "image" &&
typeof record.data === "string" &&
record.data.length > 0
);
}
function collectPreviewImages(value: unknown): PreviewImagePayload[] {
const images: PreviewImagePayload[] = [];
const stack: unknown[] = [value];
let steps = 0;
while (stack.length > 0 && steps < 200) {
steps += 1;
const current = stack.pop();
if (isPreviewImagePayload(current)) {
images.push(current);
continue;
}
if (Array.isArray(current)) {
for (const item of current) stack.push(item);
continue;
}
if (current && typeof current === "object") {
for (const nested of Object.values(current as Record<string, unknown>)) {
stack.push(nested);
}
}
}
return images;
}

18
studio/frontend/src/features/studio/sections/dataset-section.tsx
View file

@ -28,7 +28,7 @@ import {
useInfiniteScroll,
} from "@/hooks";
import { formatCompact } from "@/lib/utils";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import {
CloudUploadIcon,
Database02Icon,
@ -40,10 +40,11 @@ import {
import { HugeiconsIcon } from "@hugeicons/react";
import { useMemo, useRef, useState } from "react";
import { useShallow } from "zustand/react/shallow";
import { DatasetPreviewDialog } from "./dataset-preview-dialog";
export function DatasetSection() {
const { dataset, setDataset, datasetFormat, setDatasetFormat, hfToken } =
useWizardStore(
useTrainingConfigStore(
useShallow(
({
dataset,
@ -62,6 +63,7 @@ export function DatasetSection() {
);
const [inputValue, setInputValue] = useState("");
const [previewOpen, setPreviewOpen] = useState(false);
const selectingRef = useRef(false);
const debouncedQuery = useDebouncedValue(inputValue);
@ -305,13 +307,21 @@ export function DatasetSection() {
<Button
variant="outline"
size="sm"
className="cursor-pointer gap-1.5 text-muted-foreground"
className="cursor-pointer gap-1.5"
disabled={!dataset}
onClick={() => setPreviewOpen(true)}
>
<HugeiconsIcon icon={ViewIcon} className="size-3.5" />
Preview
View dataset
</Button>
</div>
</div>
<DatasetPreviewDialog
open={previewOpen}
onOpenChange={setPreviewOpen}
datasetName={dataset}
hfToken={hfToken}
/>
</SectionCard>
);
}

4
studio/frontend/src/features/studio/sections/model-section.tsx
View file

@ -32,7 +32,7 @@ import {
useInfiniteScroll,
} from "@/hooks";
import { formatCompact } from "@/lib/utils";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { TrainingMethod } from "@/types/training";
import {
ChipIcon,
@ -65,7 +65,7 @@ export function ModelSection() {
setTrainingMethod,
hfToken,
setHfToken,
} = useWizardStore(
} = useTrainingConfigStore(
useShallow(
({
modelType,

4
studio/frontend/src/features/studio/sections/params-section.tsx
View file

@ -21,7 +21,7 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip";
import { CONTEXT_LENGTHS, TARGET_MODULES } from "@/config/training";
import { useWizardStore } from "@/stores/training";
import { useTrainingConfigStore } from "@/features/training";
import type { GradientCheckpointing } from "@/types/training";
import {
ArrowDown01Icon,
@ -107,7 +107,7 @@ function SliderRow({
}
export function ParamsSection(): ReactElement {
const store = useWizardStore();
const store = useTrainingConfigStore();
const isLora = store.trainingMethod !== "full";
const isVision = store.modelType === "vision";
const [loraOpen, setLoraOpen] = useState(false);

263
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -5,7 +5,12 @@ import {
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
import { useWizardStore } from "@/stores/training";
import {
useTrainingConfigStore,
useTrainingActions,
useTrainingRuntimeStore,
type TrainingPhase,
} from "@/features/training";
import {
ChartAverageIcon,
DashboardSpeed01Icon,
@ -16,66 +21,160 @@ import {
ZapIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import type { ReactElement, ReactNode } from "react";
import { useEffect, useRef, useState, type ReactElement, type ReactNode } from "react";
import { useShallow } from "zustand/react/shallow";
export function ProgressSection(): ReactElement | null {
const store = useWizardStore();
const metrics = store.trainingMetrics;
if (!metrics) {
return null;
const phaseLabel: Record<TrainingPhase, string> = {
idle: "Idle",
loading_model: "Loading model",
loading_dataset: "Loading dataset",
configuring: "Configuring",
training: "Training",
completed: "Completed",
error: "Error",
stopped: "Stopped",
};
const phaseColors: Record<TrainingPhase, string> = {
idle: "bg-muted text-muted-foreground",
loading_model: "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300",
loading_dataset:
"bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300",
configuring: "bg-blue-100 text-blue-700 dark:bg-blue-900 dark:text-blue-300",
training:
"bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300",
completed:
"bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300",
error: "bg-red-100 text-red-700 dark:bg-red-900 dark:text-red-300",
stopped: "bg-muted text-muted-foreground",
};
function formatDuration(seconds: number | null): string {
if (seconds == null || seconds < 0) {
return "--";
}
const total = Math.floor(seconds);
const min = Math.floor(total / 60);
const sec = total % 60;
return `${min}m ${sec}s`;
}
const pct = Math.round((metrics.currentStep / metrics.totalSteps) * 100);
const etaSec =
metrics.totalSteps > 0
? Math.round(
((metrics.totalSteps - metrics.currentStep) /
Math.max(metrics.currentStep, 1)) *
metrics.elapsed,
function formatNumber(value: number | null | undefined, digits: number): string {
if (value == null || !Number.isFinite(value)) {
return "--";
}
return value.toFixed(digits);
}
export function ProgressSection(): ReactElement {
const runtime = useTrainingRuntimeStore(
useShallow((state) => ({
phase: state.phase,
message: state.message,
error: state.error,
currentStep: state.currentStep,
totalSteps: state.totalSteps,
currentEpoch: state.currentEpoch,
currentLoss: state.currentLoss,
currentLearningRate: state.currentLearningRate,
currentGradNorm: state.currentGradNorm,
progressPercent: state.progressPercent,
elapsedSeconds: state.elapsedSeconds,
etaSeconds: state.etaSeconds,
currentNumTokens: state.currentNumTokens,
isTrainingRunning: state.isTrainingRunning,
})),
);
const config = useTrainingConfigStore(
useShallow((state) => ({
selectedModel: state.selectedModel,
trainingMethod: state.trainingMethod,
epochs: state.epochs,
batchSize: state.batchSize,
learningRate: state.learningRate,
maxSteps: state.maxSteps,
contextLength: state.contextLength,
warmupSteps: state.warmupSteps,
loraRank: state.loraRank,
loraAlpha: state.loraAlpha,
loraDropout: state.loraDropout,
loraVariant: state.loraVariant,
})),
);
const { stopTrainingRun } = useTrainingActions();
const localStartAtRef = useRef<number | null>(null);
const [, setLocalTick] = useState(0);
const pct =
runtime.totalSteps > 0
? Math.min(
100,
Math.max(
0,
Math.round((runtime.currentStep / runtime.totalSteps) * 100),
),
)
: 0;
const fmtTime = (s: number) => {
const m = Math.floor(s / 60);
const sec = s % 60;
return `${m}m ${sec}s`;
};
: Math.round(runtime.progressPercent);
const statusColors = {
training:
"bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300",
warmup: "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300",
saving:
"bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300",
};
const statusLabels = {
training: "Training",
warmup: "Warming up",
saving: "Saving checkpoint",
};
useEffect(() => {
if (runtime.elapsedSeconds != null && runtime.elapsedSeconds >= 0) {
localStartAtRef.current = Date.now() - runtime.elapsedSeconds * 1000;
return;
}
if (runtime.currentStep > 0 && localStartAtRef.current == null) {
localStartAtRef.current = Date.now();
}
}, [runtime.currentStep, runtime.elapsedSeconds]);
const modelName = store.selectedModel ?? "—";
useEffect(() => {
if (!runtime.isTrainingRunning) {
return;
}
const timer = window.setInterval(() => {
setLocalTick((prev) => prev + 1);
}, 1000);
return () => window.clearInterval(timer);
}, [runtime.isTrainingRunning]);
const elapsed =
runtime.elapsedSeconds ??
(localStartAtRef.current == null
? null
: Math.max(0, Math.floor((Date.now() - localStartAtRef.current) / 1000)));
const derivedEta =
elapsed != null && pct > 0
? Math.round((elapsed * (100 - pct)) / Math.max(pct, 1))
: null;
const eta = runtime.etaSeconds ?? derivedEta;
const stepsPerSecond =
elapsed != null && elapsed > 0
? runtime.currentStep / elapsed
: null;
const configItems = [
{
section: "Hyperparams",
rows: [
["Epochs", store.epochs],
["Batch size", store.batchSize],
["Learning rate", store.learningRate],
["Max steps", store.maxSteps],
["Context length", store.contextLength],
["Warmup steps", store.warmupSteps],
["Epochs", config.epochs],
["Batch size", config.batchSize],
["Learning rate", config.learningRate],
["Max steps", config.maxSteps],
["Context length", config.contextLength],
["Warmup steps", config.warmupSteps],
],
},
...(store.trainingMethod !== "full"
...(config.trainingMethod !== "full"
? [
{
section: "LoRA",
rows: [
["Rank", store.loraRank],
["Alpha", store.loraAlpha],
["Dropout", store.loraDropout],
["Variant", store.loraVariant],
["Rank", config.loraRank],
["Alpha", config.loraAlpha],
["Dropout", config.loraDropout],
["Variant", config.loraVariant],
],
},
]
@ -86,7 +185,7 @@ export function ProgressSection(): ReactElement | null {
<SectionCard
icon={<HugeiconsIcon icon={ChartAverageIcon} className="size-5" />}
title="Training Progress"
description="Live training metrics"
description={runtime.message || "Live training metrics"}
accent="emerald"
className="shadow-border ring-1 ring-border"
headerAction={
@ -130,7 +229,8 @@ export function ProgressSection(): ReactElement | null {
variant="destructive"
size="sm"
className="h-7 cursor-pointer px-3 text-xs"
onClick={() => store.setIsTraining(false)}
onClick={() => void stopTrainingRun()}
disabled={!runtime.isTrainingRunning}
>
<HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
</Button>
@ -138,24 +238,22 @@ export function ProgressSection(): ReactElement | null {
}
>
<div className="grid grid-cols-1 gap-6 md:grid-cols-2">
{/* Left: Progress */}
<div className="flex flex-col gap-4">
<div className="flex items-center gap-2">
<span
className={`rounded-full px-2.5 py-0.5 text-[10px] font-semibold ${statusColors[metrics.status]}`}
className={`rounded-full px-2.5 py-0.5 text-[10px] font-semibold ${phaseColors[runtime.phase]}`}
>
{statusLabels[metrics.status]}
{phaseLabel[runtime.phase]}
</span>
<span className="text-[10px] tabular-nums text-muted-foreground">
Epoch {metrics.currentEpoch.toFixed(2)} / {metrics.totalEpochs}
Epoch {runtime.currentEpoch.toFixed(2)}
</span>
</div>
{/* Progress bar */}
<div className="flex flex-col gap-1.5">
<div className="flex justify-between text-xs text-muted-foreground">
<span>
Step {metrics.currentStep} / {metrics.totalSteps}
Step {runtime.currentStep} / {runtime.totalSteps || "--"}
</span>
<span>{pct}%</span>
</div>
@ -167,51 +265,59 @@ export function ProgressSection(): ReactElement | null {
</div>
</div>
{/* Metrics */}
<div className="flex items-baseline gap-4">
{runtime.error && (
<p className="text-xs text-red-500 leading-relaxed">{runtime.error}</p>
)}
<div className="flex flex-wrap items-baseline gap-4">
<div>
<p className="text-xs text-muted-foreground">Loss</p>
<p className="text-3xl font-bold tabular-nums tracking-tight">
{metrics.currentLoss.toFixed(4)}
{runtime.currentLoss.toFixed(4)}
</p>
</div>
<div>
<p className="text-xs text-muted-foreground">LR</p>
<p className="text-lg font-semibold tabular-nums">
{metrics.currentLR.toExponential(2)}
{runtime.currentLearningRate.toExponential(2)}
</p>
</div>
<div>
<p className="text-xs text-muted-foreground">Grad Norm</p>
<p className="text-lg font-semibold tabular-nums">
{metrics.gradNorm.toFixed(3)}
{formatNumber(runtime.currentGradNorm, 3)}
</p>
</div>
<div>
<p className="text-xs text-muted-foreground">Model</p>
<p className="text-lg font-semibold truncate max-w-[140px]">
{modelName}
{config.selectedModel ?? "--"}
</p>
</div>
<div>
<p className="text-xs text-muted-foreground">Method</p>
<p className="text-lg font-semibold">{store.trainingMethod}</p>
<p className="text-lg font-semibold">
{config.trainingMethod.toUpperCase()}
</p>
</div>
</div>
{/* Timings */}
<div className="flex gap-4 text-xs text-muted-foreground">
<span>Elapsed: {fmtTime(metrics.elapsed)}</span>
<span>ETA: {fmtTime(etaSec)}</span>
<span>{metrics.samplesPerSecond} samples/s</span>
<div className="flex flex-wrap gap-4 text-xs text-muted-foreground">
<span>Elapsed: {formatDuration(elapsed)}</span>
<span>ETA: {formatDuration(eta)}</span>
<span>
{stepsPerSecond == null
? "-- steps/s"
: `${stepsPerSecond.toFixed(2)} steps/s`}
</span>
{runtime.currentNumTokens != null && (
<span>Tokens: {runtime.currentNumTokens}</span>
)}
</div>
</div>
{/* Right: GPU */}
<div className="flex flex-col gap-3">
<p className="text-xs font-medium text-muted-foreground">
GPU Monitor
</p>
<p className="text-xs font-medium text-muted-foreground">GPU Monitor</p>
<div className="grid grid-cols-2 gap-3">
<GpuStat
label="Utilization"
@ -221,29 +327,27 @@ export function ProgressSection(): ReactElement | null {
className="size-3.5"
/>
}
value={`${metrics.gpuUtil}%`}
pct={metrics.gpuUtil}
value="--"
pct={0}
/>
<GpuStat
label="Temperature"
icon={
<HugeiconsIcon icon={TemperatureIcon} className="size-3.5" />
}
value={`${metrics.gpuTemp}°C`}
pct={metrics.gpuTemp}
icon={<HugeiconsIcon icon={TemperatureIcon} className="size-3.5" />}
value="--"
pct={0}
max={100}
/>
<GpuStat
label="VRAM"
icon={<HugeiconsIcon icon={RamMemoryIcon} className="size-3.5" />}
value={`${metrics.gpuVramUsed.toFixed(1)} / ${metrics.gpuVramTotal}GB`}
pct={(metrics.gpuVramUsed / metrics.gpuVramTotal) * 100}
value="--"
pct={0}
/>
<GpuStat
label="Power"
icon={<HugeiconsIcon icon={ZapIcon} className="size-3.5" />}
value={`${metrics.gpuPower}W`}
pct={(metrics.gpuPower / 350) * 100}
value="--"
pct={0}
/>
</div>
</div>
@ -272,6 +376,7 @@ function GpuStat({
} else if (clamped < 95) {
barColor = "bg-amber-500";
}
return (
<div className="flex flex-col gap-1.5 rounded-xl bg-muted/50 p-3">
<div className="flex items-center justify-between text-xs">

15
studio/frontend/src/features/studio/sections/training-section.tsx
View file

@ -9,7 +9,7 @@ import {
CollapsibleTrigger,
} from "@/components/ui/collapsible";
import { Input } from "@/components/ui/input";
import { useWizardStore } from "@/stores/training";
import { useTrainingActions, useTrainingConfigStore } from "@/features/training";
import {
Archive04Icon,
ArrowDown01Icon,
@ -35,7 +35,8 @@ const placeholderData = [
];
export function TrainingSection() {
const store = useWizardStore();
const store = useTrainingConfigStore();
const { isStarting, startError, startTrainingRun } = useTrainingActions();
const [logOpen, setLogOpen] = useState(false);
return (
@ -94,11 +95,15 @@ export function TrainingSection() {
{/* Start/Stop */}
<Button
className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600"
onClick={() => store.setIsTraining(true)}
onClick={() => void startTrainingRun()}
disabled={isStarting}
>
<HugeiconsIcon icon={Rocket01Icon} className="size-4" /> Start
Training
<HugeiconsIcon icon={Rocket01Icon} className="size-4" />
{isStarting ? "Starting..." : "Start Training"}
</Button>
{startError && (
<p className="text-xs text-red-500 leading-relaxed">{startError}</p>
)}
{/* Save / Clear */}
<div className="grid grid-cols-2 gap-2">

26
studio/frontend/src/features/studio/studio-page.tsx
View file

@ -1,4 +1,8 @@
import { useWizardStore } from "@/stores/training";
import {
shouldShowTrainingView,
useTrainingRuntimeLifecycle,
useTrainingRuntimeStore,
} from "@/features/training";
import type { ReactElement } from "react";
import { DatasetSection } from "./sections/dataset-section";
import { ModelSection } from "./sections/model-section";
@ -7,24 +11,32 @@ import { TrainingSection } from "./sections/training-section";
import { TrainingView } from "./training-view";
export function StudioPage(): ReactElement {
const isTraining = useWizardStore((s) => s.isTraining);
useTrainingRuntimeLifecycle();
const showTrainingView = useTrainingRuntimeStore(shouldShowTrainingView);
const runtimeMessage = useTrainingRuntimeStore((state) => state.message);
const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating);
const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated);
return (
<div className="min-h-screen bg-background">
<main className="mx-auto max-w-7xl px-6 py-8">
<main className="mx-auto max-w-7xl px-6 py-4">
{/* Header */}
<div className="mb-8 flex flex-col gap-1">
<div className="mb-8 flex flex-col gap-0.5">
<h1 className="text-2xl font-semibold tracking-tight">
Fine-tuning Studio
</h1>
<p className="text-sm text-muted-foreground">
{isTraining
? "Training in progress"
{showTrainingView
? runtimeMessage || "Training in progress"
: "Configure and start training"}
</p>
</div>
{isTraining ? (
{!hasHydratedRuntime && isHydratingRuntime ? (
<div className="rounded-xl border bg-card p-8 text-sm text-muted-foreground">
Loading training runtime...
</div>
) : showTrainingView ? (
<TrainingView />
) : (
<div className="grid grid-cols-12 items-start gap-6">

58
studio/frontend/src/features/studio/training-start-overlay.tsx Normal file
View file

@ -0,0 +1,58 @@
import {
AnimatedSpan,
Terminal,
TypingAnimation,
} from "@/components/ui/terminal"
import type { ReactElement } from "react"
type TrainingStartOverlayProps = {
message: string
currentStep: number
}
export function TrainingStartOverlay({
message,
currentStep,
}: TrainingStartOverlayProps): ReactElement {
return (
<div className="pointer-events-none absolute inset-0 z-30 flex items-center justify-center rounded-2xl bg-background/45 backdrop-blur-[1px]">
<div className="flex w-[860px] max-w-[calc(100%-2rem)] flex-col items-center gap-4">
<img
src="/Sloth emojis/large sloth wave.png"
alt="Unsloth mascot"
className="size-24 animate-bounce object-contain"
/>
<Terminal
className="w-full min-h-[390px] rounded-2xl px-7 py-6 text-left"
startOnView={false}
>
<TypingAnimation
duration={36}
className="bg-gradient-to-r from-emerald-300 via-lime-300 to-teal-300 bg-clip-text font-semibold text-transparent"
>
{"> unsloth training starts..."}
</TypingAnimation>
<AnimatedSpan className="my-2">
<pre className="whitespace-pre text-left text-muted-foreground">{`==((====))==
\\\\ /|
O^O/ \\_/ \\
\\ /
"-____-"`}</pre>
</AnimatedSpan>
<TypingAnimation duration={44}>
{"> Preparing model and dataset..."}
</TypingAnimation>
<TypingAnimation duration={44}>
{"> We are getting everything ready for your run..."}
</TypingAnimation>
<TypingAnimation duration={44}>
{"> Did you know, Mugi is actually short for \"Mugiwara\" xd"}
</TypingAnimation>
<AnimatedSpan className="mt-2 text-muted-foreground">
{`> ${message || "starting training..."} | waiting for first step... (${currentStep})`}
</AnimatedSpan>
</Terminal>
</div>
</div>
)
}

191
studio/frontend/src/features/studio/training-view.tsx
View file

@ -1,164 +1,47 @@
import { useWizardStore } from "@/stores/training";
import type { TrainingMetrics } from "@/types/training";
import { type ReactElement, useEffect, useRef } from "react";
import { cn } from "@/lib/utils";
import { useTrainingRuntimeStore } from "@/features/training";
import type { ReactElement } from "react";
import { useShallow } from "zustand/react/shallow";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
function createInitialMetrics(
totalSteps: number,
totalEpochs: number,
lr: number,
): TrainingMetrics {
return {
currentStep: 0,
totalSteps,
currentEpoch: 0,
totalEpochs,
currentLoss: 2.5,
currentLR: lr * 0.1,
gradNorm: 0,
samplesPerSecond: 0,
lossHistory: [],
lrHistory: [],
gradNormHistory: [],
gpuUtil: 0,
gpuTemp: 45,
gpuVramUsed: 0,
gpuVramTotal: 24,
gpuPower: 50,
elapsed: 0,
status: "warmup",
};
}
import { TrainingStartOverlay } from "./training-start-overlay";
export function TrainingView(): ReactElement {
const { maxSteps, epochs, learningRate, warmupSteps, setTrainingMetrics } =
useWizardStore();
const metricsRef = useRef<ReturnType<typeof setInterval> | null>(null);
const chartsRef = useRef<ReturnType<typeof setInterval> | null>(null);
const runtime = useTrainingRuntimeStore(
useShallow((state) => ({
phase: state.phase,
message: state.message,
currentStep: state.currentStep,
firstStepReceived: state.firstStepReceived,
isStarting: state.isStarting,
})),
);
useEffect(() => {
const totalSteps = maxSteps || 500;
const totalEpochs = epochs || 3;
const peakLR = learningRate;
const warmup = warmupSteps || 20;
setTrainingMetrics(createInitialMetrics(totalSteps, totalEpochs, peakLR));
let step = 0;
let elapsed = 0;
const computeStep = () => {
step++;
if (step > totalSteps) {
return null;
}
elapsed++;
let lr: number;
if (step < warmup) {
lr = peakLR * (step / warmup);
} else {
const progress = (step - warmup) / (totalSteps - warmup);
lr = peakLR * 0.5 * (1 + Math.cos(Math.PI * progress));
}
const baseLoss = 2.5 * Math.exp((-3 * step) / totalSteps) + 0.3;
const noise = (Math.random() - 0.5) * 0.08;
const loss = Math.max(0.1, baseLoss + noise);
const status =
step < warmup ? "warmup" : step % 100 === 0 ? "saving" : "training";
const gradNorm = +(
1.2 * Math.exp(-step / totalSteps) +
0.1 +
(Math.random() - 0.5) * 0.05
).toFixed(3);
return {
step,
elapsed,
lr,
loss: +loss.toFixed(4),
status: status as TrainingMetrics["status"],
gradNorm,
};
};
// Top card values — update every 1s
metricsRef.current = setInterval(() => {
const s = computeStep();
if (!s) {
if (metricsRef.current) {
clearInterval(metricsRef.current);
}
if (chartsRef.current) {
clearInterval(chartsRef.current);
}
return;
}
const prev = useWizardStore.getState().trainingMetrics;
setTrainingMetrics({
currentStep: s.step,
totalSteps,
currentEpoch:
Math.floor((s.step / totalSteps) * totalEpochs * 100) / 100,
totalEpochs,
currentLoss: s.loss,
currentLR: s.lr,
gradNorm: s.gradNorm,
samplesPerSecond: +(12 + (Math.random() - 0.5) * 2).toFixed(1),
lossHistory: prev?.lossHistory ?? [],
lrHistory: prev?.lrHistory ?? [],
gradNormHistory: prev?.gradNormHistory ?? [],
gpuUtil: Math.min(99, 85 + Math.round((Math.random() - 0.5) * 10)),
gpuTemp: Math.min(89, 68 + Math.round((Math.random() - 0.5) * 6)),
gpuVramUsed: +(18.2 + (Math.random() - 0.5) * 0.4).toFixed(1),
gpuVramTotal: 24,
gpuPower: Math.round(280 + (Math.random() - 0.5) * 30),
elapsed: s.elapsed,
status: s.status,
});
}, 1000);
// Chart history — update every 5s
chartsRef.current = setInterval(() => {
const prev = useWizardStore.getState().trainingMetrics;
if (!prev || prev.currentStep === 0) {
return;
}
setTrainingMetrics({
...prev,
lossHistory: [
...prev.lossHistory,
{ step: prev.currentStep, loss: prev.currentLoss },
],
lrHistory: [
...prev.lrHistory,
{ step: prev.currentStep, lr: prev.currentLR },
],
gradNormHistory: [
...prev.gradNormHistory,
{ step: prev.currentStep, gradNorm: prev.gradNorm },
],
});
}, 5000);
return () => {
if (metricsRef.current) {
clearInterval(metricsRef.current);
}
if (chartsRef.current) {
clearInterval(chartsRef.current);
}
};
}, [epochs, learningRate, maxSteps, setTrainingMetrics, warmupSteps]);
const isPreparingPhase =
runtime.phase === "loading_model" ||
runtime.phase === "loading_dataset" ||
runtime.phase === "configuring";
const isWaitingForFirstStep =
runtime.phase === "training" && !runtime.firstStepReceived;
const showOverlay =
runtime.isStarting ||
isPreparingPhase ||
(isWaitingForFirstStep && runtime.currentStep <= 0);
return (
<div className="flex flex-col gap-6">
<ProgressSection />
<ChartsSection />
<div className={cn("relative", showOverlay && "min-h-[72vh]")}>
<div
className={cn("flex flex-col gap-6 transition-[filter]", showOverlay && "blur")}
>
<ProgressSection />
<ChartsSection />
</div>
{showOverlay ? (
<TrainingStartOverlay
message={runtime.message}
currentStep={runtime.currentStep}
/>
) : null}
</div>
);
}

63
studio/frontend/src/features/training/api/mappers.ts Normal file
View file

@ -0,0 +1,63 @@
import type { TrainingConfigState } from "../types/config";
import type { TrainingStartRequest } from "../types/api";
const BACKEND_LORA_TYPE = "LoRA/QLoRA";
const BACKEND_FULL_TYPE = "Full Finetuning";
export function toBackendTrainingType(trainingMethod: string): string {
return trainingMethod === "full" ? BACKEND_FULL_TYPE : BACKEND_LORA_TYPE;
}
export function buildTrainingStartPayload(
config: TrainingConfigState,
): TrainingStartRequest {
const adapterMethod = config.trainingMethod !== "full";
const isQlorMethod = config.trainingMethod === "qlora";
const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
return {
model_name: config.selectedModel ?? "",
training_type: toBackendTrainingType(config.trainingMethod),
hf_token: config.hfToken.trim() || null,
load_in_4bit: adapterMethod ? isQlorMethod : false,
max_seq_length: config.contextLength,
hf_dataset: hfDataset,
local_datasets: [],
format_type: config.datasetFormat,
num_epochs: config.epochs,
learning_rate: String(config.learningRate),
batch_size: config.batchSize,
gradient_accumulation_steps: config.gradientAccumulation,
warmup_steps: config.warmupSteps,
warmup_ratio: null,
max_steps: config.maxSteps,
save_steps: config.saveSteps,
weight_decay: config.weightDecay,
random_seed: config.randomSeed,
packing: config.packing,
optim: "adamw_8bit",
lr_scheduler_type: "linear",
use_lora: adapterMethod,
lora_r: config.loraRank,
lora_alpha: config.loraAlpha,
lora_dropout: config.loraDropout,
target_modules: adapterMethod ? config.targetModules : [],
gradient_checkpointing: config.gradientCheckpointing,
use_rslora: config.loraVariant === "rslora",
use_loftq: config.loraVariant === "loftq",
train_on_completions: config.trainOnCompletions,
finetune_vision_layers: config.finetuneVisionLayers,
finetune_language_layers: config.finetuneLanguageLayers,
finetune_attention_modules: config.finetuneAttentionModules,
finetune_mlp_modules: config.finetuneMLPModules,
enable_wandb: config.enableWandb,
wandb_token: config.enableWandb ? config.wandbToken.trim() || null : null,
wandb_project: config.enableWandb
? config.wandbProject.trim() || null
: null,
enable_tensorboard: config.enableTensorboard,
tensorboard_dir: config.enableTensorboard
? config.tensorboardDir.trim() || null
: null,
};
}

173
studio/frontend/src/features/training/api/train-api.ts Normal file
View file

@ -0,0 +1,173 @@
import { authFetch } from "@/features/auth";
import type {
TrainingStartRequest,
TrainingStartResponse,
TrainingStopResponse,
} from "../types/api";
import type {
TrainingMetricsResponse,
TrainingProgressPayload,
TrainingStatusResponse,
} from "../types/runtime";
function isAbortError(error: unknown): boolean {
return error instanceof DOMException && error.name === "AbortError";
}
async function readError(response: Response): Promise<string> {
try {
const payload = (await response.json()) as { detail?: string; message?: string };
return payload.detail || payload.message || `Request failed (${response.status})`;
} catch {
return `Request failed (${response.status})`;
}
}
async function parseJson<T>(response: Response): Promise<T> {
if (!response.ok) {
throw new Error(await readError(response));
}
return (await response.json()) as T;
}
export async function startTraining(
payload: TrainingStartRequest,
): Promise<TrainingStartResponse> {
const response = await authFetch("/api/train/start", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
return parseJson<TrainingStartResponse>(response);
}
export async function stopTraining(): Promise<TrainingStopResponse> {
const response = await authFetch("/api/train/stop", { method: "POST" });
return parseJson<TrainingStopResponse>(response);
}
export async function getTrainingStatus(): Promise<TrainingStatusResponse> {
const response = await authFetch("/api/train/status");
return parseJson<TrainingStatusResponse>(response);
}
export async function getTrainingMetrics(): Promise<TrainingMetricsResponse> {
const response = await authFetch("/api/train/metrics");
return parseJson<TrainingMetricsResponse>(response);
}
type ProgressEventName = "progress" | "heartbeat" | "complete" | "error";
interface ParsedSseEvent {
event: ProgressEventName;
payload: TrainingProgressPayload;
id: number | null;
}
function parseSseEvent(rawEvent: string): ParsedSseEvent | null {
const lines = rawEvent.split(/\r?\n/);
let eventName: ProgressEventName = "progress";
let id: number | null = null;
const dataLines: string[] = [];
for (const line of lines) {
if (!line) {
continue;
}
if (line.startsWith("event:")) {
const value = line.slice(6).trim();
if (
value === "progress" ||
value === "heartbeat" ||
value === "complete" ||
value === "error"
) {
eventName = value;
}
continue;
}
if (line.startsWith("id:")) {
const value = Number(line.slice(3).trim());
id = Number.isFinite(value) ? value : null;
continue;
}
if (line.startsWith("data:")) {
dataLines.push(line.slice(5).trimStart());
}
}
if (dataLines.length === 0) {
return null;
}
const parsed = JSON.parse(dataLines.join("\n")) as TrainingProgressPayload;
return { event: eventName, payload: parsed, id };
}
export async function streamTrainingProgress(options: {
signal: AbortSignal;
lastEventId?: number | null;
onOpen?: () => void;
onEvent: (event: ParsedSseEvent) => void;
}): Promise<void> {
const headers = new Headers();
if (typeof options.lastEventId === "number") {
headers.set("Last-Event-ID", String(options.lastEventId));
}
const response = await authFetch("/api/train/progress", {
method: "GET",
headers,
signal: options.signal,
});
if (!response.ok) {
throw new Error(await readError(response));
}
if (!response.body) {
throw new Error("Progress stream unavailable");
}
options.onOpen?.();
const reader = response.body.getReader();
const decoder = new TextDecoder();
let buffer = "";
while (true) {
const { value, done } = await reader.read();
if (done) {
break;
}
buffer += decoder.decode(value, { stream: true });
let separatorIndex = buffer.search(/\r?\n\r?\n/);
while (separatorIndex >= 0) {
const rawEvent = buffer.slice(0, separatorIndex);
const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2;
buffer = buffer.slice(separatorIndex + separatorLength);
if (rawEvent.startsWith("retry:")) {
separatorIndex = buffer.search(/\r?\n\r?\n/);
continue;
}
try {
const event = parseSseEvent(rawEvent);
if (event) {
options.onEvent(event);
}
} catch (error) {
if (!isAbortError(error)) {
throw error;
}
}
separatorIndex = buffer.search(/\r?\n\r?\n/);
}
}
}
export { isAbortError };

70
studio/frontend/src/features/training/hooks/use-training-actions.ts Normal file
View file

@ -0,0 +1,70 @@
import { useCallback } from "react";
import { useTrainingConfigStore } from "../stores/training-config-store";
import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
import { startTraining, stopTraining } from "../api/train-api";
import { buildTrainingStartPayload } from "../api/mappers";
import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
import { validateTrainingConfig } from "../lib/validation";
export function useTrainingActions() {
const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
const startError = useTrainingRuntimeStore((state) => state.startError);
const startTrainingRun = useCallback(async (): Promise<boolean> => {
const config = useTrainingConfigStore.getState();
const runtimeStore = useTrainingRuntimeStore.getState();
runtimeStore.setStartError(null);
const validation = validateTrainingConfig(config);
if (!validation.ok) {
runtimeStore.setStartError(validation.message);
return false;
}
runtimeStore.setStarting(true);
try {
const payload = buildTrainingStartPayload(config);
const response = await startTraining(payload);
if (response.status === "error") {
runtimeStore.setStartError(response.error || response.message);
runtimeStore.setStarting(false);
return false;
}
runtimeStore.setStartQueued(response.job_id, response.message);
await syncTrainingRuntimeFromBackend();
return true;
} catch (error) {
const message =
error instanceof Error ? error.message : "Failed to start training";
runtimeStore.setStartError(message);
runtimeStore.setStarting(false);
return false;
}
}, []);
const stopTrainingRun = useCallback(async (): Promise<boolean> => {
const runtimeStore = useTrainingRuntimeStore.getState();
runtimeStore.setStartError(null);
try {
await stopTraining();
await syncTrainingRuntimeFromBackend();
return true;
} catch (error) {
const message =
error instanceof Error ? error.message : "Failed to stop training";
runtimeStore.setRuntimeError(message);
return false;
}
}, []);
return {
isStarting,
startError,
startTrainingRun,
stopTrainingRun,
};
}

183
studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts Normal file
View file

@ -0,0 +1,183 @@
import { useEffect } from "react";
import {
getTrainingMetrics,
getTrainingStatus,
isAbortError,
streamTrainingProgress,
} from "../api/train-api";
import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
import type { TrainingRuntimeStore } from "../types/runtime";
const STATUS_POLL_INTERVAL_MS = 3000;
const METRICS_POLL_INTERVAL_MS = 5000;
const STREAM_RECONNECT_DELAY_MS = 1500;
function shouldUseLiveSync(state: TrainingRuntimeStore): boolean {
return (
state.isTrainingRunning ||
state.phase === "loading_model" ||
state.phase === "loading_dataset" ||
state.phase === "configuring" ||
state.phase === "training"
);
}
export function useTrainingRuntimeLifecycle(): void {
useEffect(() => {
let disposed = false;
let openingStream = false;
let streamController: AbortController | null = null;
let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
const runtimeStore = useTrainingRuntimeStore;
const clearReconnect = () => {
if (reconnectTimer) {
clearTimeout(reconnectTimer);
reconnectTimer = null;
}
};
const stopStream = () => {
clearReconnect();
if (streamController) {
streamController.abort();
streamController = null;
}
runtimeStore.getState().setSseConnected(false);
};
const pollMetrics = async () => {
try {
const metrics = await getTrainingMetrics();
if (disposed) {
return;
}
runtimeStore.getState().applyMetrics(metrics);
} catch (error) {
if (!isAbortError(error) && !disposed) {
runtimeStore.getState().setSseConnected(false);
}
}
};
const pollStatus = async () => {
try {
const status = await getTrainingStatus();
if (disposed) {
return;
}
runtimeStore.getState().applyStatus(status);
const nextState = runtimeStore.getState();
if (shouldUseLiveSync(nextState)) {
void ensureStream();
} else {
stopStream();
}
} catch (error) {
if (!isAbortError(error) && !disposed) {
runtimeStore.getState().setSseConnected(false);
}
}
};
const ensureStream = async () => {
const state = runtimeStore.getState();
if (
disposed ||
openingStream ||
streamController ||
!shouldUseLiveSync(state)
) {
return;
}
clearReconnect();
openingStream = true;
const controller = new AbortController();
streamController = controller;
try {
await streamTrainingProgress({
signal: controller.signal,
lastEventId: state.lastEventId,
onOpen: () => {
runtimeStore.getState().setSseConnected(true);
},
onEvent: (event) => {
const liveStore = runtimeStore.getState();
if (typeof event.id === "number") {
liveStore.setLastEventId(event.id);
}
liveStore.applyProgress(event.payload, event.id ?? undefined);
if (event.event === "complete") {
void pollStatus();
void pollMetrics();
stopStream();
}
if (event.event === "error") {
liveStore.setRuntimeError("Training stream error");
stopStream();
}
},
});
} catch (error) {
if (!disposed && !controller.signal.aborted && !isAbortError(error)) {
runtimeStore.getState().setSseConnected(false);
}
} finally {
openingStream = false;
if (streamController === controller) {
streamController = null;
}
runtimeStore.getState().setSseConnected(false);
if (!disposed && !controller.signal.aborted) {
const liveState = runtimeStore.getState();
if (shouldUseLiveSync(liveState)) {
reconnectTimer = setTimeout(() => {
void ensureStream();
}, STREAM_RECONNECT_DELAY_MS);
}
}
}
};
const hydrate = async () => {
runtimeStore.getState().setHydrating(true);
try {
await Promise.all([pollStatus(), pollMetrics()]);
} finally {
if (!disposed) {
runtimeStore.getState().setHydrating(false);
runtimeStore.getState().setHasHydrated(true);
}
}
};
void hydrate();
const statusTimer = setInterval(() => {
void pollStatus();
}, STATUS_POLL_INTERVAL_MS);
const metricsTimer = setInterval(() => {
const state = runtimeStore.getState();
if (shouldUseLiveSync(state) || state.currentStep > 0) {
void pollMetrics();
}
}, METRICS_POLL_INTERVAL_MS);
return () => {
disposed = true;
clearInterval(statusTimer);
clearInterval(metricsTimer);
stopStream();
};
}, []);
}

8
studio/frontend/src/features/training/index.ts Normal file
View file

@ -0,0 +1,8 @@
export { useTrainingConfigStore } from "./stores/training-config-store";
export {
shouldShowTrainingView,
useTrainingRuntimeStore,
} from "./stores/training-runtime-store";
export { useTrainingActions } from "./hooks/use-training-actions";
export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
export type { TrainingPhase } from "./types/runtime";

19
studio/frontend/src/features/training/lib/sync-runtime.ts Normal file
View file

@ -0,0 +1,19 @@
import {
getTrainingMetrics,
getTrainingStatus,
} from "../api/train-api";
import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
import type { TrainingStatusResponse } from "../types/runtime";
export async function syncTrainingRuntimeFromBackend(): Promise<TrainingStatusResponse> {
const [status, metrics] = await Promise.all([
getTrainingStatus(),
getTrainingMetrics(),
]);
const runtimeStore = useTrainingRuntimeStore.getState();
runtimeStore.applyStatus(status);
runtimeStore.applyMetrics(metrics);
return status;
}

27
studio/frontend/src/features/training/lib/validation.ts Normal file
View file

@ -0,0 +1,27 @@
import type { TrainingConfigState } from "../types/config";
export interface StartValidationResult {
ok: boolean;
message: string | null;
}
export function validateTrainingConfig(
config: TrainingConfigState,
): StartValidationResult {
if (!config.selectedModel) {
return { ok: false, message: "Select a base model first." };
}
if (config.datasetSource !== "huggingface") {
return {
ok: false,
message: "Only Hugging Face dataset source is enabled right now.",
};
}
if (!config.dataset) {
return { ok: false, message: "Select a Hugging Face dataset first." };
}
return { ok: true, message: null };
}

105
studio/frontend/src/features/training/stores/training-config-store.ts Normal file
View file

@ -0,0 +1,105 @@
import { DEFAULT_HYPERPARAMS, STEPS } from "@/config/training";
import type { StepNumber } from "@/types/training";
import { create } from "zustand";
import { persist } from "zustand/middleware";
import type { TrainingConfigState, TrainingConfigStore } from "../types/config";
const MIN_STEP: StepNumber = 1;
const MAX_STEP: StepNumber = STEPS.length as StepNumber;
const initialState: TrainingConfigState = {
currentStep: MIN_STEP,
modelType: null,
selectedModel: null,
trainingMethod: "qlora",
hfToken: "",
datasetSource: "huggingface",
datasetFormat: "auto",
dataset: null,
uploadedFile: null,
...DEFAULT_HYPERPARAMS,
};
function clampStep(step: number): StepNumber {
return Math.min(MAX_STEP, Math.max(MIN_STEP, step)) as StepNumber;
}
function canProceedForStep(state: TrainingConfigState): boolean {
switch (state.currentStep) {
case 1:
return state.modelType !== null;
case 2:
return state.selectedModel !== null;
case 3:
return state.datasetSource === "upload"
? state.uploadedFile !== null
: state.dataset !== null;
case 4:
case 5:
return true;
default:
return false;
}
}
export const useTrainingConfigStore = create<TrainingConfigStore>()(
persist(
(set, get) => ({
...initialState,
setStep: (step) => set({ currentStep: step }),
nextStep: () => set({ currentStep: clampStep(get().currentStep + 1) }),
prevStep: () => set({ currentStep: clampStep(get().currentStep - 1) }),
setModelType: (modelType) => set({ modelType, selectedModel: null }),
setSelectedModel: (selectedModel) => set({ selectedModel }),
setTrainingMethod: (trainingMethod) => set({ trainingMethod }),
setHfToken: (hfToken) => set({ hfToken }),
setDatasetSource: (datasetSource) => set({ datasetSource }),
setDatasetFormat: (datasetFormat) => set({ datasetFormat }),
setDataset: (dataset) => set({ dataset }),
setUploadedFile: (uploadedFile) => set({ uploadedFile }),
setEpochs: (epochs) => set({ epochs }),
setContextLength: (contextLength) => set({ contextLength }),
setLearningRate: (learningRate) => set({ learningRate }),
setLoraRank: (loraRank) => set({ loraRank }),
setLoraAlpha: (loraAlpha) => set({ loraAlpha }),
setLoraDropout: (loraDropout) => set({ loraDropout }),
setLoraVariant: (loraVariant) => set({ loraVariant }),
setBatchSize: (batchSize) => set({ batchSize }),
setGradientAccumulation: (gradientAccumulation) =>
set({ gradientAccumulation }),
setWeightDecay: (weightDecay) => set({ weightDecay }),
setWarmupSteps: (warmupSteps) => set({ warmupSteps }),
setMaxSteps: (maxSteps) => set({ maxSteps }),
setSaveSteps: (saveSteps) => set({ saveSteps }),
setPacking: (packing) => set({ packing }),
setTrainOnCompletions: (trainOnCompletions) =>
set({ trainOnCompletions }),
setGradientCheckpointing: (gradientCheckpointing) =>
set({ gradientCheckpointing }),
setRandomSeed: (randomSeed) => set({ randomSeed }),
setEnableWandb: (enableWandb) => set({ enableWandb }),
setWandbToken: (wandbToken) => set({ wandbToken }),
setWandbProject: (wandbProject) => set({ wandbProject }),
setEnableTensorboard: (enableTensorboard) => set({ enableTensorboard }),
setTensorboardDir: (tensorboardDir) => set({ tensorboardDir }),
setLogFrequency: (logFrequency) => set({ logFrequency }),
setFinetuneVisionLayers: (finetuneVisionLayers) =>
set({ finetuneVisionLayers }),
setFinetuneLanguageLayers: (finetuneLanguageLayers) =>
set({ finetuneLanguageLayers }),
setFinetuneAttentionModules: (finetuneAttentionModules) =>
set({ finetuneAttentionModules }),
setFinetuneMLPModules: (finetuneMLPModules) => set({ finetuneMLPModules }),
setTargetModules: (targetModules) => set({ targetModules }),
canProceed: () => canProceedForStep(get()),
reset: () => set(initialState),
}),
{
name: "unsloth_training_config_v1",
partialize: (state) => {
const { modelType, ...rest } = state;
return rest;
},
},
),
);

229
studio/frontend/src/features/training/stores/training-runtime-store.ts Normal file
View file

@ -0,0 +1,229 @@
import { create } from "zustand";
import type {
TrainingMetricsResponse,
TrainingProgressPayload,
TrainingRuntimeState,
TrainingRuntimeStore,
TrainingSeriesPoint,
TrainingStatusResponse,
} from "../types/runtime";
const initialState: TrainingRuntimeState = {
jobId: null,
phase: "idle",
isTrainingRunning: false,
message: "Ready to train",
error: null,
isHydrating: false,
hasHydrated: false,
isStarting: false,
startError: null,
sseConnected: false,
firstStepReceived: false,
lastEventId: null,
currentStep: 0,
totalSteps: 0,
currentEpoch: 0,
currentLoss: 0,
currentLearningRate: 0,
progressPercent: 0,
elapsedSeconds: null,
etaSeconds: null,
currentGradNorm: null,
currentNumTokens: null,
lossHistory: [],
lrHistory: [],
gradNormHistory: [],
};
function sortSeries(points: TrainingSeriesPoint[]): TrainingSeriesPoint[] {
return [...points].sort((a, b) => a.step - b.step);
}
function toSeries(steps: number[], values: number[]): TrainingSeriesPoint[] {
const points: TrainingSeriesPoint[] = [];
for (let i = 0; i < steps.length; i += 1) {
const step = steps[i];
const value = values[i];
if (!Number.isFinite(step) || !Number.isFinite(value)) {
continue;
}
points.push({ step, value });
}
return sortSeries(points);
}
function upsertPoint(
points: TrainingSeriesPoint[],
step: number,
value: number,
): TrainingSeriesPoint[] {
const next = points.slice();
const index = next.findIndex((point) => point.step === step);
if (index >= 0) {
next[index] = { step, value };
return next;
}
next.push({ step, value });
return sortSeries(next);
}
function applyMetricHistoryFromStatus(payload: TrainingStatusResponse): {
lossHistory: TrainingSeriesPoint[] | null;
lrHistory: TrainingSeriesPoint[] | null;
} {
const history = payload.metric_history;
if (!history || !history.steps?.length) {
return { lossHistory: null, lrHistory: null };
}
const steps = history.steps;
const lossHistory = history.loss ? toSeries(steps, history.loss) : null;
const lrHistory = history.lr ? toSeries(steps, history.lr) : null;
return { lossHistory, lrHistory };
}
export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => ({
...initialState,
setHydrating: (value) => set({ isHydrating: value }),
setHasHydrated: (value) => set({ hasHydrated: value }),
setStarting: (value) => set({ isStarting: value }),
setStartError: (value) => set({ startError: value }),
setSseConnected: (value) => set({ sseConnected: value }),
setLastEventId: (value) => set({ lastEventId: value }),
resetRuntime: () =>
set({
...initialState,
lossHistory: [],
lrHistory: [],
gradNormHistory: [],
}),
setStartQueued: (jobId, message) =>
set({
jobId,
message,
error: null,
startError: null,
phase: "configuring",
isStarting: false,
}),
setRuntimeError: (message) =>
set({
error: message,
phase: "error",
isStarting: false,
startError: null,
sseConnected: false,
}),
applyStatus: (payload) =>
set((state) => {
const metricHistory = applyMetricHistoryFromStatus(payload);
const detailStep = payload.details?.step;
const detailTotal = payload.details?.total_steps;
const detailLoss = payload.details?.loss;
const detailLr = payload.details?.learning_rate;
const detailEpoch = payload.details?.epoch;
return {
...state,
jobId: payload.job_id || state.jobId,
phase: payload.phase,
isTrainingRunning: payload.is_training_running,
message: payload.message,
error: payload.error,
startError: null,
currentStep:
typeof detailStep === "number" ? Math.max(detailStep, 0) : state.currentStep,
totalSteps:
typeof detailTotal === "number"
? Math.max(detailTotal, 0)
: state.totalSteps,
currentLoss:
typeof detailLoss === "number" ? detailLoss : state.currentLoss,
currentLearningRate:
typeof detailLr === "number" ? detailLr : state.currentLearningRate,
currentEpoch:
typeof detailEpoch === "number" ? detailEpoch : state.currentEpoch,
lossHistory: metricHistory.lossHistory ?? state.lossHistory,
lrHistory: metricHistory.lrHistory ?? state.lrHistory,
};
}),
applyMetrics: (payload: TrainingMetricsResponse) =>
set((state) => {
const lossHistory = toSeries(payload.step_history, payload.loss_history);
const lrHistory = toSeries(payload.step_history, payload.lr_history);
const latestStep =
payload.current_step ??
(payload.step_history.length > 0
? payload.step_history[payload.step_history.length - 1]
: null);
return {
...state,
lossHistory: lossHistory.length > 0 ? lossHistory : state.lossHistory,
lrHistory: lrHistory.length > 0 ? lrHistory : state.lrHistory,
currentStep:
typeof latestStep === "number"
? Math.max(latestStep, state.currentStep)
: state.currentStep,
currentLoss:
typeof payload.current_loss === "number"
? payload.current_loss
: state.currentLoss,
currentLearningRate:
typeof payload.current_lr === "number"
? payload.current_lr
: state.currentLearningRate,
};
}),
applyProgress: (payload: TrainingProgressPayload, eventId?: number) =>
set((state) => {
const step = Math.max(payload.step, 0);
return {
...state,
jobId: payload.job_id || state.jobId,
currentStep: step,
totalSteps: Math.max(payload.total_steps, state.totalSteps),
currentLoss: payload.loss,
currentLearningRate: payload.learning_rate,
progressPercent: payload.progress_percent,
currentEpoch: payload.epoch ?? state.currentEpoch,
elapsedSeconds: payload.elapsed_seconds,
etaSeconds: payload.eta_seconds,
currentGradNorm: payload.grad_norm,
currentNumTokens: payload.num_tokens,
firstStepReceived: state.firstStepReceived || step > 0,
lastEventId: typeof eventId === "number" ? eventId : state.lastEventId,
lossHistory:
step > 0
? upsertPoint(state.lossHistory, step, payload.loss)
: state.lossHistory,
lrHistory:
step > 0
? upsertPoint(state.lrHistory, step, payload.learning_rate)
: state.lrHistory,
gradNormHistory:
step > 0 && typeof payload.grad_norm === "number"
? upsertPoint(state.gradNormHistory, step, payload.grad_norm)
: state.gradNormHistory,
};
}),
}));
export function shouldShowTrainingView(state: TrainingRuntimeStore): boolean {
return (
state.phase !== "idle" ||
state.isTrainingRunning ||
state.isStarting ||
state.lossHistory.length > 0 ||
state.currentStep > 0
);
}

53
studio/frontend/src/features/training/types/api.ts Normal file
View file

@ -0,0 +1,53 @@
export interface TrainingStartRequest {
model_name: string;
training_type: string;
hf_token: string | null;
load_in_4bit: boolean;
max_seq_length: number;
hf_dataset: string | null;
local_datasets: string[];
format_type: string;
num_epochs: number;
learning_rate: string;
batch_size: number;
gradient_accumulation_steps: number;
warmup_steps: number | null;
warmup_ratio: number | null;
max_steps: number | null;
save_steps: number;
weight_decay: number;
random_seed: number;
packing: boolean;
optim: string;
lr_scheduler_type: string;
use_lora: boolean;
lora_r: number;
lora_alpha: number;
lora_dropout: number;
target_modules: string[];
gradient_checkpointing: string;
use_rslora: boolean;
use_loftq: boolean;
train_on_completions: boolean;
finetune_vision_layers: boolean;
finetune_language_layers: boolean;
finetune_attention_modules: boolean;
finetune_mlp_modules: boolean;
enable_wandb: boolean;
wandb_token: string | null;
wandb_project: string | null;
enable_tensorboard: boolean;
tensorboard_dir: string | null;
}
export interface TrainingStartResponse {
job_id: string;
status: "queued" | "error";
message: string;
error: string | null;
}
export interface TrainingStopResponse {
status: "stopped" | "idle";
message: string;
}

96
studio/frontend/src/features/training/types/config.ts Normal file
View file

@ -0,0 +1,96 @@
import type {
DatasetFormat,
DatasetSource,
GradientCheckpointing,
ModelType,
StepNumber,
TrainingMethod,
} from "@/types/training";
export type LoraVariant = "lora" | "rslora" | "loftq";
export interface TrainingConfigState {
currentStep: StepNumber;
modelType: ModelType | null;
selectedModel: string | null;
trainingMethod: TrainingMethod;
hfToken: string;
datasetSource: DatasetSource;
datasetFormat: DatasetFormat;
dataset: string | null;
uploadedFile: string | null;
epochs: number;
contextLength: number;
learningRate: number;
loraRank: number;
loraAlpha: number;
loraDropout: number;
loraVariant: LoraVariant;
batchSize: number;
gradientAccumulation: number;
weightDecay: number;
warmupSteps: number;
maxSteps: number;
saveSteps: number;
packing: boolean;
trainOnCompletions: boolean;
gradientCheckpointing: GradientCheckpointing;
randomSeed: number;
enableWandb: boolean;
wandbToken: string;
wandbProject: string;
enableTensorboard: boolean;
tensorboardDir: string;
logFrequency: number;
finetuneVisionLayers: boolean;
finetuneLanguageLayers: boolean;
finetuneAttentionModules: boolean;
finetuneMLPModules: boolean;
targetModules: string[];
}
export interface TrainingConfigActions {
setStep: (step: StepNumber) => void;
nextStep: () => void;
prevStep: () => void;
setModelType: (type: ModelType) => void;
setSelectedModel: (model: string | null) => void;
setTrainingMethod: (method: TrainingMethod) => void;
setHfToken: (token: string) => void;
setDatasetSource: (source: DatasetSource) => void;
setDatasetFormat: (format: DatasetFormat) => void;
setDataset: (dataset: string | null) => void;
setUploadedFile: (file: string | null) => void;
setEpochs: (epochs: number) => void;
setContextLength: (length: number) => void;
setLearningRate: (rate: number) => void;
setLoraRank: (rank: number) => void;
setLoraAlpha: (alpha: number) => void;
setLoraDropout: (dropout: number) => void;
setLoraVariant: (variant: LoraVariant) => void;
setBatchSize: (value: number) => void;
setGradientAccumulation: (value: number) => void;
setWeightDecay: (value: number) => void;
setWarmupSteps: (value: number) => void;
setMaxSteps: (value: number) => void;
setSaveSteps: (value: number) => void;
setPacking: (value: boolean) => void;
setTrainOnCompletions: (value: boolean) => void;
setGradientCheckpointing: (value: GradientCheckpointing) => void;
setRandomSeed: (value: number) => void;
setEnableWandb: (value: boolean) => void;
setWandbToken: (value: string) => void;
setWandbProject: (value: string) => void;
setEnableTensorboard: (value: boolean) => void;
setTensorboardDir: (value: string) => void;
setLogFrequency: (value: number) => void;
setFinetuneVisionLayers: (value: boolean) => void;
setFinetuneLanguageLayers: (value: boolean) => void;
setFinetuneAttentionModules: (value: boolean) => void;
setFinetuneMLPModules: (value: boolean) => void;
setTargetModules: (value: string[]) => void;
canProceed: () => boolean;
reset: () => void;
}
export type TrainingConfigStore = TrainingConfigState & TrainingConfigActions;

102
studio/frontend/src/features/training/types/runtime.ts Normal file
View file

@ -0,0 +1,102 @@
export type TrainingPhase =
| "idle"
| "loading_model"
| "loading_dataset"
| "configuring"
| "training"
| "completed"
| "error"
| "stopped";
export interface TrainingStatusResponse {
job_id: string;
phase: TrainingPhase;
is_training_running: boolean;
message: string;
error: string | null;
details?: {
epoch?: number;
step?: number;
total_steps?: number;
loss?: number;
learning_rate?: number;
} | null;
metric_history?: {
steps?: number[];
loss?: number[];
lr?: number[];
} | null;
}
export interface TrainingMetricsResponse {
loss_history: number[];
lr_history: number[];
step_history: number[];
current_loss: number | null;
current_lr: number | null;
current_step: number | null;
}
export interface TrainingProgressPayload {
job_id: string;
step: number;
total_steps: number;
loss: number;
learning_rate: number;
progress_percent: number;
epoch: number | null;
elapsed_seconds: number | null;
eta_seconds: number | null;
grad_norm: number | null;
num_tokens: number | null;
}
export interface TrainingSeriesPoint {
step: number;
value: number;
}
export interface TrainingRuntimeState {
jobId: string | null;
phase: TrainingPhase;
isTrainingRunning: boolean;
message: string;
error: string | null;
isHydrating: boolean;
hasHydrated: boolean;
isStarting: boolean;
startError: string | null;
sseConnected: boolean;
firstStepReceived: boolean;
lastEventId: number | null;
currentStep: number;
totalSteps: number;
currentEpoch: number;
currentLoss: number;
currentLearningRate: number;
progressPercent: number;
elapsedSeconds: number | null;
etaSeconds: number | null;
currentGradNorm: number | null;
currentNumTokens: number | null;
lossHistory: TrainingSeriesPoint[];
lrHistory: TrainingSeriesPoint[];
gradNormHistory: TrainingSeriesPoint[];
}
export interface TrainingRuntimeActions {
setHydrating: (value: boolean) => void;
setHasHydrated: (value: boolean) => void;
setStarting: (value: boolean) => void;
setStartError: (value: string | null) => void;
setSseConnected: (value: boolean) => void;
setLastEventId: (value: number | null) => void;
resetRuntime: () => void;
applyStatus: (payload: TrainingStatusResponse) => void;
applyMetrics: (payload: TrainingMetricsResponse) => void;
applyProgress: (payload: TrainingProgressPayload, eventId?: number) => void;
setStartQueued: (jobId: string, message: string) => void;
setRuntimeError: (message: string) => void;
}
export type TrainingRuntimeStore = TrainingRuntimeState & TrainingRuntimeActions;

24
studio/frontend/src/hooks/use-hf-model-search.ts
View file

@ -23,6 +23,28 @@ const EXCLUDED_TAGS = new Set([
"ctranslate2",
]);
function withPopularitySort(
input: Parameters<typeof fetch>[0],
init?: Parameters<typeof fetch>[1],
): ReturnType<typeof fetch> {
const rawUrl =
typeof input === "string"
? input
: input instanceof URL
? input.toString()
: input.url;
const url = new URL(rawUrl);
if (!url.searchParams.has("sort")) {
url.searchParams.set("sort", "downloads");
}
if (!url.searchParams.has("direction")) {
url.searchParams.set("direction", "-1");
}
return fetch(url, init);
}
function mapModel(raw: unknown): HfModelResult | null {
const m = raw as {
name: string;
@ -53,10 +75,10 @@ export function useHfModelSearch(
listModels({
search: {
...(query.trim() ? { query } : { owner: "unsloth" }),
tags: ["transformers"],
...(task ? { task } : {}),
},
additionalFields: ["safetensors", "tags"],
fetch: withPopularitySort,
...(accessToken ? { credentials: { accessToken } } : {}),
}) as AsyncGenerator<unknown>,
[query, task, accessToken],

108
studio/frontend/src/stores/training.ts
View file

@ -1,106 +1,4 @@
import { DEFAULT_HYPERPARAMS } from "@/config/training";
import type { StepNumber, WizardActions, WizardState } from "@/types/training";
import { create } from "zustand";
import { useTrainingConfigStore } from "@/features/training";
const MIN_STEP: StepNumber = 1;
const MAX_STEP: StepNumber = 5;
const initialState: WizardState = {
isTraining: false,
trainingMetrics: null,
currentStep: MIN_STEP,
modelType: null,
selectedModel: null,
trainingMethod: "qlora",
hfToken: "",
datasetSource: "huggingface",
datasetFormat: "auto",
dataset: null,
uploadedFile: null,
...DEFAULT_HYPERPARAMS,
};
function clampStep(step: number): StepNumber {
return Math.min(MAX_STEP, Math.max(MIN_STEP, step)) as StepNumber;
}
function canProceedForStep(state: WizardState): boolean {
switch (state.currentStep) {
case 1:
return state.modelType !== null;
case 2:
return state.selectedModel !== null;
case 3: {
if (state.datasetSource === "upload") {
return state.uploadedFile !== null;
}
return state.dataset !== null;
}
case 4:
case 5:
return true;
default:
return false;
}
}
export const useWizardStore = create<WizardState & WizardActions>(
(set, get) => ({
...initialState,
setStep: (step) => set({ currentStep: step }),
nextStep: () => {
const { currentStep } = get();
set({ currentStep: clampStep(currentStep + 1) });
},
prevStep: () => {
const { currentStep } = get();
set({ currentStep: clampStep(currentStep - 1) });
},
setModelType: (type) => set({ modelType: type, selectedModel: null }),
setSelectedModel: (model) => set({ selectedModel: model }),
setTrainingMethod: (method) => set({ trainingMethod: method }),
setHfToken: (token) => set({ hfToken: token }),
setDatasetSource: (source) => set({ datasetSource: source }),
setDatasetFormat: (format) => set({ datasetFormat: format }),
setDataset: (dataset) => set({ dataset }),
setUploadedFile: (file) => set({ uploadedFile: file }),
setEpochs: (epochs) => set({ epochs }),
setContextLength: (length) => set({ contextLength: length }),
setLearningRate: (rate) => set({ learningRate: rate }),
setLoraRank: (rank) => set({ loraRank: rank }),
setLoraAlpha: (alpha) => set({ loraAlpha: alpha }),
setLoraDropout: (dropout) => set({ loraDropout: dropout }),
setLoraVariant: (v) => set({ loraVariant: v }),
setBatchSize: (v) => set({ batchSize: v }),
setGradientAccumulation: (v) => set({ gradientAccumulation: v }),
setWeightDecay: (v) => set({ weightDecay: v }),
setWarmupSteps: (v) => set({ warmupSteps: v }),
setMaxSteps: (v) => set({ maxSteps: v }),
setSaveSteps: (v) => set({ saveSteps: v }),
setPacking: (v) => set({ packing: v }),
setTrainOnCompletions: (v) => set({ trainOnCompletions: v }),
setGradientCheckpointing: (v) => set({ gradientCheckpointing: v }),
setRandomSeed: (v) => set({ randomSeed: v }),
setEnableWandb: (v) => set({ enableWandb: v }),
setWandbToken: (v) => set({ wandbToken: v }),
setWandbProject: (v) => set({ wandbProject: v }),
setEnableTensorboard: (v) => set({ enableTensorboard: v }),
setTensorboardDir: (v) => set({ tensorboardDir: v }),
setLogFrequency: (v) => set({ logFrequency: v }),
setFinetuneVisionLayers: (v) => set({ finetuneVisionLayers: v }),
setFinetuneLanguageLayers: (v) => set({ finetuneLanguageLayers: v }),
setFinetuneAttentionModules: (v) => set({ finetuneAttentionModules: v }),
setFinetuneMLPModules: (v) => set({ finetuneMLPModules: v }),
setTargetModules: (v) => set({ targetModules: v }),
setIsTraining: (v) => set({ isTraining: v }),
setTrainingMetrics: (v) => set({ trainingMetrics: v }),
canProceed: () => canProceedForStep(get()),
reset: () => set(initialState),
}),
);
export const useWizardStore = useTrainingConfigStore;
export { useTrainingConfigStore };

25
studio/frontend/src/types/training.ts
View file

@ -9,30 +9,7 @@ export type DatasetSource = "huggingface" | "upload";
export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt";
export type GradientCheckpointing = "none" | "true" | "unsloth";
export interface TrainingMetrics {
currentStep: number;
totalSteps: number;
currentEpoch: number;
totalEpochs: number;
currentLoss: number;
currentLR: number;
gradNorm: number;
samplesPerSecond: number;
lossHistory: { step: number; loss: number }[];
lrHistory: { step: number; lr: number }[];
gradNormHistory: { step: number; gradNorm: number }[];
gpuUtil: number;
gpuTemp: number;
gpuVramUsed: number;
gpuVramTotal: number;
gpuPower: number;
elapsed: number;
status: "training" | "warmup" | "saving";
}
export interface WizardState {
isTraining: boolean;
trainingMetrics: TrainingMetrics | null;
currentStep: StepNumber;
modelType: ModelType | null;
selectedModel: string | null;
@ -112,8 +89,6 @@ export interface WizardActions {
setFinetuneAttentionModules: (v: boolean) => void;
setFinetuneMLPModules: (v: boolean) => void;
setTargetModules: (v: string[]) => void;
setIsTraining: (v: boolean) => void;
setTrainingMetrics: (v: TrainingMetrics | null) => void;
canProceed: () => boolean;
reset: () => void;
}

0
studio/tests/__init__.py Normal file
View file

282
studio/tests/test_openai_completions.py Normal file
View file

@ -0,0 +1,282 @@
"""
Tests for the OpenAI-compatible /chat/completions endpoint.
Validates:
- Streaming: SSE chunk format matches OpenAI spec
- Non-streaming: single JSON ChatCompletion response
- System prompt extraction from messages array
- Request validation (no messages, missing model, etc.)
- Response headers for proxy compatibility
All tests mock the inference backend and bypass auth.
"""
import sys
import json
from pathlib import Path
from unittest.mock import MagicMock, patch, PropertyMock
import pytest
# ── Path setup ────────────────────────────────────────────────────
_backend_root = Path(__file__).resolve().parent.parent / "backend"
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
from fastapi.testclient import TestClient
from main import app
# ── Fixtures ──────────────────────────────────────────────────────
def _make_mock_backend(*, tokens: list[str] | None = None, active_model: str = "test-model"):
"""Build a mock InferenceBackend that yields preset tokens."""
backend = MagicMock()
backend.active_model_name = active_model
backend.models = {active_model: {"is_vision": False}}
def fake_generate(**kwargs):
for t in (tokens or ["Hello", "Hello world", "Hello world!"]):
yield t
backend.generate_chat_response = MagicMock(side_effect=fake_generate)
backend.reset_generation_state = MagicMock()
return backend
def _parse_sse_data(raw: str) -> list[dict | str]:
"""Extract `data:` payloads from raw SSE text. Returns dicts or raw strings."""
results = []
for line in raw.split("\n"):
if line.startswith("data: "):
payload = line[len("data: "):]
if payload == "[DONE]":
results.append("[DONE]")
else:
try:
results.append(json.loads(payload))
except json.JSONDecodeError:
results.append(payload)
return results
@pytest.fixture()
def client():
yield TestClient(app)
# =====================================================================
# Streaming tests
# =====================================================================
class TestStreamingChunkFormat:
"""Each SSE chunk must match the OpenAI chat.completion.chunk schema."""
def test_chunks_have_required_fields(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["Hi"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={
"messages": [{"role": "user", "content": "Hello"}],
"stream": True,
},
)
assert resp.status_code == 200
chunks = _parse_sse_data(resp.text)
# Filter to actual chunk dicts (not [DONE])
json_chunks = [c for c in chunks if isinstance(c, dict) and "choices" in c]
assert len(json_chunks) >= 2 # role chunk + content chunk(s) + final
for chunk in json_chunks:
assert "id" in chunk
assert chunk["object"] == "chat.completion.chunk"
assert "created" in chunk
assert "model" in chunk
assert len(chunk["choices"]) == 1
assert "delta" in chunk["choices"][0]
def test_first_chunk_has_role(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["Hi"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={"messages": [{"role": "user", "content": "Hello"}]},
)
chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c]
first = chunks[0]
assert first["choices"][0]["delta"].get("role") == "assistant"
def test_last_chunk_has_stop_finish_reason(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["Done"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={"messages": [{"role": "user", "content": "Hello"}]},
)
chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c]
last = chunks[-1]
assert last["choices"][0]["finish_reason"] == "stop"
# Delta should be empty on the final chunk
assert last["choices"][0]["delta"].get("content") is None
def test_stream_ends_with_done(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["x"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={"messages": [{"role": "user", "content": "Hello"}]},
)
all_data = _parse_sse_data(resp.text)
assert all_data[-1] == "[DONE]"
def test_consistent_id_across_chunks(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["a", "b", "c"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={"messages": [{"role": "user", "content": "Hello"}]},
)
chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c]
ids = set(c["id"] for c in chunks)
assert len(ids) == 1, "All chunks should share the same completion ID"
class TestStreamingHeaders:
"""Verify response headers for SSE proxy compatibility."""
def test_headers(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["x"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={"messages": [{"role": "user", "content": "Hello"}]},
)
assert resp.headers["content-type"].startswith("text/event-stream")
assert resp.headers.get("cache-control") == "no-cache"
assert resp.headers.get("x-accel-buffering") == "no"
# =====================================================================
# Non-streaming tests
# =====================================================================
class TestNonStreaming:
"""When stream=false, return a single ChatCompletion JSON object."""
def test_returns_json_object(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["Full response text"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={
"messages": [{"role": "user", "content": "Hello"}],
"stream": False,
},
)
assert resp.status_code == 200
body = resp.json()
assert body["object"] == "chat.completion"
assert body["choices"][0]["message"]["role"] == "assistant"
assert body["choices"][0]["message"]["content"] == "Full response text"
assert body["choices"][0]["finish_reason"] == "stop"
def test_non_streaming_has_model(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["x"], active_model="my-model")
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={
"messages": [{"role": "user", "content": "Hi"}],
"stream": False,
},
)
body = resp.json()
assert body["model"] == "my-model"
# =====================================================================
# System prompt extraction
# =====================================================================
class TestSystemPromptExtraction:
"""System messages should be extracted and passed as system_prompt."""
def test_system_message_extracted(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["ok"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
client.post(
"/api/inference/chat/completions",
json={
"messages": [
{"role": "system", "content": "You are a pirate."},
{"role": "user", "content": "Hello"},
],
"stream": False,
},
)
# Check that generate_chat_response was called with the correct system_prompt
call_kwargs = mock_backend.generate_chat_response.call_args[1]
assert call_kwargs["system_prompt"] == "You are a pirate."
# System message should NOT be in the chat_messages list
assert all(m["role"] != "system" for m in call_kwargs["messages"])
def test_default_system_prompt_when_none(self, client: TestClient):
mock_backend = _make_mock_backend(tokens=["ok"])
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
client.post(
"/api/inference/chat/completions",
json={
"messages": [{"role": "user", "content": "Hello"}],
"stream": False,
},
)
call_kwargs = mock_backend.generate_chat_response.call_args[1]
assert call_kwargs["system_prompt"] == "You are a helpful AI assistant."
# =====================================================================
# Error handling
# =====================================================================
class TestErrorHandling:
"""Validate error responses for bad requests."""
def test_no_model_loaded(self, client: TestClient):
mock_backend = _make_mock_backend()
mock_backend.active_model_name = None
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={"messages": [{"role": "user", "content": "Hi"}]},
)
assert resp.status_code == 400
assert "No model loaded" in resp.json()["detail"]
def test_only_system_messages_rejected(self, client: TestClient):
mock_backend = _make_mock_backend()
with patch("routes.inference.get_inference_backend", return_value=mock_backend):
resp = client.post(
"/api/inference/chat/completions",
json={
"messages": [{"role": "system", "content": "You are a bot."}],
},
)
assert resp.status_code == 400
assert "non-system message" in resp.json()["detail"]

40
studio/tests/test_remote_lora_detection.py Normal file
View file

@ -0,0 +1,40 @@
"""
Test remote LoRA adapter detection via HuggingFace Hub API.
Verifies that we can detect whether a remote HF model is a LoRA adapter
by checking for adapter_config.json in the repo file listing.
"""
import pytest
from huggingface_hub import model_info
def is_remote_lora_adapter(model_id: str, hf_token: str = None) -> bool:
"""
Check if a remote HuggingFace model is a LoRA adapter
by looking for adapter_config.json in its repo files.
"""
try:
info = model_info(model_id, token=hf_token)
filenames = [s.rfilename for s in info.siblings]
return "adapter_config.json" in filenames
except Exception:
return False
class TestRemoteLoRADetection:
"""Test remote LoRA adapter detection via HF Hub API."""
def test_lora_adapter_detected(self):
"""edbeeching/llama-se-rl-adapter is a known LoRA adapter on HF."""
result = is_remote_lora_adapter("edbeeching/llama-se-rl-adapter")
assert result is True, "Expected edbeeching/llama-se-rl-adapter to be detected as a LoRA adapter"
def test_base_model_not_detected_as_lora(self):
"""google/gemma-3-4b-it is a full base model, not a LoRA adapter."""
result = is_remote_lora_adapter("google/gemma-3-4b-it")
assert result is False, "Expected google/gemma-3-4b-it to NOT be detected as a LoRA adapter"
def test_nonexistent_model_returns_false(self):
"""A nonexistent model should return False, not raise."""
result = is_remote_lora_adapter("this-org-does-not-exist/fake-model-12345")
assert result is False, "Expected nonexistent model to return False"

320
studio/tests/test_sse_progress.py Normal file
View file

@ -0,0 +1,320 @@
"""
Tests for the SSE training progress endpoint and status fallback.
Validates:
- SSE spec compliance: `retry:`, `id:`, `event:` fields
- Named event types: progress, heartbeat, complete, error
- Last-Event-ID reconnection and history replay
- /status metric_history fallback (Option B)
All tests mock the training backend and bypass auth.
"""
import sys
from pathlib import Path
from typing import Optional
from unittest.mock import MagicMock, patch, PropertyMock
import re
import pytest
# ── Path setup ────────────────────────────────────────────────────
# Add backend root so bare `from routes…`, `from models…` etc. resolve.
_backend_root = Path(__file__).resolve().parent.parent / "backend"
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
from fastapi.testclient import TestClient
from main import app
from auth.authentication import get_current_subject
# ── Fixtures ──────────────────────────────────────────────────────
def _bypass_auth():
"""Dependency override that skips real JWT validation."""
return "test-user"
def _make_mock_backend(
*,
is_active: bool = False,
step_history: list | None = None,
loss_history: list | None = None,
lr_history: list | None = None,
total_steps: int = 100,
epoch: int | None = 1,
job_id: str = "job_test_001",
):
"""Build a lightweight mock that quacks like TrainingBackend."""
backend = MagicMock()
backend.current_job_id = job_id
backend.step_history = step_history or []
backend.loss_history = loss_history or []
backend.lr_history = lr_history or []
backend.is_training_active.return_value = is_active
backend._training_thread = None
# trainer.training_progress / get_training_progress()
tp = MagicMock()
tp.total_steps = total_steps
tp.epoch = epoch
tp.step = step_history[-1] if step_history else 0
tp.loss = loss_history[-1] if loss_history else 0.0
tp.learning_rate = lr_history[-1] if lr_history else 0.0
tp.status_message = "Training..."
tp.error = None
tp.is_completed = not is_active and bool(step_history)
backend.trainer = MagicMock()
backend.trainer.training_progress = tp
backend.trainer.get_training_progress.return_value = tp
return backend
@pytest.fixture()
def client():
"""TestClient with auth bypassed."""
app.dependency_overrides[get_current_subject] = _bypass_auth
yield TestClient(app)
app.dependency_overrides.clear()
# ── SSE Parsing Helpers ───────────────────────────────────────────
def parse_sse_events(raw: str) -> list[dict]:
"""
Parse raw SSE text into a list of event dicts.
Each dict has optional keys: 'id', 'event', 'data', 'retry'.
"""
events: list[dict] = []
current: dict = {}
for line in raw.split("\n"):
if line.startswith("retry:"):
# retry is a standalone directive, not part of a normal event
events.append({"retry": line.split(":", 1)[1].strip()})
continue
if line.startswith("id:"):
current["id"] = line.split(":", 1)[1].strip()
elif line.startswith("event:"):
current["event"] = line.split(":", 1)[1].strip()
elif line.startswith("data:"):
current["data"] = line.split(":", 1)[1].strip()
elif line == "" and current:
events.append(current)
current = {}
if current:
events.append(current)
return events
# =====================================================================
# Option A — /api/train/progress (SSE)
# =====================================================================
class TestSSERetryDirective:
"""The first thing the stream emits must be `retry: 3000`."""
def test_retry_is_first_event(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
assert resp.status_code == 200
assert resp.headers["content-type"].startswith("text/event-stream")
events = parse_sse_events(resp.text)
assert len(events) >= 1
assert events[0] == {"retry": "3000"}
class TestSSEEventFields:
"""Every non-retry event must include `id:`, `event:`, and `data:` fields."""
def test_events_have_id_and_event_type(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=False,
step_history=[1, 2, 3],
loss_history=[2.0, 1.5, 1.0],
lr_history=[1e-4, 1e-4, 1e-4],
total_steps=3,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
data_events = [e for e in events if "data" in e]
assert len(data_events) >= 1
for evt in data_events:
assert "id" in evt, f"Missing `id:` field in event: {evt}"
assert "event" in evt, f"Missing `event:` field in event: {evt}"
assert "data" in evt
class TestSSENamedEventTypes:
"""Events use the correct named types: progress, complete, heartbeat, error."""
def test_idle_sends_progress_then_complete(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=False,
step_history=[10],
loss_history=[1.5],
lr_history=[1e-4],
total_steps=10,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
data_events = [e for e in events if "event" in e and e.get("event") != "retry"]
event_types = [e["event"] for e in data_events]
assert "progress" in event_types
assert "complete" in event_types
def test_no_history_sends_complete(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
data_events = [e for e in events if "event" in e]
assert any(e["event"] == "complete" for e in data_events)
class TestSSELastEventIDResume:
"""When `Last-Event-ID` header is sent, the server replays missed steps."""
def test_replays_steps_after_last_event_id(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=False,
step_history=[1, 2, 3, 4, 5],
loss_history=[2.5, 2.0, 1.5, 1.2, 1.0],
lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4],
total_steps=5,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get(
"/api/train/progress",
headers={"Last-Event-ID": "2"},
)
events = parse_sse_events(resp.text)
# Filter to progress events (replayed ones)
progress_events = [e for e in events if e.get("event") == "progress"]
# Steps 3, 4, 5 should have been replayed
replayed_ids = [int(e["id"]) for e in progress_events]
assert 3 in replayed_ids
assert 4 in replayed_ids
assert 5 in replayed_ids
# Steps 1, 2 should NOT be replayed
assert 1 not in replayed_ids
assert 2 not in replayed_ids
def test_no_replay_without_header(self, client: TestClient):
"""Without Last-Event-ID, should start fresh (initial progress event)."""
mock_backend = _make_mock_backend(
is_active=False,
step_history=[1, 2, 3],
loss_history=[2.0, 1.5, 1.0],
lr_history=[1e-4, 1e-4, 1e-4],
total_steps=3,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
progress_events = [e for e in events if e.get("event") == "progress"]
# Should have initial step=0 progress event
assert any(e.get("id") == "0" for e in progress_events)
def test_invalid_last_event_id_treated_as_fresh(self, client: TestClient):
"""Non-integer Last-Event-ID should be ignored gracefully."""
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get(
"/api/train/progress",
headers={"Last-Event-ID": "not-a-number"},
)
assert resp.status_code == 200
events = parse_sse_events(resp.text)
# Should still work — treated as a fresh connection
assert any(e.get("event") == "progress" or e.get("event") == "complete" for e in events)
class TestSSEResponseHeaders:
"""Verify SSE response headers for proxy compatibility."""
def test_headers(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
assert resp.headers["content-type"].startswith("text/event-stream")
assert resp.headers.get("cache-control") == "no-cache"
assert resp.headers.get("x-accel-buffering") == "no"
# =====================================================================
# Option B — /api/train/status (metric_history fallback)
# =====================================================================
class TestStatusMetricHistory:
"""The /status endpoint returns metric_history for chart recovery."""
def test_metric_history_populated_when_history_exists(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=True,
step_history=[1, 2, 3, 4, 5],
loss_history=[2.5, 2.0, 1.5, 1.2, 1.0],
lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4],
total_steps=10,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/status")
assert resp.status_code == 200
body = resp.json()
assert "metric_history" in body
mh = body["metric_history"]
assert mh is not None
assert mh["steps"] == [1, 2, 3, 4, 5]
assert mh["loss"] == [2.5, 2.0, 1.5, 1.2, 1.0]
assert mh["lr"] == [1e-4, 1e-4, 1e-4, 1e-4, 1e-4]
def test_metric_history_null_when_no_history(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/status")
assert resp.status_code == 200
body = resp.json()
assert body["metric_history"] is None
def test_status_still_returns_phase_and_details(self, client: TestClient):
"""Ensure adding metric_history didn't break existing fields."""
mock_backend = _make_mock_backend(
is_active=True,
step_history=[5],
loss_history=[1.5],
lr_history=[1e-4],
total_steps=100,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/status")
body = resp.json()
assert body["phase"] == "training"
assert body["is_training_running"] is True
assert body["job_id"] == "job_test_001"
assert "details" in body