Spaces:
Running
Running
| """ | |
| ML Models — Experiments, trained models, training jobs, predictions. | |
| """ | |
| import uuid | |
| from datetime import datetime | |
| from typing import Optional | |
| from sqlalchemy import String, Boolean, Text, Integer, Float, BigInteger, Index, ForeignKey | |
| from sqlalchemy.dialects.postgresql import UUID, JSONB | |
| from sqlalchemy.orm import Mapped, mapped_column | |
| from app.models.base import Base, TimestampMixin, SoftDeleteMixin, UUIDPrimaryKeyMixin | |
| class Project(UUIDPrimaryKeyMixin, TimestampMixin, SoftDeleteMixin, Base): | |
| """Top-level project container for organizing ML work.""" | |
| __tablename__ = "projects" | |
| __table_args__ = ( | |
| Index("ix_projects_user_id", "user_id"), | |
| Index("ix_projects_name", "name"), | |
| ) | |
| user_id: Mapped[uuid.UUID] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("users.id", ondelete="CASCADE"), | |
| nullable=False, | |
| ) | |
| name: Mapped[str] = mapped_column(String(255), nullable=False) | |
| description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| status: Mapped[str] = mapped_column(String(20), default="active", nullable=False) | |
| tags: Mapped[dict] = mapped_column(JSONB, default=list, nullable=False) | |
| settings: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| def __repr__(self) -> str: | |
| return f"<Project name={self.name}>" | |
| class Dataset(UUIDPrimaryKeyMixin, TimestampMixin, SoftDeleteMixin, Base): | |
| """Dataset metadata — linked to a project.""" | |
| __tablename__ = "datasets" | |
| __table_args__ = ( | |
| Index("ix_datasets_user_id", "user_id"), | |
| Index("ix_datasets_project_id", "project_id"), | |
| ) | |
| user_id: Mapped[uuid.UUID] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("users.id", ondelete="CASCADE"), | |
| nullable=False, | |
| ) | |
| project_id: Mapped[Optional[uuid.UUID]] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("projects.id", ondelete="SET NULL"), | |
| nullable=True, | |
| ) | |
| name: Mapped[str] = mapped_column(String(255), nullable=False) | |
| description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| source_type: Mapped[str] = mapped_column(String(50), nullable=False) | |
| file_path: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| row_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) | |
| column_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) | |
| file_size_bytes: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True) | |
| schema_info: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| statistics: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| def __repr__(self) -> str: | |
| return f"<Dataset name={self.name}>" | |
| class Experiment(UUIDPrimaryKeyMixin, TimestampMixin, Base): | |
| """ML experiment — a logical grouping of training runs.""" | |
| __tablename__ = "experiments" | |
| __table_args__ = ( | |
| Index("ix_experiments_user_id", "user_id"), | |
| Index("ix_experiments_project_id", "project_id"), | |
| ) | |
| user_id: Mapped[uuid.UUID] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("users.id", ondelete="CASCADE"), | |
| nullable=False, | |
| ) | |
| project_id: Mapped[Optional[uuid.UUID]] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("projects.id", ondelete="SET NULL"), | |
| nullable=True, | |
| ) | |
| name: Mapped[str] = mapped_column(String(255), nullable=False) | |
| description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| task_type: Mapped[str] = mapped_column(String(50), nullable=False) | |
| status: Mapped[str] = mapped_column(String(20), default="created", nullable=False) | |
| config: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| best_metric_name: Mapped[Optional[str]] = mapped_column(String(50), nullable=True) | |
| best_metric_value: Mapped[Optional[float]] = mapped_column(Float, nullable=True) | |
| def __repr__(self) -> str: | |
| return f"<Experiment name={self.name} status={self.status}>" | |
| class MLModel(UUIDPrimaryKeyMixin, TimestampMixin, SoftDeleteMixin, Base): | |
| """Registered ML model in the model registry.""" | |
| __tablename__ = "ml_models" | |
| __table_args__ = ( | |
| Index("ix_ml_models_user_id", "user_id"), | |
| Index("ix_ml_models_name", "name"), | |
| ) | |
| user_id: Mapped[uuid.UUID] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("users.id", ondelete="CASCADE"), | |
| nullable=False, | |
| ) | |
| experiment_id: Mapped[Optional[uuid.UUID]] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("experiments.id", ondelete="SET NULL"), | |
| nullable=True, | |
| ) | |
| name: Mapped[str] = mapped_column(String(255), nullable=False) | |
| description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| framework: Mapped[str] = mapped_column(String(50), nullable=False) | |
| algorithm: Mapped[str] = mapped_column(String(100), nullable=False) | |
| version: Mapped[str] = mapped_column(String(50), default="1.0.0", nullable=False) | |
| stage: Mapped[str] = mapped_column( | |
| String(20), default="development", nullable=False | |
| ) | |
| metrics: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| hyperparameters: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| artifact_path: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| input_schema: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| output_schema: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| file_size_bytes: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True) | |
| def __repr__(self) -> str: | |
| return f"<MLModel name={self.name} stage={self.stage}>" | |
| class TrainingJob(UUIDPrimaryKeyMixin, TimestampMixin, Base): | |
| """Individual model training run.""" | |
| __tablename__ = "training_jobs" | |
| __table_args__ = ( | |
| Index("ix_training_jobs_user_id", "user_id"), | |
| Index("ix_training_jobs_status", "status"), | |
| ) | |
| user_id: Mapped[uuid.UUID] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("users.id", ondelete="CASCADE"), | |
| nullable=False, | |
| ) | |
| experiment_id: Mapped[Optional[uuid.UUID]] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("experiments.id", ondelete="SET NULL"), | |
| nullable=True, | |
| ) | |
| model_id: Mapped[Optional[uuid.UUID]] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("ml_models.id", ondelete="SET NULL"), | |
| nullable=True, | |
| ) | |
| dataset_id: Mapped[Optional[uuid.UUID]] = mapped_column( | |
| UUID(as_uuid=True), | |
| ForeignKey("datasets.id", ondelete="SET NULL"), | |
| nullable=True, | |
| ) | |
| status: Mapped[str] = mapped_column( | |
| String(20), default="pending", nullable=False | |
| ) | |
| config: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| metrics: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False) | |
| logs: Mapped[dict] = mapped_column(JSONB, default=list, nullable=False) | |
| error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| started_at: Mapped[Optional[datetime]] = mapped_column(nullable=True) | |
| completed_at: Mapped[Optional[datetime]] = mapped_column(nullable=True) | |
| duration_seconds: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) | |
| def __repr__(self) -> str: | |
| return f"<TrainingJob id={self.id} status={self.status}>" | |