update pydantic models for Models and Training routes
This commit is contained in:
Roland Tannous
parent
e390ca1092
commit
d7d3a5a9a5
7 changed files with 319 additions and 65 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -20,7 +20,6 @@ unsloth_compiled_cache/
|
|||
outputs/
|
||||
*.gguf
|
||||
*.safetensors
|
||||
/models/
|
||||
|
||||
# IDE / Editors
|
||||
.vscode/
|
||||
|
|
|
|||
|
|
@ -5,36 +5,8 @@ from pydantic import BaseModel, Field
|
|||
from typing import Optional, List, Dict, Any
|
||||
|
||||
|
||||
class ModelSearchRequest(BaseModel):
|
||||
"""Request schema for searching HuggingFace models"""
|
||||
query: str = Field(..., description="Search query")
|
||||
hf_token: Optional[str] = Field(None, description="HuggingFace token for authenticated searches")
|
||||
|
||||
|
||||
class ModelInfo(BaseModel):
|
||||
"""Model information"""
|
||||
id: str = Field(..., description="Model identifier")
|
||||
name: Optional[str] = Field(None, description="Display name")
|
||||
description: Optional[str] = Field(None, description="Model description")
|
||||
size: Optional[str] = Field(None, description="Model size")
|
||||
is_vision: bool = Field(False, description="Whether model is a vision model")
|
||||
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
|
||||
|
||||
|
||||
class ModelSearchResponse(BaseModel):
|
||||
"""Response schema for model search"""
|
||||
models: List[ModelInfo] = Field(default_factory=list, description="List of matching models")
|
||||
total: int = Field(0, description="Total number of results")
|
||||
|
||||
|
||||
class ModelListResponse(BaseModel):
|
||||
"""Response schema for listing available models"""
|
||||
models: List[ModelInfo] = Field(default_factory=list, description="List of available models")
|
||||
default_models: List[str] = Field(default_factory=list, description="List of default model IDs")
|
||||
|
||||
|
||||
class ModelConfigResponse(BaseModel):
|
||||
"""Response schema for model configuration"""
|
||||
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")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
Pydantic schemas for Training API
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional, List
|
||||
from typing import Optional, List, Literal
|
||||
|
||||
|
||||
class TrainingStartRequest(BaseModel):
|
||||
|
|
@ -59,38 +59,42 @@ class TrainingStartRequest(BaseModel):
|
|||
tensorboard_dir: Optional[str] = Field(None, description="TensorBoard directory")
|
||||
|
||||
|
||||
class TrainingStartResponse(BaseModel):
|
||||
"""Response schema for training start"""
|
||||
status: str = Field(..., description="Status: 'started' or 'error'")
|
||||
job_id: Optional[str] = Field(None, description="Training job ID")
|
||||
message: str = Field(..., description="Status message")
|
||||
error: Optional[str] = Field(None, description="Error message if status is 'error'")
|
||||
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 TrainingStatusResponse(BaseModel):
|
||||
"""Response schema for training status"""
|
||||
status: str = Field(..., description="Status: 'idle', 'preparing', 'training', 'stopping', 'error'")
|
||||
is_active: bool = Field(..., description="Whether training is currently active (actual training running)")
|
||||
message: str = Field(..., description="Status message")
|
||||
current_step: Optional[int] = Field(None, description="Current training step")
|
||||
total_steps: Optional[int] = Field(None, description="Total training steps")
|
||||
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 TrainingMetricsResponse(BaseModel):
|
||||
"""Response schema for training metrics"""
|
||||
loss_history: List[float] = Field(default_factory=list, description="Loss values")
|
||||
lr_history: List[float] = Field(default_factory=list, description="Learning rate values")
|
||||
step_history: List[int] = Field(default_factory=list, description="Step numbers")
|
||||
current_loss: Optional[float] = Field(None, description="Current loss value")
|
||||
current_lr: Optional[float] = Field(None, description="Current learning rate")
|
||||
current_step: Optional[int] = Field(None, description="Current step")
|
||||
|
||||
|
||||
class TrainingProgressResponse(BaseModel):
|
||||
"""Response schema for training progress updates"""
|
||||
step: int = Field(..., description="Current step")
|
||||
loss: float = Field(..., description="Current loss")
|
||||
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")
|
||||
status_message: str = Field(..., description="Status message")
|
||||
progress_percent: Optional[float] = Field(None, description="Progress percentage")
|
||||
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")
|
||||
|
||||
|
|
|
|||
27
tmp/models.py
Normal file
27
tmp/models.py
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
"""
|
||||
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")
|
||||
40
tmp/pr_description.md
Normal file
40
tmp/pr_description.md
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
# 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
|
||||
99
tmp/training.py
Normal file
99
tmp/training.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
"""
|
||||
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")
|
||||
113
tmp/training_pydantic.md
Normal file
113
tmp/training_pydantic.md
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
# 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 }
|
||||
```
|
||||
Loading…
Add table
Add a link
Reference in a new issue