Merge branch 'develop' into feature/models-tutorials

Author: Creylay

DashAI/back/api/api_v1/api.py

@@ -10,6 +10,7 @@
     generative_process,
     generative_session,
     jobs,
+    metrics,
     notebook,
     pipelines,
     plugins,
@@ -32,3 +33,4 @@
 api_router_v1.include_router(pipelines.router, prefix="/pipelines")
 api_router_v1.include_router(plugins.router, prefix="/plugin")
 api_router_v1.include_router(notebook.router, prefix="/notebook")
+api_router_v1.include_router(metrics.router, prefix="/metrics")

DashAI/back/api/api_v1/endpoints/experiments.py

@@ -215,6 +215,9 @@ async def create_experiment(
                 name=params.name,
                 input_columns=params.input_columns,
                 output_columns=params.output_columns,
+                train_metrics=params.train_metrics,
+                validation_metrics=params.validation_metrics,
+                test_metrics=params.test_metrics,
                 splits=params.splits,
             )
             db.add(experiment)

DashAI/back/api/api_v1/endpoints/metrics.py

@@ -0,0 +1,87 @@
+import asyncio
+import json
+import logging
+
+from fastapi import APIRouter, Depends, WebSocket
+from fastapi.websockets import WebSocketDisconnect
+from kink import di, inject
+from sqlalchemy.orm import sessionmaker
+
+from DashAI.back.core.enums.status import RunStatus
+from DashAI.back.dependencies.database.models import Metric, Run
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.websocket("/ws/{run_id}")
+@inject
+async def live_metrics_websocket(
+    websocket: WebSocket,
+    run_id: int,
+    session_factory: sessionmaker = Depends(lambda: di["session_factory"]),
+):
+    await websocket.accept()
+
+    last_timestamp = None
+    first_send = True
+
+    try:
+        while True:
+            with session_factory() as db:
+                query = db.query(Metric).filter(Metric.run_id == run_id)
+
+                # First send: get all metrics
+                # Subsequent sends: get only new metrics
+                if not first_send and last_timestamp is not None:
+                    query = query.filter(Metric.timestamp > last_timestamp)
+
+                # Order by step and timestamp to ensure correct sequence
+                metrics = query.order_by(
+                    Metric.step,
+                    Metric.timestamp,
+                ).all()
+
+                run = db.get(Run, run_id)
+
+            # Update last_timestamp
+            if metrics:
+                last_timestamp = metrics[-1].timestamp
+
+            # Structure payload
+            # split -> level ->
+            # metric_name -> list of {step, value, timestamp}
+            payload: dict[str, dict[str, dict[str, list]]] = {}
+            for metric in metrics:
+                split = metric.split.name
+                level = metric.level.name
+                name = metric.name
+
+                payload.setdefault(split, {}).setdefault(level, {}).setdefault(
+                    name, []
+                ).append(
+                    {
+                        "step": metric.step,
+                        "value": metric.value,
+                        "timestamp": metric.timestamp.isoformat(),
+                    }
+                )
+
+            if run:
+                payload["run_status"] = run.status.name
+
+            if payload:
+                await websocket.send_text(json.dumps(payload))
+
+            first_send = False
+
+            if run and run.status in {RunStatus.FINISHED, RunStatus.ERROR}:
+                await websocket.close(code=1000)
+                break
+
+            await asyncio.sleep(1)
+
+    except WebSocketDisconnect:
+        pass

DashAI/back/api/api_v1/endpoints/runs.py

@@ -11,10 +11,12 @@
 from sqlalchemy.orm import sessionmaker
 
 from DashAI.back.api.api_v1.schemas.runs_params import RunParams, UpdateRunParams
