Split TransitionCriterion into TransitionCriterion and GenerationCriterion (facebook#4854)

Summary:
Pull Request resolved: facebook#4854

**TLDR:** This diff splits TransitionCriterion into (1) TransitionCriterion and (2) GenerationBlockingCriterion. I think this makes sense to do because it *greatly* increases the conceptual clarity of the transition criterion. Some ways it does this include:

1. Removal of confusing dual purpose flags — block_transition_if_unmet and block_generation_if_met flags. Now transition criteria are inferred to block transition if unmet and generation criteria are inferred to raise informative errors if the criteria is met.
2. Each criterion contains less flags, and the flags are more directly intuitive.
3. With upcoming removal of special logic for online, we will need to add more generation blocking criteria (ie do we have an opt config), it is better to make this change before adding more criteria that will need to be migrated
4. It will allows the logic for transition and generation to be smoother — this diff keeps things ~= to exisiting logic as possible to minimize diff review overhead, but in subsequent diffs we can save fit time if we know we can’t generate from this node + can’t transition. It will also allow for some further clarification on generation/transition blocking logic that i think is contributing to the confusion of the file
5. i like that creating a new generation blocking criteria with a specific error to raise is easy and painless

**Cons of this change:**
- it’s a large change, sorry about that.
- There is some duplication between TrialBased transition criterion and generation criterion. I explored using a Mixin here, but i find mixins tend to add unnecessary inheritance structures to reason about.

**Most important files for review, in order of importance**
1. transition_criterion.py
2. generation_node.py
3. decoder.py
4. encoders.py
5. registery.py
6. generation_strategy_dispatch.py
7. generation_nodes.py
8. generation_strategy.py

The remaining files are mainly trivial updates to tests

**Note about backwards compatibility:**
* This diff will directly decode legacy MaxGenerationParallelism as a generation blocking criterion called MaxGenrationParallelism
* Historically, there are some instances of mintrials that have block_gen_if_met=True, this usually comes from enforce_num_trials=True. Now we call this MaxTrialsAwaitingData, and MinTrials is decoded as that. I am open to other, better names for this new criterion.

**Other notes/potential improvements:**
- we could split transition criterion, generation criterion, and utils into their own files. i kinda like them together, and if we do want to do this split i’d like to do it in a follow up to try to minimize an already v large blast radius

Reviewed By: lena-kashtelyan

Differential Revision: D92201085

fbshipit-source-id: 27a8b6c6afe81d35d1c41b7e45c384a2934c3257

Author: mgarrard

ax/api/utils/generation_strategy_dispatch.py

@@ -22,7 +22,7 @@
     GenerationStrategy,
 )
 from ax.generation_strategy.generator_spec import GeneratorSpec
-from ax.generation_strategy.transition_criterion import MinTrials
+from ax.generation_strategy.transition_criterion import MaxTrialsAwaitingData, MinTrials
 from ax.generators.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
 from botorch.acquisition.acquisition import AcquisitionFunction
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
@@ -71,21 +71,28 @@ def _get_sobol_node(
         MinTrials(  # This represents the initialization budget.
             threshold=initialization_budget,
             transition_to="MBM",
-            block_gen_if_met=(not allow_exceeding_initialization_budget),
-            block_transition_if_unmet=True,
             use_all_trials_in_exp=use_existing_trials_for_initialization,
             not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
         ),
         MinTrials(  # This represents minimum observed trials requirement.
             threshold=min_observed_initialization_trials,
             transition_to="MBM",
-            block_gen_if_met=False,
-            block_transition_if_unmet=True,
             use_all_trials_in_exp=True,
             only_in_statuses=[TrialStatus.COMPLETED],
             count_only_trials_with_data=True,
         ),
     ]
+    # If we want to enforce the initialization budget, add a pausing
+    # criterion that prevents exceeding the budget.
+    pausing_criteria = None
+    if not allow_exceeding_initialization_budget:
+        pausing_criteria = [
+            MaxTrialsAwaitingData(
+                threshold=initialization_budget,
+                not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
+                use_all_trials_in_exp=use_existing_trials_for_initialization,
+            )
+        ]
     return GenerationNode(
         name="Sobol",
         generator_specs=[
@@ -95,6 +102,7 @@ def _get_sobol_node(
             )
         ],
         transition_criteria=transition_criteria,
+        pausing_criteria=pausing_criteria,
         should_deduplicate=True,
         suggested_experiment_status=ExperimentStatus.INITIALIZATION,
     )

ax/api/utils/tests/test_generation_strategy_dispatch.py

@@ -19,7 +19,7 @@
 from ax.exceptions.core import UserInputError
 from ax.generation_strategy.center_generation_node import CenterGenerationNode
 from ax.generation_strategy.dispatch_utils import get_derelativize_config
-from ax.generation_strategy.transition_criterion import MinTrials
+from ax.generation_strategy.transition_criterion import MaxTrialsAwaitingData, MinTrials
 from ax.generators.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
@@ -99,16 +99,12 @@ def test_choose_gs_fast_with_options(self) -> None:
             MinTrials(
                 threshold=2,
                 transition_to="MBM",
-                block_gen_if_met=False,
-                block_transition_if_unmet=True,
                 use_all_trials_in_exp=False,
                 not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
             ),
             MinTrials(
                 threshold=4,
                 transition_to="MBM",
-                block_gen_if_met=False,
-                block_transition_if_unmet=True,
                 use_all_trials_in_exp=True,
                 only_in_statuses=[TrialStatus.COMPLETED],
                 count_only_trials_with_data=True,
@@ -390,7 +386,18 @@ def test_abandoned_and_failed_trials_excluded_from_initialization_budget(
             first_tc.not_in_statuses, [TrialStatus.FAILED, TrialStatus.ABANDONED]
         )
         self.assertEqual(first_tc.threshold, 5)
-        self.assertTrue(first_tc.block_gen_if_met)
+        # Verify MaxTrialsAwaitingData is in pausing_criteria
+        pausing_criteria = [
+            pc
+            for pc in sobol_node._pausing_criteria
+            if isinstance(pc, MaxTrialsAwaitingData)
+        ]
+        self.assertEqual(len(pausing_criteria), 1)
+        self.assertEqual(pausing_criteria[0].threshold, 5)
+        self.assertEqual(
+            pausing_criteria[0].not_in_statuses,
+            [TrialStatus.FAILED, TrialStatus.ABANDONED],
+        )
 
         # Test the actual behavior: Generate 5 trials, mark 3 as ABANDONED,
         # verify that Sobol can still generate more trials