From 834545f18f1a377a4ef157c4c7dc91bc81721227 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 14 Aug 2026 12:30:25 +0000 Subject: [PATCH 1/5] fix(optimizers): build a pruner instance instead of passing the schema string MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `OptunaSchema` declares `pruner` as an enum of strings, and `optimize()` forwarded that string straight to `optuna.create_study()`. The sampler right above it is resolved and instantiated; the pruner was not. `create_study()` does not validate the argument, so nothing fails at study creation — the study is simply built with `pruner` set to a `str`. Checked against Optuna 4.9.0: >>> study = optuna.create_study(direction="maximize", pruner="MedianPruner") >>> study.pruner 'MedianPruner' Two consequences follow. Today the setting is inert: `objective()` never calls `trial.report()` or `trial.should_prune()`, so the pruner is never consulted and selecting "MedianPruner" or "None" produces identical runs. And it breaks as soon as pruning is wired up, since `should_prune()` calls `self.study.pruner.prune(...)` and raises `AttributeError: 'str' object has no attribute 'prune'`. `_build_pruner()` resolves the name the same way the sampler is resolved. "None" — the string the schema sends when pruning is disabled — maps to `NopPruner` rather than to Python's `None`, because a bare `None` makes Optuna fall back to its own default (`MedianPruner`), which is not what the user picked. Pruners that cannot be built without arguments (`PatientPruner`, `PercentilePruner`, `ThresholdPruner`) raise a message naming the ones that can; exposing them means adding their configuration to the schema first. Adds `tests/back/optimizers/test_optuna_pruner.py` — 10 cases covering instance resolution, the "None" mapping, the unknown-name error and the needs-configuration error. Whether `objective()` should report intermediate values so pruning can actually take effect is a separate decision: it needs a meaningful per-step metric, which not every model exposes. That is deliberately left out of this change. --- DashAI/back/optimizers/optuna_optimizer.py | 57 ++++++++++++++++++++- tests/back/optimizers/test_optuna_pruner.py | 55 ++++++++++++++++++++ 2 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 tests/back/optimizers/test_optuna_pruner.py 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) From 537b93b5dbf9255acfeec1a6fd7ea3f24735c358 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 12:54:10 +0000 Subject: [PATCH 2/5] feat(optimizers): make pruning actually prune MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the review on #823. Resolving the pruner into an instance was necessary but not sufficient: a pruner only ever acts if the trial is told how it is doing while it still runs, and nothing was telling it. The deeper reason it could not work `optimize` called the model without validation data: self.model.train(input_dataset["train"], output_dataset["train"]) while the non-optimizer path in model_job.py passes all four: model.train(x["train"], y["train"], x["validation"], y["validation"]) Every epoch loop guards its validation metrics behind `if x_validation is not None`, so during optimization the per-epoch validation score was never computed at all. The number a pruner needs to decide did not exist — independently of whether the pruner was an instance or a string. Where the hook lives Five models train in epochs (base_torchvision, cnn, mlp, lenet5, and scikit_learn/mlp_regression) and they share no ancestor below `BaseModel`, so hooking each loop would mean five edits and would miss whatever is added next. But all five already route their per-epoch metrics through `BaseModel.calculate_metrics`, and so does the HuggingFace path via its own callback. That is the seam: one optional `_epoch_reporter` on the base class, invoked only for `level=EPOCH` and `split=VALIDATION`. Models that train in a single shot never reach that branch, so for them nothing changes. The reporter is invoked AFTER metrics are persisted, because it is allowed to raise — that is how Optuna prunes — and the epoch that triggered the stop should survive it. `TrialPruned` travels from inside the loop up to `study.optimize` with nothing in between catching it, so the trial is recorded as pruned rather than failed. A metric missing from an epoch is not an error: `calculate_metrics` drops metrics that return a non-finite value (one class present in a split, say), so the trial simply continues unpruned. Tests 11 new cases: that each epoch is reported with its step, that a rejected trial raises `TrialPruned`, that a missing metric is tolerated, that the reported metric is the one being optimized, that the hook fires for exactly epoch+validation and stays quiet for train metrics and trial-level summaries, that a model with no reporter behaves as before, and that an epoch's metrics are saved before the hook can abort. The three doubles in test_optuna_best_params.py now mirror the real `BaseModel.train` signature. They declared `train(self, x, y)`, which no actual model does — every one of them, across the torch, scikit-learn and HuggingFace families, already accepts `x_validation=None, y_validation=None`, because the normal training path has always passed them. Full suite: 806 passed. The remaining 1 failure and 6 errors reproduce identically on a clean develop checkout — the frontend build and the HuggingFace model downloads, neither reachable from this environment. --- DashAI/back/models/base_model.py | 25 ++++ DashAI/back/optimizers/optuna_optimizer.py | 49 ++++++- tests/back/models/test_epoch_reporter.py | 123 ++++++++++++++++++ .../optimizers/test_optuna_best_params.py | 6 +- tests/back/optimizers/test_optuna_pruner.py | 79 ++++++++++- 5 files changed, 277 insertions(+), 5 deletions(-) create mode 100644 tests/back/models/test_epoch_reporter.py diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index 1d051a68a..aefd28797 100644 --- a/DashAI/back/models/base_model.py +++ b/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": diff --git a/DashAI/back/optimizers/optuna_optimizer.py b/DashAI/back/optimizers/optuna_optimizer.py index 5cb66981a..bd75f4b22 100644 --- a/DashAI/back/optimizers/optuna_optimizer.py +++ b/DashAI/back/optimizers/optuna_optimizer.py @@ -170,6 +170,33 @@ def _no_arg_pruners() -> "list[str]": 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", @@ -242,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 diff --git a/tests/back/models/test_epoch_reporter.py b/tests/back/models/test_epoch_reporter.py new file mode 100644 index 000000000..dc85044a4 --- /dev/null +++ b/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" diff --git a/tests/back/optimizers/test_optuna_best_params.py b/tests/back/optimizers/test_optuna_best_params.py index 2782343a3..df09d934d 100644 --- a/tests/back/optimizers/test_optuna_best_params.py +++ b/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): diff --git a/tests/back/optimizers/test_optuna_pruner.py b/tests/back/optimizers/test_optuna_pruner.py index b24d8c8d2..2f5b7016f 100644 --- a/tests/back/optimizers/test_optuna_pruner.py +++ b/tests/back/optimizers/test_optuna_pruner.py @@ -10,7 +10,7 @@ import optuna import pytest -from DashAI.back.optimizers.optuna_optimizer import _build_pruner +from DashAI.back.optimizers.optuna_optimizer import _build_pruner, _report_epoch @pytest.mark.parametrize( @@ -53,3 +53,80 @@ def test_study_built_with_the_resolved_pruner_can_prune() -> None: # With the raw string this raised AttributeError instead of returning a bool. assert trial.should_prune() in (True, False) + + +# --- Reporting each epoch to the trial ------------------------------------- +# +# Resolving the pruner is not enough on its own: a pruner only ever acts if the +# trial is told how it is doing while it still runs. These cover that half. + + +class _Trial: + """Records what a trial was told, and answers should_prune() on cue.""" + + def __init__(self, prune_at: "int | None" = None) -> None: + self.reported: "list[tuple[float, int]]" = [] + self.prune_at = prune_at + + def report(self, value: float, step: int) -> None: + self.reported.append((value, step)) + + def should_prune(self) -> bool: + return self.prune_at is not None and len(self.reported) >= self.prune_at + + +class Accuracy: + """Stands in for a dashAI metric class. + + Only its name matters: `calculate_metrics` keys its results by + ``metric.__name__``, so that is what the reporter looks up. Assigning + ``__name__`` inside a class body does NOT rename the class — ``type.__name__`` + is a data descriptor and wins — so the double has to actually be named after + the metric. + """ + + +def test_each_epoch_is_reported_with_its_step() -> None: + trial = _Trial() + + report = _report_epoch(trial, Accuracy) + report({"Accuracy": 0.4}, 1) + report({"Accuracy": 0.6}, 2) + + assert trial.reported == [(0.4, 1), (0.6, 2)] + + +def test_a_trial_the_pruner_rejects_raises_trial_pruned() -> None: + """`TrialPruned` is how Optuna is told to abandon a trial. + + It has to travel from inside the model's epoch loop up to `study.optimize`, + so nothing in between may swallow it. + """ + trial = _Trial(prune_at=1) + + report = _report_epoch(trial, Accuracy) + with pytest.raises(optuna.TrialPruned): + report({"Accuracy": 0.1}, 1) + + +def test_a_metric_missing_from_the_epoch_is_not_an_error() -> None: + """`calculate_metrics` drops metrics that return a non-finite value. + + An epoch with only one class present in the split is a real case, so the + optimized metric can legitimately be absent. That trial keeps running. + """ + trial = _Trial(prune_at=1) + + report = _report_epoch(trial, Accuracy) + report({"F1": 0.9}, 1) + + assert trial.reported == [] + + +def test_the_reported_metric_is_the_one_being_optimized() -> None: + trial = _Trial() + + report = _report_epoch(trial, Accuracy) + report({"F1": 0.1, "Accuracy": 0.9, "Precision": 0.5}, 1) + + assert trial.reported == [(0.9, 1)] From 53c179d4b3d171fc8e82922877bca30eb0625226 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 14:26:25 +0000 Subject: [PATCH 3/5] test(optimizers): prove a trial actually gets pruned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unit tests cover each piece — the reporter reports, the hook fires only for epoch+validation, TrialPruned is raised. None of them would notice a pruner that reports faithfully and never cuts anything short, which is the failure this whole change exists to fix. This wires the real parts together and asserts the outcome Optuna records: test_a_bad_trial_is_actually_pruned at least one trial ends PRUNED test_pruning_stops_training_early pruned trials cost fewer epochs than the same study with pruning disabled test_disabled_pruning_completes_every_trial the control: with NopPruner every trial completes, so the two above are measuring pruning and not noise Real, not stubbed: OptunaOptimizer.optimize, BaseModel.calculate_metrics (where the hook lives), _report_epoch, and Optuna's own MedianPruner and trial bookkeeping. Stubbed: _save_metrics, because persistence needs a database and is not what this proves, and the training itself, replaced by a loop that improves by a fixed amount per epoch and calls calculate_metrics exactly as the five models that train in epochs do — same split, same level, same log_index. Checked against the previous commit, with the source files reverted: the two tests that assert pruning fail, and the control still passes. A test that passes before and after the change it covers is not evidence. Full suite: 809 passed. The 1 failure and 6 errors are the frontend build and the HuggingFace downloads, identical on a clean develop checkout here. --- .../test_optuna_pruning_integration.py | 141 ++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 tests/back/optimizers/test_optuna_pruning_integration.py diff --git a/tests/back/optimizers/test_optuna_pruning_integration.py b/tests/back/optimizers/test_optuna_pruning_integration.py new file mode 100644 index 000000000..9ea77c55d --- /dev/null +++ b/tests/back/optimizers/test_optuna_pruning_integration.py @@ -0,0 +1,141 @@ +"""End-to-end proof that a trial actually gets pruned. + +The unit tests cover each piece: the reporter reports, the hook fires only for +epoch+validation, `TrialPruned` is raised. This one wires the real parts +together and checks the outcome Optuna records, which is what the feature is +for — a pruner that never prunes passes every unit test in the file next door. + +Real, not stubbed: `OptunaOptimizer.optimize`, `BaseModel.calculate_metrics` +(where the hook lives), `_report_epoch`, and Optuna's own MedianPruner and +trial bookkeeping. + +Stubbed: `_save_metrics` (persistence needs a database and is not what this +proves) and the training itself, which is replaced by a loop that improves by a +fixed amount per epoch. The loop calls `calculate_metrics` exactly as the five +models that train in epochs do — same split, same level, same log_index. +""" + +import optuna +import pytest + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.models.base_model import BaseModel +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer + +EPOCHS = 12 + + +class Score: + """Metric class. Named because results are keyed by `metric.__name__`.""" + + @staticmethod + def score(y_true, y_pred): + return y_pred + + +class SteppedModel(BaseModel): + """A model whose quality is decided by `rate` and revealed epoch by epoch. + + Trials with a low rate look bad early, which is precisely the situation a + pruner exists to cut short. + """ + + def __init__(self): + self.run_id = 1 + self.rate = 1.0 + self.value = 0.0 + self.epochs_run = 0 + for split in SplitEnum: + setattr(self, f"{split.value}_metrics", [Score]) + # The optimizer also asks for trial-level metrics without passing data, + # so `calculate_metrics` falls back to what the model holds. + self.x_data = {split.value: [0] for split in SplitEnum} + self.y_data = {split.value: [0] for split in SplitEnum} + + def save(self, filename): ... + + def load(self, filename): ... + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.value = 0.0 + for epoch in range(EPOCHS): + self.value += self.rate + self.epochs_run += 1 + # Same call the real epoch loops make. + self.calculate_metrics( + split=SplitEnum.VALIDATION, + level=LevelEnum.EPOCH, + x_data=[0], + y_data=[0], + log_index=epoch + 1, + ) + + def predict(self, x_data): + return self.value + + def prepare_output(self, y_data, is_fit=False): + return y_data + + def _save_metrics(self, split, level, results, log_index=None): + """Persistence is out of scope; the database is not what this proves.""" + + +@pytest.fixture +def dataset(): + return {"train": [0], "validation": [0]} + + +def _run(pruner, n_trials=10): + model = SteppedModel() + optimizer = OptunaOptimizer( + n_trials=n_trials, sampler="RandomSampler", pruner=pruner + ) + optimizer.optimize( + model, + dataset := {"train": [0], "validation": [0]}, + dataset, + [(model, "rate", (0.01, 10.0), "number")], + {"class": Score, "metadata": {"maximize": True}}, + "TabularClassificationTask", + ) + return optimizer, model + + +def _states(optimizer): + return [t.state for t in optimizer.study.trials] + + +def test_a_bad_trial_is_actually_pruned(): + """The outcome that matters: Optuna records trials as PRUNED.""" + optimizer, _ = _run("MedianPruner") + + pruned = [s for s in _states(optimizer) if s is optuna.trial.TrialState.PRUNED] + + assert pruned, ( + "no trial was pruned. The pruner is inert: either the epoch metrics never " + "reach the trial, or TrialPruned is being swallowed before Optuna sees it." + ) + + +def test_pruning_stops_training_early(): + """A pruned trial must cost fewer epochs than a completed one. + + Pruning that reports the right verdict but keeps training saves nothing, + which is the whole point of early stopping. + """ + _, with_pruning = _run("MedianPruner") + _, without = _run("NopPruner") + + assert with_pruning.epochs_run < without.epochs_run, ( + f"pruning ran {with_pruning.epochs_run} epochs and no-pruning ran " + f"{without.epochs_run}: the trials were cut short on paper only" + ) + + +def test_disabled_pruning_completes_every_trial(): + """The control. With pruning off nothing may be cut, or the test above proves nothing.""" + optimizer, model = _run("NopPruner") + + states = _states(optimizer) + assert all(s is optuna.trial.TrialState.COMPLETE for s in states) + assert model.epochs_run == EPOCHS * (len(states) + 1) # +1: the final refit From 2762fd8461438582d0fb6a3984ba201a793b4ea8 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 16:14:22 +0000 Subject: [PATCH 4/5] test(optimizers): decide pruning by trial order, not by the sampler's draw These tests failed about one run in twenty. Measured on this branch: 2 failures in 40 runs of the file, and one in a full-suite run. The cause was that `SteppedModel`'s quality was the value Optuna sampled for `rate`, drawn by an unseeded RandomSampler. MedianPruner prunes a trial whose score falls below the median of the trials before it, so whether anything was pruned depended on where those draws landed. When trials 5 to 9 all happened to draw above the median, nothing was pruned and both pruning assertions failed -- reporting the pruner as broken when it was working. Quality now decreases with each trial (`1 / (1 + trials_run)`), so every trial after the pruner's startup window is below the median from its first epoch. The sampled `rate` is still declared as the optimizable parameter, so the optimizer's real path still runs; it just no longer decides the outcome. With the outcome deterministic the assertions can be exact, and they are: 5 trials pruned of 10, and 77 epochs against the control's 132. Before, the test only asserted that *something* was pruned and that one number was smaller than another, which is what let a flake hide. 0 failures in 30 runs after the change. Full suite: 813 passed; the 1 failure and 6 errors are the frontend build and the HuggingFace downloads, identical on a clean develop checkout here. --- .../test_optuna_pruning_integration.py | 58 ++++++++++++++----- 1 file changed, 45 insertions(+), 13 deletions(-) diff --git a/tests/back/optimizers/test_optuna_pruning_integration.py b/tests/back/optimizers/test_optuna_pruning_integration.py index 9ea77c55d..8e676c855 100644 --- a/tests/back/optimizers/test_optuna_pruning_integration.py +++ b/tests/back/optimizers/test_optuna_pruning_integration.py @@ -13,6 +13,12 @@ proves) and the training itself, which is replaced by a loop that improves by a fixed amount per epoch. The loop calls `calculate_metrics` exactly as the five models that train in epochs do — same split, same level, same log_index. + +Which trials get pruned is decided by trial order, not by the value the sampler +draws. Tying it to the draw made these tests fail about one run in twenty — the +runs where no trial happened to land below the median. A test for pruning that +only usually prunes reports the pruner as broken at random, which is worse than +not having it. """ import optuna @@ -23,6 +29,11 @@ from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer EPOCHS = 12 +N_TRIALS = 10 +# MedianPruner never prunes its first `n_startup_trials` (5 by default): with +# nothing to compare against there is no median. Every trial after those is +# below it by construction here, so the split is exact. +STARTUP_TRIALS = 5 class Score: @@ -34,10 +45,15 @@ def score(y_true, y_pred): class SteppedModel(BaseModel): - """A model whose quality is decided by `rate` and revealed epoch by epoch. + """A model that gets worse every trial, revealing it epoch by epoch. + + Each trial improves by `1 / (1 + trials already run)` per epoch, so trial 5 + onwards is below the median of everything before it from its first epoch — + exactly the situation a pruner exists to cut short, and one that does not + depend on chance. - Trials with a low rate look bad early, which is precisely the situation a - pruner exists to cut short. + `rate` is still declared as the optimizable parameter so the optimizer's + own path runs for real; it just does not decide the outcome under test. """ def __init__(self): @@ -45,6 +61,7 @@ def __init__(self): self.rate = 1.0 self.value = 0.0 self.epochs_run = 0 + self.trials_run = 0 for split in SplitEnum: setattr(self, f"{split.value}_metrics", [Score]) # The optimizer also asks for trial-level metrics without passing data, @@ -58,8 +75,10 @@ def load(self, filename): ... def train(self, x_train, y_train, x_validation=None, y_validation=None): self.value = 0.0 + quality = 1.0 / (1.0 + self.trials_run) + self.trials_run += 1 for epoch in range(EPOCHS): - self.value += self.rate + self.value += quality self.epochs_run += 1 # Same call the real epoch loops make. self.calculate_metrics( @@ -85,7 +104,7 @@ def dataset(): return {"train": [0], "validation": [0]} -def _run(pruner, n_trials=10): +def _run(pruner, n_trials=N_TRIALS): model = SteppedModel() optimizer = OptunaOptimizer( n_trials=n_trials, sampler="RandomSampler", pruner=pruner @@ -111,9 +130,11 @@ def test_a_bad_trial_is_actually_pruned(): pruned = [s for s in _states(optimizer) if s is optuna.trial.TrialState.PRUNED] - assert pruned, ( - "no trial was pruned. The pruner is inert: either the epoch metrics never " - "reach the trial, or TrialPruned is being swallowed before Optuna sees it." + assert len(pruned) == N_TRIALS - STARTUP_TRIALS, ( + f"{len(pruned)} trials were pruned, expected exactly " + f"{N_TRIALS - STARTUP_TRIALS}. None at all means the pruner is inert: " + "either the epoch metrics never reach the trial, or TrialPruned is being " + "swallowed before Optuna sees it." ) @@ -126,16 +147,27 @@ def test_pruning_stops_training_early(): _, with_pruning = _run("MedianPruner") _, without = _run("NopPruner") - assert with_pruning.epochs_run < without.epochs_run, ( - f"pruning ran {with_pruning.epochs_run} epochs and no-pruning ran " - f"{without.epochs_run}: the trials were cut short on paper only" + # Pruned trials die on their first epoch: their opening score is already + # below the median. +1 trial in both: the refit `optimize` does at the end. + completed = STARTUP_TRIALS + 1 + expected = completed * EPOCHS + (N_TRIALS - STARTUP_TRIALS) + + assert without.epochs_run == (N_TRIALS + 1) * EPOCHS + assert with_pruning.epochs_run == expected, ( + f"pruning ran {with_pruning.epochs_run} epochs, expected {expected}. " + f"Reaching {without.epochs_run} means the trials were cut short on paper " + "only and training kept going." ) def test_disabled_pruning_completes_every_trial(): - """The control. With pruning off nothing may be cut, or the test above proves nothing.""" + """The control. + + With pruning off nothing may be cut, or the test above proves nothing. + """ optimizer, model = _run("NopPruner") states = _states(optimizer) assert all(s is optuna.trial.TrialState.COMPLETE for s in states) - assert model.epochs_run == EPOCHS * (len(states) + 1) # +1: the final refit + assert len(states) == N_TRIALS + assert model.epochs_run == EPOCHS * (N_TRIALS + 1) # +1: the final refit From c958cd42514ee7dd04dc9e1b496273d2212741c8 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 16:14:38 +0000 Subject: [PATCH 5/5] test(optimizers): prove the wiring on a real DashAI model The existing integration test asserts the pruning verdict with a stand-in model whose training loop is three lines. That shape is right for asserting a verdict deterministically, but it cannot answer the question this change is actually about: does the number a pruner needs appear when a model that ships with DashAI trains? It did not. The epoch loops guard their validation metrics behind `if x_validation is not None`, and `optimize` never passed validation data, so `calculate_metrics(split=VALIDATION, level=EPOCH)` was skipped for every epoch of every trial. Nothing was reported, so nothing could be pruned, whatever the pruner was. This uses `MLPImageClassifier` unmodified -- its real `train`, its real epoch loop, its real `calculate_metrics`, and the real `Accuracy` metric -- on 24 synthetic 16x16 images. The only stub is `_save_metrics`, which needs a database. Four tests: every trial carries one intermediate value per epoch; a model trained without validation data reports nothing (the bug, as a negative control); train-split and step-level metrics never reach the hook; and the hook fires at epoch level for validation only. Pruning is disabled here on purpose. Whether a given trial deserves to be cut is the pruner's policy, asserted deterministically next door; what is under test here is that a real model produces the evidence that policy runs on. Checked with the source reverted: dropping the validation arguments from `optimize` fails the first test (0 values reported instead of 3), and removing the hook call from `calculate_metrics` fails three of the four. Costs about 10 seconds. Full suite: 813 passed; the 1 failure and 6 errors are the frontend build and the HuggingFace downloads, identical on a clean develop checkout here. --- .../back/optimizers/test_optuna_real_model.py | 216 ++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 tests/back/optimizers/test_optuna_real_model.py diff --git a/tests/back/optimizers/test_optuna_real_model.py b/tests/back/optimizers/test_optuna_real_model.py new file mode 100644 index 000000000..b35ace03c --- /dev/null +++ b/tests/back/optimizers/test_optuna_real_model.py @@ -0,0 +1,216 @@ +"""The same wiring, but against a model that ships with DashAI. + +`test_optuna_pruning_integration.py` proves the outcome — Optuna records +trials as PRUNED — using a stand-in model whose training loop is three lines. +That is the right shape for asserting a pruning verdict deterministically, but +it cannot answer the question that matters for this feature: does the number a +pruner needs actually appear when a *real* model trains? + +It did not. The epoch loops guard their validation metrics behind +`if x_validation is not None`, and `optimize` never passed validation data, so +`calculate_metrics(split=VALIDATION, level=EPOCH)` was skipped for every epoch +of every trial. Nothing was reported, so nothing could ever be pruned — no +matter what the pruner was. + +This file uses `MLPImageClassifier` unmodified: its real `train`, its real +epoch loop, its real `calculate_metrics`, and the real `Accuracy` metric, on a +small synthetic image dataset. The only stub is `_save_metrics`, which needs a +database and is not what this proves. + +The dataset is tiny (24 images of 16x16) and the models are three epochs wide +on purpose: this is a wiring test, and it should not cost a minute of CI. +""" + +import os +import tempfile +import zipfile + +import numpy as np +import pytest +from PIL import Image + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.dataloaders.classes.dashai_dataset import ( + select_columns, + split_dataset, + split_indexes, +) +from DashAI.back.dataloaders.classes.image_dataloader import ImageDataLoader +from DashAI.back.metrics.classification.accuracy import Accuracy +from DashAI.back.models.mlp_image_classifier import MLPImageClassifier +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer + +EPOCHS = 3 +N_TRIALS = 3 + + +@pytest.fixture(scope="module") +def image_splits(tmp_path_factory): + """24 synthetic images in three classes, split train/test/validation. + + Each class is a different dominant colour channel over noise, so the + problem is learnable. It does not need to be learned well — this asserts + that per-epoch scores reach Optuna, not that they are good. + """ + rng = np.random.default_rng(0) + tmp = tmp_path_factory.mktemp("images") + img_dir = tmp / "imgs" + for cls in range(3): + cls_dir = img_dir / f"class_{cls}" + cls_dir.mkdir(parents=True) + for i in range(8): + arr = rng.integers(0, 120, (16, 16, 3), dtype=np.int16) + arr[:, :, cls] += 120 + Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8)).save( + cls_dir / f"img_{i}.png" + ) + + zip_path = tmp / "imgs.zip" + with zipfile.ZipFile(zip_path, "w") as zf: + for root, _, files in os.walk(img_dir): + for f in files: + full = os.path.join(root, f) + zf.write(full, os.path.relpath(full, img_dir)) + + dataset = ImageDataLoader().load_data( + filepath_or_buffer=str(zip_path), + temp_path=tempfile.mkdtemp(), + params={}, + ) + train_idx, test_idx, val_idx = split_indexes( + total_rows=dataset.num_rows, train_size=0.5, test_size=0.25, val_size=0.25 + ) + split = split_dataset( + dataset, + train_indexes=train_idx, + test_indexes=test_idx, + val_indexes=val_idx, + ) + x, y = select_columns(split, ["image"], ["label"]) + return split_dataset(x), split_dataset(y) + + +def _model(): + model = MLPImageClassifier( + epochs=EPOCHS, learning_rate=0.01, hidden_dims=[8], image_size=16 + ) + model.run_id = 1 + model.x_data = None + model.y_data = None + model.train_metrics = None + model.validation_metrics = [Accuracy] + model.test_metrics = None + # Persistence needs a database and is not what this proves. + model._save_metrics = lambda **kwargs: None + return model + + +def test_a_real_model_reports_every_epoch_to_the_trial(image_splits): + """Each trial must carry one intermediate value per epoch. + + This is the assertion that fails without the fix: with no validation data + reaching `train`, `intermediate_values` is empty for every trial, and an + empty history is a pruner that can never fire. + + Pruning is disabled here on purpose. Whether a given trial deserves to be + cut is the pruner's policy and is asserted next door with a deterministic + stand-in; what is under test here is that a real model produces the + evidence that policy runs on. + """ + model = _model() + x, y = image_splits + + optimizer = OptunaOptimizer( + n_trials=N_TRIALS, sampler="RandomSampler", pruner="None" + ) + optimizer.optimize( + model, + x, + y, + [(model, "learning_rate", (1e-3, 5e-2), "number")], + {"class": Accuracy, "metadata": {"maximize": True}}, + "ImageClassificationTask", + ) + + reported = [len(t.intermediate_values) for t in optimizer.study.trials] + assert reported == [EPOCHS] * N_TRIALS, ( + f"trials reported {reported} epoch scores, expected {EPOCHS} each. " + "An empty history means the epoch loop skipped its validation metrics, " + "which is what left the pruner with nothing to decide on." + ) + + +def test_without_validation_data_nothing_is_reported(image_splits): + """The negative control, and the bug this feature had, in one call. + + Training a real model without validation data must leave the reporter + untouched: the epoch loop skips `calculate_metrics(VALIDATION, EPOCH)` + entirely. This is exactly what `optimize` used to do on every trial, so if + this test ever starts reporting, the assertion above stops proving that + `optimize` passes the data. + """ + model = _model() + x, y = image_splits + calls = [] + model._epoch_reporter = lambda results, step: calls.append((results, step)) + + model.train(x["train"], y["train"]) + + assert calls == [], ( + "the reporter fired without validation data, so the assertion that " + "`optimize` must pass it no longer distinguishes anything" + ) + + +def test_the_hook_only_fires_for_validation_epochs(image_splits): + """Train-split and step-level metrics must not reach the pruner. + + A pruner fed the training score prunes on how well the model memorises, + and one fed per-step noise prunes on a number that has not settled. + """ + model = _model() + x, y = image_splits + seen = [] + model._epoch_reporter = lambda results, step: seen.append(step) + + model.train(x["train"], y["train"], x["validation"], y["validation"]) + + assert seen == list(range(1, EPOCHS + 1)), ( + f"reported steps were {seen}; expected one per epoch, validation only" + ) + + +def test_calculate_metrics_reaches_the_hook_at_epoch_level(image_splits): + """Guards the hook against the levels it must ignore, on the real model.""" + model = _model() + x, y = image_splits + model.train(x["train"], y["train"]) + + seen = [] + model._epoch_reporter = lambda results, step: seen.append((results, step)) + + for split, level in ( + (SplitEnum.TRAIN, LevelEnum.EPOCH), + (SplitEnum.VALIDATION, LevelEnum.STEP), + (SplitEnum.VALIDATION, LevelEnum.TRIAL), + ): + model.calculate_metrics( + split=split, + level=level, + x_data=x["validation"], + y_data=y["validation"], + log_index=1, + ) + assert seen == [], f"the hook fired for {seen}, which the pruner must not see" + + model.calculate_metrics( + split=SplitEnum.VALIDATION, + level=LevelEnum.EPOCH, + x_data=x["validation"], + y_data=y["validation"], + log_index=1, + ) + assert len(seen) == 1 + results, step = seen[0] + assert step == 1 + assert Accuracy.__name__ in results