+from DashAI.back.core.enums.metrics import LevelEnum
 from DashAI.back.dependencies.database.models import (
     Experiment,
     GlobalExplainer,
     LocalExplainer,
+    Metric,
     Prediction,
     Run,
     RunStatus,
@@ -26,6 +28,49 @@
 router = APIRouter()
 
 
+def get_metrics_for_run(db, run_id: int):
+    """Retrieve metrics associated with a specific run.
+
+    Parameters
+    ----------
+    db : Session
+        SQLAlchemy session to interact with the database.
+    run_id : int
+        ID of the run for which to retrieve metrics.
+
+    Returns
+    -------
+    dict
+        A dictionary containing train, validation, and test metrics for the run.
+    """
+    metrics = (
+        db.query(Metric)
+        .filter(Metric.run_id == run_id, Metric.level == LevelEnum.LAST)
+        .all()
+    )
+
+    # Initialize the response structure
+    response = {
+        "train_metrics": None,
+        "validation_metrics": None,
+        "test_metrics": None,
+    }
+
+    # Group metrics by split
+    for metric in metrics:
+        # Determine the key in the response dictionary
+        split_key = f"{metric.split.name.lower()}_metrics"
+
+        if response[split_key] is None:
+            response[split_key] = {}
+
+        # In the new schema, we store 'value'.
+        # For 'LAST' level, we just want the latest name: value pair.
+        response[split_key][metric.name] = metric.value
+
+    return response
+
+
 @router.get("/")
 @inject
 async def get_runs(
@@ -67,6 +112,18 @@ async def get_runs(
                 runs = db.scalars(
                     select(Run).where(Run.experiment_id == experiment_id)
                 ).all()
+                if not runs:
+                    raise HTTPException(
+                        status_code=status.HTTP_404_NOT_FOUND,
+                        detail="Runs associated with Experiment not found",
+                    )
+
+                # Add metrics to each run
+                for run in runs:
+                    metrics = get_metrics_for_run(db, run.id)
+                    run.train_metrics = metrics["train_metrics"]
+                    run.validation_metrics = metrics["validation_metrics"]
+                    run.test_metrics = metrics["test_metrics"]
             else:
                 runs = db.query(Run).all()
         except exc.SQLAlchemyError as e:
@@ -112,6 +169,12 @@ async def get_run_by_id(
                     status_code=status.HTTP_404_NOT_FOUND,
                     detail="Run not found",
                 )
+            # Add metrics to the run
+            metrics = get_metrics_for_run(db, run_id)
+            run.train_metrics = metrics["train_metrics"]
+            run.validation_metrics = metrics["validation_metrics"]
+            run.test_metrics = metrics["test_metrics"]
+
         except exc.SQLAlchemyError as e:
             log.exception(e)
             raise HTTPException(
@@ -580,6 +643,11 @@ def reset_run(run):
     setattr(run, "delivery_time", None)
     setattr(run, "end_time", None)
 
+    # Delete metrics from DB
+    with di["session_factory"]() as db:
+        db.query(Metric).filter(Metric.run_id == run.id).delete()
+        db.commit()
+
     # Delete files
     if run.run_path and os.path.exists(run.run_path):
         remove_path(run.run_path)

DashAI/back/api/api_v1/schemas/experiments_params.py

@@ -9,6 +9,9 @@ class ExperimentParams(BaseModel):
     name: str
     input_columns: List[str]
     output_columns: List[str]
+    train_metrics: List[str]
+    validation_metrics: List[str]
+    test_metrics: List[str]
     splits: str
 
 

DashAI/back/core/enums/metrics.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class SplitEnum(Enum):
+    TRAIN = "train"
+    VALIDATION = "validation"
+    TEST = "test"
+
+
+class LevelEnum(Enum):
+    LAST = "last"
+    TRIAL = "trial"
+    STEP = "step"
+    EPOCH = "epoch"

DashAI/back/dependencies/database/models.py

@@ -3,10 +3,20 @@
 from datetime import datetime
 from typing import Any, Dict, List
 
-from sqlalchemy import JSON, Boolean, DateTime, Enum, ForeignKey, String
+from sqlalchemy import (
+    JSON,
+    Boolean,
+    DateTime,
+    Enum,
+    Float,
+    ForeignKey,
+    Integer,
+    String,
+)
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
+from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum
 from DashAI.back.core.enums.plugin_tags import PluginTag
 from DashAI.back.core.enums.status import (
     ConverterListStatus,
@@ -92,6 +102,12 @@ class Experiment(Base):
     task_name: Mapped[str] = mapped_column(String, nullable=False)
     input_columns: Mapped[str] = mapped_column(JSON, nullable=False)
     output_columns: Mapped[str] = mapped_column(JSON, nullable=False)
+
+    # Metrics per split
+    train_metrics: Mapped[list[str]] = mapped_column(JSON, nullable=True)
+    validation_metrics: Mapped[list[str]] = mapped_column(JSON, nullable=True)
+    test_metrics: Mapped[list[str]] = mapped_column(JSON, nullable=True)
+
     splits: Mapped[str] = mapped_column(JSON, nullable=False)
     created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now)
     last_modified: Mapped[DateTime] = mapped_column(
@@ -134,10 +150,6 @@ class Run(Base):
     plot_importance_path: Mapped[str] = mapped_column(String, nullable=True)
     # goal metrics
     goal_metric: Mapped[str] = mapped_column(String)
-    # metrics
-    train_metrics: Mapped[JSON] = mapped_column(JSON, nullable=True)
-    test_metrics: Mapped[JSON] = mapped_column(JSON, nullable=True)
-    validation_metrics: Mapped[JSON] = mapped_column(JSON, nullable=True)
     # artifacts
     artifacts: Mapped[str] = mapped_column(JSON, nullable=True)
     # metadata
@@ -154,6 +166,7 @@ class Run(Base):
     predictions = relationship(
         "Prediction", cascade="all, delete-orphan", back_populates="run"
     )
+    metrics = relationship("Metric", cascade="all, delete-orphan", back_populates="run")
 
     def set_status_as_delivered(self) -> None:
         """Update the status of the run to delivered and set delivery_time to now."""
@@ -223,6 +236,30 @@ def set_status_as_error(self) -> None:
         self.status = PredictionStatus.ERROR
 
 
+class Metric(Base):
+    __tablename__ = "metric"
+    """
+    Table to store all the information related to a metric
+    """
+    id: Mapped[int] = mapped_column(primary_key=True)
+    run_id: Mapped[int] = mapped_column(
+        ForeignKey("run.id", ondelete="CASCADE"), index=True
+    )
+    split: Mapped[SplitEnum] = mapped_column(Enum(SplitEnum), nullable=False)
+    level: Mapped[LevelEnum] = mapped_column(Enum(LevelEnum), nullable=False)
+
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    value: Mapped[float] = mapped_column(Float, nullable=False)
+    step: Mapped[int] = mapped_column(Integer, nullable=False)
+
+    timestamp: Mapped[datetime] = mapped_column(
+        DateTime, default=datetime.now, index=True
+    )
+
+    # Relationships
+    run: Mapped["Run"] = relationship("Run", back_populates="metrics")
+
+
 class Plugin(Base):
     __tablename__ = "plugin"
     """

DashAI/back/initial_components.py

@@ -72,7 +72,25 @@
     PipelineJob,
     PredictJob,
 )
-from DashAI.back.metrics import F1, MAE, RMSE, Accuracy, Bleu, Precision, Recall, Ter
+from DashAI.back.metrics import (
+    F1,
+    MAE,
+    MSE,
+    R2,
+    RMSE,
+    ROCAUC,
+    Accuracy,
+    Bleu,
+    Chrf,
+    CohenKappa,
+    ExplainedVariance,
+    HammingDistance,
+    LogLoss,
+    MedianAbsoluteError,
+    Precision,
+    Recall,
+    Ter,
+)
 from DashAI.back.models import (
     SVC,
     BagOfWordsTextClassificationModel,
@@ -170,8 +188,17 @@ def get_initial_components():
         Recall,
         Bleu,
         Ter,
-        MAE,
+        Chrf,
+        MSE,
         RMSE,
+        MAE,
+        R2,
+        MedianAbsoluteError,
+        ExplainedVariance,
+        ROCAUC,
+        LogLoss,
+        HammingDistance,
+        CohenKappa,
         # Optimizers
         OptunaOptimizer,
         HyperOptOptimizer,