Replace get_generator_run_limit function on GenerationStep with new m…

…ethod on GenerationNode

Summary:
This diff does the following:
Replaces the `get_generator_run_limit()` method on GenerationStep with `generator_run_limit` on GenerationNode. The new method relies on transition criterion to determine the number of generator runs, and only checks criterion that are trial based. I actually think this may not need to be expanded because the trial based criterion seem the most related to new generator run creation, but it could be expanded easily in the future if a usecase requires doing so.

upcoming:
(0) Finish removing GenerationStep methods in
(1) delete functions from GenStep that aren't needed anymore
(2) update the storage to include nodes independently (and not just as part of step)
(3) final pass on all the doc strings
(4) add transition criterion to the repr string + some of the other fields that havent made it yet on GeneratinoNode
(5) Do a final pass of the generationStrategy/GenerationNode files to see what else can be migrated/condensed
(6) rename transiton criterion to action criterion

Reviewed By: lena-kashtelyan

Differential Revision: D51169425

ax/modelbridge/generation_node.py

@@ -38,6 +38,7 @@
     MaxTrials,
     MinTrials,
     TransitionCriterion,
+    TrialBasedCriterion,
 )
 from ax.utils.common.base import Base, SortableBase
 from ax.utils.common.logger import get_logger
@@ -467,6 +468,33 @@ def should_transition_to_next_node(
             return True, transition_nodes[0]
         return False, None
 
+    def generator_run_limit(self) -> int:
+        """How many generator runs can this generation strategy generate right now,
+        assuming each one of them becomes its own trial. Only considers
+        `transition_criteria` that are TrialBasedCriterion.
+
+        Returns:
+              - the number of generator runs that can currently be produced, with -1
+                meaning unlimited generator runs,
+        """
+        gen_blocking_criterion_delta_from_threshold = [
+            criterion.num_till_threshold(
+                experiment=self.experiment, trials_from_node=self.trials_from_node
+            )
+            for criterion in self.transition_criteria
+            if criterion.block_gen_if_met and isinstance(criterion, TrialBasedCriterion)
+        ]
+
+        if len(gen_blocking_criterion_delta_from_threshold) == 0:
+            if not self.gen_unlimited_trials:
+                logger.warning(
+                    "Even though this node is not flagged for generation of unlimited "
+                    "trials, there are no generation blocking criterion, therefore, "
+                    "unlimited trials will be generated."
+                )
+            return -1
+        return min(gen_blocking_criterion_delta_from_threshold)
+
     def __repr__(self) -> str:
         "String representation of this GenerationNode"
         # add model specs

ax/modelbridge/generation_strategy.py

@@ -308,7 +308,7 @@ def current_generator_run_limit(
             return 0, True
 
         # if the generation strategy is not complete, optimization is not complete
-        return self._curr.get_generator_run_limit(), False
+        return self._curr.generator_run_limit(), False
 
     def clone_reset(self) -> GenerationStrategy:
         """Copy this generation strategy without it's state."""