fix(optimizers): build a pruner instance instead of passing the schema string

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,55 @@ 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)
+
+
 class OptunaOptimizer(BaseOptimizer):
     DISPLAY_NAME: str = MultilingualString(
         en="Optuna Optimizer",
@@ -163,14 +217,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"]

tests/back/optimizers/test_optuna_pruner.py

@@ -0,0 +1,55 @@
+"""Tests for pruner resolution in OptunaOptimizer.
+
+Regression coverage for the pruner being passed to ``optuna.create_study`` as a
+raw schema string instead of a pruner instance. ``create_study`` does not
+validate the argument, so the study was built with ``pruner`` set to a ``str``
+and any pruning call raised ``AttributeError: 'str' object has no attribute
+'prune'``.
+"""
+
+import optuna
+import pytest
+
+from DashAI.back.optimizers.optuna_optimizer import _build_pruner
+
+
+@pytest.mark.parametrize(
+    "name",
+    ["MedianPruner", "HyperbandPruner", "SuccessiveHalvingPruner", "WilcoxonPruner"],
+)
+def test_build_pruner_returns_an_instance(name: str) -> None:
+    pruner = _build_pruner(name)
+
+    assert isinstance(pruner, optuna.pruners.BasePruner)
+    assert type(pruner).__name__ == name
+
+
+@pytest.mark.parametrize("disabled", [None, "", "None"])
+def test_disabled_pruning_maps_to_nop_pruner(disabled: "str | None") -> None:
+    """The schema sends the string "None" when the user disables pruning."""
+    assert isinstance(_build_pruner(disabled), optuna.pruners.NopPruner)
+
+
+def test_pruner_needing_configuration_fails_with_a_clear_message() -> None:
+    """PatientPruner wraps another pruner, so it cannot be built from its name."""
+    with pytest.raises(ValueError, match="requires configuration"):
+        _build_pruner("PatientPruner")
+
+
+def test_unknown_pruner_lists_the_valid_options() -> None:
+    with pytest.raises(ValueError, match="Unknown pruner 'NotAPruner'") as excinfo:
+        _build_pruner("NotAPruner")
+
+    assert "MedianPruner" in str(excinfo.value)
+
+
+def test_study_built_with_the_resolved_pruner_can_prune() -> None:
+    """End to end: the failure mode this fixes was only visible when pruning."""
+    study = optuna.create_study(
+        direction="maximize", pruner=_build_pruner("MedianPruner")
+    )
+    trial = study.ask()
+    trial.report(0.1, step=0)
+
+    # With the raw string this raised AttributeError instead of returning a bool.
+    assert trial.should_prune() in (True, False)