diff --git a/DashAI/back/optimizers/optuna_optimizer.py b/DashAI/back/optimizers/optuna_optimizer.py index 8b384e09a..5cb66981a 100644 --- a/DashAI/back/optimizers/optuna_optimizer.py +++ b/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,6 +217,7 @@ 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 @@ -170,7 +225,7 @@ def optimize(self, model, input_dataset, output_dataset, parameters, metric, tas 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"] diff --git a/tests/back/optimizers/test_optuna_pruner.py b/tests/back/optimizers/test_optuna_pruner.py new file mode 100644 index 000000000..b24d8c8d2 --- /dev/null +++ b/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)