fix(optimizers): make pruning actually prune - #828
Open
Maarmapa wants to merge 5 commits into
Open
Conversation
…a string
`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.
Follow-up to the review on DashAISoftware#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.
Contributor
Author
|
Ran the full suite locally on macOS: 812 passed, 1 failed. The failure is test_app_front (DashAI/front/build/index.html missing — production frontend never built on that machine), unrelated to this change. The 6 HuggingFace errors I reported earlier were network restrictions in my other environment and pass here. |
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.
… 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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replaces #823, which was closed after review. Closes #821.
@Irozuku was right that instantiating the pruner wasn't enough. Looking into why, the reason went deeper than the schema string.
The pruner had nothing to decide with
optimize()trained the model without validation data:while the normal path in
model_job.pypasses all four: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 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_image_classifier,cnn_image_classifier,mlp_image_classifier,lenet5_image_classifierandscikit_learn/mlp_regression— and they share no ancestor belowBaseModel. 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_reporteron the base class, invoked only forlevel=EPOCHandsplit=VALIDATION.Models that train in a single shot never reach that branch, so nothing changes for them.
Two details that matter:
TrialPrunedtravels from inside the loop up tostudy.optimizewith nothing in between catching it, so the trial is recorded as pruned, not failed.A metric missing from an epoch is not an error:
calculate_metricsdrops metrics returning a non-finite value (only one class present in a split, say), so the trial simply continues unpruned.Also included
Resolving the pruner name into an instance, from #823.
create_studyaccepts any object without validating it, so the raw schema string produced a study whose pruner was astr, failing later asAttributeError: 'str' object has no attribute 'prune'. Only pruners Optuna can build with default arguments are supported;PatientPruner,PercentilePrunerandThresholdPrunerneed configuration the schema does not carry yet, and say so with a clear message.Verification
Full suite: 806 passed. The remaining 1 failure and 6 errors reproduce identically on a clean
developcheckout in the same environment — the frontend build and the HuggingFace model downloads, neither reachable here.11 new cases: each epoch reported with its step, a rejected trial raising
TrialPruned, a missing metric tolerated, the reported metric being the one optimized, the hook firing for exactly epoch+validation and staying quiet for train metrics and trial-level summaries, a model without a reporter behaving as before, and an epoch's metrics saved before the hook can abort.The three doubles in
test_optuna_best_params.pynow mirror the realBaseModel.trainsignature. They declaredtrain(self, x, y), which no actual model does — every one of them, across the torch, scikit-learn and HuggingFace families, already acceptsx_validation=None, y_validation=None.Not covered
Pruning only does something for models that train in epochs. For everything else the parameter remains inert, just visibly so rather than as a latent
AttributeError.If you would still rather drop the pruner as a parameter for now, I am happy to send that instead — it touches the schema field, the frontend config step and a few test fixtures. This PR is the other branch of the choice you offered.