diff --git a/.gitignore b/.gitignore index 17924a1137..4e816398b8 100755 --- a/.gitignore +++ b/.gitignore @@ -36,6 +36,7 @@ Thumbs.db resources/ tmp/ auth.db +studio/frontend/package-lock.json # Local working docs **/CLAUDE.md diff --git a/cli/commands/ui.py b/cli/commands/ui.py index aee2d21d08..6eaeeb793e 100644 --- a/cli/commands/ui.py +++ b/cli/commands/ui.py @@ -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, diff --git a/setup.sh b/setup.sh new file mode 100755 index 0000000000..e5a41d371f --- /dev/null +++ b/setup.sh @@ -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 < 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. diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 7b562c6cb0..a4a700d45b 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -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 diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index fe75d6a976..0fd1905ff4 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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 diff --git a/studio/backend/core/training/__init__.py b/studio/backend/core/training/__init__.py index 65bf4c3501..8fc2a6c721 100644 --- a/studio/backend/core/training/__init__.py +++ b/studio/backend/core/training/__init__.py @@ -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', ] diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index af5c3ebe01..52cf60af10 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -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 diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 9bd44400ac..fa8b3b5daf 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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 - } diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 81c2451ebe..530a9fa30b 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -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 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3917e38e36..64791b06f9 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 2afee72e89..2f55aeca7a 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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") diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 21ac5ac1eb..0ad5bb0ea8 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -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, diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 8e21acaa9e..e462bbfe18 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -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: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17d2011205..74cf37138f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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, → extract + 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)) + diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 9f6c298a4d..20a5e1e254 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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", } ) - diff --git a/studio/backend/run.py b/studio/backend/run.py index 6f93cba1d3..0959746a19 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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") diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index c16ea3fe63..a75f78d37c 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -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): diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index cadefa6ede..5882f41e2a 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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 diff --git a/studio/frontend/bun.lock b/studio/frontend/bun.lock index b27635a66a..9859b3aae5 100644 --- a/studio/frontend/bun.lock +++ b/studio/frontend/bun.lock @@ -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=="], diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 273f9e8ae9..9a10e8e639 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -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", diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index ce3af7ac54..5c2b93bcc9 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -10,7 +10,7 @@ export function AppProvider({ children }: AppProviderProps) { return ( {children} - + ); } diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 389dae424b..85bc1cacdb 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -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 && ( )} - - {currentModel?.name ?? "Select a model\u2026"} + + {currentModel?.name ?? "Select model..."} {currentModel?.description && ( - - {currentModel.description} - + {currentModel.description} )} + {children} + + ); +} + +function ModelRow({ + label, + meta, + selected, + onClick, +}: { + label: string; + meta?: string; + selected?: boolean; + onClick: () => void; +}) { + return ( + + ); +} + +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 ( +
+
+ + setQuery(event.target.value)} + placeholder="Search Hugging Face models" + className="h-9 pl-8 pr-8" + /> + {isLoading && ( + + )} +
+ +
+
+ {!showHfSection ? ( + <> + Recommended + {recommendedIds.length === 0 ? ( +
+ No default models. +
+ ) : ( + recommendedIds.map((id) => ( + onSelect(id, { source: "hub", isLora: false })} + /> + )) + )} + + ) : null} + + {showHfSection ? ( + <> + Hugging Face + {hfIds.length === 0 && !isLoading ? ( +
+ No matching models. +
+ ) : ( + hfIds.map((id) => ( + onSelect(id, { source: "hub", isLora: false })} + /> + )) + )} +
+ {isLoadingMore ? ( +
+ +
+ ) : null} + + ) : null} +
+
+
+ ); +} + +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(); + + 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 ( +
+
+ + setQuery(event.target.value)} + placeholder="Search local adapters" + className="h-9 pl-8" + /> +
+ +
+
+ {grouped.length === 0 ? ( +
No adapters found.
+ ) : ( + grouped.map(([baseModel, adapters], index) => ( +
+ {index > 0 ?
: null} + {baseModel} + {adapters.map((adapter) => ( + onSelect(adapter.id, { source: "lora", isLora: true })} + /> + ))} +
+ )) + )} +
+
+
+ ); +} + 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 ( - {models.map((model) => ( - - ))} + + + Hub models + Fine-tuned + + + + + + + + + + + + {hasSelection && onEject ? ( +
+ +
+ ) : null} ); } -function ModelSelectorItem({ - model, - isActive, - onSelect, - onEject, -}: { - model: ModelOption; - isActive: boolean; - onSelect: (id: string) => void; - onEject?: () => void; -}) { - return ( - - )} - - ); -} - -// --- 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(); + 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({ /> = ({ hideComposer, @@ -69,6 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ + !thread.isEmpty}> {!hideComposer && } @@ -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 ( +
+
+ + Warming up model... + +
+
+ ); +}; + const ThreadScrollToBottom: FC = () => { return ( @@ -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 ( - { + 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" > - + ); }; @@ -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 ( @@ -384,7 +433,9 @@ const EditComposer: FC = () => { } if (aui.thread().getState().isRunning) { + resendAfterCancelRef.current = true; aui.thread().cancelRun(); + return; } aui.composer().send(); }} diff --git a/studio/frontend/src/components/ui/data-table.tsx b/studio/frontend/src/components/ui/data-table.tsx new file mode 100644 index 0000000000..954d2ae0b7 --- /dev/null +++ b/studio/frontend/src/components/ui/data-table.tsx @@ -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 { + columns: ColumnDef[]; + data: TData[]; + className?: string; +} + +export function DataTable({ + columns, + data, + className, +}: DataTableProps) { + const [sorting, setSorting] = useState([]); + + const table = useReactTable({ + data, + columns, + getCoreRowModel: getCoreRowModel(), + getSortedRowModel: getSortedRowModel(), + onSortingChange: setSorting, + state: { sorting }, + }); + + return ( +
+ + + {table.getHeaderGroups().map((headerGroup) => ( + + {headerGroup.headers.map((header) => ( + + {header.isPlaceholder + ? null + : flexRender( + header.column.columnDef.header, + header.getContext(), + )} + + ))} + + ))} + + + {table.getRowModel().rows.length ? ( + table.getRowModel().rows.map((row, idx) => ( + + {row.getVisibleCells().map((cell) => ( + + {flexRender(cell.column.columnDef.cell, cell.getContext())} + + ))} + + )) + ) : ( + + + No results. + + + )} + +
+
+ ); +} diff --git a/studio/frontend/src/components/ui/terminal.tsx b/studio/frontend/src/components/ui/terminal.tsx new file mode 100644 index 0000000000..d9483eaca5 --- /dev/null +++ b/studio/frontend/src/components/ui/terminal.tsx @@ -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 + started: boolean +} { + const ref = useRef(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 ( +
+ {childElements.map((child, index) => + cloneElement(child, { + __sequence: sequence, + __isActive: !sequence || visibleIndex >= index, + __onDone: () => handleLineDone(index), + key: child.key ?? index, + } as InternalLineProps) + )} +
+ ) +} + +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 ( +
+ {children} +
+ ) +} + +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 ( +
+ {typed} +
+ ) +} diff --git a/studio/frontend/src/features/auth/index.ts b/studio/frontend/src/features/auth/index.ts index a6651b0c00..bd0b98f6be 100644 --- a/studio/frontend/src/features/auth/index.ts +++ b/studio/frontend/src/features/auth/index.ts @@ -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, diff --git a/studio/frontend/src/features/chat/adapter.ts b/studio/frontend/src/features/chat/adapter.ts deleted file mode 100644 index a922e834cd..0000000000 --- a/studio/frontend/src/features/chat/adapter.ts +++ /dev/null @@ -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[number]; -type RunMessages = Parameters[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(""); - 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(""); - 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("") && reasoningStart && !reasoningDuration) { - reasoningDuration = Math.round((Date.now() - reasoningStart) / 1000); - } - - if (parts.length > 0) { - yield { - content: parts, - metadata: { custom: { reasoningDuration } }, - }; - } - } - }, - }; -} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts new file mode 100644 index 0000000000..2b2356e847 --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -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[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 => + 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); + } + } + }, + }; +} diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts new file mode 100644 index 0000000000..72baf9a6f6 --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -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(response: Response): Promise { + 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 { + const response = await authFetch("/api/models/list"); + return parseJsonOrThrow(response); +} + +export async function listLoras(outputsDir = "./outputs"): Promise { + const query = new URLSearchParams({ outputs_dir: outputsDir }).toString(); + const response = await authFetch(`/api/models/loras?${query}`); + return parseJsonOrThrow(response); +} + +export async function getInferenceStatus(): Promise { + const response = await authFetch("/api/inference/status"); + return parseJsonOrThrow(response); +} + +export async function loadModel( + payload: LoadModelRequest, +): Promise { + const response = await authFetch("/api/inference/load", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + return parseJsonOrThrow(response); +} + +export async function unloadModel(payload: UnloadModelRequest): Promise { + const response = await authFetch("/api/inference/unload", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + await parseJsonOrThrow(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 { + 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/); + } + } +} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index f29d5766ae..445aa53f53 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -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 ( - +
@@ -233,33 +202,66 @@ function TopBarActions({ } export function ChatPage(): ReactElement { - const [view, setView] = useState({ mode: "single" }); + const [view, setView] = useState({ + mode: "single", + newThreadNonce: crypto.randomUUID(), + }); const [settingsOpen, setSettingsOpen] = useState(false); - const [inferenceParams, setInferenceParams] = useState( - 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( + () => + modelsFromStore.map((model) => ({ + id: model.id, + name: model.name, + description: model.description, + })), + [modelsFromStore], + ); + + const loraModels = useMemo( + () => + lorasFromStore.map((lora) => ({ + id: lora.id, + name: lora.name, + baseModel: lora.baseModel, + updatedAt: lora.updatedAt, + })), + [lorasFromStore], + ); + + useEffect(() => { + void refresh(); + }, [refresh]); return ( +
+ {modelsError && ( +
+ {modelsError} +
+ )}
); } diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 34a3511a0a..484b0d2b5e 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -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"}`} >
-
+
- - Advanced Configuration + + Configuration
-
+ {/* mt-4 matches the Playground sidebar gap (SidebarHeader py-3 + SidebarGroup pt-1) */} +
- - - - - {ENGINE_OPTIONS.map((o) => ( - - {o.label} - - ))} - - -
- - 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, + }; +} diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index c1daee9e25..b7eaaf83ae 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -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"; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index afe948387e..a94e25b67a 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -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 { + 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 { + 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 { + return Promise.resolve(); + } + + private async fileToBase64DataURL(file: File): Promise { + 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 => - useLocalRuntime(chatAdapter); +const chatAdapter = createOpenAIStreamAdapter(); + +function useRuntimeHook(): ReturnType { + 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 ( {initialThreadId && } + {!initialThreadId && newThreadNonce && ( + + )} {children} ); diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts new file mode 100644 index 0000000000..74efc57134 --- /dev/null +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -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; + 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((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: "", + }, + })), +})); diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx index b8ded1142d..e1645e77bb 100644 --- a/studio/frontend/src/features/chat/thread-sidebar.tsx +++ b/studio/frontend/src/features/chat/thread-sidebar.tsx @@ -99,11 +99,11 @@ export function ThreadSidebar({ return ( <> - - Playground + + Playground - + @@ -121,8 +121,8 @@ export function ThreadSidebar({ - - Your Chats + + Your Chats {items.map((item) => ( @@ -144,7 +144,7 @@ export function ThreadSidebar({ ))} {items.length === 0 && ( -

+

No threads yet

)} diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 45d4f9e195..b0dccab307 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -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 { diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts new file mode 100644 index 0000000000..f53bd6ca75 --- /dev/null +++ b/studio/frontend/src/features/chat/types/api.ts @@ -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[]; +} diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts new file mode 100644 index 0000000000..558f1a5464 --- /dev/null +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -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; +} diff --git a/studio/frontend/src/features/chat/utils/parse-assistant-content.ts b/studio/frontend/src/features/chat/utils/parse-assistant-content.ts new file mode 100644 index 0000000000..fdf4ce051e --- /dev/null +++ b/studio/frontend/src/features/chat/utils/parse-assistant-content.ts @@ -0,0 +1,54 @@ +import type { ChatModelRunResult } from "@assistant-ui/react"; + +type ContentPart = NonNullable[number]; + +const THINK_OPEN_TAG = ""; +const THINK_CLOSE_TAG = ""; + +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); +} diff --git a/studio/frontend/src/features/export/components/quant-picker.tsx b/studio/frontend/src/features/export/components/quant-picker.tsx index 9b6c070453..8136a46bb5 100644 --- a/studio/frontend/src/features/export/components/quant-picker.tsx +++ b/studio/frontend/src/features/export/components/quant-picker.tsx @@ -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 && ( )} {q.label} diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 376a963169..2e6b744f2d 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -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(null); const [exportMethod, setExportMethod] = useState(null); @@ -109,8 +105,8 @@ export function ExportPage() { return (
-
-
+
+

Export Model

diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 930a51ec99..65d802c77e 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -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, diff --git a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx index e7afe5f1fb..45f9c82878 100644 --- a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx @@ -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, diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index a3a2da37ec..d11ff6bac0 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -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, diff --git a/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx index a83b6486f8..8cf6e8bc07 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx @@ -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 = { 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, diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index 8dfeffba3b..ae6edd230a 100644 --- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx @@ -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, diff --git a/studio/frontend/src/features/onboarding/components/wizard-content.tsx b/studio/frontend/src/features/onboarding/components/wizard-content.tsx index 48b4371dd7..2f10be0fcd 100644 --- a/studio/frontend/src/features/onboarding/components/wizard-content.tsx +++ b/studio/frontend/src/features/onboarding/components/wizard-content.tsx @@ -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 = { }; 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]; diff --git a/studio/frontend/src/features/onboarding/components/wizard-footer.tsx b/studio/frontend/src/features/onboarding/components/wizard-footer.tsx index c167303ec6..2b3af5aa96 100644 --- a/studio/frontend/src/features/onboarding/components/wizard-footer.tsx +++ b/studio/frontend/src/features/onboarding/components/wizard-footer.tsx @@ -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, diff --git a/studio/frontend/src/features/onboarding/components/wizard-layout.tsx b/studio/frontend/src/features/onboarding/components/wizard-layout.tsx index 698b1c9052..3a38035fe1 100644 --- a/studio/frontend/src/features/onboarding/components/wizard-layout.tsx +++ b/studio/frontend/src/features/onboarding/components/wizard-layout.tsx @@ -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(null); const hasFiredRef = useRef(false); const isFinalStep = currentStep === STEPS.length; diff --git a/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx b/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx index 45390cb4e5..6cd5c36a6f 100644 --- a/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx +++ b/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx @@ -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 ( diff --git a/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx b/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx index d161c2cc5a..953b1699c7 100644 --- a/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx +++ b/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx @@ -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; diff --git a/studio/frontend/src/features/studio/sections/charts-content.tsx b/studio/frontend/src/features/studio/sections/charts-content.tsx index 3a4d2d1d29..ab1290e5eb 100644 --- a/studio/frontend/src/features/studio/sections/charts-content.tsx +++ b/studio/frontend/src/features/studio/sections/charts-content.tsx @@ -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(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 ( -
- {/* Training Loss */} +
- Training Loss + Training Loss @@ -155,12 +321,10 @@ export function ChartsContent({ - + @@ -168,19 +332,27 @@ export function ChartsContent({ formatStepTick(Number(value))} interval="preserveStartEnd" /> Number(value).toFixed(2)} /> )} {showSmoothed && ( )} @@ -232,18 +412,18 @@ export function ChartsContent({ - {/* Grad Norm */} - Gradient Norm + Gradient Norm @@ -251,19 +431,27 @@ export function ChartsContent({ formatStepTick(Number(value))} interval="preserveStartEnd" /> Number(value).toFixed(2)} /> } /> @@ -288,18 +480,15 @@ export function ChartsContent({ - {/* Learning Rate */} - Learning Rate + Learning Rate - + @@ -307,20 +496,27 @@ export function ChartsContent({ formatStepTick(Number(value))} interval="preserveStartEnd" /> v.toExponential(0)} + width={52} + tickFormatter={(value) => Number(value).toExponential(0)} /> `Step ${payload?.[0]?.payload?.step ?? ""}` } - formatter={(value) => [ - Number(value).toExponential(3), - "LR", - ]} + formatter={(value) => [Number(value).toExponential(3), "LR"]} /> } /> } /> @@ -349,10 +546,9 @@ export function ChartsContent({ - {/* Eval Loss (disabled/blurred) */} - + Eval Loss @@ -360,7 +556,7 @@ export function ChartsContent({
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 {
} > - + ); } diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx new file mode 100644 index 0000000000..80fbcb65ee --- /dev/null +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -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 | null; + detected_image_column?: string | null; + detected_text_column?: string | null; + preview_samples?: Record[] | 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 { + 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(null); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(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>[]>(() => { + if (!columns.length) return []; + return columns.map((colName) => ({ + accessorKey: colName, + header: () => ( + + {colName} + + ), + cell: ({ getValue }: { getValue: () => unknown }) => { + const value = getValue(); + const images = collectPreviewImages(value); + if (images.length > 0) { + return ( +
+ {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 ( + {`preview-${index}`} + ); + })} + {images.length > 4 && ( + + +{images.length - 4} more + + )} +
+ ); + } + + const text = formatCell(value); + if (!text) { + return ( + + -- + + ); + } + const full = + typeof value === "string" ? value : JSON.stringify(value); + return ( +

+ {text} +

+ ); + }, + })); + }, [columns]); + + return ( + + + {/* Header */} + +
+
+ +
+ + Dataset Preview + +
+ + + {/* Body */} +
+ {/* Loading */} + {loading && ( +
+
+ +
+

+ Loading preview... +

+
+ )} + + {/* Error */} + {error && ( +
+
+ +
+
+

{error}

+

+ Make sure the backend is running on port 8000. +

+
+
+ )} + + {/* Content */} + {!loading && !error && data && ( + <> + {/* Metadata card */} +
+ + + + + {columns.map((col) => ( + + {col} + + ))} + + } + /> +
+ + {/* Data table */} +
+ +
+ + {/* Footer */} +

+ Showing {rows.length} + {data.total_rows != null && + ` of ${data.total_rows.toLocaleString()}`}{" "} + rows +

+ + )} +
+ + + ); +} + +// --------------------------------------------------------------------------- +// Metadata row +// --------------------------------------------------------------------------- + +function MetaRow({ + label, + value, +}: { + label: string; + value: ReactNode; +}) { + return ( +
+ + {label}: + + {value} +
+ ); +} + +// --------------------------------------------------------------------------- +// 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; + 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)) { + stack.push(nested); + } + } + } + + return images; +} diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 6fd4d127c1..890faad2b3 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -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() {
+ ); } diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index b6a6968d64..fcaa319e1d 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -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, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index cc2953c005..8c753cde7e 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -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); diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 0656093e6a..b8656a7867 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -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 = { + idle: "Idle", + loading_model: "Loading model", + loading_dataset: "Loading dataset", + configuring: "Configuring", + training: "Training", + completed: "Completed", + error: "Error", + stopped: "Stopped", +}; + +const phaseColors: Record = { + 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(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 { } 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} > Stop @@ -138,24 +238,22 @@ export function ProgressSection(): ReactElement | null { } >
- {/* Left: Progress */}
- {statusLabels[metrics.status]} + {phaseLabel[runtime.phase]} - Epoch {metrics.currentEpoch.toFixed(2)} / {metrics.totalEpochs} + Epoch {runtime.currentEpoch.toFixed(2)}
- {/* Progress bar */}
- Step {metrics.currentStep} / {metrics.totalSteps} + Step {runtime.currentStep} / {runtime.totalSteps || "--"} {pct}%
@@ -167,51 +265,59 @@ export function ProgressSection(): ReactElement | null {
- {/* Metrics */} -
+ {runtime.error && ( +

{runtime.error}

+ )} + +

Loss

- {metrics.currentLoss.toFixed(4)} + {runtime.currentLoss.toFixed(4)}

LR

- {metrics.currentLR.toExponential(2)} + {runtime.currentLearningRate.toExponential(2)}

Grad Norm

- {metrics.gradNorm.toFixed(3)} + {formatNumber(runtime.currentGradNorm, 3)}

Model

- {modelName} + {config.selectedModel ?? "--"}

Method

-

{store.trainingMethod}

+

+ {config.trainingMethod.toUpperCase()} +

- {/* Timings */} -
- Elapsed: {fmtTime(metrics.elapsed)} - ETA: {fmtTime(etaSec)} - {metrics.samplesPerSecond} samples/s +
+ Elapsed: {formatDuration(elapsed)} + ETA: {formatDuration(eta)} + + {stepsPerSecond == null + ? "-- steps/s" + : `${stepsPerSecond.toFixed(2)} steps/s`} + + {runtime.currentNumTokens != null && ( + Tokens: {runtime.currentNumTokens} + )}
- {/* Right: GPU */}
-

- GPU Monitor -

+

GPU Monitor

} - value={`${metrics.gpuUtil}%`} - pct={metrics.gpuUtil} + value="--" + pct={0} /> - } - value={`${metrics.gpuTemp}°C`} - pct={metrics.gpuTemp} + icon={} + value="--" + pct={0} max={100} /> } - value={`${metrics.gpuVramUsed.toFixed(1)} / ${metrics.gpuVramTotal}GB`} - pct={(metrics.gpuVramUsed / metrics.gpuVramTotal) * 100} + value="--" + pct={0} /> } - value={`${metrics.gpuPower}W`} - pct={(metrics.gpuPower / 350) * 100} + value="--" + pct={0} />
@@ -272,6 +376,7 @@ function GpuStat({ } else if (clamped < 95) { barColor = "bg-amber-500"; } + return (
diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx index a624d7fd3c..fa5c4e85f5 100644 --- a/studio/frontend/src/features/studio/sections/training-section.tsx +++ b/studio/frontend/src/features/studio/sections/training-section.tsx @@ -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 */} + {startError && ( +

{startError}

+ )} {/* Save / Clear */}
diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 7478f32c36..d5e20fc194 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -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 (
-
+
{/* Header */} -
+

Fine-tuning Studio

- {isTraining - ? "Training in progress" + {showTrainingView + ? runtimeMessage || "Training in progress" : "Configure and start training"}

- {isTraining ? ( + {!hasHydratedRuntime && isHydratingRuntime ? ( +
+ Loading training runtime... +
+ ) : showTrainingView ? ( ) : (
diff --git a/studio/frontend/src/features/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx new file mode 100644 index 0000000000..b6c413b400 --- /dev/null +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -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 ( +
+
+ Unsloth mascot + + + {"> unsloth training starts..."} + + + 
{`==((====))==
+\\\\   /|
+O^O/ \\_/ \\
+\\        /
+ "-____-"`}
+ + + {"> Preparing model and dataset..."} + + + {"> We are getting everything ready for your run..."} + + + {"> Did you know, Mugi is actually short for \"Mugiwara\" xd"} + + + {`> ${message || "starting training..."} | waiting for first step... (${currentStep})`} + + +
+
+ ) +} diff --git a/studio/frontend/src/features/studio/training-view.tsx b/studio/frontend/src/features/studio/training-view.tsx index f1120b0ac2..e10faa6735 100644 --- a/studio/frontend/src/features/studio/training-view.tsx +++ b/studio/frontend/src/features/studio/training-view.tsx @@ -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 | null>(null); - const chartsRef = useRef | 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 ( -
- - +
+
+ + +
+ {showOverlay ? ( + + ) : null}
); } diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts new file mode 100644 index 0000000000..0150c62b41 --- /dev/null +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -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, + }; +} diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts new file mode 100644 index 0000000000..e3c53b0bc2 --- /dev/null +++ b/studio/frontend/src/features/training/api/train-api.ts @@ -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 { + 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(response: Response): Promise { + if (!response.ok) { + throw new Error(await readError(response)); + } + return (await response.json()) as T; +} + +export async function startTraining( + payload: TrainingStartRequest, +): Promise { + const response = await authFetch("/api/train/start", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + return parseJson(response); +} + +export async function stopTraining(): Promise { + const response = await authFetch("/api/train/stop", { method: "POST" }); + return parseJson(response); +} + +export async function getTrainingStatus(): Promise { + const response = await authFetch("/api/train/status"); + return parseJson(response); +} + +export async function getTrainingMetrics(): Promise { + const response = await authFetch("/api/train/metrics"); + return parseJson(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 { + 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 }; diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts new file mode 100644 index 0000000000..510c45ed4a --- /dev/null +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -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 => { + 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 => { + 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, + }; +} diff --git a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts new file mode 100644 index 0000000000..3baff3a62d --- /dev/null +++ b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts @@ -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 | 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(); + }; + }, []); +} diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts new file mode 100644 index 0000000000..48b9079fd3 --- /dev/null +++ b/studio/frontend/src/features/training/index.ts @@ -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"; diff --git a/studio/frontend/src/features/training/lib/sync-runtime.ts b/studio/frontend/src/features/training/lib/sync-runtime.ts new file mode 100644 index 0000000000..b5fbd0bafb --- /dev/null +++ b/studio/frontend/src/features/training/lib/sync-runtime.ts @@ -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 { + const [status, metrics] = await Promise.all([ + getTrainingStatus(), + getTrainingMetrics(), + ]); + + const runtimeStore = useTrainingRuntimeStore.getState(); + runtimeStore.applyStatus(status); + runtimeStore.applyMetrics(metrics); + + return status; +} diff --git a/studio/frontend/src/features/training/lib/validation.ts b/studio/frontend/src/features/training/lib/validation.ts new file mode 100644 index 0000000000..8e966153d3 --- /dev/null +++ b/studio/frontend/src/features/training/lib/validation.ts @@ -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 }; +} diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts new file mode 100644 index 0000000000..a9d6d37f42 --- /dev/null +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -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()( + 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; + }, + }, + ), +); diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts new file mode 100644 index 0000000000..bc40019346 --- /dev/null +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -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()((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 + ); +} diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts new file mode 100644 index 0000000000..9cf789a702 --- /dev/null +++ b/studio/frontend/src/features/training/types/api.ts @@ -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; +} diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts new file mode 100644 index 0000000000..c0d93f76a7 --- /dev/null +++ b/studio/frontend/src/features/training/types/config.ts @@ -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; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts new file mode 100644 index 0000000000..fe2afbd36d --- /dev/null +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -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; diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index e9d3b054d0..4745b4dc7b 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -23,6 +23,28 @@ const EXCLUDED_TAGS = new Set([ "ctranslate2", ]); +function withPopularitySort( + input: Parameters[0], + init?: Parameters[1], +): ReturnType { + 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, [query, task, accessToken], diff --git a/studio/frontend/src/stores/training.ts b/studio/frontend/src/stores/training.ts index 469880800c..d3bd63d648 100644 --- a/studio/frontend/src/stores/training.ts +++ b/studio/frontend/src/stores/training.ts @@ -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( - (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 }; diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index 5e206278d0..49c1dddaa3 100644 --- a/studio/frontend/src/types/training.ts +++ b/studio/frontend/src/types/training.ts @@ -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; } diff --git a/studio/tests/__init__.py b/studio/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/studio/tests/test_openai_completions.py b/studio/tests/test_openai_completions.py new file mode 100644 index 0000000000..31bbc6d51e --- /dev/null +++ b/studio/tests/test_openai_completions.py @@ -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"] diff --git a/studio/tests/test_remote_lora_detection.py b/studio/tests/test_remote_lora_detection.py new file mode 100644 index 0000000000..870758e76e --- /dev/null +++ b/studio/tests/test_remote_lora_detection.py @@ -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" diff --git a/studio/tests/test_sse_progress.py b/studio/tests/test_sse_progress.py new file mode 100644 index 0000000000..d2c007d99e --- /dev/null +++ b/studio/tests/test_sse_progress.py @@ -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