From 82dd8587e6558a2a07dff74d9a795dad8a6be24f Mon Sep 17 00:00:00 2001 From: Mia Garrard Date: Mon, 27 Nov 2023 10:33:05 -0800 Subject: [PATCH] Replace get_generator_run_limit function on GenerationStep with new method 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 | 28 +++++++++++++++++++++++++++ ax/modelbridge/generation_strategy.py | 2 +- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/ax/modelbridge/generation_node.py b/ax/modelbridge/generation_node.py index 5c3a3a541f2..c07c91af3ba 100644 --- a/ax/modelbridge/generation_node.py +++ b/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 diff --git a/ax/modelbridge/generation_strategy.py b/ax/modelbridge/generation_strategy.py index 0ccddcb3868..af2cdcfc5d2 100644 --- a/ax/modelbridge/generation_strategy.py +++ b/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."""