diff --git a/ax/modelbridge/tests/test_generation_strategy.py b/ax/modelbridge/tests/test_generation_strategy.py index 07eceb3f3de..958571e44f8 100644 --- a/ax/modelbridge/tests/test_generation_strategy.py +++ b/ax/modelbridge/tests/test_generation_strategy.py @@ -42,7 +42,7 @@ ) from ax.modelbridge.torch import TorchModelBridge from ax.modelbridge.transition_criterion import ( - AutoTransitionAfterGenCriterion, + AutoTransitionAfterGen, MaxGenerationParallelism, MaxTrials, MinTrials, @@ -206,9 +206,7 @@ def setUp(self) -> None: only_in_statuses=[TrialStatus.COMPLETED], use_all_trials_in_exp=True, ) - self.gpei_to_sobol_auto = AutoTransitionAfterGenCriterion( - transition_to="sobol_3" - ) + self.gpei_to_sobol_auto = AutoTransitionAfterGen(transition_to="sobol_3") self.competing_tc_gs = GenerationStrategy( nodes=[ GenerationNode( @@ -246,7 +244,7 @@ def setUp(self) -> None: node_name="gpei", model_specs=[self.gpei_model_spec], transition_criteria=[ - AutoTransitionAfterGenCriterion( + AutoTransitionAfterGen( transition_to="sobol_2", ) ], @@ -255,7 +253,7 @@ def setUp(self) -> None: node_name="sobol_2", model_specs=[self.sobol_model_spec], transition_criteria=[ - AutoTransitionAfterGenCriterion(transition_to="sobol_3") + AutoTransitionAfterGen(transition_to="sobol_3") ], ), GenerationNode( @@ -269,7 +267,7 @@ def setUp(self) -> None: only_in_statuses=[TrialStatus.RUNNING], use_all_trials_in_exp=True, ), - AutoTransitionAfterGenCriterion( + AutoTransitionAfterGen( transition_to="gpei", block_transition_if_unmet=True, continue_trial_generation=False, @@ -1360,7 +1358,7 @@ def test_transition_edges(self) -> None: # this gs has a single sobol node which transitions to gpei. If the MaxTrials # and MinTrials criterion are met, the transition to sobol_2 should occur, # otherwise, should transition back to sobol. - gpei_to_sobol_auto = AutoTransitionAfterGenCriterion(transition_to="sobol") + gpei_to_sobol_auto = AutoTransitionAfterGen(transition_to="sobol") gs = GenerationStrategy( nodes=[ GenerationNode( @@ -1481,21 +1479,21 @@ def test_node_gs_with_auto_transitions(self) -> None: node_name="gpei", model_specs=[self.gpei_model_spec], transition_criteria=[ - AutoTransitionAfterGenCriterion(transition_to="sobol_2") + AutoTransitionAfterGen(transition_to="sobol_2") ], ), GenerationNode( node_name="sobol_2", model_specs=[self.sobol_model_spec], transition_criteria=[ - AutoTransitionAfterGenCriterion(transition_to="sobol_3") + AutoTransitionAfterGen(transition_to="sobol_3") ], ), GenerationNode( node_name="sobol_3", model_specs=[self.sobol_model_spec], transition_criteria=[ - AutoTransitionAfterGenCriterion( + AutoTransitionAfterGen( transition_to="gpei", block_transition_if_unmet=True, continue_trial_generation=False, diff --git a/ax/modelbridge/tests/test_transition_criterion.py b/ax/modelbridge/tests/test_transition_criterion.py index 2eeef19035d..d00d7e5e400 100644 --- a/ax/modelbridge/tests/test_transition_criterion.py +++ b/ax/modelbridge/tests/test_transition_criterion.py @@ -20,7 +20,7 @@ from ax.modelbridge.model_spec import ModelSpec from ax.modelbridge.registry import Models from ax.modelbridge.transition_criterion import ( - AutoTransitionAfterGenCriterion, + AutoTransitionAfterGen, MaxGenerationParallelism, MaxTrials, MinimumPreferenceOccurances, @@ -232,7 +232,7 @@ def test_min_trials_is_met(self) -> None: self.assertTrue(min_criterion.is_met(experiment, gs._steps[0].trials_from_node)) def test_auto_transition(self) -> None: - """Very simple test to validate AutoTransitionAfterGenCriterion""" + """Very simple test to validate AutoTransitionAfterGen""" experiment = self.branin_experiment gs = GenerationStrategy( name="test", @@ -241,7 +241,7 @@ def test_auto_transition(self) -> None: node_name="sobol_1", model_specs=[self.sobol_model_spec], transition_criteria=[ - AutoTransitionAfterGenCriterion(transition_to="sobol_2") + AutoTransitionAfterGen(transition_to="sobol_2") ], ), GenerationNode( @@ -488,12 +488,10 @@ def test_repr(self) -> None: + "'use_all_trials_in_exp': False, " + "'continue_trial_generation': True})", ) - auto_transition = AutoTransitionAfterGenCriterion( - transition_to="GenerationStep_2" - ) + auto_transition = AutoTransitionAfterGen(transition_to="GenerationStep_2") self.assertEqual( str(auto_transition), - "AutoTransitionAfterGenCriterion({'transition_to': 'GenerationStep_2', " + "AutoTransitionAfterGen({'transition_to': 'GenerationStep_2', " + "'block_transition_if_unmet': True, " + "'continue_trial_generation': True})", ) diff --git a/ax/modelbridge/transition_criterion.py b/ax/modelbridge/transition_criterion.py index 06d815ed5f4..243a37d6208 100644 --- a/ax/modelbridge/transition_criterion.py +++ b/ax/modelbridge/transition_criterion.py @@ -106,7 +106,7 @@ def _unique_id(self) -> str: return str(self) -class AutoTransitionAfterGenCriterion(TransitionCriterion): +class AutoTransitionAfterGen(TransitionCriterion): """A class to designate automatic transition from one GenerationNode to another. Args: diff --git a/ax/storage/json_store/registry.py b/ax/storage/json_store/registry.py index ba0ea54eeed..676b75422a1 100644 --- a/ax/storage/json_store/registry.py +++ b/ax/storage/json_store/registry.py @@ -92,7 +92,7 @@ from ax.modelbridge.registry import ModelRegistryBase from ax.modelbridge.transforms.base import Transform from ax.modelbridge.transition_criterion import ( - AutoTransitionAfterGenCriterion, + AutoTransitionAfterGen, MaxGenerationParallelism, MaxTrials, MinimumPreferenceOccurances, @@ -184,7 +184,7 @@ AndEarlyStoppingStrategy: logical_early_stopping_strategy_to_dict, AugmentedBraninMetric: metric_to_dict, AugmentedHartmann6Metric: metric_to_dict, - AutoTransitionAfterGenCriterion: transition_criterion_to_dict, + AutoTransitionAfterGen: transition_criterion_to_dict, BatchTrial: batch_to_dict, BenchmarkMetric: metric_to_dict, BoTorchModel: botorch_model_to_dict, @@ -292,7 +292,7 @@ "AndEarlyStoppingStrategy": AndEarlyStoppingStrategy, "AugmentedBraninMetric": AugmentedBraninMetric, "AugmentedHartmann6Metric": AugmentedHartmann6Metric, - "AutoTransitionAfterGenCriterion": AutoTransitionAfterGenCriterion, + "AutoTransitionAfterGen": AutoTransitionAfterGen, "Arm": Arm, "AggregatedBenchmarkResult": AggregatedBenchmarkResult, "BatchTrial": BatchTrial, diff --git a/ax/utils/testing/modeling_stubs.py b/ax/utils/testing/modeling_stubs.py index 4c764a737d2..34aba441811 100644 --- a/ax/utils/testing/modeling_stubs.py +++ b/ax/utils/testing/modeling_stubs.py @@ -31,7 +31,7 @@ from ax.modelbridge.transforms.base import Transform from ax.modelbridge.transforms.int_to_float import IntToFloat from ax.modelbridge.transition_criterion import ( - AutoTransitionAfterGenCriterion, + AutoTransitionAfterGen, MaxGenerationParallelism, MaxTrials, MinimumPreferenceOccurances, @@ -257,7 +257,7 @@ def sobol_gpei_generation_node_gs( not_in_statuses=None, ), ] - alt_mbm_criterion = [AutoTransitionAfterGenCriterion(transition_to="MBM_node")] + alt_mbm_criterion = [AutoTransitionAfterGen(transition_to="MBM_node")] step_model_kwargs = {"silently_filter_kwargs": True} sobol_model_spec = ModelSpec( model_enum=Models.SOBOL,