From 3a2d2fd541d31f53335c63a79e644cb843a7a119 Mon Sep 17 00:00:00 2001 From: Mia Garrard Date: Tue, 21 May 2024 18:54:03 -0700 Subject: [PATCH 1/2] Update gen from multiple nodes to support generating multiple arms per node (#2469) Summary: This diff adds the functionality to generate multiple generator runs per node in a trial that uses multiple nodes. In follow up diffs we will: 1. Support fixed features -- there is ongoing discussion about fixed features on lena-kashtelyan's recent diff, and we may be able to leverage the solution there/help build it out 2. merge this gen method with our standard gen method Reviewed By: saitcakmak, lena-kashtelyan Differential Revision: D57398578 --- ax/modelbridge/generation_strategy.py | 35 +++- .../tests/test_generation_strategy.py | 161 +++++++++++++----- 2 files changed, 144 insertions(+), 52 deletions(-) diff --git a/ax/modelbridge/generation_strategy.py b/ax/modelbridge/generation_strategy.py index 96bc7881f96..170d20208d5 100644 --- a/ax/modelbridge/generation_strategy.py +++ b/ax/modelbridge/generation_strategy.py @@ -379,9 +379,8 @@ def gen_with_multiple_nodes( self, experiment: Experiment, data: Optional[Data] = None, - n: int = 1, # Total arms to generate pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, - **kwargs: Any, # TODO: @mgarrard Ensure correct dispatch to nodes + arms_per_node: Optional[Dict[str, int]] = None, ) -> List[GeneratorRun]: """Produces a List of GeneratorRuns for a single trial, either ``Trial`` or ``BatchTrial``, and if producing a ``BatchTrial`` allows for multiple @@ -398,27 +397,53 @@ def gen_with_multiple_nodes( data: Optional data to be passed to the underlying model's `gen`, which is called within this method and actually produces the resulting generator run. By default, data is all data on the `experiment`. - n: Integer representing how total arms to generate for this trial. pending_observations: A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated. + arms_per_node: A map from node name to the number of arms to generate + from that node. If not provided, the number of arms to generate + from each node defaults to one. Returns: A list of ``GeneratorRuns`` for a single trial. """ # TODO: @mgarrard merge into gen method, just starting here to derisk + # Validate `arms_per_node` if specified, otherwise construct the default + # behavior with keys being node names and values being 1 to represent + # generating a single GR from each node. + node_names = [node.node_name for node in self._nodes] + if arms_per_node is not None and not all( + node_name in arms_per_node for node_name in node_names + ): + raise UserInputError( + f""" + Each node defined in the GenerationStrategy must have + an associated number of arms to generate from that node + defined in `arms_per_node`. {arms_per_node} does not + include all of {node_names}. It may be helpful to double check + the spelling. + """ + ) + if arms_per_node is None: + arms_per_node = {node_name: 1 for node_name in node_names} grs = [] continue_gen_for_trial = True while continue_gen_for_trial: + next_node_name = self.current_node_name + should_transition, next_node = self._curr.should_transition_to_next_node( + raise_data_required_error=False + ) + if should_transition: + assert next_node is not None + next_node_name = next_node grs.extend( self._gen_multiple( experiment=experiment, num_generator_runs=1, data=data, - n=n, + n=arms_per_node[next_node_name], pending_observations=pending_observations, - **kwargs, ) ) continue_gen_for_trial = self._should_continue_gen_for_trial() diff --git a/ax/modelbridge/tests/test_generation_strategy.py b/ax/modelbridge/tests/test_generation_strategy.py index f7c0ec88ab3..0780bad4245 100644 --- a/ax/modelbridge/tests/test_generation_strategy.py +++ b/ax/modelbridge/tests/test_generation_strategy.py @@ -233,6 +233,53 @@ def setUp(self) -> None: ), ], ) + self.complex_multinode_per_trial_gs = GenerationStrategy( + nodes=[ + GenerationNode( + node_name="sobol", + model_specs=[self.sobol_model_spec], + transition_criteria=self.single_running_trial_criterion, + ), + GenerationNode( + node_name="gpei", + model_specs=[self.gpei_model_spec], + transition_criteria=[ + AutoTransitionAfterGenCriterion( + transition_to="sobol_2", + ) + ], + ), + GenerationNode( + node_name="sobol_2", + model_specs=[self.sobol_model_spec], + transition_criteria=[ + AutoTransitionAfterGenCriterion(transition_to="sobol_3") + ], + ), + GenerationNode( + node_name="sobol_3", + model_specs=[self.sobol_model_spec], + transition_criteria=[ + MaxTrials( + threshold=2, + transition_to="sobol_4", + block_transition_if_unmet=True, + only_in_statuses=[TrialStatus.RUNNING], + use_all_trials_in_exp=True, + ), + AutoTransitionAfterGenCriterion( + transition_to="gpei", + block_transition_if_unmet=True, + continue_trial_generation=False, + ), + ], + ), + GenerationNode( + node_name="sobol_4", + model_specs=[self.sobol_model_spec], + ), + ], + ) def tearDown(self) -> None: self.torch_model_bridge_patcher.stop() @@ -1453,6 +1500,70 @@ def test_transition_edges(self) -> None: }, ) + def test_multiple_arms_per_node(self) -> None: + """Test that a ``GenerationStrategy`` which expects some trials to be composed + of multiple nodes can generate multiple arms per node using `arms_per_node`. + """ + exp = get_branin_experiment() + gs = self.complex_multinode_per_trial_gs + gs.experiment = exp + # first check that arms_per node validation works + arms_per_node = { + "sobol": 3, + "sobol_2": 2, + "sobol_3": 1, + "sobol_4": 4, + } + with self.assertRaisesRegex(UserInputError, "defined in `arms_per_node`"): + gs.gen_with_multiple_nodes(exp, arms_per_node=arms_per_node) + + # now we will check that the first trial contains 3 arms, the sconed trial + # contains 6 arms (2 from gpei, 1 from sobol_2, 3 from sobol_3), and all + # remaining trials contain 4 arms + arms_per_node = { + "sobol": 3, + "gpei": 1, + "sobol_2": 2, + "sobol_3": 3, + "sobol_4": 4, + } + # for the first trial, we start on sobol, we generate the trial, but it hasn't + # been run yet, so we remain on sobol + trial0 = exp.new_batch_trial( + generator_runs=gs.gen_with_multiple_nodes(exp, arms_per_node=arms_per_node) + ) + self.assertEqual(len(trial0.arms_by_name), 3) + self.assertEqual(trial0.generator_runs[0]._generation_node_name, "sobol") + trial0.run() + + # after trial 0 is run, we create a trial with nodes gpei, sobol_2, and sobol_3 + # However, the sobol_3 criterion requires that we have two running trials. We + # don't move onto sobol_4 until we have two running trials, instead we reset + # to the last first node in a trial. + for _i in range(0, 2): + trial = exp.new_batch_trial( + generator_runs=gs.gen_with_multiple_nodes( + exp, arms_per_node=arms_per_node + ) + ) + self.assertEqual(gs.current_node_name, "sobol_3") + self.assertEqual(len(trial.arms_by_name), 6) + self.assertEqual(len(trial.generator_runs), 3) + self.assertEqual(trial.generator_runs[0]._generation_node_name, "gpei") + self.assertEqual(len(trial.generator_runs[0].arms), 1) + self.assertEqual(trial.generator_runs[1]._generation_node_name, "sobol_2") + self.assertEqual(len(trial.generator_runs[1].arms), 2) + self.assertEqual(trial.generator_runs[2]._generation_node_name, "sobol_3") + self.assertEqual(len(trial.generator_runs[2].arms), 3) + + # after running the next trial should be made from sobol 4 + trial.run() + trial = exp.new_batch_trial( + generator_runs=gs.gen_with_multiple_nodes(exp, arms_per_node=arms_per_node) + ) + self.assertEqual(trial.generator_runs[0]._generation_node_name, "sobol_4") + self.assertEqual(len(trial.generator_runs[0].arms), 4) + def test_node_gs_with_auto_transitions(self) -> None: """Test that node-based generation strategies which leverage AutoTransitionAfterGen criterion correctly transition and create trials. @@ -1496,6 +1607,7 @@ def test_node_gs_with_auto_transitions(self) -> None: ], ) exp = get_branin_experiment() + gs.experiment = exp # for the first trial, we start on sobol, we generate the trial, but it hasn't # been run yet, so we remain on sobol, after the trial is run, the subsequent @@ -1517,53 +1629,8 @@ def test_node_gs_with_auto_transitions(self) -> None: def test_node_gs_with_auto_transitions_three_phase(self) -> None: exp = get_branin_experiment() - gs_2 = GenerationStrategy( - nodes=[ - GenerationNode( - node_name="sobol", - model_specs=[self.sobol_model_spec], - transition_criteria=self.single_running_trial_criterion, - ), - GenerationNode( - node_name="gpei", - model_specs=[self.gpei_model_spec], - transition_criteria=[ - AutoTransitionAfterGenCriterion( - transition_to="sobol_2", - ) - ], - ), - GenerationNode( - node_name="sobol_2", - model_specs=[self.sobol_model_spec], - transition_criteria=[ - AutoTransitionAfterGenCriterion(transition_to="sobol_3") - ], - ), - GenerationNode( - node_name="sobol_3", - model_specs=[self.sobol_model_spec], - transition_criteria=[ - MaxTrials( - threshold=2, - transition_to="sobol_4", - block_transition_if_unmet=True, - only_in_statuses=[TrialStatus.RUNNING], - use_all_trials_in_exp=True, - ), - AutoTransitionAfterGenCriterion( - transition_to="gpei", - block_transition_if_unmet=True, - continue_trial_generation=False, - ), - ], - ), - GenerationNode( - node_name="sobol_4", - model_specs=[self.sobol_model_spec], - ), - ], - ) + gs_2 = self.complex_multinode_per_trial_gs + gs_2.experiment = exp # for the first trial, we start on sobol, we generate the trial, but it hasn't # been run yet, so we remain on sobol From c97487b1e938257934d74f240c325e56c124d4e8 Mon Sep 17 00:00:00 2001 From: Mia Garrard Date: Tue, 21 May 2024 18:54:03 -0700 Subject: [PATCH 2/2] Clean up unneccesary todos in _get_model_state_from_last_generator_run (#2476) Summary: Since we decided to keep a 1 fitted model : 1 node relationship the note about this being uncompatable with GenNodes is irrelevant now, easy compatability :) Also removed the todo about moving to gennodes - i think it's okay where it is, but if we want to move _fit_current_model and _get_model_state_from_last_generator_run to GenerationNode, that doesn't seem unreasonable and we can keep the todo Reviewed By: lena-kashtelyan Differential Revision: D57444926 --- ax/modelbridge/generation_strategy.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/ax/modelbridge/generation_strategy.py b/ax/modelbridge/generation_strategy.py index 170d20208d5..5b43f453fe8 100644 --- a/ax/modelbridge/generation_strategy.py +++ b/ax/modelbridge/generation_strategy.py @@ -887,11 +887,6 @@ def _maybe_transition_to_next_node( def _get_model_state_from_last_generator_run(self) -> Dict[str, Any]: lgr = self.last_generator_run - # NOTE: This will not be easily compatible with `GenerationNode`; - # will likely need to find last generator run per model. Not a problem - # for now though as GS only allows `GenerationStep`-s for now. - # Potential solution: store generator runs on `GenerationNode`-s and - # split them per-model there. model_state_on_lgr = {} # Need to check if model_spec_to_gen_from is none to account for # ExternalGenerationNodes which leverage models from outside Ax. @@ -910,8 +905,6 @@ def _get_model_state_from_last_generator_run(self) -> Dict[str, Any]: if grs_equal and lgr._model_state_after_gen: if self.model or isinstance(model_on_curr, ModelRegistryBase): - # TODO[drfreund]: Consider moving this to `GenerationStep` or - # `GenerationNode`. model_cls = ( self.model.model.__class__ if self.model is not None