Simplify TransitionCriteria classes by replacing MaxGenerationParallelism with MinTrials (#4819)

Summary:
Pull Request resolved: #4819

Pros of this change:
* Reduce number of transition criterion classes --> may be easier for folks to wrap their head around the options with fewer options
* In the future we want to move this type of budget to concurrency or something like that, so we can go ahead and improve the stack clarity now

Cons of this change:
* maxgenerationparallelism was a pretty explict name that defined what that criterion was for, without it you have to reason about the inputs of mintrials to understand how to configure parallelism
* some maxgen special casing need to be checked using tc.block_gen_if_met property -- this is mainly for test cases (okay imo) and gs init where we construct the transition_to property for steps

Note:
Given we only have MinTrials as subclass of TrialBasedCriterion, we could actually compress this even further to only have TrialBasedCriterion if we would like. Totally open to that, but kept it as is as it could make developing new trialbased criterion easier to have the base class.

Differential Revision: D91358715

Author: mgarrard

ax/generation_strategy/generation_node.py

@@ -42,7 +42,6 @@
 from ax.generation_strategy.generator_spec import GeneratorSpec
 from ax.generation_strategy.transition_criterion import (
     AutoTransitionAfterGen,
-    MaxGenerationParallelism,
     MinTrials,
     TransitionCriterion,
     TrialBasedCriterion,
@@ -1128,7 +1127,7 @@ def __new__(
             )
         if max_parallelism is not None:
             transition_criteria.append(
-                MaxGenerationParallelism(
+                MinTrials(
                     threshold=max_parallelism,
                     transition_to=placeholder_transition_to,
                     only_in_statuses=[TrialStatus.RUNNING],

ax/generation_strategy/generation_strategy.py

@@ -395,8 +395,8 @@ def _validate_and_set_step_sequence(self, steps: list[GenerationNode]) -> None:
                 )
             )
             for tc in step.transition_criteria:
-                if tc.criterion_class == "MaxGenerationParallelism":
-                    # MaxGenerationParallelism transitions to self (current step)
+                # Criteria that block continued generation will transition to self
+                if tc.block_gen_if_met:
                     tc._transition_to = step.name
                 else:
                     tc._transition_to = next_step_name
@@ -434,24 +434,9 @@ def _validate_and_set_node_graph(self, nodes: list[GenerationNode]) -> None:
         # Validate transition edges:
         # - All `transition_to` targets must exist in this GS
         # - All TCs on one edge must have the same `continue_trial_generation` setting
-        #  All but `MaxGenerationParallelism` TCs must have a `transition_to` set
         for node in nodes:
             for next_node, tcs in node.transition_edges.items():
-                if next_node is None:
-                    # TODO[drfreund]: Handle the case of the last generation step not
-                    # having any transition criteria.
-                    # TODO[mgarrard]: Remove MaxGenerationParallelism check when
-                    # we update TransitionCriterion always define `transition_to`
-                    # NOTE: This is done in D86066476
-                    for tc in tcs:
-                        if "MaxGenerationParallelism" not in tc.criterion_class:
-                            raise GenerationStrategyMisconfiguredException(
-                                error_info="Only MaxGenerationParallelism transition"
-                                " criterion can have a null `transition_to` argument,"
-                                f" but {tc.criterion_class} does not define "
-                                f"`transition_to` on {node.name}."
-                            )
-                elif next_node not in node_names:
+                if next_node not in node_names:
                     raise GenerationStrategyMisconfiguredException(
                         error_info=f"`transition_to` argument "
                         f"{next_node} does not correspond to any node in"

ax/generation_strategy/tests/test_dispatch_utils.py

@@ -22,10 +22,7 @@
     DEFAULT_BAYESIAN_PARALLELISM,
 )
 from ax.generation_strategy.generation_node import GenerationNode
-from ax.generation_strategy.transition_criterion import (
-    MaxGenerationParallelism,
-    MinTrials,
-)
+from ax.generation_strategy.transition_criterion import MinTrials, TrialBasedCriterion
 from ax.generators.random.sobol import SobolGenerator
 from ax.generators.winsorization_config import WinsorizationConfig
 from ax.utils.common.testutils import TestCase
@@ -621,12 +618,12 @@ def test_enforce_sequential_optimization(self) -> None:
                 sobol_gpei._nodes[0].transition_criteria[0], MinTrials
             )
             self.assertTrue(node0_min_trials.block_gen_if_met)
-            # Check that max_parallelism is set by verifying MaxGenerationParallelism
-            # criterion exists on node 1
+            # Check that max_parallelism is set by verifying a TrialBasedCriterion
+            # with block_gen_if_met=True exists on node 1
             node1_max_parallelism = [
                 tc
                 for tc in sobol_gpei._nodes[1].transition_criteria
-                if isinstance(tc, MaxGenerationParallelism)
+                if tc.block_gen_if_met and isinstance(tc, TrialBasedCriterion)
             ]
             self.assertTrue(len(node1_max_parallelism) > 0)
         with self.subTest("False"):
@@ -647,11 +644,11 @@ def test_enforce_sequential_optimization(self) -> None:
             )
             self.assertFalse(node0_min_trials.block_gen_if_met)
             # Check that max_parallelism is None by verifying no
-            # MaxGenerationParallelism criterion exists on node 1
+            # TrialBasedCriterion with block_gen_if_met=True exists on node 1
             node1_max_parallelism = [
                 tc
                 for tc in sobol_gpei._nodes[1].transition_criteria
-                if isinstance(tc, MaxGenerationParallelism)
+                if tc.block_gen_if_met and isinstance(tc, TrialBasedCriterion)
             ]
             self.assertEqual(len(node1_max_parallelism), 0)
         with self.subTest("False and max_parallelism_override"):
@@ -820,7 +817,7 @@ def test_fixed_num_initialization_trials(self) -> None:
     def _get_max_parallelism(self, node: GenerationNode) -> int | None:
         """Helper to extract max_parallelism from transition criteria."""
         for tc in node.transition_criteria:
-            if isinstance(tc, MaxGenerationParallelism):
+            if tc.block_gen_if_met and isinstance(tc, TrialBasedCriterion):
                 return tc.threshold
         return None
 

ax/generation_strategy/tests/test_generation_strategy.py

@@ -57,7 +57,6 @@
 from ax.generation_strategy.generator_spec import GeneratorSpec
 from ax.generation_strategy.transition_criterion import (
     AutoTransitionAfterGen,
-    MaxGenerationParallelism,
     MinTrials,
 )
 from ax.generators.random.sobol import SobolGenerator
@@ -1233,7 +1232,7 @@ def test_gs_setup_with_nodes(self) -> None:
                 threshold=2,
                 transition_to="node_2",
             ),
