From 15d312207647de2445f5ff5e13f52bde42cd9db6 Mon Sep 17 00:00:00 2001 From: Mia Garrard Date: Wed, 10 Jan 2024 12:37:26 -0800 Subject: [PATCH] Update `trials_as_df` method to support GenNode based GS (#2095) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/2095 This diff modifies the trials as df method to support representation of GenerationNodes as well as GenerationSteps in GSs In coming diffs: (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 (7) remove conditionals for legacy usecase ( clean up any lingering todos Reviewed By: bernardbeckerman Differential Revision: D52273582 fbshipit-source-id: 740ef0e17456e1a2a41d4472ae6432552b4b7977 --- ax/modelbridge/generation_strategy.py | 20 +++--- .../tests/test_generation_strategy.py | 63 +++++++++++++++++++ 2 files changed, 76 insertions(+), 7 deletions(-) diff --git a/ax/modelbridge/generation_strategy.py b/ax/modelbridge/generation_strategy.py index a2183e0fc7e..816dab813ea 100644 --- a/ax/modelbridge/generation_strategy.py +++ b/ax/modelbridge/generation_strategy.py @@ -276,10 +276,11 @@ def uses_non_registered_models(self) -> bool: return not self._uses_registered_models @property - @step_based_gs_only def trials_as_df(self) -> Optional[pd.DataFrame]: """Puts information on individual trials into a data frame for easy - viewing. For example: + viewing. + + For example for a GenerationStrategy composed of GenerationSteps: Gen. Step | Model | Trial Index | Trial Status | Arm Parameterizations 0 | Sobol | 0 | RUNNING | {"0_0":{"x":9.17...}} """ @@ -293,11 +294,15 @@ def trials_as_df(self) -> Optional[pd.DataFrame]: len(step.trials_from_node) == 0 for step in self._nodes ): return None + + step_or_node_col = ( + "Generation Node" if self.is_node_based else "Generation Step" + ) records = [ { - "Generation Step": step.node_name, + step_or_node_col: node.node_name, "Generation Model": self._nodes[ - step_idx + node_idx ].model_spec_to_gen_from.model_key, "Trial Index": trial_idx, "Trial Status": self.experiment.trials[trial_idx].status.name, @@ -306,12 +311,13 @@ def trials_as_df(self) -> Optional[pd.DataFrame]: for arm in self.experiment.trials[trial_idx].arms }, } - for step_idx, step in enumerate(self._nodes) - for trial_idx in step.trials_from_node + for node_idx, node in enumerate(self._nodes) + for trial_idx in node.trials_from_node ] + return pd.DataFrame.from_records(records).reindex( columns=[ - "Generation Step", + step_or_node_col, "Generation Model", "Trial Index", "Trial Status", diff --git a/ax/modelbridge/tests/test_generation_strategy.py b/ax/modelbridge/tests/test_generation_strategy.py index 5f7a31b29a9..a07069be70f 100644 --- a/ax/modelbridge/tests/test_generation_strategy.py +++ b/ax/modelbridge/tests/test_generation_strategy.py @@ -566,6 +566,69 @@ def test_trials_as_df(self) -> None: "GenerationStep_1", ) + # construct the same GS as above but directly with nodes + sobol_model_spec = ModelSpec( + model_enum=Models.SOBOL, + model_kwargs={}, + model_gen_kwargs={}, + ) + node_gs = GenerationStrategy( + nodes=[ + GenerationNode( + node_name="sobol_2_trial", + model_specs=[sobol_model_spec], + transition_criteria=[ + MaxTrials( + threshold=2, + not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED], + block_gen_if_met=True, + block_transition_if_unmet=True, + transition_to="sobol_3_trial", + ) + ], + gen_unlimited_trials=False, + ), + GenerationNode( + node_name="sobol_3_trial", + model_specs=[sobol_model_spec], + transition_criteria=[ + MaxTrials( + threshold=2, + not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED], + block_gen_if_met=True, + block_transition_if_unmet=True, + transition_to=None, + ) + ], + gen_unlimited_trials=False, + ), + ] + ) + self.assertIsNone(node_gs.trials_as_df) + # Now the trial should appear in the DF. + trial = exp.new_trial(node_gs.gen(experiment=exp)) + + self.assertFalse(node_gs.trials_as_df.empty) + self.assertEqual( + node_gs.trials_as_df.head()["Trial Status"][0], + "CANDIDATE", + ) + # Changes in trial status should be reflected in the DF. + trial._status = TrialStatus.RUNNING + self.assertEqual(node_gs.trials_as_df.head()["Trial Status"][0], "RUNNING") + # Check that rows are present for step 0 and 1 after moving to step 1 + for _i in range(3): + # attach necessary trials to fill up the Generation Strategy + trial = exp.new_trial(node_gs.gen(experiment=exp)) + self.assertEqual( + node_gs.trials_as_df.head()["Generation Node"][0], + "sobol_2_trial", + ) + self.assertEqual( + node_gs.trials_as_df.head()["Generation Node"][2], + "sobol_3_trial", + ) + def test_max_parallelism_reached(self) -> None: exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy(