fix(optimizers): make pruning actually prune

DashAI/back/models/base_model.py

@@ -31,6 +31,21 @@ class BaseModel(ConfigObject, metaclass=ABCMeta):
     COLOR: str = "#795548"
     ICON: str = "Science"
 
+    # Optional hook, set by an optimizer that wants to watch training as it goes.
+    #
+    # Signature: ``(results: dict[str, float], step: int) -> None``. It is called
+    # once per epoch with the validation metrics of that epoch, and it may raise
+    # to abort training early — that is how Optuna's pruning works.
+    #
+    # It lives here, on the base class, because every model with an epoch loop
+    # already routes its per-epoch metrics through `calculate_metrics`. Hooking
+    # the loops one by one would mean touching five files that do not share a
+    # common ancestor, and missing any model added later.
+    #
+    # Models that train in a single shot never call `calculate_metrics` with
+    # `level=EPOCH`, so for them this stays None and nothing changes.
+    _epoch_reporter = None
+
     @classmethod
     def get_metadata(cls) -> Dict[str, Any]:
         """Get metadata values for the current model.
@@ -298,6 +313,16 @@ def calculate_metrics(
             split=split, level=level, results=results, log_index=log_index
         )
 
+        # Report the epoch to whoever is watching, AFTER persisting: the reporter
+        # is allowed to raise (Optuna prunes that way), and the metrics of the
+        # epoch that triggered the stop should survive it.
+        if (
+            self._epoch_reporter is not None
+            and level is LevelEnum.EPOCH
+            and split is SplitEnum.VALIDATION
+        ):
+            self._epoch_reporter(results, log_index)
+
     def prepare_dataset(
         self, dataset: "DashAIDataset", is_fit: bool = False
     ) -> "DashAIDataset":

DashAI/back/optimizers/optuna_optimizer.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum
 from DashAI.back.core.schema_fields import (
     BaseSchema,
@@ -8,6 +10,9 @@
 from DashAI.back.core.utils import MultilingualString
 from DashAI.back.optimizers.base_optimizer import BaseOptimizer
 
+if TYPE_CHECKING:
+    import optuna
+
 
 class OptunaSchema(BaseSchema):
     n_trials: schema_field(
@@ -116,6 +121,82 @@ class OptunaSchema(BaseSchema):
     )  # type: ignore
 
 
+def _build_pruner(name: "str | None") -> "optuna.pruners.BasePruner":
+    """Resolve a pruner name from the schema into an Optuna pruner instance.
+
+    ``create_study`` accepts any object for ``pruner`` without validating it, so
+    passing the raw schema string silently produces a study whose pruner is a
+    ``str``. The failure only surfaces later, as
+    ``AttributeError: 'str' object has no attribute 'prune'``.
+
+    ``"None"`` (the string the schema sends when pruning is disabled) maps to
+    ``NopPruner``, Optuna's explicit no-op.
+
+    Only pruners that Optuna can build with default arguments are supported.
+    ``PatientPruner``, ``PercentilePruner`` and ``ThresholdPruner`` need
+    configuration (a wrapped pruner, a percentile, a threshold), so exposing them
+    means adding those fields to the schema first.
+    """
+    import optuna
+
+    if name in (None, "", "None"):
+        return optuna.pruners.NopPruner()
+
+    pruner_class = getattr(optuna.pruners, name, None)
+    if pruner_class is None:
+        raise ValueError(f"Unknown pruner '{name}'. Available: {_no_arg_pruners()}")
+    try:
+        return pruner_class()
+    except TypeError as exc:
+        raise ValueError(
+            f"Pruner '{name}' requires configuration and cannot be built from its "
+            f"name alone. Available: {_no_arg_pruners()}"
+        ) from exc
+
+
+def _no_arg_pruners() -> "list[str]":
+    """Pruner names Optuna can instantiate with no arguments."""
+    import optuna
+
+    names = []
+    for name in dir(optuna.pruners):
+        if not name.endswith("Pruner") or name == "BasePruner":
+            continue
+        try:
+            getattr(optuna.pruners, name)()
+        except TypeError:
+            continue
+        names.append(name)
+    return sorted(names)
+
+
+def _report_epoch(trial, metric):
+    """Build the per-epoch callback handed to the model during a trial.
+
+    Optuna prunes by being told how a trial is doing while it still runs:
+    `trial.report(value, step)` feeds the pruner, `trial.should_prune()` asks it
+    for a verdict, and raising `TrialPruned` is how a trial is abandoned.
+
+    Nothing between here and `study.optimize` catches that exception, so it
+    reaches Optuna and the trial is recorded as pruned rather than failed.
+
+    A missing metric is not an error: `calculate_metrics` skips any metric that
+    returns a non-finite value, so a given epoch may legitimately have nothing to
+    report. The trial simply continues unpruned.
+    """
+    import optuna
+
+    def report(results, step):
+        value = results.get(metric.__name__)
+        if value is None:
+            return
+        trial.report(value, step)
+        if trial.should_prune():
+            raise optuna.TrialPruned()
+
+    return report
+
+
 class OptunaOptimizer(BaseOptimizer):
     DISPLAY_NAME: str = MultilingualString(
         en="Optuna Optimizer",
@@ -163,14 +244,15 @@ def optimize(self, model, input_dataset, output_dataset, parameters, metric, tas
         import optuna
 
         sampler = getattr(optuna.samplers, self.sampler)
+        pruner = _build_pruner(self.pruner)
 
         self.model = model
         self.input_dataset = input_dataset
         self.output_dataset = output_dataset
         self.parameters = parameters
         direction = "maximize" if metric["metadata"]["maximize"] else "minimize"
         study = optuna.create_study(
-            direction=direction, sampler=sampler(), pruner=self.pruner
+            direction=direction, sampler=sampler(), pruner=pruner
         )
 
         self.metric = metric["class"]
@@ -187,7 +269,27 @@ def objective(trial):
                     raise ValueError(f"Unsupported parameter type for {key} : {dtype}")
                 setattr(obj, key, value)
 
-            self.model.train(self.input_dataset["train"], self.output_dataset["train"])
+            # Validation data is passed on purpose. Without it the epoch loops
+            # skip `calculate_metrics(split=VALIDATION, level=EPOCH)` entirely —
+            # they guard it behind `if x_validation is not None` — so during
+            # optimization the per-epoch validation score was never computed.
+            #
+            # That is the deeper reason pruning could not work here: the number a
+            # pruner needs to decide did not exist, independently of whether the
+            # pruner itself was an instance or a string.
+            self.model._epoch_reporter = _report_epoch(trial, self.metric)
+            try:
+                self.model.train(
+                    self.input_dataset["train"],
+                    self.output_dataset["train"],
+                    self.input_dataset["validation"],
+                    self.output_dataset["validation"],
+                )
+            finally:
+                # Cleared even when the trial is pruned: the model instance is
+                # reused across trials and by the final refit below.
+                self.model._epoch_reporter = None
+
             y_pred = self.model.predict(input_dataset["validation"])
 
             # Calculate metric for train and validation data each trial

tests/back/models/test_epoch_reporter.py

@@ -0,0 +1,123 @@
+"""Tests for the per-epoch reporting hook on BaseModel.
+
+The hook exists so an optimizer can watch a trial while it trains. It lives on
+the base class rather than inside each model's epoch loop because every model
+that trains in epochs already routes its per-epoch metrics through
+`calculate_metrics` — five loops across five files that share no common ancestor
+below `BaseModel`.
+
+What matters here is that it fires for exactly one combination (validation
+metrics, epoch level) and stays out of the way otherwise.
+"""
+
+import pytest
+
+from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum
+from DashAI.back.models.base_model import BaseModel
+
+
+class ModelStub(BaseModel):
+    """The smallest thing `calculate_metrics` will run against."""
+
+    def __init__(self) -> None:
+        self.run_id = 1
+        self.saved: "list[tuple]" = []
+        for split in SplitEnum:
+            setattr(self, f"{split.value}_metrics", [Accuracy])
+
+    def save(self, filename): ...
+
+    def load(self, filename): ...
+
+    def train(self, x_train, y_train, x_validation=None, y_validation=None): ...
+
+    def predict(self, x_data):
+        return x_data
+
+    def prepare_output(self, y_data, is_fit=False):
+        return y_data
+
+    def _save_metrics(self, split, level, results, log_index=None):
+        self.saved.append((split, level, results))
+
+
+class Accuracy:
+    @staticmethod
+    def score(y_true, y_pred):
+        return 0.75
+
+
+@pytest.fixture
+def model():
+    return ModelStub()
+
+
+def test_the_hook_fires_for_validation_metrics_of_an_epoch(model) -> None:
+    seen = []
+    model._epoch_reporter = lambda results, step: seen.append((results, step))
+
+    model.calculate_metrics(
+        split=SplitEnum.VALIDATION,
+        level=LevelEnum.EPOCH,
+        log_index=3,
+        x_data=[0],
+        y_data=[0],
+    )
+
+    assert seen == [({"Accuracy": 0.75}, 3)]
+
+
+@pytest.mark.parametrize(
+    ("split", "level"),
+    [
+        (SplitEnum.TRAIN, LevelEnum.EPOCH),
+        (SplitEnum.VALIDATION, LevelEnum.TRIAL),
+        (SplitEnum.VALIDATION, LevelEnum.LAST),
+        (SplitEnum.TEST, LevelEnum.LAST),
+    ],
+)
+def test_the_hook_stays_quiet_for_everything_else(model, split, level) -> None:
+    """Training metrics and end-of-trial summaries are not pruning signals."""
+    seen = []
+    model._epoch_reporter = lambda results, step: seen.append((results, step))
+
+    model.calculate_metrics(split=split, level=level, x_data=[0], y_data=[0])
+
+    assert seen == []
+
+
+def test_without_a_reporter_nothing_changes(model) -> None:
+    """The default. Models that no optimizer is watching must be unaffected."""
+    model.calculate_metrics(
+        split=SplitEnum.VALIDATION,
+        level=LevelEnum.EPOCH,
+        log_index=1,
+        x_data=[0],
+        y_data=[0],
+    )
+
+    assert model.saved == [(SplitEnum.VALIDATION, LevelEnum.EPOCH, {"Accuracy": 0.75})]
+
+
+def test_metrics_are_persisted_before_the_hook_can_abort(model) -> None:
+    """Optuna prunes by raising from the reporter.
+
+    The epoch that triggered the stop still happened, so its metrics have to be
+    in the database already when the exception travels up.
+    """
+
+    def prune(results, step):
+        raise RuntimeError("pruned")
+
+    model._epoch_reporter = prune
+
+    with pytest.raises(RuntimeError):
+        model.calculate_metrics(
+            split=SplitEnum.VALIDATION,
+            level=LevelEnum.EPOCH,
+            log_index=1,
+            x_data=[0],
+            y_data=[0],
+        )
+
+    assert model.saved, "the epoch's metrics were lost when the trial was pruned"

tests/back/optimizers/test_optuna_best_params.py

@@ -31,7 +31,7 @@ def __init__(self):
         self.sub = DummySubComponent()
         self.trained_with = None
 
-    def train(self, x, y):
+    def train(self, x, y, x_validation=None, y_validation=None):
         # Record what the model was actually fitted with, which is the value the
         # sub-component holds at that moment.
         self.trained_with = self.sub.C
@@ -116,7 +116,7 @@ def __init__(self):
             super().__init__()
             self.C = 1.0
 
-        def train(self, x, y):
+        def train(self, x, y, x_validation=None, y_validation=None):
             self.trained_with = self.C
 
         def predict(self, dataset):
@@ -144,7 +144,7 @@ def __init__(self):
             super().__init__()
             self.sub = IntSub()
 
-        def train(self, x, y):
+        def train(self, x, y, x_validation=None, y_validation=None):
             self.trained_with = self.sub.n
 
         def predict(self, dataset):