diff --git a/tmp/models.py b/tmp/models.py deleted file mode 100644 index fd92a169b2..0000000000 --- a/tmp/models.py +++ /dev/null @@ -1,27 +0,0 @@ -""" -Pydantic schemas for Model Management API -""" -from pydantic import BaseModel, Field -from typing import Optional, List, Dict, Any - - -class ModelDetails(BaseModel): - """Detailed model configuration and metadata""" - model_name: str = Field(..., description="Model identifier") - config: Dict[str, Any] = Field(..., description="Model configuration dictionary") - is_vision: bool = Field(False, description="Whether model is a vision model") - is_lora: bool = Field(False, description="Whether model is a LoRA adapter") - base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter") - - -class LoRAInfo(BaseModel): - """LoRA adapter information""" - display_name: str = Field(..., description="Display name for the LoRA") - adapter_path: str = Field(..., description="Path to the LoRA adapter") - base_model: Optional[str] = Field(None, description="Base model identifier") - - -class LoRAScanResponse(BaseModel): - """Response schema for scanning trained LoRA adapters""" - loras: List[LoRAInfo] = Field(default_factory=list, description="List of found LoRA adapters") - outputs_dir: str = Field(..., description="Directory that was scanned") diff --git a/tmp/pr_description.md b/tmp/pr_description.md deleted file mode 100644 index 012a3f03a4..0000000000 --- a/tmp/pr_description.md +++ /dev/null @@ -1,40 +0,0 @@ -# Refactor Pydantic API Models - -## Summary -Cleans up and restructures the Pydantic models for the training and model management APIs to reduce redundancy, improve naming clarity, and align with the frontend architecture. - ---- - -## Training Models (`studio/backend/models/training.py`) - -| Before | After | Change | -|--------|-------|--------| -| `TrainingStartResponse` | `TrainingJobResponse` | Rename - clearer that it represents a created job | -| `TrainingStatusResponse` | `TrainingStatus` | Rename + add `phase` field with explicit pipeline stages | -| `TrainingProgressResponse` | `TrainingProgress` | Rename + add `epoch`, `elapsed_seconds`, `eta_seconds` | -| `TrainingMetricsResponse` | — | **Removed** - frontend stores history in IndexedDB | - -**Key improvements:** -- `TrainingStatus` unifies status polling and streaming with explicit `phase` literals: `idle`, `loading_model`, `loading_dataset`, `configuring`, `training`, `completed`, `error`, `stopped` -- Renamed `is_active` → `is_training_running` for clarity -- `TrainingProgress` now includes timing info (`elapsed_seconds`, `eta_seconds`) for better UX - ---- - -## Model Management (`studio/backend/models/models.py`) - -| Before | After | Change | -|--------|-------|--------| -| `ModelSearchRequest` | — | **Removed** - search handled via query params | -| `ModelSearchResponse` | — | **Removed** | -| `ModelListResponse` | — | **Removed** | -| `ModelInfo` | — | **Removed** - redundant with ModelDetails | -| `ModelConfigResponse` | `ModelDetails` | Rename | - -**Remaining models:** `ModelDetails`, `LoRAInfo`, `LoRAScanResponse` - ---- - -## Breaking Changes -- All renamed/removed models will require updates in any code referencing them -- Frontend TypeScript types should be regenerated diff --git a/tmp/training.py b/tmp/training.py deleted file mode 100644 index fed2076802..0000000000 --- a/tmp/training.py +++ /dev/null @@ -1,99 +0,0 @@ -""" -Pydantic schemas for Training API -""" -from pydantic import BaseModel, Field -from typing import Optional, List, Literal - - -class TrainingStartRequest(BaseModel): - """Request schema for starting training""" - # Model parameters - model_name: str = Field(..., description="Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')") - training_type: str = Field(..., description="Training type: 'LoRA/QLoRA' or 'Full Finetuning'") - hf_token: Optional[str] = Field(None, description="HuggingFace token") - load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization") - max_seq_length: int = Field(2048, description="Maximum sequence length") - - # Dataset parameters - 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") - - # Training parameters - num_epochs: int = Field(1, description="Number of training epochs") - learning_rate: str = Field("2e-4", description="Learning rate") - batch_size: int = Field(1, description="Batch size") - gradient_accumulation_steps: int = Field(1, description="Gradient accumulation steps") - warmup_steps: Optional[int] = Field(None, description="Warmup steps") - warmup_ratio: Optional[float] = Field(None, description="Warmup ratio") - max_steps: Optional[int] = Field(None, description="Maximum training steps") - save_steps: int = Field(100, description="Steps between checkpoints") - weight_decay: float = Field(0.01, description="Weight decay") - random_seed: int = Field(42, description="Random seed") - packing: bool = Field(False, description="Enable sequence packing") - optim: str = Field("adamw_8bit", description="Optimizer") - lr_scheduler_type: str = Field("linear", description="Learning rate scheduler type") - - # LoRA parameters - use_lora: bool = Field(True, description="Use LoRA (derived from training_type)") - lora_r: int = Field(16, description="LoRA rank") - lora_alpha: int = Field(16, description="LoRA alpha") - lora_dropout: float = Field(0.0, description="LoRA dropout") - target_modules: List[str] = Field(default_factory=list, description="Target modules for LoRA") - gradient_checkpointing: str = Field("", description="Gradient checkpointing setting") - use_rslora: bool = Field(False, description="Use RSLoRA") - use_loftq: bool = Field(False, description="Use LoftQ") - train_on_completions: bool = Field(False, description="Train on completions only") - - # Vision-specific LoRA parameters - finetune_vision_layers: bool = Field(False, description="Finetune vision layers") - finetune_language_layers: bool = Field(False, description="Finetune language layers") - finetune_attention_modules: bool = Field(False, description="Finetune attention modules") - finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules") - - # Logging parameters - enable_wandb: bool = Field(False, description="Enable Weights & Biases logging") - wandb_token: Optional[str] = Field(None, description="W&B token") - wandb_project: Optional[str] = Field(None, description="W&B project name") - enable_tensorboard: bool = Field(False, description="Enable TensorBoard logging") - tensorboard_dir: Optional[str] = Field(None, description="TensorBoard directory") - - -class TrainingJobResponse(BaseModel): - """Immediate response when training is initiated""" - job_id: str = Field(..., description="Unique training job identifier") - status: Literal["queued", "error"] = Field(..., description="Initial job status") - message: str = Field(..., description="Human-readable status message") - error: Optional[str] = Field(None, description="Error details if status is 'error'") - - -class TrainingStatus(BaseModel): - """Current training job status - works for streaming or polling""" - job_id: str = Field(..., description="Training job identifier") - phase: Literal[ - "idle", - "loading_model", - "loading_dataset", - "configuring", - "training", - "completed", - "error", - "stopped" - ] = Field(..., description="Current phase of training pipeline") - is_training_running: bool = Field(..., description="True if training loop is actively running") - 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'}") - - -class TrainingProgress(BaseModel): - """Training progress metrics - for streaming or polling""" - job_id: str = Field(..., description="Training job identifier") - step: int = Field(..., description="Current training step") - total_steps: int = Field(..., description="Total training steps") - 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") - elapsed_seconds: Optional[float] = Field(None, description="Time elapsed since training started") - eta_seconds: Optional[float] = Field(None, description="Estimated time remaining") diff --git a/tmp/training_pydantic.md b/tmp/training_pydantic.md deleted file mode 100644 index 63e5289ab0..0000000000 --- a/tmp/training_pydantic.md +++ /dev/null @@ -1,113 +0,0 @@ -# Training Pydantic Models Recommendation - -## Current State - -| Model | Purpose | -|-------|---------| -| `TrainingStartRequest` | Request payload to start training | -| `TrainingStartResponse` | Immediate response after `/start` | -| `TrainingStatusResponse` | Response for `/status` polling | -| `TrainingMetricsResponse` | Historical metrics | -| `TrainingProgressResponse` | Per-step progress data | - ---- - -## Proposed Models - -### 1. `TrainingStartRequest` — Keep as-is -Well-structured, no changes needed. - ---- - -### 2. `TrainingJobResponse` (rename from `TrainingStartResponse`) -Returned **once** when `/train/start` is called. - -```python -class TrainingJobResponse(BaseModel): - """Immediate response when training is initiated""" - job_id: str - status: Literal["queued", "error"] - message: str - error: Optional[str] = None -``` - ---- - -### 3. `TrainingStatus` (unifies `TrainingStatusResponse` + `TrainingPhaseUpdate`) -Single model for **both streaming and polling**. - -```python -class TrainingStatus(BaseModel): - """Current training job status - works for streaming or polling""" - job_id: str - phase: Literal[ - "idle", - "loading_model", - "loading_dataset", - "configuring", - "training", - "completed", - "error", - "stopped" - ] - is_training_running: bool # True if training loop is actively running - message: str - error: Optional[str] = None - details: Optional[dict] = None # Phase-specific info, e.g. {"model_size": "8B"} -``` - -**Usage:** -- **Streaming**: Push when phase changes -- **Polling**: Return from `GET /train/status/{job_id}` - ---- - -### 4. `TrainingProgress` (rename from `TrainingProgressResponse`) -Per-step metrics during active training. - -```python -class TrainingProgress(BaseModel): - """Training progress metrics - for streaming or polling""" - step: int - total_steps: int - loss: float - learning_rate: float - progress_percent: float - epoch: Optional[int] = None - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None -``` - ---- - -### 5. `TrainingMetricsResponse` — Remove -Frontend stores history in IndexedDB. - ---- - -## Summary - -| Current | Proposed | Action | -|---------|----------|--------| -| `TrainingStartRequest` | `TrainingStartRequest` | Keep | -| `TrainingStartResponse` | `TrainingJobResponse` | Rename | -| `TrainingStatusResponse` | `TrainingStatus` | Rename + enhance | -| `TrainingMetricsResponse` | — | **Remove** | -| `TrainingProgressResponse` | `TrainingProgress` | Rename + enhance | - ---- - -## API Flow - -``` -POST /train/start - └─► TrainingJobResponse { job_id, status: "queued" } - -StreamingResponse / Polling: - └─► TrainingStatus { phase: "loading_model", is_training_running: false } - └─► TrainingStatus { phase: "loading_dataset", is_training_running: false } - └─► TrainingStatus { phase: "training", is_training_running: true } - └─► TrainingProgress { step: 1, loss: 2.5, ... } - └─► TrainingProgress { step: 2, loss: 2.3, ... } - └─► TrainingStatus { phase: "completed", is_training_running: false } -```