-            MaxGenerationParallelism(
+            MinTrials(
                 threshold=1,
                 only_in_statuses=[TrialStatus.RUNNING],
                 block_gen_if_met=True,

ax/generation_strategy/tests/test_transition_criterion.py

@@ -22,7 +22,6 @@
     AutoTransitionAfterGen,
     AuxiliaryExperimentCheck,
     IsSingleObjective,
-    MaxGenerationParallelism,
     MinTrials,
 )
 from ax.utils.common.logger import get_logger
@@ -186,7 +185,7 @@ def test_default_step_criterion_setup(self) -> None:
                 threshold=2,
                 transition_to="GenerationStep_2_BoTorch",
             ),
-            MaxGenerationParallelism(
+            MinTrials(
                 threshold=1,
                 only_in_statuses=[TrialStatus.RUNNING],
                 block_gen_if_met=True,
@@ -466,7 +465,7 @@ def test_repr(self) -> None:
             + "'continue_trial_generation': False, "
             + "'count_only_trials_with_data': False})",
         )
-        max_parallelism = MaxGenerationParallelism(
+        max_parallelism_criterion = MinTrials(
             only_in_statuses=[TrialStatus.EARLY_STOPPED],
             threshold=3,
             transition_to="GenerationStep_2",
@@ -475,16 +474,17 @@ def test_repr(self) -> None:
             not_in_statuses=[TrialStatus.FAILED],
         )
         self.assertEqual(
-            str(max_parallelism),
-            "MaxGenerationParallelism({'threshold': 3, "
+            str(max_parallelism_criterion),
+            "MinTrials({'threshold': 3, "
             + "'transition_to': 'GenerationStep_2', "
             + "'only_in_statuses': "
             + "[<enum 'TrialStatus'>.EARLY_STOPPED], "
             + "'not_in_statuses': [<enum 'TrialStatus'>.FAILED], "
             + "'block_transition_if_unmet': False, "
             + "'block_gen_if_met': True, "
             + "'use_all_trials_in_exp': False, "
-            + "'continue_trial_generation': False})",
+            + "'continue_trial_generation': False, "
+            + "'count_only_trials_with_data': False})",
         )
         auto_transition = AutoTransitionAfterGen(transition_to="GenerationStep_2")
         self.assertEqual(

ax/generation_strategy/transition_criterion.py

@@ -29,7 +29,7 @@
 DATA_REQUIRED_MSG = (
     "All trials for current node {node_name} have been generated, "
     "but not enough data has been observed to proceed to the next "
-    "Generation node. Try again when more is are available."
+    "Generation node. Try again when more data is available."
 )
 
 
@@ -48,9 +48,7 @@ class TransitionCriterion(SortableBase):
             we will raise an error, otherwise we will continue to generate trials
             until ``MinTrials`` is met (thus overriding MinTrials).
         block_transition_if_unmet: A flag to prevent the node from completing and
-            being able to transition to another node. Ex: MaxGenerationParallelism
-            defaults to setting this to False since we can complete and move on from
-            this node without ever reaching its threshold.
+            being able to transition to another node.
         continue_trial_generation: A flag to indicate that all generation for a given
             trial is not completed, and thus even after transition, the next node will
             continue to generate arms for the same trial. Example usage: in
@@ -234,9 +232,7 @@ class TrialBasedCriterion(TransitionCriterion):
         threshold: The threshold as an integer for this criterion. Ex: If we want to
             generate at most 3 trials, then the threshold is 3.
         block_transition_if_unmet: A flag to prevent the node from completing and
-            being able to transition to another node. Ex: MaxGenerationParallelism
-            defaults to setting this to False since we can complete and move on from
-            this node without ever reaching its threshold.
+            being able to transition to another node.
         block_gen_if_met: A flag to prevent continued generation from the
             associated GenerationNode if this criterion is met but other criterion
             remain unmet. Ex: ``MinTrials`` has not been met yet, but
@@ -383,93 +379,17 @@ def is_met(
         )
 
 
-class MaxGenerationParallelism(TrialBasedCriterion):
-    """Specific TransitionCriterion implementation which defines the maximum number
-    of trials that can simultaneously be in the designated trial statuses. The
-    default behavior is to block generation from the associated GenerationNode if the
-    threshold is met. This is configured via the `block_gen_if_met` flag being set to
-    True. This criterion defaults to not blocking transition to another node via the
-    `block_transition_if_unmet` flag being set to False.
-
-    Args:
-        threshold: The threshold as an integer for this criterion. Ex: If we want to
-            generate at most 3 trials, then the threshold is 3.
-        only_in_statuses: A list of trial statuses to filter on when checking the
-            criterion threshold.
-        not_in_statuses: A list of trial statuses to exclude when checking the
-            criterion threshold.
-        transition_to: The name of the GenerationNode the GenerationStrategy should
-            transition to when this criterion is met, if it exists.
-        block_transition_if_unmet: A flag to prevent the node from completing and
-            being able to transition to another node. Ex: MaxGenerationParallelism
-            defaults to setting this to False since we can complete and move on from
-            this node without ever reaching its threshold.
-        block_gen_if_met: A flag to prevent continued generation from the
-            associated GenerationNode if this criterion is met but other criterion
-            remain unmet. Ex: ``MinTrials`` has not been met yet, but
-            MinTrials has been reached. If this flag is set to true on MinTrials then
-            we will raise an error, otherwise we will continue to generate trials
-            until ``MinTrials`` is met (thus overriding MinTrials).
-        use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
-            only those generated by the current GenerationNode.
-        continue_trial_generation: A flag to indicate that all generation for a given
-            trial is not completed, and thus even after transition, the next node will
-            continue to generate arms for the same trial. Example usage: in
-            ``BatchTrial``s we may  enable generation of arms within a batch from
-            different ``GenerationNodes`` by setting this flag to True. Defaults to
-            False for MaxGenerationParallelism since this criterion isn't currently
-            used for node -> node or trial -> trial transition.
-        count_only_trials_with_data: If set to True, only trials with data will be
-            counted towards the ``threshold``. Defaults to False.
+class MinTrials(TrialBasedCriterion):
     """
+    Simple class to enforce a threshold for the number of trials with the
+    designated statuses being generated by a specific GenerationNode.
 
-    def __init__(
-        self,
-        threshold: int,
-        transition_to: str,
-        only_in_statuses: list[TrialStatus] | None = None,
-        not_in_statuses: list[TrialStatus] | None = None,
-        block_transition_if_unmet: bool | None = False,
-        block_gen_if_met: bool | None = True,
-        use_all_trials_in_exp: bool | None = False,
-        continue_trial_generation: bool | None = False,
-    ) -> None:
-        super().__init__(
-            threshold=threshold,
-            only_in_statuses=only_in_statuses,
-            not_in_statuses=not_in_statuses,
-            transition_to=transition_to,
-            block_gen_if_met=block_gen_if_met,
-            block_transition_if_unmet=block_transition_if_unmet,
-            use_all_trials_in_exp=use_all_trials_in_exp,
-            continue_trial_generation=continue_trial_generation,
-        )
+    This class can be configured to behave as either:
+    - A minimum trials criterion (default): blocks transition to next node until
+        the threshold is met; block_transition_if_unmet=True
+    - A maximum parallelism criterion: blocks further generation from the current node
+        when threshold is reached; block_gen_if_met=True
 
-    def block_continued_generation_error(
-        self,
-        node_name: str,
-        experiment: Experiment,
-        trials_from_node: set[int],
-    ) -> None:
-        """Raises the appropriate error (should only be called when the
-        ``GenerationNode`` is blocked from continued generation). For this
-        class, the exception is ``MaxParallelismReachedException``.
-        """
-        assert self.block_gen_if_met  # Sanity check.
-        raise MaxParallelismReachedException(
-            node_name=node_name,
-            num_running=self.num_contributing_to_threshold(
-                experiment=experiment, trials_from_node=trials_from_node
-            ),
-        )
-
-
-class MinTrials(TrialBasedCriterion):
-    """
-    Simple class to enforce a minimum threshold for the number of trials with the
-    designated statuses being generated by a specific GenerationNode. The default
-    behavior is to block transition to the next node if the threshold is unmet, but
-    not affect continued generation.
 
     Args:
         threshold: The threshold as an integer for this criterion. Ex: If we want to
@@ -481,15 +401,12 @@ class MinTrials(TrialBasedCriterion):
         transition_to: The name of the GenerationNode the GenerationStrategy should
             transition to when this criterion is met.
         block_transition_if_unmet: A flag to prevent the node from completing and
-            being able to transition to another node. Ex: MaxGenerationParallelism
-            defaults to setting this to False since we can complete and move on from
-            this node without ever reaching its threshold.
+            being able to transition to another node. Defaults to True for minimum
+            trials behavior. Set to False for maximum parallelism behavior.
         block_gen_if_met: A flag to prevent continued generation from the
             associated GenerationNode if this criterion is met but other criterion
-            remain unmet. Ex: ``MinTrials`` has not been met yet, but
-            MinTrials has been reached. If this flag is set to true on MinTrials then
-            we will raise an error, otherwise we will continue to generate trials
-            until ``MinTrials`` is met (thus overriding MinTrials).
+            remain unmet. Defaults to False for minimum trials behavior. Set to True
+            for maximum parallelism behavior.
         use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
             only those generated by the current GenerationNode.
         continue_trial_generation: A flag to indicate that all generation for a given
@@ -531,12 +448,18 @@ def block_continued_generation_error(
         experiment: Experiment,
         trials_from_node: set[int],
     ) -> None:
-        """Raises the appropriate error (should only be called when the
-        ``GenerationNode`` is blocked from continued generation). For this
-        class, the exception is ``DataRequiredError``.
+        """Raises MaxParallelismReachedException when generation is blocked.
+
+        This method is only called when block_gen_if_met=True (max parallelism
+        behavior), which always has block_transition_if_unmet=False.
         """
         assert self.block_gen_if_met  # Sanity check.
-        raise DataRequiredError(DATA_REQUIRED_MSG.format(node_name=node_name))
+        raise MaxParallelismReachedException(
+            node_name=node_name,
+            num_running=self.num_contributing_to_threshold(
+                experiment=experiment, trials_from_node=trials_from_node
+            ),
+        )
 
 
 class AuxiliaryExperimentCheck(TransitionCriterion):
@@ -571,9 +494,7 @@ class AuxiliaryExperimentCheck(TransitionCriterion):
             we will raise an error, otherwise we will continue to generate trials
             until ``MinTrials`` is met (thus overriding MinTrials).
         block_transition_if_unmet: A flag to prevent the node from completing and
-            being able to transition to another node. Ex: MaxGenerationParallelism
-            defaults to setting this to False since we can complete and move on from
-            this node without ever reaching its threshold.
+            being able to transition to another node.
         continue_trial_generation: A flag to indicate that all generation for a given
             trial is not completed, and thus even after transition, the next node will
             continue to generate arms for the same trial. Example usage: in

ax/orchestration/tests/test_orchestrator.py

@@ -49,7 +49,7 @@
     GenerationStep,
     GenerationStrategy,
 )
-from ax.generation_strategy.transition_criterion import MaxGenerationParallelism
+from ax.generation_strategy.transition_criterion import TrialBasedCriterion
 from ax.metrics.branin import BraninMetric
 from ax.metrics.branin_map import BraninTimestampMapMetric
 from ax.orchestration.orchestrator import (
@@ -1171,7 +1171,7 @@ def test_run_trials_and_yield_results_with_early_stopper(self) -> None:
             # Extract max_parallelism from transition criteria
             node0_max_parallelism = None
             for tc in self.two_sobol_steps_GS._nodes[0].transition_criteria:
-                if isinstance(tc, MaxGenerationParallelism):
+                if tc.block_gen_if_met and isinstance(tc, TrialBasedCriterion):
                     node0_max_parallelism = tc.threshold
                     break
             self.assertEqual(