Drop critertion from AutoTransitionAfterGenCriterion name (#2625)

Summary: Pull Request resolved: #2625 Lena pointed out in a previous diff there is some inconsistency in the naming of transitioncriterion. This diff creates the following pattern: if a class is a base class for other criterion (transitioncriterion, trialbasedcriterion) it keeps criterion at the end of the name since it is a bucket of criterion, and isn't directly used (except for comparisions to see if a single critierion is of that type: aka maxtrials is of trialbasedcriterion). If a criterion is a leaf of the inheritance tree, it does not have criterion at the end of the name. Differential Revision: D60526927 fbshipit-source-id: 40e908226f6a819c2e53f3c25fd42f87f3b84aa2
facebook · Aug 7, 2024 · 0b6b160 · 0b6b160
1 parent 8ba04b7
commit 0b6b160
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 24 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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,