Add database module with models

app/database.py

@@ -0,0 +1,241 @@
+"""
+Database models and initialization for Universal Model Trainer
+"""
+
+from sqlalchemy import (
+    Column, Integer, String, Text, Float, Boolean, DateTime, 
+    ForeignKey, JSON, Enum, create_engine
+)
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker, relationship, declarative_base
+from datetime import datetime
+import enum
+import os
+
+from app.config import settings
+
+# Create async engine
+DATABASE_PATH = settings.DATABASE_URL.replace("sqlite:///./", "").replace("sqlite://", "")
+os.makedirs(os.path.dirname(DATABASE_PATH) if os.path.dirname(DATABASE_PATH) else ".", exist_ok=True)
+
+# Use async engine for SQLite
+ASYNC_DB_URL = settings.DATABASE_URL.replace("sqlite://", "sqlite+aiosqlite://")
+engine = create_async_engine(ASYNC_DB_URL, echo=settings.DEBUG)
+
+AsyncSessionLocal = sessionmaker(
+    engine, class_=AsyncSession, expire_on_commit=False
+)
+
+Base = declarative_base()
+
+
+class JobStatus(str, enum.Enum):
+    """Training job status enum."""
+    PENDING = "pending"
+    QUEUED = "queued"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    PAUSED = "paused"
+
+
+class TaskType(str, enum.Enum):
+    """Supported task types."""
+    CAUSAL_LM = "causal-lm"
+    SEQ2SEQ = "seq2seq"
+    TOKEN_CLASSIFICATION = "token-classification"
+    SEQUENCE_CLASSIFICATION = "sequence-classification"
+    QUESTION_ANSWERING = "question-answering"
+    SUMMARIZATION = "summarization"
+    TRANSLATION = "translation"
+    TEXT_CLASSIFICATION = "text-classification"
+    MASKED_LM = "masked-lm"
+    VISION_CLASSIFICATION = "vision-classification"
+    VISION_SEGMENTATION = "vision-segmentation"
+    AUDIO_CLASSIFICATION = "audio-classification"
+    AUDIO_TRANSCRIPTION = "audio-transcription"
+
+
+class TrainingJob(Base):
+    """Model for training jobs."""
+    __tablename__ = "training_jobs"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    job_id = Column(String(36), unique=True, index=True, nullable=False)
+    name = Column(String(255), nullable=False)
+    description = Column(Text, nullable=True)
+    
+    # Task configuration
+    task_type = Column(String(50), nullable=False)
+    base_model = Column(String(255), nullable=False)
+    output_model_name = Column(String(255), nullable=True)
+    
+    # Dataset configuration
+    dataset_source = Column(String(50), default="huggingface")
+    dataset_name = Column(String(255), nullable=True)
+    dataset_config = Column(String(100), nullable=True)
+    dataset_split = Column(String(50), default="train")
+    custom_dataset_path = Column(String(512), nullable=True)
+    
+    # Training arguments
+    training_args = Column(JSON, default=dict)
+    peft_config = Column(JSON, nullable=True)
+    deepspeed_config = Column(JSON, nullable=True)
+    
+    # Status and progress
+    status = Column(String(20), default=JobStatus.PENDING.value)
+    progress = Column(Float, default=0.0)
+    current_epoch = Column(Integer, default=0)
+    total_epochs = Column(Integer, default=3)
+    current_step = Column(Integer, default=0)
+    total_steps = Column(Integer, default=0)
+    
+    # Metrics
+    train_loss = Column(Float, nullable=True)
+    eval_loss = Column(Float, nullable=True)
+    learning_rate = Column(Float, nullable=True)
+    metrics = Column(JSON, default=dict)
+    
+    # Output
+    output_path = Column(String(512), nullable=True)
+    hub_model_id = Column(String(255), nullable=True)
+    model_card = Column(Text, nullable=True)
+    
+    # Error handling
+    error_message = Column(Text, nullable=True)
+    traceback = Column(Text, nullable=True)
+    retry_count = Column(Integer, default=0)
+    max_retries = Column(Integer, default=3)
+    
+    # Timestamps
+    created_at = Column(DateTime, default=datetime.utcnow)
+    started_at = Column(DateTime, nullable=True)
+    completed_at = Column(DateTime, nullable=True)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+    
+    # User info
+    created_by = Column(String(100), nullable=True)
+    tags = Column(JSON, default=list)
+    
+    # Relationships
+    checkpoints = relationship("Checkpoint", back_populates="job", cascade="all, delete-orphan")
+    logs = relationship("TrainingLog", back_populates="job", cascade="all, delete-orphan")
+    
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "job_id": self.job_id,
+            "name": self.name,
+            "description": self.description,
+            "task_type": self.task_type,
+            "base_model": self.base_model,
+            "output_model_name": self.output_model_name,
+            "dataset_name": self.dataset_name,
+            "status": self.status,
+            "progress": self.progress,
+            "current_epoch": self.current_epoch,
+            "total_epochs": self.total_epochs,
+            "current_step": self.current_step,
+            "total_steps": self.total_steps,
+            "train_loss": self.train_loss,
+            "eval_loss": self.eval_loss,
+            "metrics": self.metrics,
+            "output_path": self.output_path,
+            "hub_model_id": self.hub_model_id,
+            "error_message": self.error_message,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "started_at": self.started_at.isoformat() if self.started_at else None,
+            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
+            "tags": self.tags
+        }
+
+
+class Checkpoint(Base):
+    """Model for training checkpoints."""
+    __tablename__ = "checkpoints"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    job_id = Column(Integer, ForeignKey("training_jobs.id"), nullable=False)
+    checkpoint_name = Column(String(255), nullable=False)
+    checkpoint_path = Column(String(512), nullable=False)
+    step = Column(Integer, nullable=False)
+    epoch = Column(Float, nullable=False)
+    loss = Column(Float, nullable=True)
+    metrics = Column(JSON, default=dict)
+    is_best = Column(Boolean, default=False)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    size_mb = Column(Float, nullable=True)
+    
+    # Relationship
+    job = relationship("TrainingJob", back_populates="checkpoints")
+
+
+class TrainingLog(Base):
+    """Model for training logs."""
+    __tablename__ = "training_logs"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    job_id = Column(Integer, ForeignKey("training_jobs.id"), nullable=False)
+    level = Column(String(10), default="INFO")
+    message = Column(Text, nullable=False)
+    step = Column(Integer, nullable=True)
+    epoch = Column(Float, nullable=True)
+    loss = Column(Float, nullable=True)
+    learning_rate = Column(Float, nullable=True)
+    metrics = Column(JSON, nullable=True)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    
+    # Relationship
+    job = relationship("TrainingJob", back_populates="logs")
+
+
+class ModelRegistry(Base):
+    """Registry of trained and available models."""
+    __tablename__ = "model_registry"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    name = Column(String(255), unique=True, nullable=False)
+    model_id = Column(String(255), nullable=False)
+    task_type = Column(String(50), nullable=False)
+    description = Column(Text, nullable=True)
+    tags = Column(JSON, default=list)
+    parameters = Column(String(20), nullable=True)
+    is_local = Column(Boolean, default=False)
+    local_path = Column(String(512), nullable=True)
+    hub_url = Column(String(512), nullable=True)
+    is_trained = Column(Boolean, default=False)
+    training_job_id = Column(Integer, ForeignKey("training_jobs.id"), nullable=True)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    last_used = Column(DateTime, nullable=True)
+
+
+class DatasetCache(Base):
+    """Cache for downloaded datasets."""
+    __tablename__ = "dataset_cache"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    name = Column(String(255), unique=True, nullable=False)
+    config = Column(String(100), nullable=True)
+    split = Column(String(50), nullable=True)
+    local_path = Column(String(512), nullable=False)
+    size_mb = Column(Float, nullable=True)
+    num_samples = Column(Integer, nullable=True)
+    features = Column(JSON, nullable=True)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    last_accessed = Column(DateTime, default=datetime.utcnow)
+
+
+async def init_db():
+    """Initialize database tables."""
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+
+async def get_db():
+    """Get database session."""
+    async with AsyncSessionLocal() as session:
+        try:
+            yield session
+        finally:
+            await session.close()