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 8b384e09a..bd75f4b22 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,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,6 +244,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 +252,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"] @@ -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 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 new file mode 100644 index 000000000..2f5b7016f --- /dev/null +++ b/tests/back/optimizers/test_optuna_pruner.py @@ -0,0 +1,132 @@ +"""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, _report_epoch + + +@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) + + +# --- 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)] 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..8e676c855 --- /dev/null +++ b/tests/back/optimizers/test_optuna_pruning_integration.py @@ -0,0 +1,173 @@ +"""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. + +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 +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 +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: + """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 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. + + `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): + self.run_id = 1 + 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, + # 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 + quality = 1.0 / (1.0 + self.trials_run) + self.trials_run += 1 + for epoch in range(EPOCHS): + self.value += quality + 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=N_TRIALS): + 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 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." + ) + + +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") + + # 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. + """ + optimizer, model = _run("NopPruner") + + states = _states(optimizer) + assert all(s is optuna.trial.TrialState.COMPLETE for s in states) + assert len(states) == N_TRIALS + assert model.epochs_run == EPOCHS * (N_TRIALS + 1) # +1: the final refit 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