From 7532c6f0d14e91b13f81c55097d2b0bfa80f6b97 Mon Sep 17 00:00:00 2001 From: Elizabeth Santorella Date: Thu, 28 Mar 2024 08:45:27 -0700 Subject: [PATCH] Ax `Trial`: Bug fix for error message that references nonexistent function (#2304) Summary: `BaseTrial._validate_can_attach_data` raises the following exception: f"Trial {self.index} has already been completed with data. To add more data to it (for example, for a different metric), use `Trial.update_trial_data()` or BatchTrial.update_batch_trial_data()." This is not right because there is no `BatchTrial.update_batch_trial_data`. Also, it's an anti-pattern for a superclass to reference its subclasses like this instead of delegating to them. This PR: - Makes the behavior in the base class make sense for clases that have no method like "update trial data" - Overrides that error message for `Trial` which does have such an update method Differential Revision: D55439108 --- ax/core/base_trial.py | 18 ++++++++++++------ ax/core/tests/test_batch_trial.py | 11 +++++++++++ ax/core/tests/test_trial.py | 12 +++++++++++- ax/core/trial.py | 13 +++++++++++++ 4 files changed, 47 insertions(+), 7 deletions(-) diff --git a/ax/core/base_trial.py b/ax/core/base_trial.py index 905f140bd1c..69667f341cb 100644 --- a/ax/core/base_trial.py +++ b/ax/core/base_trial.py @@ -852,15 +852,21 @@ def _make_evaluations_and_data( ) return evaluations, data + def _raise_cant_attach_if_completed(self) -> None: + """ + Helper method used by `validate_can_attach_data` to raise an error if + the user tries to attach data to a completed trial. Subclasses such as + `Trial` override this by suggesting a remediation. + """ + raise UnsupportedError( + f"Trial {self.index} already has status 'completed', so data cannot " + "be attached." + ) + def _validate_can_attach_data(self) -> None: """Determines whether a trial is in a state that can be attached data.""" if self.status.is_completed: - raise UnsupportedError( - f"Trial {self.index} has already been completed with data." - "To add more data to it (for example, for a different metric), " - "use `Trial.update_trial_data()` or " - "BatchTrial.update_batch_trial_data()." - ) + self._raise_cant_attach_if_completed() if self.status.is_abandoned or self.status.is_failed: raise UnsupportedError( f"Trial {self.index} has been marked {self.status.name}, so it " diff --git a/ax/core/tests/test_batch_trial.py b/ax/core/tests/test_batch_trial.py index 6f5c105cfea..9cd6d110bf3 100644 --- a/ax/core/tests/test_batch_trial.py +++ b/ax/core/tests/test_batch_trial.py @@ -19,6 +19,7 @@ from ax.core.generator_run import GeneratorRun, GeneratorRunType from ax.core.parameter import FixedParameter, ParameterType from ax.core.search_space import SearchSpace +from ax.exceptions.core import UnsupportedError from ax.runners.synthetic import SyntheticRunner from ax.utils.common.testutils import TestCase from ax.utils.common.typeutils import checked_cast @@ -47,6 +48,16 @@ def setUp(self) -> None: self.weights = weights[1:] self.batch.add_arms_and_weights(arms=self.arms, weights=self.weights) + def test__validate_can_attach_data(self) -> None: + self.batch.mark_running(no_runner_required=True) + self.batch.mark_completed() + + expected_msg = ( + "Trial 0 already has status 'completed', so data cannot " "be attached." + ) + with self.assertRaisesRegex(UnsupportedError, expected_msg): + self.batch._validate_can_attach_data() + def test_Eq(self) -> None: new_batch_trial = self.experiment.new_batch_trial() self.assertNotEqual(self.batch, new_batch_trial) diff --git a/ax/core/tests/test_trial.py b/ax/core/tests/test_trial.py index 36aaf3bbd00..93a051ed4ba 100644 --- a/ax/core/tests/test_trial.py +++ b/ax/core/tests/test_trial.py @@ -16,7 +16,7 @@ from ax.core.data import Data from ax.core.generator_run import GeneratorRun, GeneratorRunType from ax.core.runner import Runner -from ax.exceptions.core import UserInputError +from ax.exceptions.core import UnsupportedError, UserInputError from ax.runners.synthetic import SyntheticRunner from ax.utils.common.result import Ok from ax.utils.common.testutils import TestCase @@ -60,6 +60,16 @@ def setUp(self) -> None: def tearDown(self) -> None: self.mock_supports_trial_type.stop() + def test__validate_can_attach_data(self) -> None: + self.trial.mark_running(no_runner_required=True) + self.trial.mark_completed() + + expected_msg = ( + "Trial 0 has already been completed with data. To add more data to " + ) + with self.assertRaisesRegex(UnsupportedError, expected_msg): + self.trial._validate_can_attach_data() + def test_eq(self) -> None: new_trial = self.experiment.new_trial() self.assertNotEqual(self.trial, new_trial) diff --git a/ax/core/trial.py b/ax/core/trial.py index 74a49bd9b10..200e5fff8e0 100644 --- a/ax/core/trial.py +++ b/ax/core/trial.py @@ -19,6 +19,7 @@ from ax.core.data import Data from ax.core.generator_run import GeneratorRun, GeneratorRunType from ax.core.types import TCandidateMetadata, TEvaluationOutcome +from ax.exceptions.core import UnsupportedError from ax.utils.common.docutils import copy_doc from ax.utils.common.logger import _round_floats_for_logging, get_logger from ax.utils.common.typeutils import not_none @@ -351,3 +352,15 @@ def clone_to( new_trial.add_generator_run(self.generator_run.clone()) self._update_trial_attrs_on_clone(new_trial=new_trial) return new_trial + + def _raise_cant_attach_if_completed(self) -> None: + """ + Helper method used by `validate_can_attach_data` to raise an error if + the user tries to attach data to a completed trial. Subclasses such as + `Trial` override this by suggesting a remediation. + """ + raise UnsupportedError( + f"Trial {self.index} has already been completed with data. " + "To add more data to it (for example, for a different metric), " + f"use `{self.__class__.__name__}.update_trial_data()`." + )